Multi-step agent safety stops checking each step alone: TAS caps tail risk with one shared budget, cutting average attack success to 45.15%

Safety for multi-step LLM agents shifts from independent per-step verification to tracing risk propagation along the whole workflow and constraining tail risk with one shared budget. Earlier prompt engineering and post-training alignment mainly improve the safety of single-step LLM generation, while external-verification defenses check each step separately and ignore risk propagation between steps. On Agent Security Bench, TAS reaches average attack success rates of 45.15% under direct prompt injection and 29.73% under indirect prompt injection, below no defense and the eleven defense baselines it compares against.

The approach treats the agent workflow's safety states as a chain of dependent states, so each step's action and output are conditioned on the current safety state. TAS models this sequence as a Markov decision process, making risk propagation traceable instead of treating steps as independent checkpoints. Across the Llama3-70B, Qwen2-72B, GPT-4o and GPT-5.2 backbones, removing MDP modeling raises attack success under direct prompt injection from 43.50% to 50.75% (Llama3-70B) and under indirect prompt injection from 13.95% to 18.50%.

The risk measure weights high-cost tail outcomes instead of looking only at overall average cost. An average-based measure is diluted by the many low-cost normal steps and can hide rare but high-cost failures. Replacing the tail-sensitive measure with an overall average raises Llama3-70B's attack success from 43.50% to 69.75% under direct prompt injection and from 13.95% to 22.25% under indirect prompt injection.

Each step's decision is constrained by a safety budget shared across steps; when the budget is insufficient, the step is rolled back and an alternative action is resampled. Independent per-step thresholds cannot trace risk accumulated over the whole workflow, and an earlier high-risk step consumes budget available to later steps. Removing the shared budget constraint raises Llama3-70B's attack success under direct prompt injection from 43.50% to 87.75%, the largest safety degradation in the ablation.

Evaluation covers three benchmarks, Agent Security Bench, Agent-SafetyBench and AgentDojo, with Llama3-70B, Qwen2-72B, GPT-4o and GPT-5.2 as the four main-experiment backbones; results for the smaller Llama3-8B and Qwen2-7B appear in supplementary material. The method targets safety-critical multi-step agent workflows where reducing high-risk behavior is the priority.

Each step's safety cost comes from a rule-based score combined with an LLM scorer, so scorer bias can affect the accuracy of risk estimation; TAS may still slightly compromise benign-task utility and adds computation from sampling-based estimation, so users with strict utility or latency requirements need to tune the safety-related hyperparameters carefully.