Source-linked AI summary
The Task Shield: Enforcing Task Alignment to Defend Against Indirect Prompt Injection in LLM Agents
Feiran Jia, Tong Wu, Xin Qin, Anna Squicciarini
TL;DR
Indirect prompt injections hide malicious instructions in external data and can divert tool-integrated agents from user intentions. Task Shield checks whether instructions and tool calls contribute to user goals; on AgentDojo it reports a 2.07% attack success rate and 69.79% utility under attack on GPT-4o.
Problem
Indirect prompt injections embed malicious instructions in external data that can divert tool-integrated agents, while existing defenses face practical limitations in maintaining security and task functionality.
Source: §1 · p. 1Method
Task Shield is a test-time defense that checks whether agent and tool directives align with user goals and provides feedback or prevents misaligned tool calls.
Source: §1 · p. 1, §3.1 · p. 3, §4.2 · p. 5Results
2.07% attack success rate and 69.79% utility under attack (U) on GPT-4o; Task Shield outperformed baseline defenses across task suites.
Source: §5.2 · p. 7Takeaways & Limitations
The paper concludes that task alignment can guide defense against indirect prompt injection while preserving agent capabilities across simulated real-world tasks in AgentDojo.
Source: §7 · p. 8Reliance on LLMs for instruction extraction and ContributeTo scoring risks degradation with weaker models and adaptive attacks; evaluation was limited to one benchmark and one model family.
Source: §7 · p. 8
Abstract
from arXiv · showhide
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.
1 Introduction
Tool-integrated LLM agents can act on sensitive data and systems, making indirect prompt injection a serious security concern. Task Shield reframes defense as checking whether directives serve user goals, and the paper reports reduced injection vulnerability while preserving task utility.
1 Introduction
- Indirect prompt injections hide malicious directives in external data, and more capable LLMs are described as increasingly vulnerable.The agents considered can access sensitive data, perform financial transactions, and interact with critical systems.Source: §1 · p. 1
- Task alignment asks whether each directive serves user objectives rather than whether it appears harmful, addressing limits of rule, source-spotlighting, and authentication defenses.The authors note that detailed rules are challenging and malicious directives can appear benign.Source: §1 · p. 1
- Task Shield checks directives from agents and tools against user goals and intervenes to prevent unrelated actions while preserving task completion.The overview presents it as a guardian for tool-integrated LLM agents.Source: §1 · pp. 1–2, Figure 1 · p. 2
- On AgentDojo, the paper reports that Task Shield reduces prompt-injection vulnerabilities while preserving user-task utility.The introduction summarizes experiments on the benchmark without giving numerical results here.Source: §1 · p. 2
2 Preliminary
The preliminary section describes the message roles and privilege hierarchy in tool-integrated conversations, then frames indirect prompt injection as instructions hidden in external data processed during a task.
2 Preliminary
- Conversation messages include system, user, assistant, and tool-output roles, with system instructions ranked above user, assistant, and tool instructions.The hierarchy assigns privilege levels and specifies the precedence order system > user > assistant > tool.Source: §2 · p. 2
- The restaurant-search example illustrates a user request followed by assistant planning and an external API query for data.The example situates tool use within the message hierarchy.Source: §2 · p. 2
- Indirect prompt injection embeds instructions in the environment, such as a hidden webpage directive that redirects an agent summarizing the page.The passage characterizes these attacks as stealthy because the directives are concealed in legitimate external data.Source: §2 · p. 2
3 Task Alignment
Task alignment treats whether an instruction advances a user goal as the basis for judging agent behavior. The section defines task instructions, explains how tools can support or specify subtasks, and formalizes alignment against user instructions.
3 Task Alignment
- Task alignment evaluates whether actionable instructions contribute to user objectives, allowing the framework to capture injected prompts even when they appear benign.The section frames attacks as succeeding when agents execute directives that deviate from user or conversational goals.Source: §3 · pp. 2–3
3.1 Task Instructions
- A task instruction is an actionable directive extracted from a message, including user requests, assistant plans or tool calls, and tool-generated suggestions.The analysis primarily focuses on user-level directives, while noting that system messages can define high-level tasks in specialized agents.Source: §3.1 · p. 3
3.2 Task Interactions
- Tool outputs may supply information for a clear user request or define concrete subtasks, and conversation history provides context for judging their relation to user goals.The examples contrast retrieving dentist contact details with a to-do tool that supplies specific tasks.Source: §3.2 · p. 3
3.3 Formalization of Task Alignment
- An instruction contributes to a prior task when it helps achieve that task; alignment requires each lower-level instruction to contribute to at least one user task.A conversation is fully aligned when all assistant-level task instructions meet that condition, keeping plans and tool calls in service of user goals.Source: §3.3 · pp. 3–4
4 The Task Shield Framework
Task Shield continuously monitors instructions against user objectives, operationalizing task alignment through instruction extraction, alignment checks, and feedback. It applies these checks across message types in agent conversations.
4 The Task Shield Framework
- The framework combines instruction extraction, alignment checks, and feedback generation to enforce alignment with user objectives throughout a conversation.Instruction extraction and ContributesTo score calculation use an LLM.Source: §4 · p. 4
4.1 Task Shield Components
- Task Shield extracts user objectives into a User Task Set and identifies directives from other sources for alignment checks.Its conservative strategy targets implicit, nested, and complex instructions, rewriting information-seeking queries as explicit instructions and preserving task dependencies.Source: §4.1 · p. 4
- A fuzzy score from 0 to 1 represents how each instruction contributes to user goals; an instruction is misaligned when its total contribution score is 0.The score allows direct contributions, intermediate steps, and reasonable attempts at resolution to be represented.Source: §4.1 · p. 4
- When it detects misalignment, Task Shield alerts the agent to the instruction, explains potential risks, and reminds it of current user objectives.Source: §4.1 · p. 4
4.2 Interaction with the LLM Agent System
- Each message passes an alignment check before proceeding, with tool-output directives flagged and potentially harmful assistant responses caught for refinement.The system describes these checks as multiple barriers against attacks while maintaining overall conversation alignment.Source: §4.2 · pp. 4–5
- Newly extracted user instructions update the User Task Set and become alignment targets for subsequent message processing.Source: §4.2 · p. 5
- For assistant messages, both response content and tool calls are checked; misaligned content prompts feedback, while misaligned tool calls are blocked.The agent may receive several rounds of feedback to refine its queries.Source: §4.2 · p. 5
- Tool outputs are checked with their source and arguments attached to each instruction; detected misalignment returns the original output alongside feedback to the assistant.Source: §4.2 · p. 5
5 Experiments
On AgentDojo, Task Shield sharply reduced attack success while retaining utility, and its security–utility balance compared favorably with other defenses. Experiments examined GPT-4o and GPT-4o-mini across benchmark task suites and indirect-attack settings.
5 Experiments
- AgentDojo evaluation covered multi-turn tool interactions in Travel, Workspace, Banking, and Slack, with one trial per task; metrics were CU, utility under attack (U), and target ASR.Source: §5.1 · pp. 5–6, §5 · p. 5
5.1 Settings
- The primary evaluation used GPT-4o, with GPT-4o-mini as an additional model, and compared Task Shield with Data Delimiting, PI Detector, Repeat Prompt, and Tool Filter.The models used default configurations, and the same model implemented the protective Task Shield.Source: §5.1 · pp. 5–6
5.2 Results
- Against Important Instructions, Task Shield reduced GPT-4o ASR from 47.69% without defense to 2.07% while preserving 69.79% utility; this was the strongest of three tested attacks.The subsequent experiments focused on Important Instructions; Table 1 compares U and ASR with and without Task Shield across task suites.Source: Table 1 · p. 6, §5.2 · p. 7
- Task Shield consistently approached the Pareto front in benign and attacked conditions, occupying the desirable lower-right region across task suites.Figure 3 compares ASR with benign and under-attack utility; its arrows show utility change under attack.Source: Figure 3 · p. 6, §5.2 · p. 7
- Other defenses split the trade-off: PI Detector achieved low ASR with severe utility degradation, Tool Filter was moderate on both metrics, and Repeat Prompt retained utility but offered inadequate defense.Source: §5.2 · p. 7
- 2.07% ASR and 69.79% utility under attack (U) on GPT-4o; Task Shield also achieved 2.23% ASR and 54.53% U on GPT-4o-mini.Across task suites, it showed near-optimal or optimal performance on clean utility (CU), U, and ASR.Source: §5.2 · p. 7, Table 2 · p. 7
- GPT-4o-mini had lower undefended ASR than GPT-4o (27.19% versus 47.69%); Task Shield remained effective across both models, unlike Repeat Prompt.Task Shield was the only defense that maintained or improved GPT-4o-mini clean utility; most defenses improved GPT-4o clean utility but hurt GPT-4o-mini’s.Source: §5.2 · pp. 7–8
6 Related Work
Prior work develops LLM agents through tool integration and studies indirect prompt injection and defenses against it. The paper situates Task Shield among test-time defenses, emphasizing task alignment as its alternative approach.
LLM Agent and Tool Integration
- LLM agents commonly integrate external tools for tasks such as searching and decision making, and approaches to autonomous tool use support broad integration into agent systems.The cited work includes Ouyang et al. (2022), Sharma et al. (2024), and Schick et al. (2024).Source: §6 · p. 8
Indirect Prompt Injection Attacks
- Indirect prompt injection embeds malicious content in external data and can lead agents to unsafe actions, motivating benchmarks such as Injecagent and AgentDojo.Source: §6 · p. 8
Defense Methods
- Training-time defenses can be costly and require internal model access, while test-time defenses include authentication tags and delimiters; Task Shield instead targets task alignment for robustness and utility.Source: §6 · p. 8
7 Conclusion
The conclusion presents task alignment as a guiding principle for defending against indirect prompt injection and Task Shield as a test-time mechanism that enforces it. The authors report state-of-the-art defense while preserving agent capabilities, and identify model and evaluation constraints.
7 Conclusion
- Task Shield enforces task alignment by checking instruction alignment with user goals, with state-of-the-art defense and preserved agent capabilities reported across simulated real-world tasks in AgentDoJo.The authors describe it as a test-time mechanism.Source: §7 · p. 8
- The framework relies on LLMs for instruction extraction and ContributeTo scoring, leaving it vulnerable to weaker-model performance degradation and adaptive attacks.The authors also say resource constraints limited evaluation to one benchmark and one model family.Source: §7 · p. 8
- Future work includes improving efficiency and robustness, addressing threats beyond prompt injection, adapting to domain-specific business contexts, and generating synthetic training data for task dependencies and misalignment.Source: §7 · p. 8
A Appendix: Detailed Discussion on Task Alignment
The appendix explains that seemingly benign external directives can exceed a user’s request, exposing unnecessary information and consuming resources, while also undermining trust and predictability. It notes that system-level objectives create a hierarchical alignment case not included in the user-level experiments.
A.1 Why Task Alignment Matters: Beyond Overtly Harmful Instructions
- A menu-footer directive to provide details about every dish goes beyond a request about one dish, potentially exposing unrelated information and increasing token use and processing.Source: §A.1 · p. 11
- Following unrelated external suggestions undermines user trust and makes assistant behavior and information flow less predictable.Source: §A.1 · p. 11
A.2 Discussion on System-Level Instructions
- When system messages specify conversation objectives, assistants must align with both system objectives and user goals, creating a hierarchical alignment scenario.Source: §A.2 · p. 11
- The experiments focus on user-level tasks and do not include system-level task objectives in the main analysis.Source: §A.2 · p. 11
A.3 Examples of Task Misalignments
The examples illustrate misalignment through unrequested actions and tool use that fails to serve the user’s stated goal.
A.3 Examples of Task Misalignments
- Posting a requested article summary to social media is an unrequested action that does not align with the user’s original intent.Source: §A.3 · p. 11
- Sending a meeting-confirmation email to Bob instead of Alice misuses the tool and does not contribute to the user’s goal.Source: §A.3 · p. 11
- The examples characterize task-instruction misalignment as proposing tool misuse or actions that do not contribute to the user’s original goals.Source: §A.3 · p. 11
B Appendix: Detials in Task Shield Frameworks Design
The appendix illustrates fuzzy contribution scoring with intermediate steps and reasonable attempts, then outlines Task Shield’s core processing: extract instructions, score their contributions to user tasks, and return feedback.
B.1 Examples of Fuzzy-logic Based Contribution Scoring
- A room-availability lookup receives a high contribution score because it supplies essential information for booking, despite not completing the task directly.Source: §B.1 · p. 11
- A recent-files search for a project budget receives a medium score: it is a reasonable attempt to resolve file-location ambiguity but does not guarantee success.Source: §B.1 · p. 12
B.2 Task Shield Core Processing Algorithm
- The algorithm takes the current message, conversation history, threshold, and user task instructions as input and returns a feedback message.Source: §B.2 · p. 12
- The algorithm extracts potential instructions from each message, adds instructions from user-level messages to the task set, and skips further processing for those messages.Source: §B.2 · p. 12
- For extracted instructions, the algorithm computes contribution scores relative to each user task, totals the scores, and checks the total against threshold ϵ.Source: §B.2 · p. 12
C.1 Results on GPT-3.5-turbo
The GPT-3.5-turbo experiments report defense performance against an instruction-injection attack, with comparisons to GPT-4o and GPT-4o-mini requiring caution because of context-length limitations.
C.1 Results on GPT-3.5-turbo
- For GPT-3.5-turbo, context-length constraints affect benchmark evaluations, so its results should be interpreted cautiously relative to GPT-4o and GPT-4o-mini.The experiments are presented as supplementary insight into Task Shield on another model architecture.Source: §C.1 · p. 12
- Table 3 reports Task Shield and baseline-defense performance on GPT-3.5-turbo across task suites, using Clean Utility (CU), Utility under Attack (U), and Attack Success Rate (ASR) metrics.Because context-length constraints affect benchmark evaluation, these results warrant caution when compared with GPT-4o and GPT-4o-mini.Source: Table 3 · p. 12, §C.1 · p. 12
C.2 Omitted Details in Experiments
- The experimental details specify temperature 0.0; for tool calls, conversation history is truncated so each call follows its corresponding tool output.ContributesTo uses a significant portion of conversation history to capture context.Source: §C.2 · p. 13
D Prompts
- The appendix prompts extract actionable task instructions, check instruction alignment using task relevance and privilege, and verify tool-call alignment.A separate feedback prompt combines content and tool-call misalignment with user-intention reminders.Source: Figure 4 · p. 14, Figure 5 · p. 15, Figure 6 · p. 16