PaperlayerView exact revision on arXiv

Source-linked AI summary

Reassembling Distributed Risk: Trajectory-Conditioned Action Generation for Multi-Turn Agent Safety

Yanbo Dai, Zhenlan Ji, Zongjie Li, Shuai Wang

arXiv:2608.25711v1cs.CR

TL;DR

Tool-using agents can accumulate harmful intent across plausible multi-turn interactions, while existing defenses rely on auxiliary online reasoning or post-generation action assessment. ReDiR injects a compact, trajectory-conditioned safety representation into a frozen base model before each action, achieving below 8% attack success rates, transfer to unseen domains, and preserved benign fidelity.

  • Problem

    Multi-turn agent attacks distribute harmful objectives across plausible requests and tool calls, making isolated-step safety judgments inadequate for actions that affect external systems.

  • Method

    ReDiR compresses each current trajectory into a latent safety representation learned through same-model, cross-view supervision and injects it into the frozen base model before action generation.

  • Results

    Below 8% attack success rate, ReDiR transfers across eight held-out tool domains while preserving benign fidelity and low computational overhead.

  • Takeaways & Limitations

    Trajectory-conditioned safety representations can shape action generation directly while avoiding a separate action-level safety module and transferring beyond the supervised tool domain.

  • Takeaways & Limitations

    Deployment assumes white-box access to the base model and may incur backbone-specific latency, particularly for Gemma-4-E4B.

Abstract

from arXiv · show

Tool-using LLM agents extend security risks beyond generated text to actions that affect external systems. Under multi-turn decomposition attacks, a harmful objective can be distributed across individually plausible requests and tool calls, becoming apparent only from the accumulated trajectory. Existing defenses either rely on auxiliary online reasoning to recover long-horizon security evidence or assess actions after generation, often incurring additional inference cost or depending on runtime-specific action representations. We propose \emph{Reassembling Distributed Risk} (ReDiR), a generation-time defense that conditions action generation on trajectory-level security evidence. Before each action, ReDiR compresses the current trajectory into a compact latent safety representation and injects it into the frozen base model. The representation is learned through same-model, cross-view supervision, where safe behavior from an explicit task view provides supervision for recovering distributed safety evidence from the original multi-turn trajectory. This design enables ReDiR to integrate cross-turn security information directly within the generation process without relying on a separate action-level safety module. We evaluate ReDiR on two agent-safety benchmarks across three model families and eight held-out tool domains. ReDiR reduces attack success rates to below 8\%, transfers to unseen tool domains, and preserves benign fidelity with low computational overhead.

1 Introduction

Tool-using agents face trajectory-level safety risks because harmful objectives can emerge from individually plausible turns and culminate in consequential actions. ReDiR addresses this by conditioning action generation on compact trajectory-level safety evidence and achieves strong attack reduction while preserving benign behavior.

  • Motivation: Tool-using agents can modify external systems, so safety must constrain executable actions before they take effect.Dispatched actions may create difficult-to-reverse external state changes, unlike harmful text that can sometimes be filtered before reaching downstream systems.
  • Motivation: Multi-turn decomposition distributes harmful intent across plausible requests, tool observations, and state changes, making isolated-step judgments insufficient.A destructive command can be assembled from helper functions across turns before being executed in a final tool call.
  • ReDiR: ReDiR conditions action generation on trajectory-level safety evidence by compressing the current trajectory into a compact latent representation before each action.The representation is learned through same-model, cross-view supervision and injected into the base model to shape the next action.
  • Results: Below 8% attack success rate, ReDiR outperforms competing defenses near 20%–34% and remains below 7% across eight held-out tool domains.The evaluation covers two benchmarks and three model families; benign fidelity remains high, while harmful-task runtime decreases by up to 89% relative to base agents.
  • ReDiR: ReDiR aggregates distributed safety evidence before action generation through a trajectory-conditioned representation abstraction.Its training combines audited same-model, cross-view action targets with selective corrective training and benign-behavior retention.

2 Preliminaries and Related Work

The paper models tool use as an agent–runtime–environment interaction whose full trajectory determines later actions. Related defenses either assess generated actions after the fact or use auxiliary trajectory-risk mechanisms, motivating direct trajectory-conditioned action generation.

  • Tool-Using Agent Execution: A tool-using agent combines a language model with runtime R, which mediates tool interactions and stateful environment updates.The runtime formats the trajectory into model input, interprets outputs as valid actions, and incorporates observations into the next trajectory prefix.
  • Tool-Using Agent Execution: The model-visible trajectory includes prior instructions, responses, actions, tool observations, and exposed tool specifications rather than only the current turn.These accumulated interactions form the execution trajectory governing multi-turn behavior.
  • Multi-Turn Agent Security: Multi-turn harmful intent may emerge only from accumulated evidence, because individually plausible actions can compose into unsafe outcomes.Safe-looking text can still accompany a policy-violating tool call, so isolated textual or stepwise assessment is insufficient.
  • Related Work: Existing defenses use auxiliary long-horizon reasoning or post-generation action inspection, which can add online decisions or depend on runtime-specific action formats.MAGE uses shadow memory and an online judge, while runtime guards inspect generated tool calls before dispatch.
  • Research Gap: The open design goal is to integrate trajectory-level safety evidence directly into action generation while transferring across tool schemas and preserving benign behavior.The desired defense should avoid reliance on runtime-specific post-generation guards and maintain practical online efficiency.

3 Threat Model and Design Requirements

The threat model gives attackers control over multi-turn inputs while limiting them from modifying the agent stack, and gives defenders white-box control before each action. The design requirements emphasize effectiveness, transferability, fidelity, and efficiency across deployment settings.

  • 3.1 Threat Model: Attackers control the sequence of user inputs and seek policy-violating actions, while unable to modify the model, runtime, tools, credentials, or defense.The harmful objective may remain distributed across turns until the trajectory is considered as a whole.
  • 3.1 Threat Model: Before each action, defenders observe the visible trajectory prefix and use white-box access to influence generation without future observations or hidden task information.The defender aims to prevent policy violations while preserving legitimate agent behavior.
  • 3.2 Design Requirements: Effectiveness requires preventing multi-turn harmful objectives before they produce policy-violating tool effects, even when individual requests appear benign.Transferability requires working across tool schemas and runtime formats without safety-specific retraining, rules, or per-deployment target construction.
  • 3.2 Design Requirements: The framework is evaluated under effectiveness, transferability, fidelity, and efficiency requirements, with strong support indicating explicit design and reported evaluation evidence.Table 1 uses one, two, and three filled bars for limited, partial, and strong support.
  • 3.2 Design Requirements: MAGE preserves benign utility but relies on external memory and an online judge, whereas TRACE produces auxiliary risk judgments and ToolShield uses tool-specific exploration for new tools.These trade-offs motivate a defense that directly controls action generation across deployments.

4 Reassembling Distributed Risk

ReDiR constructs audited safe action supervision from explicit task views, then learns trajectory-conditioned latent safety representations that guide generation while preserving the base model’s action interface.

  • 4 Reassembling Distributed Risk: ReDiR addresses safe-action supervision, trajectory-conditioned learning, and action-level control through offline target construction, safety encoding, and latent-guided generation.The design pairs original trajectory prefixes with audited safe targets and injects the learned representation before the next action is generated.
  • 4.1 Offline Action Supervision Construction: The supervision pipeline decomposes harmful objectives into benign-looking multi-turn requests, executes them to obtain trajectories, and filters prefixes before harm occurs.Prefixes whose earlier actions already caused the harmful outcome are discarded because the next action can no longer prevent that harm.
  • 4.1 Offline Action Supervision Construction: Same-model cross-view target generation makes distributed safety evidence explicit, samples candidate actions, and retains targets that are safe, executable, and terminate harmful execution.The collapsed view preserves deployed instructions and tool specifications while exposing the canonical objective more clearly.
  • 4.1 Offline Action Supervision Construction: Corrective supervision is routed as NONE, TERMINATE, or FULL according to whether the student output already safely terminates, needs only termination repair, or requires full correction.Already-safe behavior is left unchanged, while supervision targets the termination span or full audited target as appropriate.
  • 4.2 Trajectory-Conditioned Safety Training: The safety encoder uses learned latent queries to aggregate trajectory evidence into a fixed-size representation, normalizes it, and appends it to the trajectory-prefix embeddings.The encoder is adapted with LoRA while the action-generating base model remains frozen.
  • 4.2 Trajectory-Conditioned Safety Training: The action-grounded objective applies weighted token-level corrective supervision only at routed positions, emphasizing reasoning entry tokens and executable decision tokens.Examples routed to NONE contribute no corrective loss; entry and decision weights are combined through the larger applicable weight.

5 Experimental Setup

The experiments evaluate ReDiR across multiple agent-safety benchmarks, model backbones, threat settings, and efficiency measures using execution-oriented safety metrics.

  • Models: ReDiR is evaluated with Qwen3.5-9B, Ministral-3-8B-Instruct, and Gemma-4-E4B, generally sharing the pretrained backbone between policy and safety encoder.The safety encoder uses an additional LoRA module, while a smaller encoder is also considered as an alternative.
  • Benchmarks: MT-AgentRisk covers multi-turn decomposition attacks across Addition and Decomposition formats, Data Files and Environment States, and eight attack subcategories.The benchmark spans 365 tasks.
  • Benchmarks: AgentDojo evaluates indirect prompt injection through untrusted tool observations, providing an out-of-distribution test beyond MT-AgentRisk’s multi-turn decomposition attacks.The evaluations use native tool environments and agent pipelines.
  • Metrics and Efficiency: ASR measures whether harmful outcomes are realized, while SSR measures valid safety-preserving termination before realization; both distinguish outcome prevention from explicit intervention.The study also reports benign utility for AgentDojo and multiple efficiency measures.
  • Results Reporting: Table 2 reports MT-AgentRisk Strict ASR and SSR by transformation format and risk carrier, and AgentDojo ASR and benign utility.The table uses percentage metrics with lower ASR and higher SSR or benign utility preferred.

6 Evaluation

ReDiR consistently reduces multi-turn attack success, transfers safety across unseen tool domains, preserves benign behavior, and operates with low online and moderate offline costs. Ablations further show that backbone and target compatibility, sufficient latent capacity, and supervised token coverage materially affect performance.

  • Multi-Turn Security Effectiveness: 0.0–7.9% ASR across three model families, versus 20.3–34.2% for MAGE, with SSRs of 82.7–98.6%.ReDiR also achieves the lowest ASR across every tested transformation format and risk carrier.
  • Multi-Turn Security Effectiveness: 5.5%, 0.0%, and 7.9% overall ASR on Qwen3.5-9B, Ministral-3-8B, and Gemma-4-E4B, respectively.With Ministral-3-8B, ReDiR eliminates successful attacks across all four MT-AgentRisk constructions.
  • Cross-Tool Transferability: A maximum 6.1-percentage-point difference separates held-out-domain and Filesystem ASR across model families.Evaluation uses Filesystem supervision and eight held-out environments without target-domain adaptation.
  • Benign Behavior Fidelity: 0.0% FPR across all four evaluated benign domains, matching the unmodified base agent and ToolShield.MAGE introduces unnecessary refusals on benign tasks in several domains.
  • Efficiency: 1.328–3.250 MiB of additional GPU memory and a 53.3% average end-to-end latency reduction across the three models.The memory savings come from reusing one shared backbone; ReDiR also incurs 38.14 GPU-hours of offline compute, 37.0% below MAGE’s 60.55 GPU-hours.

7 Resilience to Adaptive Attackers

ReDiR is evaluated against benign-history dilution, an adaptive attack that prepends read-only benign interactions to weaken safety-relevant trajectory evidence. Across the tested conditions, this manipulation does not systematically degrade ReDiR’s safety performance.

  • Adaptive attack: The attack targets ReDiR’s fixed latent safety embeddings by adding benign context before the harmful request.
  • Adaptive attack: Benign-history dilution prepends N benign requests and read-only tool calls before a harmful final request to dilute safety-relevant evidence.The evaluation varies N ∈ {0,1,2,4,8} across 30 harmful tasks in Filesystem, Terminal, and PostgreSQL.

N Filesystem Terminal PostgreSQL Overall

ReDiR remains stable under increasing benign-history dilution, with attack and safety-success rates staying within narrow ranges across the tested conditions.

  • ASR remains within 6.7–10.0% as benign-history length increases across the evaluated conditions.
  • SSR stays between 86.7% and 90.0% across all benign-history conditions.

8 Conclusion

ReDiR addresses multi-turn tool attacks by learning trajectory-conditioned safety representations that guide the frozen base model before action generation. Across models and tool domains, it lowers attack success, transfers to unseen domains, and preserves benign behavior with low overhead.

  • Multi-turn tool attacks distribute harmful intent across individually plausible actions, so safety depends on evidence accumulated throughout the interaction.
  • ReDiR uses same-model, cross-view supervision to learn trajectory-conditioned safety representations that guide the frozen base model before action generation.
  • ReDiR reduces attack success rates to below 8% across multiple model families and tool domains, generalizes to unseen tool domains, and preserves benign behavior with low overhead.

Ethical Considerations

The work evaluates a defensive method in controlled benchmark environments while recognizing dual-use risks and limitations. The authors therefore pair restrained artifact release with a recommendation that ReDiR complement broader execution safeguards.

  • The study evaluates harmful actions only in benchmark-controlled environments using test data and credentials, not production services or real user accounts.
  • ReDiR’s attack transformations and defense-aware evaluation may be dual use, motivating omission of unnecessary operational detail and safety review of released artifacts.
  • The authors state that ReDiR may miss unseen attacks, interrupt benign tasks, or motivate adaptive bypasses.
  • The paper recommends complementing ReDiR with least privilege, sandboxing, and execution monitoring.

A Training Data Construction

The training-data pipeline collects eligible pre-action states, generates task-level refusal candidates, audits and filters them, then routes and selects corrective state–target pairs. It constructs supervision from 55 source tasks while holding out 15 tasks entirely.

  • 55 of 70 MT-AgentRisk Filesystem tasks provide safety supervision, while 15 tasks are held out from collection, generation, routing, and optimization.
  • 76 eligible final-turn states remain after prefix auditing removes trajectories that already realized the harmful objective.The initial pool contained 78 states spanning 43 tasks.
  • 440 teacher candidates are sampled by generating eight native-protocol completions for each of the 55 source tasks.Generation occurs at the task level, regardless of whether a task contributes an eligible trajectory state.
  • Candidates are structurally filtered for a well-formed, non-empty FINISH action with no other native action, while recovering exact routing token spans.The parser extracts the action header, FINISH function name, and message argument spans.
  • Semantic auditing accepts only VALID-REFUSAL candidates with confidence at least 0.8, requiring risk identification, an unambiguous refusal, and no harmful completion or nontermination proposal.The judge evaluates the canonical task, complete generation, and extracted refusal message, using labels including VALID-REFUSAL, WRONG-RISK, NOT-COMMITTED, UNSAFE, and AMBIGUOUS.
  • 52 audited targets across 28 tasks yield a final safety set of 36 trajectory–target pairs from 36 states and 21 tasks.The final set contains 34 full-target examples and two termination-only examples; one corrective target is retained per state to avoid overrepresentation.

B Additional Training Details

The remaining optimization settings for the main experiments are reported in Table 16, while the action-generating base model remains frozen during training.

  • The main experiments use the encoder architecture, LoRA placement, supervision weighting, and batch composition described in Section 5.
  • Table 16 reports the remaining optimization settings used in the main experiments.
  • The action-generating base model remains frozen throughout training.
  • Training is conducted on NVIDIA A100 GPUs.
Loading 2608.25711v1…