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

cs.CLMar 3, 2026

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

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

Organizations: Microsoft Research

Abstract

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.

Figures & tables

Figure 1: MOSAIC overview. Agents follow a Plan → Check → Act/Refuse loop, where an explicit safety check determines whether to execute a tool action or refuse and replan, making safety decisions modular, explicit, and learnable.

Figure 2: MOSAIC System design.

Model	AgentHarm (Harmful)			AgentHarm (Benign)		ASB (DPI)		ASB (IPI)		ASB (–)
	Harm $\downarrow$	Refusal $\uparrow$	NR Harm $\downarrow$	Refusal $\downarrow$	NR Score $\uparrow$	ASR $\downarrow$	RR $\uparrow$	ASR $\downarrow$	RR $\uparrow$	CR $\uparrow$
GPT-4o (base $\rightarrow$ MOSAIC)	0.31 $\rightarrow$ 0.07	0.00 $\rightarrow$ 0.92	0.31 $\rightarrow$ 0.67	0.00 $\rightarrow$ 0.19	0.80 $\rightarrow$ 0.75	0.76 $\rightarrow$ 0.21	0.00 $\rightarrow$ 0.79	0.67 $\rightarrow$ 0.27	0.00 $\rightarrow$ 0.63	0.89 $\rightarrow$ 0.93
	+0.24	+0.92	-0.36	-0.19	-0.05	+0.55	+0.79	+0.40	+0.63	+0.04
GPT-5 (base $\rightarrow$ MOSAIC)	0.11 $\rightarrow$ 0.06	0.00 $\rightarrow$ 0.91	0.11 $\rightarrow$ 0.57	0.00 $\rightarrow$ 0.24	0.68 $\rightarrow$ 0.73	0.42 $\rightarrow$ 0.26	0.00 $\rightarrow$ 0.67	0.03 $\rightarrow$ 0.02	0.00 $\rightarrow$ 0.65	0.98 $\rightarrow$ 0.99
	+0.05	+0.91	-0.46	-0.24	+0.05	+0.16	+0.67	+0.01	+0.65	+0.01
Qwen2.5-7B (base $\rightarrow$ MOSAIC)	0.18 $\rightarrow$ 0.09	0.74 $\rightarrow$ 0.87	0.58 $\rightarrow$ 0.52	0.13 $\rightarrow$ 0.15	0.51 $\rightarrow$ 0.61	0.55 $\rightarrow$ 0.42	0.42 $\rightarrow$ 0.58	0.40 $\rightarrow$ 0.33	0.44 $\rightarrow$ 0.61	0.90 $\rightarrow$ 0.84

Table 1: Impact of MOSAIC safety scaffolding on frontier models. Each column reports baseline

\rightarrow

MOSAIC, with

\Delta

colored by whether MOSAIC improves (green) or degrades (red) the metric, according to the metric direction (

\uparrow/\downarrow

Model	Mis. Func.	Mis. Param.	Base	LC
Qwen2.5-base	17.0	12.0	21.0	12.5
Qwen2.5- MOSAIC	18.0	14.0	28.5	13.5

Table 2: BFCLv3 benign agentic performance. Higher is better.

Model	LR $\downarrow$	ALR $\downarrow$	Help. $\uparrow$	Help. $\uparrow$
			(0–1)	(0–3)
Qwen2.5 (base)	0.32	0.48	0.59	1.79
Qwen2.5 ( MOSAIC )	0.26	0.37	0.63	1.90
Phi-4 (base)	0.38	0.42	0.87	2.61
Phi-4 ( MOSAIC )	0.32	0.41	0.76	2.27

Table 3: PrivacyLens results for base and MOSAIC -trained models.

Panel of Figure 3 — Figure 3: Mean response length over training.

Model	Benign			Harmful
Model	Safety	Think	Total	Safety	Think	Total
Qwen2.5 (base)	9	24	124	17	36	122
Qwen2.5 (MOSAIC)	36	62	182	37	81	188
Qwen3 (base)	15	726	1172	88	593	983
Qwen3 (MOSAIC)	0	181	262	0	187	258
Phi-4 (base)	37	105	230	57	85	226

Table 4: Average per-turn token usage on AH, decomposed into safety, reasoning, and total tokens for benign and harmful tasks.

Model	AgentHarm			ASB (DPI)
	H $\downarrow$	R $\uparrow$	NR $\uparrow$	ASR $\downarrow$	RR $\uparrow$	CR $\uparrow$
Qwen2.5 (only think )	0.12	0.59	0.42	0.49	0.34	0.94
Qwen2.5 (ptwise reward)	0.14	0.79	0.54	0.51	0.48	0.89
Qwen2.5 ( MOSAIC )	0.09	0.87	0.61	0.42	0.58	0.84

Table 5: Ablation: Qwen2.5-7B without/with safety thoughts (row 1 and row 3); with pointwise and pairwise rewards (row 2 and row 3).

Figure 4: Judge agreement over training.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

R(\tau_{i})=r_{\text{align}}(\tau_{i})+r_{\text{fmt}}(\tau_{i})-p_{\text{len}}(\tau_{i})

Appendix

Algorithm 1 Training MOSAIC with GRPO and Preference-Based Rewards

Judge Preference	vs. $T_{1}$	vs. $T_{2}$	vs. $T_{3}$	vs. $T_{4}$	Summed Reward
Trajectory $T_{1}$	—	Win	Win	Lose	2
Trajectory $T_{2}$	Lose	—	Tie	Tie	1
Trajectory $T_{3}$	Lose	Tie	—	Lose	0.5
Trajectory $T_{4}$	Win	Tie	Win	—	2.5

Appendix

Table 6: Pairwise Preference Matrix for Trajectories

T_{1}

–

T_{4}

Model	AgentHarm					Agent Security Bench
	Harmful			Benign		DPI		IPI		CR
	Harm $\downarrow$	Ref. $\uparrow$	NR $\downarrow$	Ref. $\downarrow$	NR $\uparrow$	ASR $\downarrow$	RR $\uparrow$	ASR $\downarrow$	RR $\uparrow$	$\uparrow$
Qwen2.5-7B (only think )	0.12	0.59	0.28	0.07	0.42	0.49	0.34	0.24	0.28	0.94
Qwen2.5-7B (Pointwise reward)	0.14	0.79	0.61	0.1	0.54	0.51	0.48	0.44	0.52	0.89
Qwen2.5-7B (MOSAIC)	0.09	0.87	0.52	0.15	0.61	0.42	0.58	0.33	0.61	0.84

Appendix

Table 7: Ablation on Qwen2.5-7B with and without explicit safety checks, evaluated on AgentHarm and Agent Security Bench.

Training
Hyperparameter	Value
Train batch size	4
PPO mini-batch size	16
Max prompt length	2000
Max response length	5000
Rollout samples ( $n$ )	4

Appendix

Table 8: Training and evaluation hyperparameters.

Figure 5: Judge agreement over training.

LLM Agent Safety

Large language model agent safety studies how to keep tool-using, persistent, and multi-agent systems aligned with user intent and authorized limits.

See all 321 papers

AgentBoundary: Counterfactual Evaluation of Safety in Tool-Using LLM Agents
Sep 27, 2026Tianzhuo Yang, Zirui Mi, Yantao Huang +4LLM Agent Evaluation LLM Agents
Peking University · Beijing Academy of Artificial Intelligence
Agent Safety Is Action Alignment
Jun 27, 2026Shawn Li, Yue ZhaoLanguage Model Safety Evaluation AI Alignment
University of Southern California · Los Angeles, California, USA
MOSAIC: Composable Safety Alignment with Modular Control Tokens
Mar 17, 2026Jingyu Peng, Hongyu Chen, Jiancheng Dong +5LLM Safety Alignment Inference-Time Language Model Alignment
University of Science and Technology of China · City University of Hong Kong · Baidu Inc +1

Sep 27, 2026J/K move · Enter open · S save

AgentBoundary: Counterfactual Evaluation of Safety in Tool-Using LLM Agents

Tianzhuo Yang, Zirui Mi, Yantao Huang, Guoxi Zhang+3

Peking University · Beijing Academy of Artificial Intelligence

LLM Agent Evaluation LLM Agents Over-Refusal AI Agent Safety Benchmarks LLM Agent Safety

Open paper AskPDF arXiv

Safety alignment for large language models (LLMs) in conversational settings is largely framed around whether to answer or refuse a request. In agentic settings, however, the same models must decide whether to act as permission-critical evidence emerges during execution. This creates a distinct challenge: apparent risk, action permissibility, and task competence are easily confounded, making agentic over-refusal difficult to distinguish from ordinary task failure. To address this, we introduce AgentBound, the first four-way counterfactual generation-and-evaluation framework for tool-using agent safety. AgentBound transforms the same executable workflow by independently varying apparent risk and action permissibility, enabling controlled comparisons of risky-looking but authorized tasks and routine-looking but unauthorized tasks. These comparisons jointly diagnose over-refusal and unsafe compliance while controlling for task competence. We instantiate AgentBound as a human-validated 4,000-task evaluation suite with trajectory-based and post-state-based judgments. Across 17 model and harness configurations, high safety frequently coexists with poor authorized-task completion: GPT-5.5 blocks 99.5% of routine-looking unauthorized actions yet completes only 28.7% of risky-looking authorized tasks. We further train a lightweight runtime calibration module that improves authorized-task completion by 18.2% on average across 10 evaluated configurations, while improving unsafe-action blocking by 5.4% on average. These show that effective agentic alignment requires action decisions to track permission-relevant execution evidence, rather than refusal strength alone.

Abstract

LLM Agent Safety

Explore similar work

AgentBoundary: Counterfactual Evaluation of Safety in Tool-Using LLM Agents

Agent Safety Is Action Alignment

MOSAIC: Composable Safety Alignment with Modular Control Tokens