PaperlayerView exact revision on arXiv

Source-linked AI summary

RTBAS: Defending LLM Agents Against Prompt Injection and Privacy Leakage

Peter Yong Zhong, Siyuan Chen, Ruiqi Wang, McKenna McCall, Ben L. Titzer, Heather Miller, Phillip B. Gibbons

arXiv:2502.08966v2cs.CRcs.AI

TL;DR

Tool-based agent systems expand language-model capabilities but expose users to prompt injection and privacy leakage, while per-call confirmation burdens users. RTBAS adapts information flow control with LM-Judge and attention-based dependency screeners to automate safe tool calls selectively. It thwarts all policy-violating attacks with less than 2% utility degradation and reaches near-oracle privacy-leakage performance.

  • Problem

    TBAS create risks of prompt injection and inadvertent confidentiality leakage because external tools and interaction histories can influence agent behavior and expose sensitive data.

  • Method

    RTBAS adapts information flow control for TBAS and selectively propagates security metadata using LM-Judge and attention-based dependency screeners with selective masking.

  • Results

    RTBAS thwarts all policy-violating exploits with less than 2% task-utility degradation under attack and achieves near-oracle privacy-leakage performance.

  • Takeaways & Limitations

    RTBAS safeguards TBAS against prompt injection and inadvertent privacy leakage while selectively handling security labels.

  • Takeaways & Limitations

    RTBAS assumes tools comply with the defined information-flow policy, and screener errors can degrade performance through over-tainting or under-tainting.

Abstract

from arXiv · show

Tool-Based Agent Systems (TBAS) allow Language Models (LMs) to use external tools for tasks beyond their standalone capabilities, such as searching websites, booking flights, or making financial transactions. However, these tools greatly increase the risks of prompt injection attacks, where malicious content hijacks the LM agent to leak confidential data or trigger harmful actions. Existing defenses (OpenAI GPTs) require user confirmation before every tool call, placing onerous burdens on users. We introduce Robust TBAS (RTBAS), which automatically detects and executes tool calls that preserve integrity and confidentiality, requiring user confirmation only when these safeguards cannot be ensured. RTBAS adapts Information Flow Control to the unique challenges presented by TBAS. We present two novel dependency screeners, using LM-as-a-judge and attention-based saliency, to overcome these challenges. Experimental results on the AgentDojo Prompt Injection benchmark show RTBAS prevents all targeted attacks with only a 2% loss of task utility when under attack, and further tests confirm its ability to obtain near-oracle performance on detecting both subtle and direct privacy leaks.

1 Introduction

TBAS extend language models with external tools, but prompt injections and confidentiality risks can trigger harmful actions or disclose sensitive data. RTBAS adapts information flow control with dependency screening and selective masking to automate safe tool calls while limiting user confirmations.

  • TBAS let language models use external tools for tasks such as website search, flight booking, and financial transactions.
  • Prompt injections can manipulate tool-using agents into unauthorized transfers, altered itineraries, or disclosure of personally identifiable and financial data.
  • Existing per-tool-call confirmation safeguards burden users and can increase mindless approvals or abandonment through fatigue.
  • RTBAS adapts information flow control to dynamic, opaque TBAS environments where dependencies between inputs, model processing, and tool calls are implicit.
  • RTBAS selectively propagates security metadata by masking irrelevant history regions, using LM-Judge and attention-based dependency screeners.
  • RTBAS prevents 100% of AgentDojo attacks violating security policies with less than 2% task-utility degradation and achieves near-oracle confidentiality protection across almost all leakage tasks.

2 Background and Related Work

Related work covers prompt-injection detection, impact mitigation, privacy protection, selective propagation, and attention-based saliency. RTBAS is positioned as a policy-aware alternative whose dependency screening reduces the overhead of prior exponential enumeration.

  • Prompt injections introduce malicious inputs into agent history to alter behavior, including through data retrieved from untrusted websites and public reviews.
  • Existing defenses detect anomalies or mitigate impacts through sanitization, delimiters, fine-tuning, and tool restrictions.
  • Most existing techniques are heuristic, lack explicit security-policy support, or remain vulnerable to unrestricted tools and adversarial jailbreaking.
  • Inference-time privacy techniques commonly detect or desensitize personally identifiable information but are typically not designed for tool-based environments.
  • Prior selective propagation can require exponential enumeration, including 16-64 additional language-model calls for typical four-to-six-level lattices.
  • RTBAS uses attention scores as dependency-screening inputs because they capture how output tokens are influenced by input regions.

3 Motivation

Tool responses can carry attacker-controlled instructions or confidential-data dependencies, making integrity and privacy difficult to protect in TBAS. Attention scores provide evidence of these dependencies and motivate attention-based screening.

  • Prompt injection: Tool responses can contain attacker-controlled prompts that influence integrity-sensitive actions without compromising the underlying tool.Examples include malicious website content and transaction descriptions.
  • Privacy leakage: Confidential data may be used explicitly, implicitly, or during intermediate reasoning, complicating leakage tracking.Tool descriptions and system instructions can also subtly influence its flow.
  • Attention-based dependency: Attention scores measure how strongly output tokens depend on input tokens by comparing model loss before and after masking them.Scores are averaged across attention heads and layers.
  • Attention-based dependency: 74% to 86% of attention mass for non-dependent data falls below 0.2, versus 14% and 44% for dependent data across GPT-4o and Claude.The case study uses 3,424 argument–tool-call pairs and attention scores from OPT-125m.
  • Attention-based dependency: When the LM follows injected instructions, attention shifts toward the tool response by a large margin.This supports attention scores as a signal of dependency between instructions and outputs.

4 Tool-based Agent Systems

A TBAS combines a language-model agent, conversation history, and developer-provided tools to act through external APIs. The runtime appends tool results to history before the LM produces its next response or tool call.

  • System structure: TBAS agents use internal knowledge, prior interactions, and tool calls to fulfill user requests through external environments.The user interacts through a chat interface similar to ChatGPT or Gemini.
  • System structure: The developer supplies a system message and tools with names, signatures, and descriptions corresponding to callable runtime APIs.These components initialize the agent’s role, tone, and legal invocation formats.
  • Execution cycle: The application history concatenates system messages, user requests, tool outputs, and previous assistant responses into the LM’s input state.The LM uses this state to decide whether to respond or invoke a tool.
  • Execution cycle: The runtime generates tool calls, executes their APIs, appends returned messages to history, and then obtains the LM’s response.The cycle repeats when the user continues the conversation.

5 Attack Model for Prompt Injection

The attack model assumes an informed attacker can influence externally sourced tool outputs to manipulate tool calls, causing confidential-data leakage or harmful side effects. The attacker cannot directly observe or alter the agent’s internal execution.

  • Attacker goals: The attacker seeks to manipulate TBAS tool calls so they leak confidential information or cause harmful side effects.Examples include sending credit-card information or exploiting money-transfer tools.
  • Attacker capabilities: The attacker knows the system instructions, available tools, and tool-specific instructions but lacks timing knowledge and access to internal agent state.The attacker also cannot directly observe tool arguments or responses.
  • Execution model: The TBAS execution model initializes history, accepts user requests, generates tool calls, runs APIs, appends responses, and produces LM responses.These steps define where manipulated tool outputs can enter subsequent decisions.
  • Attacker capabilities: The attacker can influence outputs of tools that depend on external inputs without compromising tool implementations or intercepting API communication.The underlying data source can only be affected as feasible for an untrusted legitimate user.

6 Robust TBAS Objectives and Assumptions

RTBAS aims to prevent privacy leakage and prompt-injection side effects while preserving utility and minimizing overhead. Its information-flow policy relies on developer-provided security labels and assumes tools themselves obey that policy.

  • Objectives: The primary objectives are preventing private data from leaving without confirmation and preventing attacker instructions from causing unwanted side effects.Secondary objectives are maintaining utility under attack and minimizing compute or confirmation overhead.
  • Information-flow assumptions: The system assumes developers provide a finite security-label lattice with a flows-to order and join operation.Labels pair integrity and confidentiality levels, with a four-point trusted/untrusted and private/public example.
  • Information-flow assumptions: Developers and users jointly specify a policy that allows a tool call only when its calling environment’s label flows to the tool’s policy label.The technique can generalize to finer-grained lattices.
  • Information-flow assumptions: The system assumes user and tool messages contain regions labeled as low-integrity external data or high-confidentiality data.These labels are assumed to be supplied to the system, which the paper identifies as an open IFC problem and developer burden.
  • Information-flow assumptions: Tools must comply with the information-flow policy because otherwise one tool can launder private or low-integrity information into another tool’s public or high-integrity output.The paper gives sending a private message and reading sent messages as an example.

7 Approach

RTBAS performs selective information-flow control for TBAS by screening which history regions influence the next output, propagating their security labels, and redacting more restrictive content. It supports both LM-judge and attention-based dependency screeners while allowing finer-grained message regions and lightweight runtime classification.

  • Dependency screening: The classifier maps input text, output text, and input regions to Boolean predictions indicating which regions influence the output.It formulates dependency analysis as sequential binary classification over region-level dependencies.
  • Implementation and deployment: The attention-based classifier achieved 85% training accuracy and 81% test accuracy on 40 labeled AgentDojo test cases with minimal runtime overhead.A local feature extractor and lightweight classifier run whenever the LM generates an output to track dependencies.
  • Security metadata: RTBAS supports multiple independently labeled regions within a message, although the simplified presentation treats each message as one uniformly labeled region.Messages carry integrity and confidentiality labels, and the implementation supports finer granularity for user messages and tool responses.
  • Selective information-flow enforcement: RTBAS screens history regions for relevance, propagates their confidentiality and integrity labels, and redacts content more restrictive than the resulting label.The screener aggregates labels from relevant regions; the redactor preserves permissible messages and replaces more restrictive ones with redaction markers.
  • Dependency screening: The dependency screener can be instantiated with either an LM-as-a-judge method or an attention-based method.The attention approach extracts regional attention features and maps them to dependency predictions with a lightweight two-layer LSTM.
  • Limitations: Screener errors cannot compromise security under selective masking, but they can reduce performance through over-tainting or under-tainting.Over-tainting may cause excessive user confirmations, while under-tainting may leave insufficient content available for task completion.

8 Evaluation

RTBAS is evaluated on prompt-injection and privacy-leakage benchmarks using utility, integrity, confirmation burden, and detection outcomes. It prevents policy-violating attacks while generally preserving utility, but performance depends on task domain, labeling, and confirmation decisions.

  • Setup: AgentDojo contains 79 realistic multi-round tasks across banking, travel, workspace, and Slack suites, with tools such as send_money and reserve_restaurant.The benchmark assigns integrity labels to regions and tools according to application requirements.
  • Prompt Injection Results: Without attacks, utility degrades by 10% for LM-Judge and 7.4% for Attention-based detectors, while Tool Filter slightly increases utility.The authors speculate that Tool Filter benefits from implicitly excluding irrelevant tools during planning.
  • Prompt Injection Results: Travel and workspace achieve the highest utility, exceeding the no-defense baseline by 5% on average, whereas Slack utility falls to 33% and 22% for Attention-based and LM-Judge detectors.Slack tasks often combine untrusted website content with actions, making them inherently unsafe and likely to require user confirmation.
  • Privacy Leakage Results: When users allow every tool call, utility reaches 83% and 75% for LM-Judge and Attention-based approaches, matching GPT’s 83% in the reported comparison.Users can permit potentially risky calls with confirmation, unlike Confirm Never, which provides no such autonomy.

8.3 Analysis

The analysis evaluates selective privacy-label propagation and compares the two dependency screeners. It finds improved label accuracy over conservative baselines, while detector overhead remains a practical cost.

  • Label Accuracy: More private categories than actually propagated can cause erroneous rejection, while fewer categories can mask information needed for the task.The label quality therefore affects both policy decisions and task functionality.
  • Privacy Leakage Analysis: More than 70% and 57% of selective-propagation decisions exactly match ground-truth privacy labels for the prompting and attention approaches, respectively.Both approaches outperform baselines that redact all sensitive regions or assume every tool call leaks every secret.
  • Screener Comparison: LM Judge performs slightly better overall under prompt injection and much better on privacy-leakage detection accuracy than Attention-based screening.The authors conjecture that explicit dependency reasoning and chain-of-thought output may help generalize across unseen tasks.
  • Runtime Overhead: Detector overhead increases runtime costs because Attention-based screening uses two LLM runs and LM Judge runs before each new agent message.The comparison is reported for executing banking-suite AgentDojo tasks, with metrics averaged across test cases and prices calculated against OpenAI pricing.

9 Discussion

The discussion identifies labeling requirements and resource cost as important deployment constraints. It also describes smaller-model screening as a possible efficiency improvement, though preliminary tests found general-purpose small models inadequate.

  • Labeling: RTBAS requires labeled tool and user messages plus an information-flow policy understandable to users.Region-based labeling can map fields such as email content, sender identity, and account balances to integrity or confidentiality policies.
  • Cost: Operating both dependency screeners is resource-intensive because Attention-based screening generates a preliminary message before producing the final masked output.The LM-Judge screener also adds computation before each new agent message.
  • Cost: Smaller models may reduce screening costs, but preliminary experiments found that small local models struggle as general-purpose screeners.Task-specific fine-tuning or prompt-tuning is proposed as a possible way to improve their efficiency and performance.

10 Conclusion

RTBAS applies fine-grained dynamic information-flow control to tool-based LLM agents through selective security-label propagation and redaction. On AgentDojo and a privacy-leakage benchmark, it blocked policy-violating attacks with under 2% utility degradation and achieved near-oracle leakage detection.

  • Framework: RTBAS selectively propagates relevant security labels through LM-Judge and Attention-based dependency screeners while redacting unused data.The mechanism targets both prompt injection and inadvertent privacy leaks.
  • Results: Under prompt injection attacks, RTBAS thwarts all policy-violating exploits with less than 2% degradation to agent task utility.The result is reported on the AgentDojo benchmark.
  • Results: RTBAS achieves near-oracle performance on detecting undesirable confidential data disclosures in the privacy-leakage benchmark.The evaluation also reports that the framework curbs malicious manipulations.

11 Ethics considerations

The study used public benchmarks and constructed datasets without real-world user data or personally identifiable information. Its attack evaluations simulated documented threats and potentially harmful tasks.

  • Data and Evaluation: Experiments used publicly available benchmarks and constructed datasets, without involving real-world user data or personally identifiable information.The attacks studied were documented within the research community and required no additional disclosure.
  • Data and Evaluation: The evaluation subjected language models to potentially harmful tasks and simulated attacks as part of the study.

A Appendix A: RTBAS Example Walk Through

The walkthrough shows RTBAS screening transaction data before allowing a potentially harmful money-transfer tool call. It identifies the malicious instruction, redacts private content, checks the proposed call against policy, and stops it pending user confirmation.

  • Screening: The screener identifies relevant regions in the transaction data and merges their labels for further processing.The example states that Regions 1 and 3 are determined relevant and their labels are returned together.
  • Prompt injection: A transaction description contains an untrusted instruction asking the agent to send Mallory $100 after answering the user’s question.The injected instruction is embedded in a transaction retrieved while answering a 31-day payment query.
  • Information handling: RTBAS redacts the private transaction description while retaining the public transaction information and the untrusted instruction for analysis.The example output replaces the private description with <<REDACTED>> and preserves the other displayed fields.
  • Policy enforcement: The proposed send_money call is evaluated by the runtime against policy before execution.The walkthrough explicitly identifies this as Step 4.
  • Policy enforcement: Because the tool call violates policy, RTBAS stops it pending user confirmation.The example therefore prevents the injected transfer from executing automatically.
Loading 2502.08966v2…