TaskTracker Generalization to Agentic Benchmarks

Evaluating whether TaskTracker (Abdelnabi et al., 2024) generalizes as a prompt injection detector when tested on agentic benchmarks it was never trained on.

Motivation

TaskTracker trains linear probes on LLM hidden-state activations to detect prompt injection attacks. The original paper reports strong results across several datasets (TensorTrust, Ignore-This, the authors' own injected samples, etc.). But the evaluation setup is entirely static: single-turn inputs, no tool calls, no multi-step agent traces. It is unclear whether the features the probe learns transfer to more realistic agentic settings where injections show up inside tool outputs or multi-turn conversations.

Two benchmarks are good candidates for testing this:

• InjecAgent (Zhan et al., 2024) - prompt injections embedded in tool outputs within agent pipelines. Covers both direct and indirect injection across multiple tool-use scenarios.
• AgentDojo (Debenedetti et al., 2024) - a framework for evaluating injection robustness in agentic systems with function calling. Includes benign tasks alongside adversarial injections.

Neither was part of TaskTracker's training or evaluation data.

Approach

We take the pretrained TaskTracker probe and apply it directly to inputs drawn from InjecAgent and AgentDojo. No retraining, no adaptation. Just: does the probe, as-is, detect injections in these new settings?

Input formatting

TaskTracker expects a specific three-part input: a system prompt, a user instruction, and a data segment (retrieved content or tool output where injections might appear). Both InjecAgent and AgentDojo have richer structures (tool call sequences, multi-turn conversations) that need to be flattened into this format. The details are in src/data_loader.py.

This is a real limitation. Flattening multi-step interactions into a single data segment throws away ordering and context that a real agent would have. We note this but proceed with the flattened format since that is what the probe was trained on.

Project structure

src/
  data_loader.py   - loading and formatting InjecAgent + AgentDojo
  eval.py          - run probe, compute ROC AUC / FPR / FNR
data/              - raw benchmark data (not checked in)
results/           - evaluation outputs

Current status

Work in progress. Data loaders are done. Evaluation pipeline runs end-to-end. Still collecting results and working through some formatting edge cases in AgentDojo's export format.

References

• Abdelnabi et al. (2024). "You Are What You Read: TaskTracker - An AI-Assisted Framework for Detecting Prompt Injection Attacks." Microsoft Research.
• Zhan et al. (2024). "InjecAgent: Benchmarking Indirect Prompt Injections in Tool-Integrated LLM Agents."
• Debenedetti et al. (2024). "AgentDojo: A Dynamic Environment to Evaluate Attacks and Defenses for LLM Agents."

Reproducing

pip install -r requirements.txt
# see data/README.md for data setup
python src/eval.py --dataset injecagent --probe-path /path/to/tasktracker/probe.pt
python src/eval.py --dataset agentdojo --probe-path /path/to/tasktracker/probe.pt