use std::sync::{Arc, RwLock};
use async_trait::async_trait;
use serde_json::{json, Value};
use crate::error::{Error, Result};
use crate::llm::ToolSpec;
use super::{Tool, ToolSideEffect};
pub const TOOL_SEARCH_TOOL_NAME: &str = "ToolSearchTool";
pub type DeferredCatalog = Arc<RwLock<Vec<ToolSpec>>>;
#[derive(Clone)]
pub struct ToolSearchTool {
catalog: DeferredCatalog,
}
impl ToolSearchTool {
pub fn new(catalog: DeferredCatalog) -> Self {
Self { catalog }
}
}
#[async_trait]
impl Tool for ToolSearchTool {
fn spec(&self) -> ToolSpec {
ToolSpec {
name: TOOL_SEARCH_TOOL_NAME.to_string(),
description: "Fetches full schema definitions for deferred tools so they can be \
called. Until fetched, only the name is known — there is no parameter schema, \
so the tool cannot be invoked. This tool takes a query, matches it against the \
deferred tool list, and returns the matched tools' complete JSONSchema \
definitions inside a <functions> block. Once a tool's schema appears in that \
result, it is callable exactly like any tool defined at the top of the prompt.\n\n\
Result format: each matched tool appears as one \
<function>{\"description\": \"...\", \"name\": \"...\", \"parameters\": {...}}\
</function> line inside the <functions> block.\n\n\
Query forms:\n\
- \"select:Read,Edit,Grep\" — fetch these exact tools by name\n\
- \"notebook jupyter\" — keyword search, up to max_results best matches\n\
- \"+slack send\" — require \"slack\" in the name, rank by remaining terms"
.to_string(),
parameters: json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Query to find deferred tools. Use \"select:<tool_name>\" \
for direct selection, or keywords to search."
},
"max_results": {
"type": "integer",
"description": "Maximum number of results to return (default: 5)",
"default": 5
}
},
"required": ["query"]
}),
}
}
async fn execute(&self, arguments: Value) -> Result<String> {
let query = arguments
.get("query")
.and_then(|v| v.as_str())
.unwrap_or("")
.trim();
let max_results = arguments
.get("max_results")
.and_then(|v| v.as_u64())
.map(|n| n as usize)
.unwrap_or(5)
.clamp(1, 20);
if query.is_empty() {
return Err(Error::BadToolArgs {
name: TOOL_SEARCH_TOOL_NAME.to_string(),
message: "\"query\" must be a non-empty string".to_string(),
});
}
let catalog = self.catalog.read().map_err(|_| Error::BadToolArgs {
name: TOOL_SEARCH_TOOL_NAME.to_string(),
message: "deferred catalog lock poisoned".to_string(),
})?;
let matched = resolve(query, &catalog, max_results);
if matched.is_empty() {
return Ok(format!(
"No matching deferred tools found for query \"{query}\". \
Available: {}",
catalog
.iter()
.map(|s| s.name.as_str())
.collect::<Vec<_>>()
.join(", ")
));
}
let names: Vec<&str> = matched.iter().map(|s| s.name.as_str()).collect();
Ok(serde_json::to_string(&names).unwrap_or_default())
}
fn side_effect_class(&self) -> ToolSideEffect {
ToolSideEffect::ReadOnly
}
fn is_deferred(&self) -> bool {
false }
}
pub fn resolve<'a>(query: &str, catalog: &'a [ToolSpec], max_results: usize) -> Vec<&'a ToolSpec> {
let q = query.trim().to_lowercase();
if let Some(names_str) = q.strip_prefix("select:") {
let mut result = Vec::new();
for name in names_str.split(',').map(|s| s.trim()) {
if let Some(spec) = catalog.iter().find(|s| s.name.to_lowercase() == name) {
result.push(spec);
}
}
result.truncate(max_results);
return result;
}
if let Some(spec) = catalog.iter().find(|s| s.name.to_lowercase() == q) {
return vec![spec];
}
let terms: Vec<&str> = q.split_whitespace().collect();
let mut scored: Vec<(usize, &ToolSpec)> = catalog
.iter()
.filter_map(|spec| {
let haystack = format!(
"{} {}",
spec.name.to_lowercase(),
spec.description.to_lowercase()
);
let score = terms.iter().filter(|&&t| haystack.contains(t)).count();
if score > 0 {
Some((score, spec))
} else {
None
}
})
.collect();
scored.sort_by(|a, b| b.0.cmp(&a.0).then(a.1.name.cmp(&b.1.name)));
scored.truncate(max_results);
scored.into_iter().map(|(_, s)| s).collect()
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn make_spec(name: &str, desc: &str) -> ToolSpec {
ToolSpec {
name: name.to_string(),
description: desc.to_string(),
parameters: json!({"type": "object"}),
}
}
#[test]
fn select_prefix_returns_exact_matches() {
let catalog = vec![
make_spec("Read", "Read a file"),
make_spec("WebFetch", "Fetch a URL"),
make_spec("Write", "Write a file"),
];
let result = resolve("select:WebFetch,Read", &catalog, 10);
assert_eq!(result.len(), 2);
assert_eq!(result[0].name, "WebFetch");
assert_eq!(result[1].name, "Read");
}
#[test]
fn select_skips_unknown_names() {
let catalog = vec![make_spec("Read", "Read a file")];
let result = resolve("select:Read,Unknown", &catalog, 10);
assert_eq!(result.len(), 1);
assert_eq!(result[0].name, "Read");
}
#[test]
fn exact_name_match_case_insensitive() {
let catalog = vec![make_spec("WebFetch", "Fetch a URL")];
let result = resolve("webfetch", &catalog, 10);
assert_eq!(result.len(), 1);
assert_eq!(result[0].name, "WebFetch");
}
#[test]
fn keyword_search_scores_by_term_matches() {
let catalog = vec![
make_spec("Read", "Read a UTF-8 file from disk"),
make_spec("WebFetch", "Fetch content from a URL via HTTP"),
make_spec("Write", "Write content to a file"),
];
let result = resolve("file", &catalog, 10);
assert!(result.iter().any(|s| s.name == "Read"));
assert!(result.iter().any(|s| s.name == "Write"));
assert!(!result.iter().any(|s| s.name == "WebFetch"));
}
#[test]
fn max_results_is_respected() {
let catalog = (0..10)
.map(|i| make_spec(&format!("Tool{i}"), "do something useful"))
.collect::<Vec<_>>();
let result = resolve("useful", &catalog, 3);
assert_eq!(result.len(), 3);
}
#[tokio::test]
async fn execute_returns_json_array_of_names() {
let catalog: DeferredCatalog =
Arc::new(RwLock::new(vec![make_spec("WebFetch", "Fetch a URL")]));
let tool = ToolSearchTool::new(catalog);
let out = tool
.execute(json!({"query": "select:WebFetch"}))
.await
.unwrap();
let names: Vec<String> = serde_json::from_str(&out).expect("valid JSON array");
assert!(
names.contains(&"WebFetch".to_string()),
"missing WebFetch: {out}"
);
}
#[tokio::test]
async fn execute_no_match_returns_helpful_message() {
let catalog: DeferredCatalog =
Arc::new(RwLock::new(vec![make_spec("Read", "Read a file")]));
let tool = ToolSearchTool::new(catalog);
let out = tool
.execute(json!({"query": "select:Unknown"}))
.await
.unwrap();
assert!(
out.contains("No matching"),
"expected no-match message: {out}"
);
}
#[test]
fn select_prefix_truncates_at_max_results() {
let catalog = vec![
make_spec("A", "tool A"),
make_spec("B", "tool B"),
make_spec("C", "tool C"),
];
let result = resolve("select:A,B,C", &catalog, 2);
assert_eq!(
result.len(),
2,
"select: results must be truncated to max_results"
);
}
#[test]
fn keyword_no_match_returns_empty() {
let catalog = vec![make_spec("Read", "Read a file")];
let result = resolve("xyzzy_nonexistent_term", &catalog, 10);
assert!(result.is_empty(), "zero-score tools must not appear");
}
#[test]
fn keyword_search_ranks_higher_score_first() {
let catalog = vec![
make_spec("AlphaRead", "Read files"),
make_spec("BetaReadWrite", "Read and write files"),
];
let result = resolve("read write", &catalog, 10);
assert!(!result.is_empty());
assert_eq!(
result[0].name, "BetaReadWrite",
"higher score tool must rank first"
);
}
#[test]
fn keyword_search_stable_sort_by_name_asc_when_equal_score() {
let catalog = vec![
make_spec("ZTool", "do something useful"),
make_spec("ATool", "do something useful"),
];
let result = resolve("something", &catalog, 10);
assert_eq!(result.len(), 2);
assert_eq!(
result[0].name, "ATool",
"equal-score tools must sort by name asc"
);
assert_eq!(result[1].name, "ZTool");
}
#[test]
fn exact_name_match_returns_exactly_one() {
let catalog = vec![
make_spec("Read", "Read a file"),
make_spec("ReadAll", "Read all files"),
];
let result = resolve("Read", &catalog, 10);
assert_eq!(result.len(), 1, "exact match must return exactly 1 tool");
assert_eq!(result[0].name, "Read");
}
#[test]
fn is_deferred_returns_false() {
let catalog: DeferredCatalog = Arc::new(RwLock::new(vec![]));
let tool = ToolSearchTool::new(catalog);
assert!(
!tool.is_deferred(),
"ToolSearchTool must always be eager (not deferred)"
);
}
#[tokio::test]
async fn execute_max_results_clamps_to_min_1() {
let catalog: DeferredCatalog = Arc::new(RwLock::new(
(0..5)
.map(|i| make_spec(&format!("T{i}"), "useful thing"))
.collect(),
));
let tool = ToolSearchTool::new(catalog);
let out = tool
.execute(json!({"query": "useful", "max_results": 0}))
.await
.unwrap();
let names: Vec<String> = serde_json::from_str(&out).expect("valid JSON array");
assert_eq!(names.len(), 1, "max_results=0 must be clamped to 1");
}
#[tokio::test]
async fn execute_empty_query_returns_error() {
let catalog: DeferredCatalog = Arc::new(RwLock::new(vec![make_spec("Read", "file")]));
let tool = ToolSearchTool::new(catalog);
let result = tool.execute(json!({"query": ""})).await;
assert!(result.is_err(), "empty query must return an error");
}
}