DriftNet detects and locates prompt injection attacks in AI agents
DriftNet: A Dual-Head Trajectory Transformer for Detecting and Localizing Prompt Injection in LLM Agents
Cryptography and SecurityArtificial IntelligenceMachine Learning
Summary
When AI agents are tricked by sneaky instructions hidden in their tasks, these attacks show up in their actions. The authors present DriftNet, a tool that reads the agent’s sequence of steps and tells if an attack happened, where it started, which steps were affected, and if the agent resisted it. DriftNet works without seeing inside the AI's brain, using only the logged step data, and it identifies attacks very accurately on a test set. This helps operators understand and respond to hidden attacks on AI agents more precisely.
What this means in practice
- •For ai security teams: Monitor AI agents’ behavior logs to detect and pinpoint where prompt injection attacks occur during operation.
- •For software engineers integrating ai tools: Improve reliability by identifying and isolating steps in AI workflows corrupted by malicious inputs without accessing the AI model internals.
Authors
Asif Pinjari, Mithun Paul Saint-Germain
Abstract
When an indirect prompt injection succeeds against an LLM agent, the compromise is visible in the agent's own behavior: a benign prefix of tool calls, a poisoned observation, and a suffix of actions that serve the attacker. An operator needs three facts: where the attack entered, which steps it corrupted, and whether apparent poison was resisted. Existing systems return either a whole-trace verdict or a single unsafe index. We present DriftNet, a dual-head trajectory Transformer that reads a logged tool-call trajectory and answers all three questions in one forward pass: one head classifies the trajectory as compromised or not, and a second assigns every step one of four labels (benign, injection point, hijacked, failed injection). To our knowledge it is the first supervised detector to produce this joint output. A frozen sentence encoder and four identity-free world features embed each step; the trained trunk, under two million parameters and optimized with a class-weighted joint objective over both heads, needs no access to the agent's model. On the task-disjoint split of the AgentDrift benchmark (12,536 trajectories, 71,024 labeled steps), with a 20-configuration sweep bounding hyperparameter sensitivity to 0.011 F1 and the test part evaluated exactly once, DriftNet reaches trajectory-level F1 of 0.983, exact injection-point recovery on 98.7% of attacked trajectories, hijacked-span IoU of 0.979, zero flags on 218 resisted attacks, and 2.9% flags on hard negatives. A surface baseline retrained on the identical split recovers 11.1% of partial hijacks and 17.1% of delayed executions; DriftNet reaches 98.6% and 93.2% while lowering every false-alarm rate. Reading all 26 residual errors shows that most misses trace to trajectories whose labeled injection observation carries no legible instruction, and we report the benchmark's measured world-identity regularity alongside the results.