AnShi-HW/tool-call-anomaly-mmbert-small-v1

Tool-call anomaly detector (mmBERT-small)

Binary classifier over an agent session: the user task plus the sequence of tool calls the agent made. Output is the probability that the session contains an unsafe, unauthorized or task-deviating tool call.

Input

A session is serialized as one string:

TASK: <task text> CALLS: tool_a({"arg": "value"}) tool_b({...}) ...

Arguments are JSON, cut at 200 characters per call. The string is truncated to 1024 tokens. score_sessions.py in this repo does the serialization; use it rather than reimplementing it.

Usage

pip install torch transformers numpy tqdm
python score_sessions.py --model AnShi-HW/tool-call-anomaly-mmbert-small \
    --file stepshield_eval.jsonl agentdyn_eval.jsonl --threshold 0.5299

Input files: one JSON record per line with task and tool_calls ([{"name": ..., "args": {...}}, ...]). If records carry label (1 attack, 0 benign) and source, the script prints recall, false-positive rate and AUROC per source at the given threshold. Evaluation files are in AnShi-HW/tool-call-anomaly-data.

Threshold

Recommended: 0.5299. This is the score at which 5% of the bfcl benign set (7,927 sessions the model never trained on) is flagged. The value is also stored in classifier_meta.json as recommended_threshold. The threshold field in the same file is the training-set calibration and is not the one to use. Different traffic will need a different threshold; set it on a benign sample of your own.

Results at threshold 0.5299

source	seen in training	benign	attack	AUROC	recall	FPR
harm_actions	yes	23	32	0.990	0.906	0.000
agent_trust	yes	44	52	0.930	1.000	0.386
stepshield	yes	2,562	108	0.841	0.759	0.215
agentdyn	no	9,806	1,310	0.880	0.651	0.087
agentdojo	no	6,413	1,217	0.751	0.064	0.021
bfcl (benign only)	no	7,927	0			0.050
injecagent (benign only)	no	1,054	0			0.000

Eval rows for the seen sources are held out and session-disjoint from training. FPR is measured on each source's own benign rows.

Known weaknesses: high false-positive rate on agent_trust and stepshield benign traffic; low recall on agentdojo (prompt-injection attacks, a family absent from training); results on unseen sources vary across training seeds.

Training

• Student: jhu-clsp/mmBERT-small, sequence classification head, max length 1024, 4 epochs, lr 2e-5, batch 16, warmup 0.1, seed 42.
• Data: stepshield, harm_actions, agent_trust (attack sources) plus seven benign tool-call corpora (agent_align, funcdex_mt, glaive_fc_v2, apigen_mt, toolace, hermes_fc, toolmind_toolace_query). 227,157 benign and 1,102 attack sessions.
• Loss: class-weighted cross-entropy on hard labels, plus a soft loss (alpha 3.0) matching Qwen2.5-7B-Instruct zero-shot scores on a 49,000-row benign sample and all attacks, excluding stepshield rows from the soft term.

Files

• model weights and tokenizer
• classifier_meta.json: training settings and thresholds
• score_sessions.py: scorer and evaluator