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

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, Phi-4 and Qwen3-30B-A3B-Thinking, 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.

AI is rapidly shifting from simply answering questions to acting as "agents" that take real actions for us—browsing the web, editing files, even sending money. Unlike a chatbot's words, these actions can cause lasting harm, and a single misstep is hard to undo. Yet today's safety training was built for conversation, so agents often slip: they follow malicious instructions hidden in a webpage or document, or they timidly refuse perfectly harmless requests. We built MOSAIC, which reorganizes how an agent thinks before each action into a simple loop: plan, check for safety, then act or refuse. Refusing becomes a deliberate choice rather than an afterthought. To teach this, instead of scoring each attempt on its own, we let an AI "judge" compare pairs of attempts and pick the safer one—because catching a problem early beats stopping only after the damage is done, a distinction that single scores miss. MOSAIC cut harmful actions by up to half, reduced privacy leaks and vulnerability to hidden attacks, and still finished legitimate tasks—often using fewer words. Remarkably, small open models trained this way matched far larger commercial ones, suggesting safe AI agents come from better training, not just bigger size.