Learning When to Act or Refuse: Guarding Agentic Reasoning Models for Safe Multi-Step Tool Use
On the map: LLM reasoning → Reasoning for agents, planning and code
also in 1 more
Abstract (arXiv)
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.
Key facts. Cited by 7 papers in the map (data through 7 Oct 2026). Placed in Reasoning for agents, planning and code, one of 17 directions in LLM reasoning. Builds on 30 of its 52 references that are in the map. First posted on arXiv 3 Mar 2026.
Builds on
30 of 52 references.
Closest ancestors · 4
Same direction, newest first
- 2025AgentAuditor: Human-Level Safety and Security Evaluation for LLM Agents31 May
- 2025Think Twice Before You Act: Enhancing Agent Behavioral Safety with Thought Correction16 May
- 2025Agentic Reasoning and Tool Integration for LLMs via Reinforcement Learning28 Apr
- 2025ShieldAgent: Shielding Agents via Verifiable Safety Policy Reasoning26 Mar
Other references · 26
- 2025DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learningcited by 6,146
- 2024GPT-4o System Cardcited by 4,467
- 2024Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phonecited by 1,883
- 2024A Survey on LLM-as-a-Judgecited by 1,384
- 2023Not what you've signed up for: Compromising Real-World LLM-Integrated Applications with Indirect Prompt Injectioncited by 1,339
- 2025OpenAI GPT-5 System Cardcited by 910
- 2024DeepSeek-V3 Technical Reportcited by 724
- 2024Phi-4 Technical Reportcited by 675
- 2023Prompt Injection attack against LLM-integrated Applicationscited by 568
- 2024AgentHarm: A Benchmark for Measuring Harmfulness of LLM Agentscited by 406
- 2024Agent Security Bench (ASB): Formalizing and Benchmarking Attacks and Defenses in LLM-based Agentscited by 341
- 2025Small Language Models are the Future of Agentic AIcited by 254
14 more
- 2024The Landscape of Emerging AI Agent Architectures for Reasoning, Planning, and Tool Calling: A Surveycited by 164
- 2024PrivacyLens: Evaluating Privacy Norm Awareness of Language Models in Actioncited by 161
- 2024Aligning with Human Judgement: The Role of Pairwise Preference in Large Language Model Evaluatorscited by 146
- 2024AutoDefense: Multi-Agent LLM Defense against Jailbreak Attackscited by 135
- 2025Qwen3Guard Technical Reportcited by 125
- 2024SafeAgentBench: A Benchmark for Safe Task Planning of Embodied LLM Agentscited by 102
- 2025Semantic-Space Exploration and Exploitation in RLVR for LLM Reasoningcited by 17
- 2025Towards Improving Reward Design in RL: A Reward Alignment Metric for RL Practitionerscited by 10
- 2025Exposing Weak Links in Multi-Agent Systems under Adversarial Promptingcited by 8
- 2025SafeAgent: Safeguarding LLM Agents via an Automated Risk Simulatorcited by 6
- 2025ToolRegistry: A Protocol-Agnostic Tool Management Library for Function-Calling LLMscited by 5
- 2025STAC: When Innocent Tools Form Dangerous Chains for LLM Agentscited by 5
- 2024Challenges in Guardrailing Large Language Models for Sciencecited by 4
- 2025Guardians of the Agentic System: Preventing Many Shots Jailbreak with Agentic Systemcited by 4
Built on by
Citing papers per year: 2026: 7.
7 papers.
- 2026From Question Answering to Task Completion: A Survey on Agent System and Harness Designcited by 9
- 2026Trust No Tool: Evaluating and Defending LLM Agents under Untrusted Tool Feedbackcited by 2
- 2026MIRROR: A Hierarchical Benchmark for Metacognitive Calibration in Large Language Modelscited by 1
- 2026Don't Offer What Can't Be Done: Deterministic Executability Gating for LLM Skill Selection at Scalecited by 1
- 2026Learning When to Act: Interval-Aware Reinforcement Learning with Predictive Temporal Structurecited by 0
- 2026Confuse the Model, Control the Flow: Understanding and Mitigating Privacy Leakage from LLM Agents with Information Flow Controlcited by 0
- 2026Jailbreak Context Lingers: Divergent Safety Routing and Its Cross-Task Predictability in Tool Agentscited by 0
- 2026Jailbreak Context Lingers: Divergent Safety Routing and Its Cross-Task Predictability in Tool Agents28 Sep 2026
- 2026Confuse the Model, Control the Flow: Understanding and Mitigating Privacy Leakage from LLM Agents with Information Flow Control12 Sep 2026
- 2026Don't Offer What Can't Be Done: Deterministic Executability Gating for LLM Skill Selection at Scale2 Aug 2026
- 2026From Question Answering to Task Completion: A Survey on Agent System and Harness Design14 Jun 2026
- 2026Trust No Tool: Evaluating and Defending LLM Agents under Untrusted Tool Feedback17 May 2026
- 2026MIRROR: A Hierarchical Benchmark for Metacognitive Calibration in Large Language Models15 Apr 2026
- 2026Learning When to Act: Interval-Aware Reinforcement Learning with Predictive Temporal Structure23 Mar 2026
Similar, not cited
Close in wording, no citation either way
- 2025Agent Safety Alignment via Reinforcement Learning
- 2026AgentBoundary: Counterfactual Evaluation of Safety in Tool-Using LLM Agents
- 2025AlphaAlign: Incentivizing Safety Alignment with Extremely Simplified Reinforcement Learning
- 2025AgentAlign: Navigating Safety Alignment in the Shift from Informative to Agentic Large Language Models
- 2025InvThink: Premortem Reasoning for Safer Language Models
- 2026Agent Safety Is Action Alignment
- 2026MOSAIC: Composable Safety Alignment with Modular Control Tokens
- 2025Agentic Moderation: Multi-Agent Design for Safer Vision-Language Models