Risk-Aware Reranking for Agentic Tool Retrieval

• URL: http://arxiv.org/abs/2608.22751v1
• Date: Mon, 24 Aug 2026 03:20:54 GMT
• Title: Risk-Aware Reranking for Agentic Tool Retrieval
• Abstract summary: We study risk-aware tool retrieval, where the goal is to retrieve useful tools while reducing exposure to higher-risk tools.<n>We propose a lightweight reranking framework on top of a frozen first-stage retriever.<n>We show that our approach improves the relevance--safety tradeoff over relevance--only retrievers and reranking baselines.
• Score: 13.019844558764099
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Tool retrieval determines which external tools are exposed to an LLM agent for a user query or task, making retrieval a critical pre-execution safety boundary. Unlike document retrieval, tool retrieval exposes executable actions: a tool that is useful for one task may be unnecessary or risky for another. However, existing tool-retrieval methods primarily optimize semantic relevance, and safety evaluations often focus on failures after tool execution rather than risks introduced during retrieval. We study risk-aware tool retrieval, where the goal is to retrieve useful tools while reducing exposure to higher-risk tools. We propose a lightweight reranking framework on top of a frozen first-stage retriever. The framework models query-conditioned relevance and tool-level exposure risk separately, combines them through an explicit parameter controlling the tradeoff between safety and utility, smooths scores over a ToolGraph, and optionally applies rule-based safety constraints. To support retrieval-time safety evaluation, we annotate 6,108 tools across UltraTool and Seal-Tools with five ordinal risk levels and define metrics that measure risky-tool exposure in the top-$k$ results. Experiments on UltraTool and Seal-Tools show that our approach improves the relevance--safety tradeoff over relevance-only retrievers and reranking baselines, with the rule-filtered variant providing a conservative operating point for safety-critical deployments. These findings indicate that retrieval-stage filtering can reduce the candidate action space exposed to agents before execution, complementing downstream tool-use safeguards. The code and supplementary materials are available at: https://github.com/qli447/risk-aware-tool-retrieval-release.

Related papers

• ROPE: Routed Origin Policy Enforcement against Indirect Prompt Injection [61.4148196085256]
  We present ROPE (Routed Origin Policy Enforcement), which is anchored in a structural notion of trust.<n>Our approach admits two provable guarantees: 1) at every step of a trajectory, no value whose only origin is attacker-writable content reaches an origin-guarded parameter, and 2) no rewording of an injection changes an admission decision.
  arXiv  Detail & Related papers  (2026-08-27T01:22:14Z)
• Forgotten in Weights, Recovered by Tools: Agentic Tool Unlearning for LLM Agents [54.050506620820784]
  Agentic Tool Unlearning (ATU) aims to reduce both parametric recall and tool-mediated recovery while preserving normal tool use for retained knowledge.<n>ATU achieves a better balance between target forgetting and retained utility, making unlearning more robust under tool-augmented agent deployment.
  arXiv  Detail & Related papers  (2026-08-21T18:34:48Z)
• Breadcrumbing Search Agents [118.20284823674382]
  LLM-based search agents are widely used for information-seeking tasks, but their reliance on external tool returns introduces a critical security risk.<n>We show that the channel delivering search and page observations is a fragile security boundary.<n>Under a constrained tool-intermediary threat model, appending only one controlled result per query can substantially increase attack success.
  arXiv  Detail & Related papers  (2026-08-05T07:57:27Z)
• ToolChain-CRC: Conformal Risk Control for Agentic AI Under Retrieval and Tool-Use Drift [0.0]
  We propose ToolChain-CRC, a conformal risk-control method for retrieval-augmented and tool-using agents under drift.<n>The method treats each agent run as a full trajectory of actions, observations, and final output.<n>It builds step-level risk scores, combines them into a trajectory risk score, calibrates an accept-or-intervene rule, and adds an anytime alarm.
  arXiv  Detail & Related papers  (2026-06-16T20:27:37Z)
• RUBAS: Rubric-Based Reinforcement Learning for Agent Safety [93.80584563695646]
  RUBAS is a rubric-based reinforcement learning framework for agent safety.<n>It decomposes agent behavior into four dimensions: tool-use safety, argument safety, response safety, and helpfulness.<n>Experiments show that RUBAS improves safety over standard alignment baselines, reduces tool-grounded hallucinations, and maintains competitive utility.
  arXiv  Detail & Related papers  (2026-06-02T09:02:14Z)
• Trust No Tool: Evaluating and Defending LLM Agents under Untrusted Tool Feedback [38.251599309409]
  We study a different failure mode, cognitive poisoning, in which a malicious tool behaves plausibly during exploration.<n>We present VISTA-Guard, a backbone-agnostic framework for final-action risk scoring.
  arXiv  Detail & Related papers  (2026-05-17T13:51:34Z)
• Unsafer in Many Turns: Benchmarking and Defending Multi-Turn Safety Risks in Tool-Using Agents [68.20752678837377]
  We propose a principled taxonomy that transforms single-turn harmful tasks into multi-turn attack sequences.<n>Using this taxonomy, we construct MT-AgentRisk, the first benchmark to evaluate multi-turn tool-using agent safety.<n>We propose ToolShield, a training-free, tool-agnostic, self-exploration defense.
  arXiv  Detail & Related papers  (2026-02-13T18:38:18Z)
• AgenTRIM: Tool Risk Mitigation for Agentic AI [5.4672006013914975]
  We introduce AgenTRIM, a framework for detecting and mitigating tool-driven agency risks.<n>AgenTRIM addresses these risks through complementary offline and online phases.<n>AgenTRIM substantially reduces attack success while maintaining high task performance.
  arXiv  Detail & Related papers  (2026-01-18T15:10:18Z)
• ToolSafe: Enhancing Tool Invocation Safety of LLM-based agents via Proactive Step-level Guardrail and Feedback [53.2744585868162]
  Monitoring step-level tool invocation behaviors in real time is critical for agent deployment.<n>We first construct TS-Bench, a novel benchmark for step-level tool invocation safety detection in LLM agents.<n>We then develop a guardrail model, TS-Guard, using multi-task reinforcement learning.
  arXiv  Detail & Related papers  (2026-01-15T07:54:32Z)
• Quantifying Distributional Robustness of Agentic Tool-Selection [8.457056023589951]
  We introduce ToolCert, the first statistical framework that formally certifies tool selection robustness.<n>We show that ToolCert produces a high-confidence lower bound on accuracy, formally quantifying the agent's worst-case performance.<n>Our evaluation with ToolCert uncovers the severe fragility: under attacks injecting deceptive tools or saturating retrieval, the certified accuracy bound drops near zero.
  arXiv  Detail & Related papers  (2025-10-05T01:50:34Z)
• ToolTweak: An Attack on Tool Selection in LLM-based Agents [52.17181489286236]
  We show that adversaries can systematically bias agents toward selecting specific tools, gaining unfair advantage over equally capable alternatives.<n>We present ToolTweak, a lightweight automatic attack that increases selection rates from a baseline of around 20% to as high as 81%.<n>To mitigate these risks, we evaluate two defenses: paraphrasing and perplexity filtering, which reduce bias and lead agents to select functionally similar tools more equally.
  arXiv  Detail & Related papers  (2025-10-02T20:44:44Z)
• SafeToolBench: Pioneering a Prospective Benchmark to Evaluating Tool Utilization Safety in LLMs [35.180946816997164]
  Large Language Models (LLMs) have exhibited great performance in autonomously calling various tools in external environments.<n>This paper focuses on prospective ways to assess the safety of LLM tool utilization, aiming to avoid irreversible harm caused by directly executing tools.<n>We propose SafeToolBench, the first benchmark to comprehensively assess tool utilization security in a prospective manner.<n>We also propose a novel framework, SafeInstructTool, which aims to enhance LLMs' awareness of tool utilization security from three perspectives.
  arXiv  Detail & Related papers  (2025-09-09T01:31:25Z)
• Prompt Injection Attack to Tool Selection in LLM Agents [60.95349602772112]
  A popular approach follows a two-step process - emphretrieval and emphselection - to pick the most appropriate tool from a tool library for a given task.<n>In this work, we introduce textitToolHijacker, a novel prompt injection attack targeting tool selection in no-box scenarios.
  arXiv  Detail & Related papers  (2025-04-28T13:36:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.