Skip to content

DRIFT: Dynamic Rule-Based Defense with Injection Isolation for Securing LLM Agents

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

DRIFT is a system-level defense framework that protects LLM agents from prompt injection attacks by combining pre-execution constraint planning, runtime dynamic policy validation, and post-execution memory sanitization. On AgentDojo with GPT-4o-mini, it reduces attack success rate from 30.7% to 1.3% while outperforming the strongest security baseline (CaMeL) by 21.8% on task utility.

Problem

LLM agents that interact with untrusted external environments (tools, web pages, emails) are vulnerable to prompt injection attacks, where adversarial content in tool responses hijacks agent behavior. Existing system-level defenses either enforce static, manually-crafted security policies (sacrificing utility on complex/long-trajectory tasks) or fail to eliminate injected content from the memory stream, leaving residual risk during prolonged interactions.

Method

DRIFT enforces security through three layered components operating in sequence:

  1. Secure Planner (pre-interaction): Before any tool call, an LLM analyzes the user query and generates (a) a minimal function trajectory (ordered list of expected function calls) as control-level constraints and (b) a JSON-schema-style parameter checklist per function node as data-level constraints.

  2. Dynamic Validator (per tool call): After each tool-call request is generated, validates it against both control and data constraints. Deviations are classified by OS-inspired privilege levels (Read / Write / Execute). Read-privilege deviations are auto-approved; Write/Execute deviations trigger an LLM intent-alignment check against the original query. Approved deviations are folded back into the running trajectory, allowing the policy to adapt without breaking legitimate multi-step tasks.

  3. Injection Isolator (post tool response): Inspects each tool response for instructions conflicting with the original user query using an LLM classifier. Detected injections are externally masked before being written to the agent's memory stream, preventing accumulated injection risk over long interactions.

A trainable variant fine-tunes Qwen2.5-7B-Instruct (LoRA) on 1,000 Planner samples and 1,000 Isolator samples derived from ToolBench, with tool environments augmented to 10,000+ unique tools to simulate realistic scale.

Key Contributions

  • DRIFT framework integrating control-level trajectory constraints, data-level parameter checklists, dynamic privilege-based policy updating, and memory-stream injection masking into a single cohesive pipeline.
  • Demonstration that static policies collapse on tasks with trajectory length ≥ 3, motivating the dynamic constraint-update mechanism.
  • Trainable policy via LoRA fine-tuning that drives ASR to 0.0% while improving utility relative to the base tuned model.
  • Empirical validation across five LLMs (GPT-4o, GPT-4o-mini, Claude-3.5-Sonnet, Claude-3-Haiku, Qwen2.5-7B) and two benchmarks (AgentDojo, ASB).

Results

AgentDojo (GPT-4o-mini, important_instruction attack): - ASR: 30.7% (undefended) → 1.3% (DRIFT); CaMeL achieves 0.0% ASR but at 7× token cost - Utility (no attack): DRIFT 58.5% vs. CaMeL 35.4% (+21.8% pp); vs. undefended 63.6% - Utility (under attack): DRIFT 47.9% vs. CaMeL 37.1% (+10.8% pp)

ASB (GPT-4o-mini, OPI attack): - ASR: DRIFT 4.8% vs. Progent 15.8% (best competing baseline) - Utility maintained near undefended levels

Cross-model generalization (AgentDojo): - GPT-4o: ASR 51.7% → 1.5%; Claude-3.5-Sonnet: ASR ~11% → 4.4% - All five tested models reduced to single-digit ASR

Qwen2.5-7B fine-tuned policy: - ASR: 15.1% → 0.0%; Utility (no attack): 26.6% → 32.2%; Utility (under attack): 19.1% → 22.2%

Adaptive attacks (Table 2): Combined isolator + validator adaptive attack raises ASR by only 0.81%; PAIR attack raises ASR by 0.31%.

Token cost: DRIFT uses 2.37M tokens vs. 6.09M for CaMeL; efficiency metric (Utility−ASR / Total Tokens) = 19.7 vs. 5.8 for CaMeL.

Limitations

  • Dynamic Validator and Injection Isolator each require an additional LLM call per tool invocation, increasing latency and token cost (~1.89× over undefended baseline).
  • Secure Planner must accurately decompose the user query before any environment interaction; ambiguous or underspecified queries may produce incomplete initial trajectories.
  • The Injection Isolator relies on an LLM to detect conflicting instructions and cannot directly modify tool outputs, requiring an external masking layer — potential gap if masking is bypassed.
  • Training data for the fine-tuned variant is generated by GPT-4o-mini rewriting ToolBench conversations, which may introduce distributional noise or inherit GPT-4o-mini's own biases.
  • Evaluation is restricted to text-based agentic settings; applicability to multimodal or computer-use agents is not assessed.

Relevance to Agentic AI / LLM Agents

Prompt injection is one of the most pressing unsolved security problems for deployed LLM agents, and DRIFT directly targets the failure mode of tool-augmented agents operating in adversarial environments. The framework's design — plan-before-act constraint generation plus runtime dynamic policy update — mirrors standard operating-system privilege models, offering a principled architecture pattern for agent sandboxing that could generalize beyond injection defense. The memory isolation component addresses a subtle but critical risk for long-horizon agents: accumulated context poisoning that bypasses action-level controls. For researchers building or evaluating agentic systems, DRIFT provides both a practical defense and a structured decomposition (Planner / Validator / Isolator) that clarifies where in an agent's execution loop different attack surfaces reside.