Tail Risk-Aware Agent Safety (TAS): Constraining Tail Risk in Multi-Step LLM Agent Workflows

1. 01What problem does this paper attempt to solve?

   Existing agent defenses mostly check each step independently; the authors argue that cross-step risk propagation should be modeled and tail risk constrained.

   Core question: how to defend against attacks targeting individual steps in multi-step LLM agent workflows, especially rare but costly tail risks, without harming utility on normal tasks.

   • Problem setting: LLM agents extend the LLM with tool calls, memory, planning and retrieval, and attackers can launch attacks against different steps such as input, retrieval, planning and execution, causing task failure, malicious content or harmful actions (Introduction).
   • Shortcomings of existing methods (per the authors): prompt engineering and post-training alignment only improve single-step LLM safety and ignore non-LLM components such as retrieval, planning and execution; external-verification methods (e.g., GuardAgent) check each step independently and ignore risk propagation between steps.
   • Core observation: most agent behaviors are safe and near-zero cost, and only a small fraction are costly; measuring risk by an overall average masks these tail risks, while using an average constraint to cover tail risks may over-penalize normal cases and harm utility (§3.3).
   • Proposed approach: the paper proposes Tail-risk-aware Agent Safety (TAS), which models workflow safety as a sequence of safety-dependent states (MDP), measures tail risk with a spectral risk measure, and constrains it step by step with a budget shared across the entire workflow, deriving a step-wise safety policy.
   • Threat model:
     • Attacker goal: make the agent task fail, produce malicious content, or execute harmful actions.
     • Attack surface: the input, retrieval, planning and execution steps of the workflow; experiments cover direct prompt injection (DPI), indirect prompt injection (IPI) and POT backdoor attacks.
     • Victim system: agent workflows with open-source and closed-source LLMs as backbones (ASB, Agent-SafetyBench, AgentDojo environments).
     • Information available to the defender: at each step, the tool name, parameters, PII/secret features, external target checks and current workflow context, used for rule-based and API semantic scoring (§3.3).

   Collapse this answer

2. 02What related research exists?

   Related work is divided into two groups, LLM safety (prompt engineering, post-training alignment) and agent safety (internal alignment, external verification); the authors argue the latter ignores cross-step risk propagation.

   LLM Safety: Prompt Engineering

   • Xiong et al. (2025), Zhou, Li, and Wang (2024), Li et al. (2025b)—respectively design adaptive suffix prompts, optimization-based robust prompts, and prompt compression to defend against jailbreaks; the authors point out that such prompts need to be carefully designed for the task, limiting generalization.
   • Xu et al. (2024), Zhang et al. (2024b)—SafeDecoding amplifies the probability of safe tokens at decoding time; the goal-priority method reconciles helpfulness and safety during training and inference; the authors point out that they are vulnerable to adversarial prompt attacks.

   LLM Safety: Post-training Alignment

   • Chaudhary et al. (2024), Mu et al. (2024)—risk-averse RLHF and Rule-Based Rewards (RBR), which reduce harmful content and unnecessary refusals during fine-tuning.
   • Gu et al. (2025), Maeng et al. (2025), Hsu et al. (2024), Ji et al. (2024)—fine-tuning or output-repair methods such as Mask-DPO, GC-DPO, Safe LoRA, Aligner. The authors summarize that SFT relies on human annotation, and RL requires carefully designed rewards and may exhibit reward hacking.

   Agent Safety: Internal Alignment and External Verification

   • Internal alignment: Tennant, Hailes, and Musolesi (2025) design a moral alignment reward function for LLM agents; Qiao et al. (2024) propose a World Knowledge Model to provide global priors and reduce hallucinated actions. The authors cite Huang et al. (2024) to point out that LLMs struggle to self-correct reasoning errors without external feedback.
   • External verification: GuardAgent (Xiang et al. 2024) uses a guardrail agent to verify behavior; Progent (Shi et al. 2025) performs programmable privilege control; DRIFT (Li et al. 2026) performs dynamic rule isolation; Bhagwatkar et al. (2025) study tool input/output firewalls; AgentSentinel (Hu et al. 2025) intercepts sensitive operations and audits execution traces. The authors point out that these methods usually handle each step in isolation or rely on local tool-level checks, ignoring risk propagation across multi-step processes.

   Baselines and Benchmarks

   • Baselines: under DPI/IPI, compare Deli., Para., Instr., Rewr., Sand.; under POT, compare Para. and Shuffle; additionally compare 11 defense baselines in total: Progent, DRIFT, Firewall, GuardAgent, Sentinel.
   • Benchmarks: Agent Security Bench (ASB), Agent-SafetyBench, AgentDojo; the appendix additionally uses the native protocols of GuardAgent and AgentSentinel (EHRAgent/EICU-AC, SeeAct/Mind2Web-SC).

   The authors position this paper as follows: unlike external verification that checks each step independently, TAS models dependencies between steps and uses a shared budget to constrain the tail risk of the entire workflow.

   Collapse this answer

3. 03How does the paper address this problem?

   TAS has three steps: MDP modeling of safety state transitions, SRM measuring tail risk, and a stepwise safety policy using a shared budget plus sampling estimation.

   Overall idea: TAS formalizes the agent workflow as a sequential decision process over safety states, estimates the risk distribution of candidate actions step by step, constrains only tail (rare high-cost) risk, and uses a budget shared across steps to decide whether to allow, stop, or roll back.

   Safety states modeled as a sequential decision process (§3.2)

   • Safety state: the observable safety state at step t, ht ≜ (x, st, Ct), comprising the user query x, the cumulative workflow observations st (memory, retrieval, planning, LLM generation, tool call results, environment feedback), and the cumulative safety cost Ct.
   • Two-stage decision: given ht, first select an action according to Pθ(at+1 | ht) (which module/tool to execute), then generate the in-step output according to Pθ(s̃t+1 | ht, at+1); the joint distribution is the product of the two (Eq. (1)).
   • State transition: st+1 = concat(st, s̃t+1), Ct+1 = Ct + c(ht, at+1, s̃t+1), yielding ht+1; a T-step workflow induces a safety state trajectory τ, whose distribution is expanded by the chain rule (Eq. (2)). The authors state that they use an MDP rather than ordinary sequential transitions so that transitions are conditioned on action decisions, tracking risk propagation.

   Tail safety measure based on spectral risk (§3.3)

   • Stepwise safety cost: ct ≜ c(ht, at+1, s̃t+1), composed of two parts: rule scoring (first a risk pre-classification by tool name, parameters, PII/secret patterns, and external target checks, then computing a nonlinear feature cost) and API semantic scoring (inputting the tool name, truncated parameters, rule features, and current context, returning a semantic risk score); the two are combined into a final ct in budget units; the full rules and prompts are in Appendix C.
   • Cumulative cost: C(\tau) \triangleq \sum_{t=0}^{T-1} c_t, i.e., the sum of step costs over the entire trajectory.
   • Tail risk: the u-quantile weighted integral of C(τ) using a spectral risk measure (SRM), M_\sigma(C(\tau)) \triangleq \int_\alpha^1 \sigma(u) F_C^{-1}(u)\,du, where α is the quantile cutoff of the tail region, and σ is non-negative, monotonically increasing, and normalized, so that high-cost outcomes receive greater weight. The authors state that expectation-type measures would mask tail risk under a long-tailed cost distribution.

   Shared budget constraint and empirical estimation (§3.4)

   • Shared budget: the constraint Mσ(C(τ)) ≤ B; its online form checks whether the tail risk of the incremental cost ΔCt+1 does not exceed the remaining budget Bt ≜ B − Ct (Eq. (5)). The authors state that independent single-step thresholds cannot track cumulative risk, and that the shared budget couples the safety decisions of each step.
   • Empirical estimation: the true risk distribution is unknown; K candidate outputs are sampled under the same (ht, at+1) (Eq. (6)), sorted, and the tail samples after ⌈αK⌉ are weighted and summed to obtain the empirical tail risk (Eq. (7)).
   • Stepwise policy: maximize utility Jutil(π) under the constraint; according to the budget check, there are three cases: M < Bt 放行并优先效用；M = Bt 放行该步但随后停止工作流；M > Bt rejects the step, rolls back to ht, blocks the action and resamples; if retries (e.g., 3 times) all fail, return a deterministic safe fallback response and terminate the workflow.

   Key parameters

   • Main setting: α = 0.95, B = 20, K = 200, σ(·) = e^(·), sampling temperature 0.7, top-p 0.95, API scoring model gpt-5-mini (Table 5).
   • Parameter sweeps: B ∈ {5, 10, 20, 40, 80, 160, 10000} (α = 0.95 fixed); α ∈ {0.99, 0.95, 0.85, 0.70, 0.50, 0.25} (B = 20 fixed); K ∈ {1, 10, 100, 200, 500, 1000} (Appendices A, B.5).

   Collapse this answer

4. 04What experiments did the paper conduct?

   Compared with 11 baselines on three benchmarks and four backbone models, the authors state that TAS has the lowest average ASR and utility close to no defense, with improvements passing the McNemar test at p < 0.01.

   Experimental Setup

   • Models tested: four main backbones Llama3-70B-Instruct, Qwen2-72B-Instruct, GPT-4o, GPT-5.2; Appendix B.1 additionally reports LlaMA3-8B-Instruct and Qwen2-7B-Instruct. Only the backbone is changed; the rest of the workflow stays the same.
   • Benchmarks: ASB (DPI, IPI, POT attacks), Agent-SafetyBench, AgentDojo; Appendix B.3 evaluates EHRAgent/EICU-AC and SeeAct/Mind2Web-SC under the original protocols. The paper does not state the sample size of each benchmark.
   • Baselines: under DPI/IPI, Deli., Para., Instr., Rewr., Sand.; under POT, Para., Shuffle; additionally Progent, DRIFT, Firewall, GuardAgent, Sentinel.
   • Metrics: ASR↓ (ASB, AgentDojo), Safety Score↑ (Agent-SafetyBench), PNA↑, NRP = PNA × (1 − ASR), Utility↑ (AgentDojo).
   • Scoring method: rule-based scoring plus API semantic scoring, with gpt-5-mini as the API scoring model (Table 5); the paper does not report agreement between the judge and human annotators.
   • Statistical test: two-sided exact McNemar test, p < 0.01 (Table 15, App. B.8); the paper does not report the number of repetitions.

   Main Results

   Average ASR, NRP, PNA for DPI/IPI attacks on ASB (%, from Table 1):

   Swipe horizontally to read the table

   Attack	Metric	No Def.	Strongest baseline	Ours
   DPI	ASR ↓	76.89	63.91 (Para.)	45.15
   DPI	NRP ↑	14.29	22.37 (Para.)	34.90
   DPI	PNA ↑	61.81	62.00 (Para.)	63.63
   IPI	ASR ↓	48.60	46.03 (Deli.)	29.73
   IPI	NRP ↑	31.77	33.50 (Sand.)	44.71
   IPI	PNA ↑	61.81	62.19 (Sand.)	63.63

   • By model, under DPI the ASR of TAS is 43.50 for Llama3-70B, 48.35 for Qwen2-72B, 40.35 for GPT-4o, 48.40 for GPT-5.2; under IPI it is 13.95, 17.05, 45.65, 42.25 (Table 1).
   • The authors state that TAS improves average safety over No Defense by 47.95% (DPI) and 70.25% (IPI), and exceeds the previous SOTA baselines by 24.28% and 20.53%; PNA is close to No Defense.
   • Exception: in Table 1, a few baselines are not lower than TAS on individual items; e.g., under DPI, Instr. on Llama3-70B has PNA 61.75, higher than TAS's 58.25.

   POT Backdoor Attacks and Other Benchmarks

   • POT (Table 2, 8): across the four backbones, TAS's ASR is 15.00 (Llama3-70B), 43.25 (Qwen2-72B), 42.75 (GPT-4o), 39.75 (GPT-5.2), averaging 35.19, lower than No Defense's 85.19 and Para.'s 63.56; average NRP 41.24 is the highest, and PNA 63.63 is higher than No Defense's 61.81.
   • Agent-SafetyBench (Table 3): TAS's Safety Score is 0.508, 0.483, 0.758, 0.908 (four backbones); the authors state it improves on average by 61.36% over No Defense and by 11.41% over the strongest baseline; on Qwen2-72B, TAS ties GuardAgent at 0.483.
   • AgentDojo (Table 3): TAS's ASR is 0.00, 0.00, 3.93, 1.36, dropping from No Defense's 47.51 to 3.93 on GPT-4o; Utility is 11.11, 9.15, 63.32, 76.43, higher than No Defense's 56.74 on GPT-4o, while on GPT-5.2 76.43 is slightly lower than No Defense's 76.91.

   Ablations and Analysis

   • Component ablation (Table 4): averaged over the four backbones under DPI, removing the shared budget constraint (w/o constr.) gives the highest ASR (e.g., Qwen2-72B 99.75, Llama3-70B 87.75); the authors state the budget is most critical for controlling risk propagation; w/o tail, w/o sampling, w/o mdp, w/o srm also all give higher ASR than the full method.
   • Tail risk visualization (Figure 2): the horizontal axis is tail probability and the vertical axis is the risk cost threshold; the authors state that on the four backbones the baseline curves rise toward the upper left, while the TAS curve stays in the lower left, indicating smaller tail risk (the figure is a line plot; only the authors' statement and the axes are relayed here).
   • Utility-safety trade-off (Figure 3, Appendix A): the authors state that increasing B or decreasing α improves utility but lowers safety, and that B ≈ 20, α ≈ 0.95 is a practical balance point.

   Other Ablations and Analysis

   • Number of samples K (Table 12, GPT-4o, DPI): as K goes from 1 to 1000, ASR drops from 68.58 to 31.09, and relative computational cost rises from 1.00 to 12.27; the main experiments use K = 200 (ASR 40.35, cost 2.34).
   • Computational overhead (Table 11): TAS's total overhead is 2.34× that of No Defense, while baselines are 0.98×–1.14×; in the overhead breakdown, sampling accounts for 1.19× and scorer calls 0.15×, and MDP modeling introduces no additional inference latency.
   • Scorer robustness (Table 13): after artificially introducing 10% scoring misclassification, average ASR changes from 45.15 to 45.61, and PNA drops from 63.63 to 59.37.
   • FPR/FNR (Table 14, GPT-4o, AgentDojo): TAS's FPR is 36.68 and FNR 3.93; Sentinel's FPR is 89.81, and DRIFT's FPR is 69.40.
   • Small models (Table 6): on LlaMA3-8B and Qwen2-7B, TAS's ASR is also lower than each baseline (e.g., Qwen2-7B's IPI ASR is 1.65).
   • GuardAgent/AgentSentinel original protocols (Table 9, 10): under the AgentSentinel protocol, TAS averages DPI DSR 0.825, IPI DSR 0.775, FT 0.129, all better than AgentSentinel's 0.658, 0.460, 0.390.

   Collapse this answer

5. 05What are some points that can be further explored?

   In Appendix D.2, the authors point out two limitations, scorer bias and utility and computational overhead, and propose directions such as a more unbiased scorer and dynamic budget allocation.

   Limitations and future directions noted by the authors

   • Scorer dependence and bias (App. D.2): TAS relies on rule and LLM scorers to estimate stepwise safety cost, which may introduce bias and in some cases harm the accuracy of risk estimation; the authors state that designing more unbiased and more generalizable risk scorers is an important challenge for future work.
   • Richer environment abstractions and dynamic budgets (App. D.2): the authors state that another promising direction is to introduce richer environment abstractions and perform more dynamic risk budget allocation within the workflow.
   • Utility and computational cost (App. D.2): the authors state that TAS may slightly harm benign task utility and incur additional computational cost, because stronger safety constraints and larger sampling scale reject or modify some useful executions; users with strict utility or latency requirements need to carefully tune the safety-related hyperparameters.

   Experimental coverage

   • The models under test in the main experiments are 4 backbones in total: Llama3-70B-Instruct, Qwen2-72B-Instruct, GPT-4o, GPT-5.2; results for LlaMA3-8B-Instruct and Qwen2-7B-Instruct are in Appendix B.1 (Tables 6, 7).
   • The evaluation benchmarks are 3 in total: ASB (three attacks: DPI, IPI, POT), Agent-SafetyBench, AgentDojo, plus supplementary experiments on the native protocols of GuardAgent and AgentSentinel (App. B.3).
   • The defense comparison covers 11 baselines; the attack types cover DPI, IPI, and POT backdoor attacks.
   • The paper does not report the sample size of each benchmark, the number of experimental repetitions, or the agreement between the LLM scorer (gpt-5-mini) and human judgment.
   • The main hyperparameter setting is α = 0.95, B = 20, K = 200 (Table 5); the sweep analysis is conducted under GPT-4o and the DPI setting (App. B.5).

   Collapse this answer

6. 06Summarize the main content of the paper

   TAS uses an MDP plus a shared budget to constrain the tail risk of agent workflows; the authors state that on three benchmarks it outperforms 11 baselines in safety and overall performance.

   An agent safety defense paper addressing the problem that in multi-step LLM agent workflows attacks can occur at each step of input, retrieval, planning, and execution, and that existing defenses mostly check each step independently.

   • Method: TAS models the workflow as an MDP over safety states (query, cumulative observations, cumulative cost) to track risk propagation; uses a spectral risk measure (SRM) to weight rare high-cost tail risk; checks step by step with a budget B shared across the entire workflow, allowing, stopping, or rolling back and resampling according to whether the budget is satisfied, at the boundary, or exceeded; the risk distribution is empirically estimated with K = 200 samples, and scoring is given by rules and a gpt-5-mini API scorer.
   • Main results (Table 1): averaged over the four backbones on ASB, DPI attack ASR drops from 76.89% for No Defense to 45.15%, IPI from 48.60% to 29.73%, and PNA is close to No Defense (63.63 vs. 61.81).
   • POT and other benchmarks: POT attack average ASR is 35.19%, lower than 85.19% for No Defense (Table 2); on Agent-SafetyBench, GPT-5.2's Safety Score reaches 0.908; on AgentDojo, GPT-4o's ASR drops from 47.51% to 3.93% (Table 3). The authors state that the improvements pass the two-sided exact McNemar test at p < 0.01.
   • Ablation (Table 4): removing the shared budget constraint raises ASR the most, and the authors state that budget coupling is key to controlling risk propagation; removing the tail measure, the SRM weighting, or the sampling estimation also worsens safety.
   • Cost: computational overhead is 2.34× that of No Defense (Table 11), mainly from sampling estimation.
   • Authors' conclusion: constraining tail risk rather than the overall average can suppress rare high-cost failures while preserving normal utility; the authors also note that the scorer may introduce bias and that the method may slightly harm utility and increase computational cost (App. D.2).

   Collapse this answer

Related papers