# ToolSafe: Enhancing Tool Invocation Safety of LLM-based agents via Proactive Step-level Guardrail and Feedback

arXiv: https://arxiv.org/abs/2601.10156 · HTML: https://mlsanity.com/paper/2601.10156/

## Key facts

- Cited by 72 papers in the map (data through 14 Oct 2026).
- Placed in Agent threats, guardrails and governance, one of 16 directions in LLM agents and tool use.
- Builds on 29 of its 39 references that are in the map.
- First posted on arXiv 15 Jan 2026.

## Placed in

- LLM agents and tool use / [Agent threats, guardrails and governance](https://mlsanity.com/agents/agent-threats-guardrails-and-governance/)
- LLM safety and security / [Prompt injection and agent security](https://mlsanity.com/safety/prompt-injection-and-agent-security/)

## Builds on (most cited first, 29 of 29)

- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [τ-bench: A Benchmark for Tool-Agent-User Interaction in Real-World Domains](https://mlsanity.com/paper/2406.12045/) (2024)
- [SFT Memorizes, RL Generalizes: A Comparative Study of Foundation Model Post-training](https://mlsanity.com/paper/2501.17161/) (2025)
- [WildGuard: Open One-Stop Moderation Tools for Safety Risks, Jailbreaks, and Refusals of LLMs](https://mlsanity.com/paper/2406.18495/) (2024)
- [GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models](https://mlsanity.com/paper/2508.06471/) (2025)
- [Agent-SafetyBench: Evaluating the Safety of LLM Agents](https://mlsanity.com/paper/2412.14470/) (2024)
- [ShieldGemma: Generative AI Content Moderation Based on Gemma](https://mlsanity.com/paper/2407.21772/) (2024)
- [SALAD-Bench: A Hierarchical and Comprehensive Safety Benchmark for Large Language Models](https://mlsanity.com/paper/2402.05044/) (2024)
- [WASP: Benchmarking Web Agent Security Against Prompt Injection Attacks](https://mlsanity.com/paper/2504.18575/) (2025)
- [ShieldAgent: Shielding Agents via Verifiable Safety Policy Reasoning](https://mlsanity.com/paper/2503.22738/) (2025)
- [AgentDojo: A Dynamic Environment to Evaluate Prompt Injection Attacks and Defenses for LLM Agents](https://mlsanity.com/paper/2406.13352/) (2024)
- [Watch Out for Your Agents! Investigating Backdoor Threats to LLM-Based Agents](https://mlsanity.com/paper/2402.11208/) (2024)
- [BadAgent: Inserting and Activating Backdoor Attacks in LLM Agents](https://mlsanity.com/paper/2406.03007/) (2024)
- [AgentAuditor: Human-Level Safety and Security Evaluation for LLM Agents](https://mlsanity.com/paper/2506.00641/) (2025)
- [AGrail: A Lifelong Agent Guardrail with Effective and Adaptive Safety Detection](https://mlsanity.com/paper/2502.11448/) (2025)
- [ToolSword: Unveiling Safety Issues of Large Language Models in Tool Learning Across Three Stages](https://mlsanity.com/paper/2402.10753/) (2024)
- [A-MemGuard: A Proactive Defense Framework for LLM-Based Agent Memory](https://mlsanity.com/paper/2510.02373/) (2025)
- [AdvAgent: Controllable Blackbox Red-teaming on Web Agents](https://mlsanity.com/paper/2410.17401/) (2024)
- [All Roads Lead to Likelihood: The Value of Reinforcement Learning in Fine-Tuning](https://mlsanity.com/paper/2503.01067/) (2025)
- [PolyGuard: A Multilingual Safety Moderation Tool for 17 Languages](https://mlsanity.com/paper/2504.04377/) (2025)
- [IPIGuard: A Novel Tool Dependency Graph-Based Defense Against Indirect Prompt Injection in LLM Agents](https://mlsanity.com/paper/2508.15310/) (2025)
- [InjecGuard: Benchmarking and Mitigating Over-defense in Prompt Injection Guardrail Models](https://mlsanity.com/paper/2410.22770/) (2024)
- [LessLeak-Bench: A First Investigation of Data Leakage in LLMs Across 83 Software Engineering Benchmarks](https://mlsanity.com/paper/2502.06215/) (2025)
- [Building a Foundational Guardrail for General Agentic Systems via Synthetic Data](https://mlsanity.com/paper/2510.09781/) (2025)
- [OS-Sentinel: Towards Safety-Enhanced Mobile GUI Agents via Hybrid Validation in Realistic Workflows](https://mlsanity.com/paper/2510.24411/) (2025)
- [PromptGuard: Soft Prompt-Guided Unsafe Content Moderation for Text-to-Image Models](https://mlsanity.com/paper/2501.03544/) (2025)
- [Efficient Universal Goal Hijacking with Semantics-guided Prompt Organization](https://mlsanity.com/paper/2405.14189/) (2024)
- [Quantifying Misalignment Between Agents: Towards a Sociotechnical Understanding of Alignment](https://mlsanity.com/paper/2406.04231/) (2024)

## Built on by (most cited first, 30 of 72)

- [SkillLearnBench: Benchmarking Continual Learning Methods for Agent Skill Generation on Real-World Tasks](https://mlsanity.com/paper/2604.20087/) (2026)
- [Agent-World: Scaling Real-World Environment Synthesis for Evolving General Agent Intelligence](https://mlsanity.com/paper/2604.18292/) (2026)
- [AgenTRIM: Tool Risk Mitigation for Agentic AI](https://mlsanity.com/paper/2601.12449/) (2026)
- [HINTBench: Horizon-agent Intrinsic Non-attack Trajectory Benchmark](https://mlsanity.com/paper/2604.13954/) (2026)
- [AgentDoG 1.5: A Lightweight and Scalable Alignment Framework for AI Agent Safety and Security](https://mlsanity.com/paper/2605.29801/) (2026)
- [StepShield: When, Not Whether to Intervene on Rogue Agents](https://mlsanity.com/paper/2601.22136/) (2026)
- [Visual Confused Deputy: Exploiting and Defending Perception Failures in Computer-Using Agents](https://mlsanity.com/paper/2603.14707/) (2026)
- [From Governance Norms to Enforceable Controls: A Layered Translation Method for Runtime Guardrails in Agentic AI](https://mlsanity.com/paper/2604.05229/) (2026)
- [Toward Secure LLM Agents: Threat Surfaces, Attacks, Defenses, and Evaluation](https://mlsanity.com/paper/2606.10749/) (2026)
- [From Risk Classification to Action Plan Remediation: A Guardrail Feedback Driven Framework for LLM Agents](https://mlsanity.com/paper/2606.05805/) (2026)
- [PolicyGuard: A Dialogue-Grounded Sub-Agent Verifier for Policy Adherence in LLM Agents](https://mlsanity.com/paper/2606.29225/) (2026)
- [SHE: Trajectory-driven Safety Harness Evolution for LLM Agents](https://mlsanity.com/paper/2608.09885/) (2026)
- [StepGuard: Learning Step-Level Guardrails with Scalable Supervision and Safety-Utility Balancing](https://mlsanity.com/paper/2608.24777/) (2026)
- [OpenClawBench: Benchmarking Process-side Anomalies in Real-world Agent Execution Trajectories](https://mlsanity.com/paper/2605.29253/) (2026)
- [Oversight Has a Capacity: Calibrating Agent Guards to a Subjective, Fatiguing Human](https://mlsanity.com/paper/2606.08919/) (2026)
- [What Can Be Enforced? A Theory of Certified Runtime Safety for Tool-Using Agents](https://mlsanity.com/paper/2607.22868/) (2026)
- [MCPShield: A Security Cognition Layer for Adaptive Trust Calibration in Model Context Protocol Agents](https://mlsanity.com/paper/2602.14281/) (2026)
- [ProMAS: Proactive Error Forecasting for Multi-Agent Systems Using Markov Transition Dynamics](https://mlsanity.com/paper/2603.20260/) (2026)
- [Plant, Persist, Trigger: Sleeper Attack on Large Language Model Agents](https://mlsanity.com/paper/2605.28201/) (2026)
- [PROJECTMEM: A Local-First, Event-Sourced Memory and Judgment Layer for AI Coding Agents](https://mlsanity.com/paper/2606.12329/) (2026)
- [AgentLens: Interpretable Safety Steering via Mechanistic Subspaces for Multi-Turn Coding Agent](https://mlsanity.com/paper/2606.22673/) (2026)
- [ScopeJudge: Cost-Aware Pre-Execution Gating for Offensive Security Agents](https://mlsanity.com/paper/2607.07774/) (2026)
- [DreamGuard: Efficient Runtime Guardrail for LLM Agents via Risk-Aware World Model](https://mlsanity.com/paper/2608.05695/) (2026)
- [LLM-enabled Applications Require System-Level Threat Monitoring](https://mlsanity.com/paper/2602.19844/) (2026)
- [Who Tests the Testers? Systematic Enumeration and Coverage Audit of LLM Agent Tool Call Safety](https://mlsanity.com/paper/2603.18245/) (2026)
- [Beyond Task Success: An Evidence-Synthesis Framework for Evaluating, Governing, and Orchestrating Agentic AI](https://mlsanity.com/paper/2604.19818/) (2026)
- [Trust No Tool: Evaluating and Defending LLM Agents under Untrusted Tool Feedback](https://mlsanity.com/paper/2605.17453/) (2026)
- [HARP: Measuring Harm Amplification in Multi-Agent LLM Systems](https://mlsanity.com/paper/2605.27489/) (2026)
- [RUBAS: Rubric-Based Reinforcement Learning for Agent Safety](https://mlsanity.com/paper/2606.04051/) (2026)
- [Outcome Monitors: Recovery Affordances for Silent Tool Failures](https://mlsanity.com/paper/2608.19303/) (2026)

## Abstract (arXiv)

While LLM-based agents can interact with environments via invoking external tools, their expanded capabilities also amplify security risks. Monitoring step-level tool invocation behaviors in real time and proactively intervening before unsafe execution is critical for agent deployment, yet remains under-explored. In this work, we first construct TS-Bench, a novel benchmark for step-level tool invocation safety detection in LLM agents. We then develop a guardrail model, TS-Guard, using multi-task reinforcement learning. The model proactively detects unsafe tool invocation actions before execution by reasoning over the interaction history. It assesses request harmfulness and action-attack correlations, producing interpretable and generalizable safety judgments and feedback. Furthermore, we introduce TS-Flow, a guardrail-feedback-driven reasoning framework for LLM agents, which reduces harmful tool invocations of ReAct-style agents by 65 percent on average and improves benign task completion by approximately 10 percent under prompt injection attacks.