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

TL;DR

This work addresses the safety challenges of tool invocation in LLM-based agents by introducing TS-Bench, the first step-level benchmark for pre-execution safety, and two complementary mechanisms: TS-Guard, a multi-task RL-trained guardrail that assesses harmfulness, attack linkage, and safety of proposed actions, and TS-Flow, a guardrail-driven feedback system that guides agents toward safe tool use rather than aborting tasks. TS-Guard outputs interpretable reasoning and safety signals, enabling proactive intervention, while TS-Flow enhances safety without sacrificing task performance, achieving up to 65% reductions in harmful invocations and about 10% gains in benign task completion under prompt-injection conditions. The framework leverages GRPO and multi-task rewards to improve generalization across diverse unsafe patterns, and analyses reveal guardrail feedback increases agent entropy in risky steps to promote exploration of safe trajectories. Overall, TS-Bench, TS-Guard, and TS-Flow offer a practical, low-latency solution for deploying safer, more reliable LLM-based agents in open-ended environments.

Abstract

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.

Figures (25)

• Figure 1: Illustration of two categories of tool invocation security risks considered in this study. (a) Malicious user requests that directly induce unsafe tool invocation. (b) Prompt injection attacks occurring during benign task execution, leading to unintended tool use.
• Figure 2: Illustration of our proactive step-level guardrail and feedback framework for LLM agents. (a) Input and output format of TS-Guard. (b) TS-Flow feeds guardrail feedback to the agent, enabling safe tool invocation reasoning rather than aborting execution.
• Figure 3: Ablation results on training methods and reward designs. (a) Comparision of SFT, SFT+RL and RL only (b) Comparision of multi-task rewards and single-task rewards.
• Figure 4: Entropy comparison of guardrails. (a) Specialized models show lower entropy than general LLMs. (b) TS-Guard lowers final-decision entropy while preserving reasoning-step entropy to facilitate exploration.
• Figure 5: Token-wise entropy of a ReAct-style agent (Qwen2.5-14B-IT). Without guardrails, entropy decreases as the agent grows overconfident; with TS-Flow, TS-Guard feedback raises entropy in risky steps, maintaining uncertainty and guiding safe exploration..