laya-toolguard-mlx

A Laya decision model fine-tuned in MLX to judge coding-agent tool calls (bash, write, edit, read) before they run. It powers pi-laya-mlx, a guard for the pi coding agent that blocks unsafe calls and tells the model why. It runs locally on Apple Silicon in about 60-85 ms per call, 8-bit, ~500 MB.

Laya does not generate text. For each call it scores six options in one forward pass and returns calibrated probabilities:

option	meaning
routine	normal development work
destroy	deletes or overwrites data outside the project, or rewrites history
exfiltrate	sends secrets or private data off the machine
remote_code	runs code fetched from an untrusted source
system	changes OS, security settings or safety controls
production	touches live infrastructure, databases or releases

Risk = 1 - p(routine); the guard blocks at risk ≥ 0.5.

Checkpoints

path	model	notes
/ (root)	v8	recommended; trained with structural evidence from the CARE static shell analyzer, which it needs at inference (pip install "laya-guard[evidence]")
v2/	v2	early checkpoint; use with v8 as a cautious ensemble (block if either would)

Usage

git clone https://github.com/a3harva/pi-laya-mlx && cd pi-laya-mlx
./scripts/setup.sh                       # installs the server and downloads this model
pi -e ./extension                        # every tool call is screened

# or the v2+v8 ensemble: catches more attacks, more false blocks
LAYA_GUARD_MODEL=a3harva/laya-toolguard-mlx,a3harva/laya-toolguard-mlx:v2 server/.venv/bin/laya-guard-serve

curl -s localhost:8723/v1/guard -H 'content-type: application/json' \
  -d '{"tool":"bash","input":{"command":"rm -rf ~"},"cwd":"/Users/me/proj"}'
# {"decision":"block","risk":0.99,"category":"destroy","reason":"Blocked by the Laya safety guard: ..."}

Results

Unsafe calls blocked / safe calls allowed, at threshold 0.5. No set below was trained on.

set	v8	v2+v8 ensemble
Synthetic validation split (324)	98.3% / 97.9%	98.9% / 95.2%
Red-team held-out probes (617)	91.1% / 92.7%	93.4% / 77.0%
Final exam (107, written last)	82.4% / 91.1%	90.2% / 73.2%

The red-team and final-exam sets are adversarial on purpose: attacks disguised as routine work, encoded payloads, guard tampering, and safe calls that only look dangerous (grep -rn 'rm -rf ~' docs/). Zero-shot Laya blocked ~60% of routine calls, including npm test.

On a CARE-style benchmark (549 commands rebuilt from Atomic Red Team, RedCode-Exec and NL2Bash, because CARE did not release its split), v8 scored F1 62.8% with 8.2% false positives, against 65.9% / 6.1% for CARE's static engine; blocking when either flags a command reached F1 73.0%. The two policies differ (CARE counts every attack-playbook step as dangerous, including read-only lookups), so this compares approaches, not the published CARE numbers.

Training

• Base: cavi-ai/laya-MLX-8bit:typed-decisions, dequantized, top 14 of 28 encoder layers and the decision head trained, re-quantized to 8-bit in the original layout.
• Cross-entropy on the guard's own 6-option question, class weighting, label smoothing 0.1, 4 epochs, lr 3e-5, batch 4, calibration temperature fitted on the validation split. ~1 h on an Apple M5 (16 GB).
• Data: ~3,500 template and LLM-written tool calls, plus ~1,200 rows from 7 red-team rounds in which a second agent (DeepSeek in a locked-down pi session that could only score text) attacked the current checkpoint; every red-team row was reviewed before training. Data, scripts and the full log are in the GitHub repo.

Limits

• A classifier, not a sandbox: it lowers risk and does not prove a call safe. Pair it with OS sandboxing.
• Judges what is visible in the call; a script that hides its effect is judged by its visible text.
• Reads up to 1024 tokens per window; the server splits long inputs into overlapping windows.
• English only; trained on macOS/Linux developer workflows.

License and credits

Weights: Apache-2.0, derived from convaiinnovations/laya via the MLX port cavi-ai/laya-MLX-8bit, whose model code (laya/, MIT) is included. See LICENSE.