PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Hao Li, Xiaogeng Liu, Hung-Chun Chiu, Dianqi Li, Ning Zhang, Chaowei Xiao

arXiv:2506.12104v3cs.CRcs.AI

TL;DR

Prompt injection threatens LLM agents that use tools in untrusted environments, and static defenses do not fully address dynamic policy updates or memory-stream risks. DRIFT combines query-derived planning, privilege-aware validation, and injection isolation; on AgentDojo, it reduced ASR from 30.7% to 1.4% while improving utility over CaMeL and generalizing across models.

  • Problem

    LLM agents can be misled by prompt injections from untrusted environments, while existing static system-level defenses leave dynamic policy and memory-isolation gaps.

  • Method

    DRIFT derives a minimal function trajectory and parameter checklist from the user query, validates deviations by privilege and intent, and masks conflicting instructions in memory.

  • Results

    30.7% to 1.4%: DRIFT reduced Attack Success Rate on AgentDojo with GPT-4o-mini, while utility exceeded CaMeL by 18.9% without attack and 15.5% under attack.

  • Takeaways & Limitations

    DRIFT provides a fully automatic system-level defense with strong security and utility across diverse scenarios and models.

  • Takeaways & Limitations

    AgentDojo’s domains do not fully cover the diverse tasks and attack scenarios encountered in real-world agentic systems.

Abstract

from arXiv · show

Large Language Models (LLMs) are increasingly central to agentic systems due to their strong reasoning and planning capabilities. By interacting with external environments through predefined tools, these agents can carry out complex user tasks. Nonetheless, this interaction also introduces the risk of prompt injection attacks, where malicious inputs from external sources can mislead the agent's behavior, potentially resulting in economic loss, privacy leakage, or system compromise. System-level defenses have recently shown promise by enforcing static or predefined policies, but they still face two key challenges: the ability to dynamically update security rules and the need for memory stream isolation. To address these challenges, we propose Dynamic Rule-based Isolation Framework for Trustworthy agentic systems (DRIFT), which enforces the dynamic security policy and injection isolation for securing LLM agents against prompt injection attacks. A Secure Planner first constructs a minimal function trajectory and a JSON-schema-style parameter checklist for each function node based on the user query. A Dynamic Validator then monitors deviations from the original plan, assessing whether changes comply with privilege limitations and the user's intent. Finally, an Injection Isolator detects and masks any instructions that may conflict with the user query from the memory stream to mitigate long-term risks. We empirically validate the effectiveness of DRIFT on the AgentDojo, ASB, and AgentDyn benchmark, demonstrating its strong security performance while maintaining high utility across diverse models, showcasing both its robustness and adaptability. The project website is available at https://safo-lab.github.io/DRIFT.

1 Introduction

LLM agents face prompt injection risks when interacting with untrusted environments, while existing system-level defenses rely largely on static policies. DRIFT addresses dynamic policy enforcement and memory isolation, achieving strong security and utility across models.

  • Motivation: Prompt injection attacks exploit untrusted external inputs to mislead LLM-agent workflows during tool-mediated interactions.Such attacks can arise from malicious instructions embedded in third-party content.
  • Motivation: Existing defenses include model-level guardrails and system-level static policies, but static restrictions can limit agent flexibility.System-level methods improve reliability against unseen attacks by constraining action spaces before interaction.
  • Motivation: Dynamic policy enforcement must limit validator exposure to untrusted inputs while adapting permissions to the risks of different actions.The policy mechanism itself can become an attack surface if adversaries influence policy decisions.
  • Approach: DRIFT combines a Secure Planner, Dynamic Validator, and Injection Isolator to enforce adaptive constraints and isolate conflicting memory instructions.The planner derives a minimal trajectory and parameter checklist; validation checks deviations against user intent and privilege limits; isolation masks conflicting instructions.
  • Evaluation: 30.7% to 1.4%: DRIFT reduced Attack Success Rate on AgentDojo with GPT-4o-mini while outperforming CaMeL in utility under and without attack.Utility outperformed CaMeL by 18.9% under no attack and 15.5% under attack; performance generalized across four advanced agents.

2 Related Works

Research on securing LLM agents against prompt injection spans model-level and system-level defenses. System-level approaches constrain actions or information flow, but earlier methods often rely on predefined policies or face transfer challenges.

  • Defense Taxonomy: Prompt-injection defenses are commonly classified into model-level and system-level approaches.This classification frames the related work surveyed for tool-integrated agents.
  • Model-Level Defenses: Model-level methods strengthen intrinsic model resistance through structured queries, preference optimization, or external injection detectors.Examples include StruQ, preference optimization, LlamaGuard, and InjecGuard.
  • System-Level Defenses: System-level defenses constrain the model’s action space with predefined security policies, but coding-focused methods can transfer poorly to tool-integrated agents.These approaches address security outside the model itself.
  • System-Level Defenses: Recent system-level approaches isolate application execution, propagate untrusted-data labels, or construct control and data flows from the original query.The cited examples include IsolateGPT, f-secure, RTBAS, and CaMeL.

3 DRIFT: Dynamic Rule-based Isolation Framework

DRIFT protects LLM agents with control- and data-level constraints that can adapt during execution, while isolating conflicting instructions from memory. Its workflow combines planning, deviation validation, and tool-output sanitization.

  • Framework Overview: DRIFT is a rule-based system-level defense that enforces control- and data-level constraints with dynamic permissions and memory isolation.Its components are Secure Planner, Dynamic Validator, and Injection Isolator.
  • Secure Planner: The Secure Planner operates before environmental interaction, using the user query to establish foundational control- and data-level policies.This timing avoids exposure to injection attacks while initial constraints are created.
  • Secure Planner: The planner decomposes the query into subtasks, creates a minimal function trajectory, and generates JSON parameter checklists with value dependencies.The trajectory constrains control flow, while checklists constrain data flow and can detect discrepancies that control-only policies miss.
  • Dynamic Validator: The Dynamic Validator checks each tool request against planned trajectory and parameter constraints, handling deviations to preserve task utility.It compares executed functions and parameters with the planner’s trajectory and checklists.
  • Dynamic Validator: Dynamic policy updates classify deviating functions as Read, Write, or Execute and apply privilege-aware approval based on user intent.Read operations are generally approved; Write and Execute deviations require intent alignment checks.
  • Injection Isolator: Memory-resident injections create long-term, tool-independent risks because reused conversations and tool responses expose conflicting instructions to agents and security modules.Some such attacks produce no tool-call deviation and therefore evade control- or data-based constraints.
  • Injection Isolator: The Injection Isolator analyzes tool outputs, masks instructions conflicting with the original query, and stores cleaned responses for future steps.This extends isolation beyond tool-call trajectory checks into the memory stream.

4 Experiments

Across AgentDojo, ASB, and AgentDyn, DRIFT generally improves security while preserving utility, adapts across models, and outperforms competing dynamic or static defenses in key comparisons. Ablations and trajectory analyses attribute this balance to dynamic policy updates and injection isolation, while efficiency results show higher cost than undefended agents but lower cost than major policy-based alternatives.

  • Defense Techniques Comparison: On ASB, DRIFT achieved 4.8% ASR versus 15.8% for Progent while maintaining robust utility under no-attack and attack conditions.The undefended agent had higher utility, but DRIFT retained robust functionality while providing stronger security.
  • Defense Techniques Comparison: On AgentDyn, DRIFT provided the best utility–security balance among dynamic defenses, whereas prompt-based methods exceeded 35% ASR and CaMeL had zero utility.DRIFT was more stable than Progent in both utility and security on the open-ended benchmark.
  • DRIFT Adaptation: Across GPT-4o, GPT-4o-mini, Claude-3.5-Haiku, and Claude-3.5-Sonnet, DRIFT reduced ASR to single-digit levels while keeping utility stable or sometimes improving it.For GPT-4o, ASR fell from 51.7% to 1.7% after deploying DRIFT.
  • Ablation Studies: Ablations show that the Secure Planner reduced ASR from 30.67% to 2.14% but sharply lowered utility, while the Dynamic Validator restored utility with ASR at 5.41%.Benign Utility increased to 58.48% and Utility Under Attack to 45.82% after adding the Dynamic Validator.
  • Ablation Studies: Adding the Injection Isolator further reduced ASR to 1.35% with only a slight utility drop, and dynamic policy remained stable as trajectory length reached or exceeded three.Static-policy success rates sharply decreased at longer trajectories, while dynamic-policy success remained stable.
  • Efficiency: DRIFT consumed approximately 1.89× more tokens than the undefended agent but fewer than most defenses, including CaMeL and Progent.CaMeL incurred roughly 7× the token cost, while DRIFT reduced ASR to 1.4% compared with tool_filter’s 7.6%.

5 Conclusion

The paper presents DRIFT as a system-level defense for prompt injection attacks, combining dynamic action constraints with injection isolation to support secure and functional agentic systems.

  • DRIFT generates dynamic policies to constrain agent actions while maintaining utility.
  • DRIFT includes injection isolation that removes injected content from the memory stream to preserve long-term security.
  • The framework comprises a Secure Planner, a Dynamic Validator, and an Injection Isolator.

Limitations

Although DRIFT improves utility and security on AgentDojo, the benchmark’s domains do not fully represent the diverse tasks and attack scenarios of real-world agentic systems.

  • AgentDojo’s limited domains do not fully cover the diverse tasks and attack scenarios encountered in real-world agentic systems.
  • Future work will evaluate DRIFT in more realistic and diverse environments.

A.1 Dynamic Mechanism Performance on Open-ended Task

DRIFT retains substantial capability on open-ended AgentDojo tasks despite a minor completion-rate reduction, whereas CaMeL’s fixed constraints prevent completion.

  • Task setup: The evaluation manually selected six open-ended AgentDojo tasks across Banking, Slack, and Workspace.
  • Task setup: DRIFT was compared with the base agent and CaMeL using GPT-4o-mini to reduce bias from base-model capability.
  • Results: Approximately 70% of the base agent’s capability was retained by DRIFT on unpredictable open-ended tasks.DRIFT slightly reduced completion rate, but the decrease was described as minor.
  • Results: CaMeL achieved a zero completion rate on these open-ended tasks because of its fixed constraints.

A.2 Further Analysis of DRIFT and Progent

DRIFT and Progent show comparable utility and security with GPT-4o, but DRIFT is substantially more secure with GPT-4o-mini and maintains robustness across models.

  • Cross-model comparison: With GPT-4o, DRIFT and Progent achieve comparable levels of utility and security.
  • Cross-model comparison: With GPT-4o-mini, DRIFT records 1.35% ASR versus Progent’s 9.39% on AgentDojo and 4.75% versus 15.75% on ASB.
  • Cross-model comparison: DRIFT maintains robustness similar to its GPT-4o performance when switching to GPT-4o-mini, unlike Progent’s substantial security-performance drop.
  • Mechanistic analysis: Progent’s dynamic policy updates require the model to decide when and how to update policies, whereas DRIFT identifies tool privileges through simpler subtasks.
  • Mechanistic analysis: The comparison suggests decomposing dynamic-module tasks into simpler subtasks whenever possible.

A.3 Case Study for Injection Isolator

The case study shows that correct tool usage and parameter control do not prevent injection instructions from influencing an agent through its memory stream. Injection Isolator blocks this persistent risk and prevents malicious content from affecting the final response.

  • The Figure 8 case study is presented to illustrate the effectiveness of Injection Isolator against prompt injection attacks.
  • An AgentDojo attack succeeds even when the tool trajectory and parameters remain correct, because injected instructions in returned messages influence the final answer.The case demonstrates that control and data constraints alone are insufficient against all injection attacks.
  • Because the injection enters during the first tool call and historical conversations are re-input later, harmful content can continue influencing subsequent reasoning and the final output.The example includes later interactions such as reserve_hotels before final answer generation.
  • Injection Isolator successfully defends against the attack by avoiding storage of malicious content in the memory stream.This also reduces the possibility that the content is exposed to other modules or subsequent interactions.

B Detailed Results on AgentDojo

This section presents AgentDojo utility measurements without attack and under attack, together with attack-success-rate measurements. It also identifies prompts for constraint generation, privilege assignment, intent-alignment validation, and prompt-injection detection.

  • Table 6 reports utility on the AgentDojo benchmark without attack, measured as percentages.
  • Table 7 reports utility on the AgentDojo benchmark under attack, measured as percentages.
  • Table 8 reports attack success rate on the AgentDojo benchmark under attack, measured as percentages.
  • Figure 9 presents the prompt of constraints generation.
  • Figure 10 presents the prompt of privilege assignment, Figure 11 the prompt of intent alignment validation, and Figure 12 the prompt of prompt injection detection.
Loading 2506.12104v3…