Table of Contents
Fetching ...

AgentArmor: Enforcing Program Analysis on Agent Runtime Trace to Defend Against Prompt Injection

Peiran Wang, Yang Liu, Yunfei Lu, Yifeng Cai, Hongbo Chen, Qingyou Yang, Jie Zhang, Jue Hong, Ye Wu

TL;DR

This work tackles prompt-injection risks in autonomous LLM agents by reframing agent runtime traces as structured programs and enforcing security through Program Dependence Graphs (PDGs). AgentArmor builds a three-part pipeline—graph construction (CFG/DFG/PDG), graph annotation with a security-type system, and graph-level inspection—to perform fine-grained, dependency-aware checks before tool execution. Empirical results on AgentDojo and ASB show that AgentArmor reduces attack success rate to as low as $0$–$3 ext{ extpercent}$ with only about $1 ext{ extpercent}$ utility loss, outperforming many prompt-level and some system-level defenses. The approach highlights the value of program-analysis techniques for ensuring robust, verifiable security in dynamic, tool-enabled AI agents, with future work aimed at dynamic rule generation and multi-agent extension.

Abstract

Large Language Model (LLM) agents offer a powerful new paradigm for solving various problems by combining natural language reasoning with the execution of external tools. However, their dynamic and non-transparent behavior introduces critical security risks, particularly in the presence of prompt injection attacks. In this work, we propose a novel insight that treats the agent runtime traces as structured programs with analyzable semantics. Thus, we present AgentArmor, a program analysis framework that converts agent traces into graph intermediate representation-based structured program dependency representations (e.g., CFG, DFG, and PDG) and enforces security policies via a type system. AgentArmor consists of three key components: (1) a graph constructor that reconstructs the agent's runtime traces as graph-based intermediate representations with control and data flow described within; (2) a property registry that attaches security-relevant metadata of interacted tools \& data, and (3) a type system that performs static inference and checking over the intermediate representation. By representing agent behavior as structured programs, AgentArmor enables program analysis for sensitive data flow, trust boundaries, and policy violations. We evaluate AgentArmor on the AgentDojo benchmark, the results show that AgentArmor can reduce the ASR to 3\%, with the utility drop only 1\%.

AgentArmor: Enforcing Program Analysis on Agent Runtime Trace to Defend Against Prompt Injection

TL;DR

This work tackles prompt-injection risks in autonomous LLM agents by reframing agent runtime traces as structured programs and enforcing security through Program Dependence Graphs (PDGs). AgentArmor builds a three-part pipeline—graph construction (CFG/DFG/PDG), graph annotation with a security-type system, and graph-level inspection—to perform fine-grained, dependency-aware checks before tool execution. Empirical results on AgentDojo and ASB show that AgentArmor reduces attack success rate to as low as with only about utility loss, outperforming many prompt-level and some system-level defenses. The approach highlights the value of program-analysis techniques for ensuring robust, verifiable security in dynamic, tool-enabled AI agents, with future work aimed at dynamic rule generation and multi-agent extension.

Abstract

Large Language Model (LLM) agents offer a powerful new paradigm for solving various problems by combining natural language reasoning with the execution of external tools. However, their dynamic and non-transparent behavior introduces critical security risks, particularly in the presence of prompt injection attacks. In this work, we propose a novel insight that treats the agent runtime traces as structured programs with analyzable semantics. Thus, we present AgentArmor, a program analysis framework that converts agent traces into graph intermediate representation-based structured program dependency representations (e.g., CFG, DFG, and PDG) and enforces security policies via a type system. AgentArmor consists of three key components: (1) a graph constructor that reconstructs the agent's runtime traces as graph-based intermediate representations with control and data flow described within; (2) a property registry that attaches security-relevant metadata of interacted tools \& data, and (3) a type system that performs static inference and checking over the intermediate representation. By representing agent behavior as structured programs, AgentArmor enables program analysis for sensitive data flow, trust boundaries, and policy violations. We evaluate AgentArmor on the AgentDojo benchmark, the results show that AgentArmor can reduce the ASR to 3\%, with the utility drop only 1\%.

Paper Structure

This paper contains 21 sections, 1 equation, 15 figures, 3 tables.

Figures (15)

  • Figure 1: Methodology overview for implementing AgentArmor on the LLM agent runtime: AgentArmor hooks the agent runtime to get the runtime trace, consisting of dozens of messages. Then, the graph constructor transforms the hooked agent runtime trace into graph-based abstraction representations; Next, the graph annotator adds the security semantics upon the constructed graph-based abstraction representations; At last, AgentArmor enforces the graph inspector to ensure the security of agent runtime.
  • Figure 2: The graph constructor and the property registry (tool registry plus data registry) construct the graph in 8 steps: First, the graph constructor converts the agent runtime trace into a control flow graph by composing messages from the trace into nodes and constructing control flow edges. Then, the graph constructor calls the dependency analyzer to get the control dependency edges and adds them to the graph. Next, the data flow graph is built by first filtering nodes from CFG, then constructing the data flow edges. The data dependency edges are inferred using the dependency analyzer. Furthermore, the graph constructor complements the graph based on the metadata in the tool registry. At last, the program dependency graph is constructed with essential information from the control and data flow graphs.
  • Figure 3: We provide 2 reasoning pattern examples: direct execution and parameterized execution.
  • Figure 4: AgentArmor's graph annotator works as follows: The annotator first assigns predefined types to some nodes in the input program dependence graph, by retrieving metadata from the data registry. Then, the annotator infers the rest of the nodes' types based on lattice propagation.
  • Figure 5: The graph inspector first extracts the rule type from the node , then it evaluates the constraints of the rule type , and resolves the violation at last.
  • ...and 10 more figures