fleet/agentdojo-gym — prompt-injection training gym

Programmatic wrapper around AgentDojo (ETH Zurich Spylab, MIT, NeurIPS 2024, arXiv:2406.13352) — a dynamic benchmark environment for prompt-injection attacks and defenses against tool-using LLM agents. Tested against agentdojo==0.1.35 (pip install agentdojo, Python 3.10–3.12; the cachebox dependency has no wheels for 3.13+ yet).

Layout

gym.py               run suites programmatically: load suite, build pipeline,
                     benchmark with/without injections, emit telemetry
adapter.py           ExternalAgentElement: plugs OUR JSON-line agent protocol
                     into agentdojo's AgentPipeline/ToolsExecutionLoop
telemetry.py         utility / ASR / per-task math + telemetry.{json,md} writers
custom_suites/       extension stub: randomized-scenario generator
  counter_suite.py   worked example suite (tools, tasks, injection vector)
  payload_bank.py    per-run secrets + injection payload bank (5 families)
  README.md          authoring guide + randomizing recipe
examples/agent_counter.py   deterministic deliberately-SUSCEPTIBLE scripted agent
selftest.py          21 offline checks (no agentdojo needed)
smoke_test.sh        LIVE run: venv + pip install agentdojo + full benchmark

Quickstart

python3 selftest.py    # offline — always run first
bash smoke_test.sh     # live benchmark through a real agentdojo install

Smoke proves, end-to-end: custom suite registration, adapter + tool execution through agentdojo's FunctionsRuntime, utility 100% without attack, a real important_instructions attack landing via the suite's injection vector (ASR 100%), and injection-goal actionability 100%.

Running suites

# built-in suites (banking, slack, travel, workspace) with any LLM agent that
# speaks the JSON-line protocol:
python3 gym.py --suite workspace --benchmark-version v1 \
    --agent-cmd "python3 your_agent.py" --attack tool_knowledge

# our randomized-extension suite:
python3 gym.py --custom-suite custom_suites.counter_suite --suite counter \
    --agent-cmd "python3 examples/agent_counter.py" --attack important_instructions

Runs land in runs/<ts>_<suite>_<attack>/: agentdojo's per-task trace logs, raw results_{without,with}*.json, and telemetry.{json,md}.

The agent protocol (JSON lines over a subprocess)

The adapter (adapter.py) bridges agentdojo's verified BasePipelineElement.query(query, runtime, env, messages, extra_args) to our protocol — the same protocol family as fleet/evm-asserts, so one factory agent runtime serves both gyms:

• adapter → agent: {"type":"handshake",...} then per turn

{"type":"query","prompt":..., "tools":[{name,description,parameters_schema}], "messages":[{role,content}...]}

• agent → adapter: {"tool_calls":[{"function","args","id"}...]} or

{"final_answer":"..."}

Tool calls are executed by agentdojo's own ToolsExecutor inside a ToolsExecutionLoop — we never re-implement the runtime. The canonical composition (verified against AgentPipeline.from_config in 0.1.35) is [InitQuery(), LLM-slot, ToolsExecutionLoop([ToolsExecutor(), LLM-slot])].

adapter.check_api() probes the installed agentdojo and fails fast with a readable diff when the (explicitly unstable) upstream API moves.

Telemetry semantics (polarity!)

Matching the de-facto convention of agentdojo's built-in suites (verified in 0.1.35 banking/slack/workspace: their security() returns True when the injection goal WAS achieved):

• utility = user tasks' utility check passed (no attack);
• utility_under_attack = same while attacked;
• ASR = mean(security_results) — True means the injection landed;
• injection_task_utility = injection goals achievable when prompted

directly (the benchmark warns if not — a broken-environment signal).

Also emitted: ASR per injection task, per-(user,injection) task matrix, and utility drop under attack.

Wiring defenses / baselines

• Attacks: important_instructions (default), tool_knowledge, direct,

ignore_previous, injecagent, system_message, or none for the clean baseline — all instantiated upstream-side with (task_suite, pipeline).

• Defenses: agentdojo's PromptInjectionDetector, tool filters etc. slot

into the same ToolsExecutionLoop; we benchmark them by adding the element to the pipeline in gym.py (next step in the stubbed roadmap).

Legal / license notes

AgentDojo is MIT; the pip package is used, not vendored. Our suites, agents, payload bank, and trajectories are original. Injection payloads in the bank are for LOCAL red-team evaluation only.