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

1. Problem Statement

The deployment of Large Language Models (LLMs) as autonomous agents introduces a fundamentally different safety regime compared to static chat models. While chat models generate text, agents plan, call tools, and execute long-horizon actions in real-world environments.

• The Core Challenge: A single misstep (e.g., accessing sensitive files, executing payment commands, or entering credentials) can cause irreversible harm.
• Limitations of Current Methods:
  • Static Alignment: Existing safety alignment methods are optimized for single-turn generation and fail in sequential decision-making settings.
  • Implicit Safety: Agents often treat safety as an implicit byproduct of reasoning rather than an explicit decision, leading to "cascading failures" where an agent proceeds with unsafe intermediate steps before realizing the risk.
  • Scalar Reward Failure: Standard reinforcement learning (RL) using scalar rewards (e.g., "task completed" vs. "failed") collapses multi-step safety distinctions. For instance, a trajectory that refuses immediately is often scored similarly to one that follows an injected instruction and aborts late, despite the latter being significantly more dangerous.
  • Small Language Models (SLMs): SLMs, favored for cost and privacy, are particularly vulnerable to prompt injection and adversarial tool feedback due to compressed world models and tighter context budgets.

2. Methodology: The MOSAIC Framework

The authors introduce MOSAIC (Modular, Organized, Safety-Aware, Integrated Control), a post-training framework designed to make safety decisions explicit, modular, and learnable.

A. Structured Inference Loop

MOSAIC restructures the agent's reasoning process into a Plan → Check → Act/Refuse loop:

1. Plan: The agent generates a plan and candidate tool action using standard <thought> reasoning.
2. Check (Explicit Safety Reasoning): The agent may invoke an explicit <safety thoughts> block. This is a first-class action, not an implicit part of the thought process. It performs structured self-reflection on:
   • Potential harm and irreversibility.
   • Permission changes and sensitive data handling.
   • Risks revealed by recent tool feedback.
3. Act or Refuse: Based on the check, the agent chooses one of three paths:
   • Execute: Proceed with the tool call.
   • Refuse: Invoke a dedicated refuse_unsafe_task tool with a justification (terminal action).
   • Replan/Halt: Request user clarification or verification.

A learned gating mechanism determines whether to invoke the <safety thoughts> block, ensuring safety reasoning is applied only when necessary to maintain token efficiency.

B. Training: Preference-Based Reinforcement Learning

To train agents without requiring massive annotated corpora of "safe" vs. "unsafe" trajectories, MOSAIC employs Group Relative Policy Optimization (GRPO) with Pairwise Preference Learning.

• LLM-as-a-Judge: Instead of scalar rewards, an LLM judge compares pairs of trajectories for the same task.
• Relative Ordering: The judge selects the "safer" trajectory. This captures critical temporal distinctions, such as preferring early refusal over late abort (after unsafe actions have occurred).
• Composite Reward Function: The final reward $R(\tau)$ combines:
  • Alignment Reward ($r_{align}$): Derived from pairwise preferences (0–3 scale).
  • Format Reward ($r_{fmt}$): Penalizes malformed outputs (invalid tags, nesting errors).
  • Length Penalty ($p_{len}$): Discourages unnecessary verbosity while allowing flexibility for complex safety checks.

3. Key Contributions

1. Explicit Safety as a First-Class Action: MOSAIC decouples safety reasoning from general planning, making "checking" and "refusing" explicit, learnable decisions within the agent's action space.
2. Pairwise Preference Optimization: The paper demonstrates that pairwise trajectory comparisons are superior to scalar rewards for agentic safety, as they preserve the temporal ordering of safety decisions (e.g., distinguishing between a safe refusal and a dangerous execution followed by a late stop).
3. Model-Adaptive Generalization: The framework is shown to generalize across different model families (Qwen2.5, Qwen3, Phi-4) and scales, adapting to each model's specific failure modes (e.g., over-refusal vs. under-refusal).

4. Experimental Results

The authors evaluated MOSAIC on three open-weight models (Qwen2.5-7B, Qwen3-4B-Thinking, Phi-4) and compared them against frontier closed models (GPT-4o, GPT-5) across four benchmarks: AgentHarm, Agent Security Bench (ASB), BFCL v3, and PrivacyLens.

• Safety Improvements:

  • Harm Reduction: MOSAIC reduced harmful behavior by up to 50% (e.g., Qwen2.5 harmful score dropped from 0.18 to 0.09).
  • Refusal Rates: Harmful-task refusal increased significantly (e.g., Qwen2.5 from 74% to 87%).
  • Injection Robustness: On ASB, MOSAIC reduced Direct Prompt Injection (DPI) success rates by ~20% and Indirect Prompt Injection (IPI) by ~10%.
  • Privacy: On PrivacyLens, MOSAIC reduced information leakage by up to 23% while maintaining helpfulness.

• Utility Preservation:

  • Unlike rigid safety filters that cause over-refusal, MOSAIC improved benign task performance. For example, Qwen3-4B-Thinking saw benign task completion jump from 44% to 85% by avoiding endless reasoning loops.
  • Phi-4 reduced benign over-refusal by 56%, recovering utility without sacrificing safety.

• Efficiency:

  • Token Efficiency: Safety reasoning was invoked dynamically. Safety tokens remained below 20% of total usage.
  • Length Reduction: For verbose models like Qwen3, the length penalty combined with selective safety checks reduced total token usage by 75% (from ~1000 to ~262 tokens per turn) without degrading performance.

• Comparison to Frontier Models:

  • Without explicit safety scaffolding, GPT-4o and GPT-5 failed significantly on agentic safety (0% refusal on harmful tasks).
  • MOSAIC-trained open models outperformed unscaffolded frontier models and became broadly comparable to them once the frontier models were also equipped with explicit safety scaffolding.

5. Significance and Conclusion

The paper establishes that agentic safety is not merely a function of model scale but of structured inference and training signals.

• Paradigm Shift: It moves safety from being an external filter or an implicit byproduct of reasoning to an explicit, modular control mechanism integrated into the agent's decision loop.
• Scalability: The approach is effective for Small Language Models (SLMs), which are critical for edge deployment and privacy, proving that robust agentic safety can be achieved through post-training alignment rather than relying solely on massive pre-trained models.
• Generalizability: The framework successfully handles diverse threats, including prompt injection, privacy leakage, and harmful task execution, while preserving the agent's ability to complete benign, complex multi-step tasks.

In summary, MOSAIC provides a principled path to building reliable AI agents by teaching them when to act, when to verify, and when to refuse, ensuring that safety is a learnable, dynamic capability rather than a static constraint.