Source-linked AI summary
ToolSafe: Enhancing Tool Invocation Safety of LLM-based agents via Proactive Step-level Guardrail and Feedback
Yutao Mou, Zhangchi Xue, Lijun Li, Peiyang Liu, Shikun Zhang, Wei Ye, Jing Shao
TL;DR
LLM-based agents’ external tool access creates security risks, but step-level pre-execution safety monitoring remains under-explored. The paper constructs TS-Bench, trains TS-Guard with multi-task reinforcement learning, and integrates it into TS-Flow, reducing harmful invocations by 65% on average while improving benign task completion by approximately 10%.
Problem
LLM-based agents can induce unsafe tool invocations through malicious requests or prompt injection, motivating under-explored step-level monitoring before execution.
Method
The paper constructs TS-Bench and develops TS-Guard and TS-Flow for interpretable, proactive, pre-execution tool-invocation safety detection and feedback.
Results
65% average reduction in harmful tool invocations accompanies approximately 10% improvement in benign task performance.
Takeaways & Limitations
The proposed framework improves agent safety while preserving utility for reliable deployment in open-ended environments.
Takeaways & Limitations
TS-Flow appends guardrail feedback externally, and independently trained agents and guards may fail to incorporate feedback or become misaligned.
Abstract
from arXiv · showhide
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.
1 Introduction
LLM-based agents’ tool access creates security risks that require step-level, pre-execution monitoring. This work introduces TS-Bench, TS-Guard, and TS-Flow to detect and mitigate unsafe invocations while preserving benign performance.
- Motivation: Agent tool access amplifies security risks, including malicious requests and prompt injections that induce unsafe or unintended tool use.Because agents act directly on external environments, tool invocation safety is critical for reliable deployment.
- Motivation: Existing guardrails primarily moderate static inputs and outputs, while agent safety methods often analyze complete plans or trajectories rather than individual invocation steps.The paper motivates dynamic monitoring that can intervene against emerging risks before execution.
- Research Questions: The paper asks which step-level signals indicate unsafe invocations, how to train a generalizable detector, and how to improve safety without compromising benign task performance.These questions organize the proposed benchmark, guardrail model, and agent integration framework.
- Contributions: TS-Bench benchmarks step-level tool invocation safety detection, while TS-Guard uses multi-task reinforcement learning to identify harmful requests, attack vectors, and unsafe actions.TS-Guard also provides interpretable analysis and reasoning before execution.
- Contributions: TS-Flow monitors each tool invocation and provides pre-execution feedback that guides ReAct-style agents toward safer reasoning instead of terminating detected unsafe tasks.The framework is designed to improve safety while preserving benign task performance.
- Results: 65% average reduction in harmful tool invocations accompanies approximately 10% improvement in benign task performance.The reported results concern the proposed guardrail-feedback approach under the evaluated settings.
2 Related Work
Prior agent guardrails address several safety stages, but existing approaches have limited coverage, operate on complete trajectories, or incur latency that hinders step-level monitoring. TS-Bench targets fine-grained safety detection for general tool calls.
- Existing Guardrails: Existing LLM guardrails commonly classify prompts and responses, while newer agent guardrails extend protection to memory, planning, and tool invocation risks.These approaches broaden safety coverage beyond static input and output moderation.
- Benchmark Gap: Trajectory-level methods evaluate complete execution traces, whereas TS-Bench provides fine-grained step-level labels for pre-execution evaluation of general tool invocation safety.The benchmark is contrasted with post-hoc trajectory annotations and domain-specific safety benchmarks.
- Risk Coverage: TS-Bench organizes unsafe invocation risks by malicious user requests or prompt injections and by harmful tools or benign tools with risky arguments.The resulting four patterns are MUR, PI, HT, and BTRA.
- Limitations of Prior Work: Rule-based guardrails have limited coverage and generalization, while complex reasoning and verification pipelines can incur high latency for step-level monitoring.These limitations motivate efficient and generalizable step-level safety monitoring.
3 TS-Bench: Step-Level Tool Invocation Safety Detection Benchmark
TS-Bench formalizes pre-execution, step-level safety detection for tool invocations using interaction context. It captures four unsafe patterns across malicious requests, prompt injections, harmful tools, and risky arguments.
- Task formulation: Step-level detection determines whether a candidate tool action is unsafe from the preceding interaction history and user request before execution.The interaction history contains earlier action-observation pairs in the agent’s iterative reasoning-action loop.
- Risk patterns: Four unsafe invocation patterns combine two causes—malicious requests or third-party prompt injections—with two manifestations—harmful tools or risky arguments to benign tools.These patterns were identified by analyzing interaction logs from four representative agent-safety datasets.
- Task formulation: TS-Bench represents each sample with available tools, the user request, interaction history, a candidate tool action, and a three-class safety label.Labels are safe, controversial, or unsafe, indicating whether executing the candidate action would create potential risk.
- Benchmark construction: TS-Bench provides fine-grained, step-level labels for general unsafe tool-call risks, enabling pre-execution evaluation beyond post-hoc trajectory annotations and domain-specific benchmarks.Its source logs cover benign and malicious requests, prompt injections, and harmful or benign tool sets.
- Annotation: The benchmark annotates safety, prompt-injection association, and user-request harmfulness for tool-invocation steps sampled from complete execution trajectories.Annotations use safe, controversial or potentially unsafe, and significantly unsafe categories.
- Train-test split: TS-Bench-train and TS-Bench-eval use largely disjoint trajectories and domains, with the evaluation split combining remaining ASB domains with AgentHarm and AgentDojo.The split is designed to minimize overlap and indicate the absence of data leakage.
4 Method
The method combines TS-Guard, a step-level guardrail that produces interpretable multi-task safety judgments, with TS-Flow, which feeds pre-execution feedback back to the agent. This interaction-based design aims to improve safety while preserving benign task performance.
- Framework overview: The framework comprises TS-Guard for interpretable safety feedback and TS-Flow for feedback-driven reasoning during tool invocation.TS-Flow uses guardrail feedback rather than terminating the agent workflow.
- TS-Guard: TS-Guard sequentially predicts request harmfulness, candidate-action links to third-party attacks, and the safety of the current tool invocation.The model performs these three tasks within a single inference and reasons over preceding interaction logs.
- Optimization: TS-Guard is optimized with GRPO using a multi-task reward that compares predictions for all three tasks with their ground-truth labels.The three reward weights are uniformly set to w1 = w2 = w3 = 1/3.
- TS-Guard: TS-Guard outputs brief reasoning, request harmfulness, action-attack association, and a final safety rating to provide fine-grained supervision and step-level feedback.The multiple signals support causal analysis and more informative feedback to the agent.
- TS-Flow: TS-Flow provides pre-execution feedback for potentially unsafe actions, enabling agents to correct behavior instead of being terminated.The framework is presented as an agent–guardrail interaction paradigm intended to improve safety while largely preserving performance.
5 Experiments
The experiments evaluate guardrail models on step-level detection and guarded agents across safety and utility benchmarks. TS-Guard performs strongly across unsafe patterns, while TS-Flow improves the safety–utility trade-off through dynamic feedback.
- Evaluation setup: Experiments evaluate guardrail models on TS-Bench and guarded agents on AgentDojo, ASB, and AgentHarm.Guardrail models are assessed with accuracy, F1, and recall; agent evaluations use benchmark-specific safety and utility metrics.
- Guardrail model evaluation: Most guardrail models lose effectiveness under prompt injection, with GPT-4o’s F1 and recall dropping from 84.8 and 96.19 to 63.03 and 60.11.The comparison is between AgentHarm-Traj and ASB-Traj.
- Guardrail model evaluation: Prompt injection can cause benign tool invocations to be misclassified as unsafe, producing over-defensiveness and degraded agent utility.This pattern appears in benign-tool-with-risky-argument cases on AgentDojo-Traj.
- Guardrail model evaluation: TS-Guard performs strongly across all four unsafe patterns while substantially reducing over-defensiveness.The four patterns vary by malicious request versus prompt injection and harmful versus benign tools with risky arguments.
- Guarded agent evaluation: Detect-and-abort defenses improve safety but reduce benign task completion by terminating execution after prompt injection detection.Using TS-Guard in LlamaFirewall produces the most significant safety improvement among the described detect-and-abort defenses.
- Guarded agent evaluation: TS-Flow’s dynamic agent–guardrail interaction avoids unsafe calls while using feedback to guide benign task completion and achieve a better safety–utility trade-off.Compared with other models, TS-Guard within TS-Flow more effectively reduces attack success rates while preserving, and sometimes improving, benign task completion.
6 Analyses
Analyses attribute TS-Guard’s performance to reinforcement-learning and multi-task reward choices, while entropy studies show how guardrail feedback preserves exploration during risky reasoning. Richer feedback improves both agent safety and utility.
- Training methods: RL-only training generalizes better to diverse agent trajectories than SFT or SFT+RL.The analysis compares SFT, SFT+RL, and RL-only training methods.
- Training methods: SFT reduces Qwen2.5-7B-IT output entropy from 0.74 to 0.61, indicating reduced output diversity and potentially limiting subsequent RL gains.This reduction helps explain the inferior performance of SFT+RL.
- Reward design: Multi-task rewards improve F1 and reduce false positives compared with single-task safety-rating rewards.Multi-task rewards supervise request harmfulness and attack correlation in addition to safety ratings.
- Entropy analysis: TS-Guard lowers entropy mainly at the final judgment token while preserving higher entropy during intermediate reasoning.This differs from ShieldAgent-THU, which reduces entropy throughout its outputs, and is described as encouraging exploration for more reliable judgments.
- Entropy analysis: TS-Flow raises agent entropy at risky steps, countering overconfident execution and encouraging safety-aware exploration.Without guardrails, entropy steadily decreases during Qwen2.5-14B-IT reasoning and tool invocation.
- Feedback richness: Agents receiving full TS-Guard output achieve better safety and utility than agents receiving only the current-action safety rating.Full feedback includes interaction-history analysis, request harmfulness, and action–attack correlation.
7 Conclusion
The paper introduces TS-Bench and a proactive framework consisting of TS-Guard and TS-Flow for real-time, pre-execution tool-invocation safety. Experiments show improved agent safety while preserving utility.
- Conclusion: TS-Bench is introduced as the first benchmark for evaluating step-level tool-invocation safety.The benchmark supports the paper’s investigation of safety before tool execution.
- Conclusion: TS-Guard and TS-Flow provide proactive monitoring and feedback-driven reasoning for safer tool invocation in open-ended environments.TS-Guard supplies guardrail judgments, while TS-Flow integrates feedback into agent reasoning.
Limitations
The paper identifies limitations in how TS-Flow integrates guardrail feedback and coordinates agent and guardrail training. Future work is directed toward tighter integration and joint optimization.
- Limitations: Agents may fail to fully incorporate TS-Flow’s externally appended guardrail feedback, limiting step-level intervention effectiveness.The feedback interface is simple and model-agnostic but depends on the agent using the added signal.
- Limitations: Independent training of the agent and guardrail can misalign the agent’s reasoning process with the guard’s safety judgments.The authors propose joint training or tighter coupling as future directions.
Ethics Statement
TS-Bench is constructed from diverse agent-safety trajectories with step-level safety labels and auxiliary annotations, while access to harmful content is restricted to authorized researchers. Its split is evaluated for leakage and found to have negligible overlap under a conservative similarity criterion.
- Ethics and data governance: Access to harmful-content data is restricted to authorized researchers following strict ethical guidelines.
- Source datasets: AgentHarm-Traj contains benign and malicious instructions with simulated tools, while its sampled ReAct-style trajectories form the evaluation split.
- Source datasets: ASB-Traj covers direct and indirect prompt injection across application scenarios, with separate training and testing domains.
- Benchmark construction: TS-Bench combines trajectories from AgentAlign, AgentHarm, ASB, and AgentDojo to cover diverse agent environments and unsafe tool-invocation patterns.
- Annotation: Each tool invocation receives a three-level safety label, alongside annotations for user-request harmfulness and prompt-injection association.
- Data leakage analysis: TS-Bench-train and TS-Bench-eval show negligible overlap under the study’s conservative no-leakage criterion of similarity below 0.9.
G Supplementary results of Guardrail Model Evaluation
TS-Guard performs strongly across benchmark settings, ranking competitively in loose mode and achieving the best overall performance in exact mode. The results support fine-grained step-level safety judgments.
- Loose mode: TS-Guard ranks second on AgentHarm-Traj, first on ASB-Traj, and third on AgentDojo-Traj in loose mode.
- Exact mode: TS-Guard achieves the best overall performance in exact mode among the evaluated models.
- Overall assessment: The evaluation concludes that TS-Guard provides the strongest comprehensive performance for step-level tool-invocation safety detection.
H Overhead Analysis of TS-Flow
TS-Flow reduces harmful or injected behavior while improving benign-task utility and often shortening interaction sequences. It increases input-token use, but the reported expansion remains below twofold, while guardrail agents can incur much larger latency overhead.
- Task efficiency: TS-Flow reduces attack success rate, improves benign-task utility, and decreases interaction steps on AgentDojo.
- Token overhead: TS-Flow increases LLM input-token length, but the increase remains below a twofold expansion.
- Feedback frequency: GPT-4o- and Qwen2.5-14B-Instruct-driven agents trigger guardrail feedback an average of 0.97 and 1.41 times per user task, respectively.
- Comparison with guardrail agents: AGrail performs substantially worse than TS-Guard in general-purpose tool-invocation safety detection and requires approximately 8× more time per safety check.
- Auxiliary safety tasks: TS-Guard significantly outperforms other evaluated models on user-request harmfulness and prompt-injection association detection auxiliary tasks.