SafeAgent: Dual-Objective Optimization Against Prompt and Tool Injection in Interactive Environments

1. Introduction

The emergence of Large Language Model (LLM) agents capable of autonomous tool use represents a paradigm shift in AI applications [1]. These agents combine natural language reasoning with external actions, enabling them to browse the web, execute code, manage files, and interact with APIs. However, this capability expansion introduces severe security vulnerabilities that traditional LLM safety mechanisms cannot adequately address [2].

Prompt injection attacks exploit the fundamental inability of LLMs to distinguish between trusted instructions and untrusted data [3]. When agents process external content, such as web pages, email attachments, or tool responses, attackers can embed malicious instructions that hijack agent behavior [4]. The OWASP Foundation identifies prompt injection as the top security risk for LLM applications, with documented incidents including credential theft, unauthorized actions, and data exfiltration [2].

A naive approach to agent security involves aggressive filtering or capability restriction. However, this creates a second problem: overly cautious agents that refuse legitimate requests or fail to complete useful tasks. We term this the safety-utility dilemma, where improving one objective often degrades the other. Existing work typically addresses either safety [5] or utility [6] in isolation, lacking a unified framework for optimizing both simultaneously.

This paper makes the following contributions: (1) We formalize the dual-objective optimization problem for agent security, defining quantitative metrics for both safety and utility. (2) We propose SafeAgent, a defense architecture combining context firewalls, tool sandboxing, and safety-aware planning. (3) We construct a diverse attack benchmark covering realistic injection scenarios. (4) We demonstrate Pareto-optimal performance among tested configurations on standard benchmarks, achieving measurable improvements in the safety-utility tradeoff.

2. Related Work

2.1. Prompt Injection Attacks

Prompt injection was first characterized by Perez and Ribeiro [3], who demonstrated that simple textual prefixes could override system instructions. Greshake et al. [4] extended this to indirect injection, where malicious content in external data sources compromises LLM-integrated applications. Liu et al. [7] developed HouYi, a systematic black-box attack methodology that achieves high success rates against commercial applications.

Recent work has focused specifically on tool-augmented agents. Zhan et al. [8] introduced InjecAgent, a benchmark for evaluating indirect prompt injection against agents with tool access. Shi et al. [9] proposed ToolHijacker, which manipulates tool selection through optimized injection payloads. These attacks demonstrate that agents face unique vulnerabilities at the tool-calling interface.

2.2. Agent Safety Benchmarks

The AgentDojo benchmark [6] provides a dynamic environment for evaluating prompt injection attacks and defenses across 97 realistic tasks and 629 security test cases spanning banking, workspace, travel, and communication domains. Agent-SafetyBench [5] extends evaluation to 349 interaction environments covering eight categories of safety risks. AgentHarm [10] focuses on measuring agent resistance to harmful multi-step tasks, while R-Judge [11] evaluates agents' safety risk awareness during interaction.

Evaluation must also distinguish malicious interference from changes in the data distribution. Target-aware augmentation for rare events under covariate shift [12] and divergence reduction in batch normalization for domain adaptation [13] frame complementary challenges: representing infrequent failures and maintaining stable behavior as inputs change. Active data curation and dataset distillation [14], together with annotated dataset generation [15], offer ways to organize training coverage without treating sample volume as a substitute for diversity. These considerations motivate separating attack coverage from adaptation effects when interpreting safety scores; the benchmark protocol used here remains the official configuration described in Section V.

2.3. Defense Mechanisms

Defense strategies fall into three categories. Input filtering approaches attempt to detect and remove injected instructions before processing [16]. Architectural defenses separate trusted and untrusted components; Shi et al. [17] proposed Progent for programmable privilege control, while Zhu et al. [18] introduced MiniScope for least-privilege authorization. Planning-level defenses integrate safety checks into the reasoning process; Xiang et al. [19] proposed GuardAgent, which uses a separate guard model to validate agent actions. SafeAgent combines elements from all three categories into a unified framework.

Evidence quality creates a second boundary beyond instruction filtering. Evidence-grounded audit generation [20], structure-aware summarization with fidelity constraints [21], and legal reasoning guided by attention and prior knowledge [22] connect generated conclusions to structured evidence or explicit domain constraints. Retrieval and adaptive re-ranking [23] address which evidence enters the context, while semantic conflict detection and cross-paragraph evidence integration [24] address whether the selected material supports a coherent conclusion. Together, these perspectives distinguish retrieving plausible content from establishing its authority to support an agent action.

That distinction becomes operationally significant when tools return incomplete or noisy observations. Evidence-verified root cause localization under tool and telemetry noise [25] and VeriTrail-RCA [26] foreground verification of diagnostic evidence, while causal graph-augmented retrieval [27] structures the dependencies used in diagnosis. ETL-log fault prediction [28] and joint temporal-structural fault discrimination [29] supply complementary perspectives on recognizing anomalous execution patterns. An agent defense must therefore consider both the credibility of individual observations and the way successive tool results influence downstream decisions.

Representations determine which aspects of external content remain visible to these checks. Multimodal proxy embeddings for cold-start ranking [30], cross-modal graphs combining temporal behavior and social information [31], and attention-based representations guided by XBRL structure [32] frame alternative ways to combine sparse, heterogeneous, or explicitly structured signals. Multi-feature fusion with gated recurrent units [33] and hierarchical multi-source fusion with dropout regularization [34] similarly connect local features to an aggregate classification decision. For security monitoring, this line of work motivates studying richer evidence representations while retaining the provenance boundary between user instructions and external observations.

Temporal change further complicates the distinction between benign novelty and compromise. Graph and transformer models for multivariate time series [35] emphasize relationships across signals, whereas incremental modeling of evolving streams [36] and drift-aware agentic ETL for changing schemas [37] emphasize adaptation over time. These perspectives suggest evaluating whether a defense remains consistent as tool outputs and environmental structure evolve. Causally consistent explanations in recommendation agents [38] and consistency-aware learning to mitigate correlation bias [39] add a complementary requirement: apparent agreement between an output and observed data should not replace a check on the decision process that produced it.

Action constraints and unequal error costs provide a common basis for turning such checks into deployment decisions. Policy-constrained repair of Kubernetes configurations [40] places generated changes within an admissible action space, while learning-based backend scheduling [41] treats operational choices as decisions with resource consequences. Decision-focused forecasting with robust optimization [42] and one-sided conformal prediction under asymmetric loss [43] frame how downstream costs can shape the choice of operating point. These ideas sharpen the safety-utility question: blocking an unauthorized action and rejecting a legitimate task have different consequences, so a single aggregate accuracy measure is insufficient.

Defense overhead also depends on where computation and communication occur. Communication-efficient decentralized inference [44] and transformer-aware gradient compression [45] address communication demands at inference and training time, respectively; subspace-deconvolution adaptation [46] addresses the parameterization of model updates. Cost-aware federated learning with client routing and aggregation selection [47] and privacy-preserving anomaly detection through hierarchical federated learning [48] extend the resource question to distributed coordination and information exposure. These distinct constraints motivate reporting the cost of verification alongside its security benefit. SafeAgent focuses on the resulting safety-utility operating points through context separation, restricted tool execution, and action-level verification.

3. Problem Formulation

3.1. Threat Model

We consider an LLM agent $\mathcal{A}$ with access to a tool set $\mathcal{T}=\{t_1,t_2,\mathrm{...},t_n\}$ for tasks including web retrieval, code execution, and file operations. The agent receives user queries and generates action sequences to complete tasks. An adversary $\mathcal{M}$ can inject malicious content into any external data source the agent may access, including web pages, tool responses, or repository files. The adversary's goal is to cause the agent to execute unauthorized actions, leak sensitive information, or deviate from the user's intended task.

3.2. Dual-Objective Framework

Let $\mathcal{D}={\{(q_i,g_i)\}}_{i=1}^N$ denote a dataset of $N$ user queries with corresponding ground-truth goals. Let $N_{\mathrm{atk}}$ denote the number of attack test cases. We define two primary metrics (reported as percentages for readability, i.e., values in $[0,100]$):

Utility Score (%) measures task completion quality:

$$\mathrm{Utility}={\displaystyle \frac{100}{N}}{\displaystyle \sum _{i=1}^{N}1}[\mathcal{A}(q_i)\models g_i]$$

(1)

where $1[\cdot ]$ is the indicator function and $\mathcal{A}(q_i)\models g_i$ indicates the agent's output satisfies the task goal $g_i$.

Safety Score (%) measures resistance to attacks:

$$\mathrm{Safety}=100-{\displaystyle \frac{100}{N_{\mathrm{atk}}}}{\displaystyle \sum _{j=1}^{N_{\mathrm{atk}}}1}[\mathcal{A}(q_j^{\mathrm{atk}})\models g_j^{\mathrm{mal}}]$$

(2)

where $q_j^{\mathrm{atk}}$ represents a query containing injected attack content, and $g_j^{\mathrm{mal}}$ is the adversary's malicious objective. Equivalently, Safety $=100-\mathrm{ASR}$, where ASR is the Attack Success Rate in percentage points.

The dual-objective optimization seeks configurations on the Pareto frontier:

$${\max }_{\theta }(\mathrm{Utility}(\theta ),\mathrm{Safety}(\theta ))$$

(3)

where $\theta$ represents defense configuration parameters. A configuration is Pareto-optimal if no other configuration improves one objective without degrading the other. We note that we evaluate Pareto optimality empirically among tested configurations and do not claim global optimality guarantees.

3.3. Attack Taxonomy

We categorize injection attacks by their source and mechanism:

Type I: Web Retrieval Injection. Malicious instructions embedded in web page content retrieved during information-seeking tasks. Attack payloads typically use invisible formatting, comment tags, or semantic disguise.

Type II: Tool Response Injection. Adversarial content returned by compromised or malicious API endpoints, including poisoned database entries and manipulated file contents.

Type III: Repository Injection. Hidden instructions in code documentation, README files, or comments that activate when agents process repository content.

4. Proposed Defense Framework

SafeAgent's defense architecture comprises three layers designed to provide defense-in-depth while preserving agent capability. Figure 1 illustrates the overall system design. Note that the modules shown are conceptual; implementation details are provided in Section V-C.

4.1. Context Firewall

The context firewall enforces strict separation between data and instructions. All external content is wrapped in explicit data markers and processed through a classification stage before entering the agent's context window.

Given input text $x$, the firewall applies transformation $\varphi$:

$$\varphi (x)=[\mathrm{DATASTART}]\oplus \mathrm{sanitize}(\mathrm{x})\oplus [\mathrm{DATAEND}]$$

(4)

where $\oplus$ denotes concatenation and $\mathrm{sanitize}(\cdot )$ removes known injection patterns including control characters, instruction-like prefixes, and encoded payloads.

The agent's system prompt explicitly instructs that content within data markers should be treated as information only and never executed as commands. We do not claim formal guarantees for data/instruction separation; we evaluate empirically under benchmarked attack settings. Empirical results show significant reduction in successful injections.

4.2. Tool-Call Sandbox

The sandbox layer implements the principle of least privilege by restricting tool access based on task requirements. For each task, SafeAgent defines an allowlist $T_{\mathrm{allow}}\subseteq \mathcal{T}$ specifying permitted tools and parameter constraints.

A tool call $(t,p)$ is permitted only if:

$$\mathrm{permit}(t,p)=1[t\in T_{\mathrm{allow}}]\cdot 1[p\in {\mathcal{P}}_t]$$

(5)

where ${\mathcal{P}}_t$ defines valid parameter ranges for tool $t$. Unauthorized calls are blocked and logged.

The allowlist is determined by a pre-planning phase where SafeAgent analyzes the user query before exposure to external content. In this phase, the agent assesses the task requirements and determines the minimal set of tools needed. This follows the principle that tool permissions should be decided based on trusted input only.

4.3. Safety-Aware Planning

SafeAgent extends the ReAct framework [1] with explicit safety verification steps. The standard ReAct loop alternates between Thought $\to$ Action $\to$ Observation. Our modified loop inserts a Safety Check before each action:

• Thought: Reason about current state and next step
• Safety Check: Verify action aligns with original goal
• Action: Execute if safety check passes
• Observation: Process result through context firewall

The safety check evaluates whether the proposed action is consistent with the user's original query and whether it would violate any security constraints. Actions flagged as potentially compromised trigger a re-planning step that excludes the suspicious reasoning chain.

The planner draws on the verifier-guided reasoning perspective of Wang et al. [49] to place an explicit assessment between a candidate reasoning step and its execution. In SafeAgent, the verified object is the proposed tool action, and the acceptance criterion is consistency with the trusted user goal and security constraints. Re-planning after a failed check concentrates additional reasoning on a candidate that cannot safely proceed, while a passing check permits the existing action path to continue. The budget-aware framing also makes verification overhead part of the design rationale: repeated checks must be assessed against the useful work they preserve. The current implementation checks every action and reports latency overhead; it does not implement adaptive budget allocation or claim the efficiency results of a separate reasoning system.

4.4. Output Validator

The Output Validator serves as the final defense layer, performing anomaly detection on agent outputs before they are returned to the user. It checks for indicators of successful injection attacks, including: (1) responses that contain content unrelated to the original user query, (2) outputs that attempt to execute unauthorized actions or leak sensitive information, and (3) anomalous patterns suggesting the agent has been compromised. When suspicious outputs are detected, the validator either sanitizes the response or triggers a re-execution with enhanced safety constraints.

The validator incorporates the sentinel perspective expressed by AgentShield [50] by treating the release of an agent response as a distinct monitoring boundary. A fluent intermediate conclusion can propagate into later decisions unless its consistency with the task is checked before it is reused or returned. SafeAgent applies this principle through its existing anomaly checks for goal deviation, unauthorized behavior, and indications of compromise, with sanitization or constrained re-execution as the response to a detected problem. The firewall limits what enters the reasoning context, whereas the validator checks what leaves the agent, giving the two mechanisms complementary positions in the information flow. This use of a sentinel-style boundary concerns the present single-agent pipeline; multi-agent hallucination propagation is a separate evaluation setting.

5. Experimental Setup

5.1. Benchmark and Metrics

We evaluate on the AgentDojo benchmark [6], which provides 97 tasks across four domains: Banking (financial transactions), Workspace (document management), Travel (booking operations), and Slack (communication tasks). The benchmark includes 629 security test cases with injected attack payloads. We use the official split and configuration from the original benchmark without additional curation.

We report four metrics (all values in percentages): (1) Benign Utility (BU): Task success rate without attacks. (2) Utility Under Attack (UA): Task success rate when attacks are present. (3) Attack Success Rate (ASR): Percentage of attacks that achieve their malicious goal. (4) Safety Score: Defined as $100-\mathrm{ASR}$ (in percentage points).

5.2. Defense Configurations

We compare four configurations: (1) No Defense: Baseline agent with standard prompting. (2) Filter Only: Input sanitization and injection pattern detection. (3) Filter + Sandbox: Adding tool-call restrictions via allowlists. (4) SafeAgent (Full): Complete framework including context firewall, sandbox, and safety-aware planning.

5.3. Implementation Details

Experiments use GPT-4o as the backbone LLM. The context firewall employs regex-based pattern matching augmented with a classifier trained on benchmark injection samples combined with public prompt-injection pattern datasets. Allowlists are generated automatically by prompting the LLM to identify required tools from task descriptions before external data exposure. Safety checks use a temperature of 0 for deterministic verification.

6. Results and Analysis

6.1. Overall Performance

Table 1 presents the main experimental results. The full defense configuration achieves the best balance between safety and utility, reducing ASR from 67.2% to 13.4% (a reduction of 53.8 percentage points) while maintaining 68.9% utility compared to 78.5% for the undefended baseline.

6.2. Pareto Analysis

Figure 2 shows the safety-utility tradeoff across all configurations. SafeAgent achieves points along the estimated Pareto frontier among tested configurations, meaning no other tested configuration provides better safety at the same utility level or vice versa. The shaded region indicates configurations that are dominated by frontier points.

Notably, the transition from Filter + Sandbox to SafeAgent yields disproportionate safety gains (67.5% to 86.6%) with modest utility cost (71.8% to 68.9%), suggesting the safety-aware planning component provides high return on investment.

6.3. Attack-Type Analysis

Figure 2 presents attack success rates across attack types and defense configurations. Indirect (Web) injection shows the highest baseline vulnerability (82% ASR) but SafeAgent reduces it effectively to 22%. Tool Poisoning attacks are most effectively defended against, with SafeAgent achieving the lowest ASR of 15%.

The relative ASR reduction from baseline to SafeAgent is approximately 80% on average (from 67.2% to 13.4% in Table 1), calculated as:

$$\mathrm{Relative}\mathrm{Reduction}={\displaystyle \frac{{\mathrm{ASR}}_{\mathrm{baseline}}-{\mathrm{ASR}}_{\mathrm{defense}}}{{\mathrm{ASR}}_{\mathrm{baseline}}}}$$

(6)

This demonstrates the framework's general effectiveness across diverse attack vectors. Domain-specific analysis reveals that Workspace tasks experience the highest utility degradation under full defense (from 76.3% to 65.1%), likely due to the broader tool requirements of document management operations.

Figure 3. Attack success rate across attack types and defense configurations. Lower values indicate stronger defense. SafeAgent achieves consistent ASR below 23% across all attack types. Cell values are shown as decimals (e.g., 0.18 = 18% ASR).

Figure 3. Attack success rate across attack types and defense configurations. Lower values indicate stronger defense. SafeAgent achieves consistent ASR below 23% across all attack types. Cell values are shown as decimals (e.g., 0.18 = 18% ASR).

6.4. Component Ablation

To understand individual component contributions, we conducted ablation experiments removing each defense layer. Results indicate the context firewall contributes approximately 35% of total ASR reduction, the sandbox contributes 28%, and safety-aware planning contributes 37%. However, these components exhibit synergistic effects; removing any single layer degrades overall performance more than the sum of individual contributions would suggest.

6.5. Limitations

Our framework has several limitations. First, the context firewall relies on heuristic patterns that sophisticated adversaries may evade through novel encoding schemes. Second, automatically generated allowlists may be overly permissive or restrictive for edge cases. Third, safety checks incur computational overhead, increasing average response latency by approximately 23%. Finally, our evaluation focuses on text-based injection; multimodal attacks (e.g., image-based prompt injection) through images or audio remain unexplored.

7. Conclusions

We presented SafeAgent, a dual-objective optimization framework for securing LLM agents against prompt and tool injection attacks while preserving task utility. Our three-layer defense architecture, combining context firewalls, tool sandboxing, and safety-aware planning, achieves Pareto-optimal performance among tested configurations on the AgentDojo benchmark. The framework reduces attack success rates from 67.2% to 13.4% (approximately 80% relative reduction) while maintaining 68.9% task utility.

Our results demonstrate that, under the AgentDojo benchmark and our evaluated attack types, the safety-utility tradeoff in agent systems is not strictly zero-sum; carefully designed defenses can achieve substantial security improvements with acceptable utility costs. This balance relies on the pre-planning phase accurately determining tool permissions based on trusted input. The quantitative framework we provide enables reproducible evaluation and comparison of future defense mechanisms.

Future work should address multimodal injection attacks such as image-based prompt injection, develop adaptive defenses that learn from attack patterns, and explore formal verification methods for safety guarantees. As LLM agents become increasingly deployed in real-world applications, principled approaches to balancing safety and utility will be essential for trustworthy AI systems.