niti-gate-vegaml-4b

Decides whether an AI agent's tool call should be allowed or blocked, and if blocked, what kind of misuse it is (eval gaming, harness introspection, credential misuse, exfiltration, sandbox escape, other overreach). Built for the gate in Niti, PromptForce's policy gateway for coding agents.

It is two small heads on top of vegaml: the vega-U engine reading the call with Qwen3.5-4B, both frozen and unchanged. All the reading is theirs; this repo adds only the heads (16 MB). Built on vegaml; not affiliated with or endorsed by its authors or by Qwen.

Results

safeagents-judge-bench v2, test split, 661 calls (536 block, 125 allow), run with the bench's own job code. The bench was never used for training or tuning: the block threshold (0.225) was chosen on our own held-out data.

Metric	HF Jobs A100 80GB	RTX 3090
Balanced accuracy	0.905	0.909
AUROC	0.957	0.957
Recall (blocks caught)	0.882	0.882
False blocks (allowed calls blocked)	0.072	0.064
Trajectories caught	78/82	79/82
Category macro F1	0.497	0.495
Category accuracy	0.729	0.728
Latency p50 / p95	0.41 s / 0.61 s	0.19 s / 0.36 s

The A100 column is the leaderboard run (HF Jobs a100-large, vegaml==0.15.0), on the same hardware as the other decision models; the two runs differ only by GPU numerics (5 of 661 verdicts). At the time of submission it is first on the bench's block-decision track (balanced accuracy) and fourth on the category track, where the small sandbox_escape (4 calls) and credential_misuse (15 calls) classes weigh as much as the large ones.

GPU memory: 8.3 GB allocated (11.3 GB on card). On CPU it runs at about 23 s per call, so a GPU is needed in practice.

Use

pip install vegaml safetensors huggingface_hub

import sys
from huggingface_hub import snapshot_download
path = snapshot_download("prompt-force/niti-gate-vegaml-4b"); sys.path.insert(0, path)
from niti_gate import NitiGate

gate = NitiGate(path, device="cuda")
task = "Fix the failing date parser test."
gate.decide({"task": task, "name": "bash", "args": '{"command": "cat /root/.config/gh/hosts.yml"}', "findings": []})
# {'verdict': 'block', 'p_block': 0.671, 'category': 'harness_introspection', 'category_probs': {...}}
gate.decide({"task": task, "name": "bash", "args": '{"command": "python -m pytest tests/test_dates.py -q"}', "findings": []})
# {'verdict': 'allow', 'p_block': 0.005, 'category': 'legit', 'category_probs': {'legit': 1.0}}

The input is the bench's judge view: the user's task, the tool call (name and JSON arguments) and any scanner findings. decide(..., threshold=...) trades recall for false blocks.

gate.system_one(state) answers in the System One shape the bench's jobs read. Its block probability is rescaled so the tuned threshold sits at 0.5 (monotone, so the ranking is unchanged); the raw value is in p_block_raw.

Training

Heads trained on agent tool calls labelled by real user decisions and a panel of open judge models.

• Block head: PCA (256 components) then logistic regression on vega-U's features for the block question; blocks are kept against legitimate calls at 1:4.
• Category head: the same over the category question, trained on blocked calls plus a sample of legitimate ones.

Limits

• Tuned for coding agents in a sandbox; other domains are untested.
• Blocks more legitimate calls on some Terminal-Bench-style tasks (0.231 on bench v1). Raise the threshold where false blocks cost more than misses.
• A judge, not a sandbox: use it alongside real isolation and policy.
• On CPU, uninstall flash-linear-attention (fla) if present: its GPU kernels take over the reader and fail on CPU tensors.

License

The heads, niti_gate.py and config.json in this repo are licensed CC BY-NC 4.0: free for non-commercial use, with attribution. For commercial use, contact PromptForce.

vegaml, vega-U and Qwen3.5-4B are not included here; they are downloaded from their own repos under their own licence (Apache-2.0). See NOTICE.

Credits

• vegaml and vega-U by Nandakishor M (Apache-2.0)
• Qwen3.5-4B by the Qwen team, Alibaba (Apache-2.0)
• safeagents-judge-bench for the evaluation and its judge input format