The Task Shield: Enforcing Task Alignment to Defend Against Indirect Prompt Injection in LLM Agents

TL;DR

This work reframes LLM agent security around task alignment, arguing that all agent actions should serve user objectives to counter indirect prompt injection. It introduces Task Shield, a test-time defense that extracts task directives, assesses their contribution to user goals via a fuzzy ContributesTo score, and provides feedback to maintain alignment across user, assistant, and tool messages. Empirical results on AgentDojo with GPT-4o and GPT-4o-mini show Task Shield significantly reduces attack success rates while preserving substantial task utility, outperforming baseline defenses in most scenarios. The approach demonstrates a practical, modular defense that improves security-utility trade-offs in real-world, tool-augmented AI agents, with clear avenues for future enhancements and broader threat coverage.

Abstract

Large Language Model (LLM) agents are increasingly being deployed as conversational assistants capable of performing complex real-world tasks through tool integration. This enhanced ability to interact with external systems and process various data sources, while powerful, introduces significant security vulnerabilities. In particular, indirect prompt injection attacks pose a critical threat, where malicious instructions embedded within external data sources can manipulate agents to deviate from user intentions. While existing defenses based on rule constraints, source spotlighting, and authentication protocols show promise, they struggle to maintain robust security while preserving task functionality. We propose a novel and orthogonal perspective that reframes agent security from preventing harmful actions to ensuring task alignment, requiring every agent action to serve user objectives. Based on this insight, we develop Task Shield, a test-time defense mechanism that systematically verifies whether each instruction and tool call contributes to user-specified goals. Through experiments on the AgentDojo benchmark, we demonstrate that Task Shield reduces attack success rates (2.07\%) while maintaining high task utility (69.79\%) on GPT-4o.

Figures (7)

• Figure 1: Overview of the Task Shield interacting with a tool-integrated LLM agent. The framework enforces task alignment and defends against indirect prompt injection attacks.
• Figure 2: This diagram illustrates how the Task Shield framework processes different message types from the conversational flow through task instruction extraction, alignment checks, and feedback generation.
• Figure 3: GPT-4o: Comparison of Attack Success Rate (ASR) versus Utility. Solid markers represent ASR versus benign utility, while hollow markers represent ASR versus utility under attack. Arrows indicate the change in utility due to the attack, with their direction showing the impact of the attack on model performance. The green circles highlight the Pareto front in benign conditions, and the orange circles highlight the Pareto front under attack. Numbers along the arrows indicate the magnitude of the utility change when an attack is introduced (positive values show improvement, and negative values indicate degradation).
• Figure 4: Task Extraction Prompt: This prompt outlines the methodology for extracting actionable task instructions from the conversation content.
• Figure 5: Content Checker Prompt: This prompt evaluates the alignment of new actionable instructions with user task instructions based on task relevance and privilege level.

Theorems & Definitions (4)

• Definition 1: Task Instruction
• Definition 2: $\mathrm{ContributesTo}$ Relation
• Definition 3: Task Instruction Alignment Condition
• Definition 4: Task Alignment