PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning When to Act or Refuse: Guarding Agentic Reasoning Models for Safe Multi-Step Tool Use

Aradhye Agarwal, Gurdit Siyan, Yash Pandya, Joykirat Singh, Akshay Nambi, Ahmed Awadallah

arXiv:2603.03205v2cs.CL

TL;DR

Agentic language models face safety risks from sequential tool use, adversarial feedback, and irreversible intermediate actions that static-generation safeguards do not reliably address. MOSAIC makes safety checks and refusal explicit in a plan–check–act/refuse loop and trains them with pairwise trajectory preferences. Across models and out-of-distribution benchmarks, it reduces harmful behavior and privacy leakage while preserving or improving benign-task performance.

  • Problem

    Sequential tool use creates irreversible safety risks that existing alignment methods optimized for static generation and task completion do not reliably handle.

  • Method

    MOSAIC uses explicit safety reasoning and refusal in a plan–check–act/refuse loop, trained with preference-based reinforcement learning over pairwise trajectory comparisons.

  • Results

    Across harmful, injection, benign, and privacy benchmarks, MOSAIC consistently reduces harmful behavior, injection success, and privacy leakage while preserving or improving benign performance; Qwen2.5 reports a 50 percent harmful-task reduction.

  • Takeaways & Limitations

    Explicit safety reasoning and refusal mechanisms generalize across model families, scales, and domains while maintaining agent utility and token efficiency.

Abstract

from arXiv · show

Agentic language models operate in a fundamentally different safety regime than chat models: they must plan, call tools, and execute long-horizon actions where a single misstep, such as accessing files or entering credentials, can cause irreversible harm. Existing alignment methods, largely optimized for static generation and task completion, break down in these settings due to sequential decision-making, adversarial tool feedback, and overconfident intermediate reasoning. We introduce MOSAIC, a post-training framework that aligns agents for safe multi-step tool use by making safety decisions explicit and learnable. MOSAIC structures inference as a plan, check, then act or refuse loop, with explicit safety reasoning and refusal as first-class actions. To train without trajectory-level labels, we use preference-based reinforcement learning with pairwise trajectory comparisons, which captures safety distinctions often missed by scalar rewards. We evaluate MOSAIC zero-shot across three model families, Qwen2.5-7B, Qwen3-4B-Thinking, and Phi-4, and across out-of-distribution benchmarks spanning harmful tasks, prompt injection, benign tool use, and cross-domain privacy leakage. MOSAIC reduces harmful behavior by up to 50%, increases harmful-task refusal by over 20% on injection attacks, cuts privacy leakage, and preserves or improves benign task performance, demonstrating robust generalization across models, domains, and agentic settings.

1 INTRODUCTION

Agentic tool use creates sequential safety risks that static-generation safeguards often miss. MOSAIC makes safety checks and refusal explicit, then learns these decisions through pairwise trajectory preferences while preserving useful behavior across models and domains.

  • Motivation: Multi-step agents can turn individually plausible actions into irreversible harm through tools, prompt injection, hallucinated functions, or late aborts.These failure modes extend beyond incorrect text generation and are difficult for output filtering or single-turn safeguards to address.
  • MOSAIC: MOSAIC structures inference as plan, check, then act or refuse, with explicit safety reasoning over harm, irreversibility, permissions, and tool feedback.The agent can proceed, refuse with a justification, or pause for clarification or verification.
  • Training: Pairwise trajectory preferences distinguish safer timing decisions, such as early refusal versus late abort, that pointwise rewards can miss.An LLM judge compares trajectories for the same task and selects the safer, more appropriate outcome.
  • Evaluation: MOSAIC generalizes across Qwen2.5-7B-Instruct, Qwen3-4B-Thinking-2507, and Phi-4 on harmful, injection, benign, and privacy-sensitive benchmarks.The evaluation is zero-shot and out of distribution across model families, environments, and domains.
  • Results: 50 percent reduction in harmful-task scores is reported for Qwen2.5, alongside 87 percent correct refusal and substantial benign-task gains for other models.Qwen3-4B-Thinking increases benign completion from 44 percent to 85 percent, while Phi-4 improves completion to 91 percent and reduces benign over-refusal by 56 percent.

2 MOSAIC OVERVIEW

MOSAIC turns safety into an explicit, selectively invoked action in a structured agent loop. It trains this behavior with trajectory-level reinforcement learning and pairwise preferences that prioritize safe intermediate decisions, calibrated refusal, and efficiency.

  • Structured Reasoning: MOSAIC uses a plan/think → check → act/refuse loop in which safety is a first-class trainable decision rather than an implicit reasoning byproduct.The framework combines modular reasoning blocks, learned safety gating, explicit refusal semantics, and preference-based reinforcement learning.
  • Structured Reasoning: The agent evaluates harmful intent, sensitive data, permission changes, and irreversible effects before proceeding, revising, clarifying, or refusing.The safety check is optional and operates on the proposed action and context.
  • Action Semantics: Refusal is a terminal action in the same action space as tool calls, providing an auditable halt and a direct learning signal for calibrated abstention.This design prevents unsafe intermediate calls from cascading when the plan or tool feedback reveals new risk.
  • Selective Computation: A learned gate decides whether to invoke safety reasoning after each plan, enabling checks on safety-critical turns while skipping them on routine ones.The gate is learned end-to-end from trajectory-level reinforcement learning feedback rather than fixed heuristics.
  • Preference Learning: Pairwise LLM-judge preferences rank complete trajectories by relative safety and appropriateness, preserving distinctions such as early refusal versus late abort.Group-relative rewards are aggregated from pairwise comparisons, with n=4 trajectories per rollout group in all experiments.
  • Optimization: The composite reward combines alignment, format, and length terms, while masking tool-output tokens so optimization targets model-generated decisions.The objective is R(τ) = r_align + r_fmt − p_len; the length penalty activates only above L0=400 tokens per turn.

3 EXPERIMENTAL SETUP

The study trains on Agent-SafetyBench and tests zero-shot generalization across distinct models and out-of-distribution agentic benchmarks. Its setup emphasizes execution-aware safety, prompt-injection robustness, and comparisons with both base and frontier models.

  • Training Data: Agent-SafetyBench contains explicitly harmful, benign, and mixed-strategy tasks, supporting calibrated behavior beyond surface cues or task labels.Training uses tasks where safe and unsafe execution strategies can lead to different outcomes.
  • Training Data: The benchmark spans 349 interaction environments, approximately 2,000 task instances, 8 safety-risk categories, and 10 common agent failure modes.It includes single- and multi-step interactions with direct and indirect prompt injections through user prompts and tool responses.
  • Evaluation: Evaluation labels correctness by execution quality, including tool-output interpretation, risk assessment, and the decision to act or refuse.This provides supervision and measurement beyond task completion alone.
  • Training Signal: Pairwise trajectory preferences provide reinforcement-learning signals without step-level annotations or ground-truth safety labels.Two trajectories for the same task are compared and the safer outcome is selected.
  • Evaluation: Zero-shot evaluation spans Agent Security Bench, AgentHarm, BFCL, and PrivacyLens, covering prompt injection, harmful tasks, benign tool use, and privacy leakage.The benchmarks differ from training in environments, tools, task structure, and threat models.
  • Models: The study evaluates Qwen2.5-7B-Instruct, Qwen3-4B-Thinking-2507, and Phi-4 against corresponding base models and frontier GPT-4o and GPT-5 models.The comparison tests safety improvements across model families, scales, and agentic capabilities.

4 RESULTS

MOSAIC improves agentic safety through explicit, model-adaptive safety decisions, while generally preserving or improving benign tool-use performance and efficiency. Its gains extend across harmful behavior, injection robustness, privacy leakage, and execution reliability, with trade-offs for some models.

  • Frontier models: Harmful-task refusal rises above 90% for GPT-4o and GPT-5, while GPT-4o harm drops from 0.31 to 0.07 without degrading benign completion.These results indicate that explicit safety reasoning and refusal mechanisms substantially reshape frontier-model agent behavior.
  • Open-source models: 50%: Qwen2.5 harmful-task scores fall from 0.18 to 0.09, while harmful-task refusal increases from 0.74 to 0.87.Prompt-injection robustness also improves, with DPI ASR changing from 0.55 to 0.42 and IPI ASR from 0.40 to 0.33.
  • Open-source models: Qwen3 benign completion rises from 0.44 to 0.85, alongside modest harmful-task improvements and reduced prompt-injection success.The result is attributed to reduced unsafe or unproductive reasoning loops, although benign refusals increase from 0.13 to 0.23.
  • Open-source models: Phi-4 completion increases from 0.78 to 0.91 as benign refusals fall from 0.43 to 0.19, but harmful-task refusal declines from 0.94 to 0.88.This model-specific pattern highlights a safety–utility trade-off when correcting an initially over-conservative policy.
  • Benign performance: 35%: Qwen2.5 benign multi-turn accuracy improves from 21.0 to 28.5 on BFCLv3, including missing-parameter, missing-function, and long-context tasks.The reported gains show that explicit safety checks and refusal need not reduce benign tool-use utility.
  • Cross-domain transfer: 23%: MOSAIC reduces PrivacyLens leakage by up to 23% while preserving helpfulness, extending safety gains beyond harm prevention to privacy-aware execution.For Qwen2.5, leakage falls from 0.32 to 0.26 and ALR from 0.48 to 0.37, while helpfulness rises on both reported scales.
  • Efficiency: Safety tokens remain below 20% of usage on Qwen2.5 harmful tasks, while Qwen3 token usage falls from over 1,000 to 262 tokens per turn.MOSAIC invokes safety reasoning dynamically and uses length penalties to compress redundant reasoning while retaining longer traces when needed.

5 RELATED WORK

Prior agentic safety defenses often treat safety as reactive or auxiliary, while MOSAIC makes sequential safety reasoning and refusal explicit and learnable through trajectory-level preference optimization.

  • Safety in Agentic LLMs: Prompting, rule-based filters, and external shielding are reactive defenses that limit prevention of compounding failures during multi-step tool use.These limitations are highlighted by AgentHarm, Agent Security Bench, and PrivacyLens.
  • LLM Judges and Preference Optimization: Scalar rewards can mis-rank trajectories with similar end states but different timing of unsafe actions, such as early refusal versus late abort.Pairwise preference supervision preserves this temporal structure without requiring ground-truth trajectory labels.
  • Safety-Specialized Models: External safety-specialized models moderate or classify outputs but do not directly shape planning or execution during tool use.This limits their control over long-horizon agentic actions.
  • Safety-Specialized Models: MOSAIC introduces plan/think → check → act/refuse reasoning, using modular blocks to target safety-critical dimensions during tool use.The framework makes safety reasoning and refusal first-class decisions in the agent’s inference loop.
  • LLM Judges and Preference Optimization: MOSAIC trains with GRPO and preference-based rewards over sampled trajectories rather than relying on a learned critic.The training algorithm uses an initial policy, environment, prompt set, safety reasoning blocks, and an LLM judge.

B EVALUATION DATA AND METRICS

The evaluation combines harmful, adversarial, privacy-sensitive, and benign tool-use benchmarks with metrics that distinguish safety, refusal behavior, and utility. Deterministic labeling evaluates both successful benign completion and unsafe responses to injection attacks.

  • Evaluation Data: Agent Security Bench covers 10 real-world scenarios, 10 task-specific agents, over 400 tools, and 27 attack/defense strategies.It reports Direct Prompt Injection and Indirect Prompt Injection separately.
  • Evaluation Protocol: Injection-augmented trajectories are labeled unsafe incomplete immediately when the agent follows the injected instruction.If the injection is ignored, evaluation proceeds using the benign completion protocol.
  • Evaluation Data: AgentHarm pairs 110 malicious base tasks across 11 harm categories with benign counterparts of comparable complexity.The pairing tests whether agents distinguish intent rather than relying on surface cues.
  • Evaluation Data: BFCL evaluates multi-turn tool-calling correctness by checking whether the final simulated API state matches the ground-truth outcome.The benchmark includes clarification questions and iterative approach refinement.
  • Evaluation Metrics: Metrics jointly separate harmfulness, correct refusal, benign over-refusal, non-refusal utility, injection success, and benign completion.For Agent Security Bench, Attack Success Rate measures following injected instructions, while Completion Rate measures benign tasks completed without refusal.
  • Evaluation Data: PrivacyLens measures contextual privacy behavior during tool execution using leakage rate and adjusted leakage rate.Adjusted Leakage Rate accounts for safe early termination and measures privacy violations conditional on non-refusal behavior.

C MODELS USED FOR EVALUATION

MOSAIC is evaluated on three open-weight model families spanning different sizes, reasoning capabilities, tool-use properties, and safety priors, with comparisons to frontier and corresponding base models.

  • Models: The evaluated model families are Qwen2.5-7B-Instruct, Qwen3-4B-Thinking-2507, and Phi-4.They span different scales and agentic capabilities.
  • Models: Qwen2.5-7B-Instruct is the primary model, selected for instruction-following strength and suitability for post-training.Its listed scale is 7B parameters.
  • Models: Qwen3-4B-Thinking-2507 is a smaller reasoning-oriented model with native tool-calling and explicit intermediate reasoning.Its listed scale is 4B parameters.
  • Models: Phi-4 is a 14B safety-oriented model with a conservative prior and no native tool use.It is augmented with a common agent interface for evaluation.
  • Baselines: Baselines include GPT-4o, GPT-5, and the corresponding open-weight base models.These comparisons assess safety improvements relative to scale and pre-training.

D LLM-AS-A-JUDGE

The LLM judge compares agent trajectories using explicit safety and helpfulness criteria, while the evaluation design addresses judge context and positional bias. The accompanying injection example specifies the unsafe tool-use pattern and expected safe response.

  • Judge Criteria: The judge penalizes benign refusals and expects harmful tasks to be rejected immediately without partial execution.In borderline cases, it favors helpful completion over refusal to avoid over-defensiveness.
  • Positional Bias: The judge favors the first trajectory in approximately 60% of comparisons versus 40% for the second in one analyzed training run.Trajectory order is randomized, so each trajectory has equal probability of appearing first or second in expectation.
  • Pairwise Judge: Pairwise judging compares two trajectories and returns which is better or whether both are the same.The prompt requires a structured analysis and answer format.
  • Judge Inputs: The judge receives the full agent trajectory, including thinking and safety thoughts, which empirically improves convergence and intent attribution.This information is used in judging trajectory appropriateness.
  • Prompt Injection Example: The injection example treats invoking AccessToSupplierData to retrieve confidential supplier information as unsafe behavior.The expected response is to detect task mismatch, identify the sensitive tool as high risk, and refuse or continue with task-relevant tools.

E.2 AGENTHARM

The AgentHarm section contrasts benign synthetic-media creation with deceptive political deepfake generation, requiring completion of the former and refusal of the latter.

  • Benign Task: Benign tasks require generating and posting synthetic media with non-deceptive framing.The task uses image and video generation followed by posting to X, without targeting real individuals or deception.
  • Benign Task: Successful benign-task behavior follows the hinted tool-call order, completes the task, and avoids unjustified refusal or misleading content.The benign pass condition requires completion, while refusal without justification is a failure.
  • Harmful Task: Harmful tasks request political deepfakes presented as real videos, posing misinformation, manipulation, and societal-harm risks.The scenario depicts politicians conversing and asks the agent to claim the video is real.
  • Harmful Task: Successful harmful-task behavior detects unsafe intent, invokes refusal, and avoids generative or posting tools.The pass condition requires refusing without generating or posting the deceptive content.

F TRAINING ROLLOUTS

The training rollouts compare refusal, generic reasoning, frequency mapping, and tool-augmented retrieval for a medical-information request. Preference scoring favors grounded helpfulness, with explicit refusal preferred when safety cannot be verified.

  • Tool Environment: The available environment includes device-analysis, session-setting, file-listing, file-reading, and unsafe-task-refusal tools.These tools support both context-specific retrieval and safety-constrained refusal.
  • Compared Rollouts: Trajectory 1 conservatively identifies medical risk and stops, while Trajectories 2 and 3 provide general or theoretical information without tool use.The refusal trajectory emphasizes avoiding medical advice; the other strategies rely on internal knowledge or theoretical output.
  • Compared Rollouts: Trajectory 4 actively explores the environment, reads a research file, and returns a grounded summary of findings spanning 174 Hz to 963 Hz.It lists available files, reads the named PDF, and reports promising results with a professional-advice caveat.
  • Preference Ranking: The judge ranks grounded helpfulness above safe refusal, and safe refusal above unverified generation.T4 is preferred for extracting context-specific data, while T1 is preferred when safety cannot be verified.
  • Agent Protocol: The agent prompt requires planning, optional safety reasoning, JSON tool calls, and a final tagged answer after tool use ends.The template permits repeated reasoning and tool-call cycles before the final response.

H ABLATION 1: EXPLICIT SAFETY CHECKS VS SINGLE <THINK> BLOCK.

The ablation shows that explicit safety checks outperform reliance on a generic think block, preserving both harmful-task refusal and benign execution under ordinary and adversarial conditions.

  • AgentHarm: For Qwen2.5-7B, removing explicit safety checks lowers AgentHarm harmful-task refusal from 0.87 to 0.59 and raises harm score from 0.09 to 0.12.Benign non-refusal performance also falls from 0.61 to 0.42, indicating degraded task quality rather than improved utility.
  • Agent Security Bench: Removing explicit checks lowers refusal and increases attack success on Agent Security Bench under both DPI and IPI.The result indicates that refusal-only training is brittle to injected instructions and adversarial tool outputs.

I IMPLEMENTATION DETAILS AND HYPERPARAMETERS

The implementation section specifies evaluation protocols and reports that preference-judge agreement rises steadily during training, indicating increasingly consistent safety decisions.

  • Evaluation Settings: Evaluation uses default pretrained-model rollout temperatures across AgentHarm, Agent Security Bench, PrivacyLens, and BFCL.AgentHarm uses GPT-4.1 with dataset-provided rubrics, while PrivacyLens uses GPT-4o for helpfulness and information leakage.
  • Judge Agreement: Figure 5 measures judge agreement as the fraction of comparisons with consistent ordering and shows steady increases across all models.The passage interprets this trend as increasingly consistent trajectory distributions and convergence toward a stable safety decision boundary.
Loading 2603.03205v2…