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

2025-06-13

AI Summary

To address the vulnerability of LLM-based agents to prompt injection attacks during external tool invocation—and the lack of dynamic rule updating and memory-flow isolation in existing defenses—this paper proposes a synergistic defense framework combining *dynamic rule-driven enforcement* and *instruction-level memory isolation*. The framework integrates three novel components: (1) a security planner for JSON Schema–based parameter validation; (2) a dynamic validator performing runtime deviation detection and privilege auditing; and (3) an injection isolator identifying and masking memory-flow conflicts. Together, they jointly constrain both control flow and data flow while enabling user-intent consistency verification. Evaluated on the AgentDojo benchmark, our approach achieves a 92% attack mitigation rate while maintaining a task completion rate above 98%, and demonstrates compatibility with mainstream open- and closed-source LLMs.

Technology Category

• Multiagent Systems: Adversarial Agents
• Natural Language Processing: Safety and Robustness
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

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 DRIFT, a Dynamic Rule-based Isolation Framework for Trustworthy agentic systems, which enforces both control- and data-level constraints. 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 benchmark, demonstrating its strong security performance while maintaining high utility across diverse models -- showcasing both its robustness and adaptability.

Problem

• Prevent prompt injection attacks on LLM agents
• Dynamically update security rules for agent safety
• Isolate malicious inputs to protect system integrity

Innovation

• Dynamic rule-based security framework
• Secure Planner for minimal function trajectory
• Injection Isolator for memory stream protection