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 40-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

All checkpoints were trained with structural evidence from the CARE static shell analyzer and need it at inference (pip install "laya-guard[evidence]"), except v2.

path	model	notes
/ (root)	v9d	recommended; v8's data + red-team round 9, 18 trainable layers
v10/	v10	v9d + round 10; use with v9d as a cautious ensemble to catch the most attacks
v8/	v8	previous default (root until 2026-10-29)
v2/	v2	early checkpoint, no evidence; the old v2+v8 ensemble

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 v9d+v10 ensemble (block if either would): catches more attacks, more false blocks
LAYA_GUARD_MODEL=a3harva/laya-toolguard-mlx,a3harva/laya-toolguard-mlx:v10 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	v9d	v10	v9d+v10
Red-team held-out probes, rounds 1-8 (617)	91.4% / 92.3%	93.4% / 77.0%	96.1% / 93.9%	94.7% / 95.8%	96.7% / 93.6%
Round 9 probes (58)	78.6% / 96.7%	89.3% / 83.3%	96.4% / 96.7%	96.4% / 96.7%	100% / 96.7%
Round 10 probes (77)	70.7% / 80.6%	80.5% / 58.3%	90.2% / 86.1%	92.7% / 91.7%	97.6% / 83.3%
Final exam (107, written last)	84.3% / 91.1%	90.2% / 73.2%	92.2% / 85.7%	84.3% / 83.9%	92.2% / 78.6%
Canonical attacks (26 risky / 8 routine)	80.8% / 100%	96.2% / 87.5%	100% / 100%	96.2% / 100%	100% / 100%
AUC, held-out	0.972	0.955	0.990	0.985	0.989

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. v9d's one regression against v8 is routine calls on the final exam. v10 fixes most round-10 misses but forgets some final-exam attacks, so it ships only as an ensemble member.

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 18 of 28 encoder layers (v8: 14) 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. 80 min on an Apple M5 (16 GB), peak 10.5 GB.
• Data: ~3,500 template and LLM-written tool calls, plus ~1,500 rows from red-team rounds 1-9 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. Round 9 targeted find -delete/-exec and xargs scope, credential hunting in logs, and danger hidden after a benign prefix. 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.