| 1 | //! Deferred tool catalog and built-in advanced tool helpers. |
| 2 | //! |
| 3 | //! The streaming turn loop owns when tools are offered or executed. This module |
| 4 | //! owns the catalog-level policy around deferred loading, tool search, missing |
| 5 | //! tool suggestions, and the small set of built-in advanced tools that are not |
| 6 | //! registered by the normal runtime tool registry. |
| 7 | |
| 8 | use std::collections::HashSet; |
| 9 | use std::path::Path; |
| 10 | use std::time::Duration; |
| 11 | |
| 12 | use serde_json::json; |
| 13 | |
| 14 | use crate::models::Tool; |
| 15 | use crate::tools::spec::{ToolError, ToolResult, required_str}; |
| 16 | use crate::tui::app::AppMode; |
| 17 | |
| 18 | pub(super) const MULTI_TOOL_PARALLEL_NAME: &str = "multi_tool_use.parallel"; |
| 19 | pub(super) const REQUEST_USER_INPUT_NAME: &str = "request_user_input"; |
| 20 | pub(super) const CODE_EXECUTION_TOOL_NAME: &str = "code_execution"; |
| 21 | const CODE_EXECUTION_TOOL_TYPE: &str = "code_execution_20250825"; |
| 22 | const TOOL_SEARCH_REGEX_NAME: &str = "tool_search_tool_regex"; |
| 23 | const TOOL_SEARCH_REGEX_TYPE: &str = "tool_search_tool_regex_20251119"; |
| 24 | pub(super) const TOOL_SEARCH_BM25_NAME: &str = "tool_search_tool_bm25"; |
| 25 | const TOOL_SEARCH_BM25_TYPE: &str = "tool_search_tool_bm25_20251119"; |
| 26 | |
| 27 | pub(super) fn is_tool_search_tool(name: &str) -> bool { |
| 28 | matches!(name, TOOL_SEARCH_REGEX_NAME | TOOL_SEARCH_BM25_NAME) |
| 29 | } |
| 30 | |
| 31 | pub(super) fn should_default_defer_tool(name: &str, mode: AppMode) -> bool { |
| 32 | if mode == AppMode::Yolo { |
| 33 | return false; |
| 34 | } |
| 35 | |
| 36 | // Shell tools are kept active in Agent so the model can run verification |
| 37 | // commands (build/test/git/cargo) without first having to discover the |
| 38 | // tool through ToolSearch. Plan mode never registers shell tools. |
| 39 | let always_loaded_in_action_modes = matches!(mode, AppMode::Agent) |
| 40 | && matches!( |
| 41 | name, |
| 42 | "exec_shell" |
| 43 | | "exec_shell_wait" |
| 44 | | "exec_shell_interact" |
| 45 | | "exec_wait" |
| 46 | | "exec_interact" |
| 47 | ); |
| 48 | if always_loaded_in_action_modes { |
| 49 | return false; |
| 50 | } |
| 51 | |
| 52 | !matches!( |
| 53 | name, |
| 54 | "read_file" |
| 55 | | "list_dir" |
| 56 | | "grep_files" |
| 57 | | "file_search" |
| 58 | | "diagnostics" |
| 59 | | "rlm" |
| 60 | | "recall_archive" |
| 61 | | MULTI_TOOL_PARALLEL_NAME |
| 62 | | "update_plan" |
| 63 | | "checklist_write" |
| 64 | | "todo_write" |
| 65 | | "task_create" |
| 66 | | "task_list" |
| 67 | | "task_read" |
| 68 | | "task_gate_run" |
| 69 | | "task_shell_start" |
| 70 | | "task_shell_wait" |
| 71 | | "github_issue_context" |
| 72 | | "github_pr_context" |
| 73 | | REQUEST_USER_INPUT_NAME |
| 74 | ) |
| 75 | } |
| 76 | |
| 77 | pub(super) fn apply_native_tool_deferral(catalog: &mut [Tool], mode: AppMode) { |
| 78 | for tool in catalog { |
| 79 | tool.defer_loading = Some(should_default_defer_tool(&tool.name, mode)); |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | fn should_keep_mcp_tool_loaded(name: &str) -> bool { |
| 84 | matches!( |
| 85 | name, |
| 86 | "list_mcp_resources" |
| 87 | | "list_mcp_resource_templates" |
| 88 | | "mcp_read_resource" |
| 89 | | "read_mcp_resource" |
| 90 | | "mcp_get_prompt" |
| 91 | ) |
| 92 | } |
| 93 | |
| 94 | pub(super) fn apply_mcp_tool_deferral(catalog: &mut [Tool], mode: AppMode) { |
| 95 | for tool in catalog { |
| 96 | tool.defer_loading = |
| 97 | Some(mode != AppMode::Yolo && !should_keep_mcp_tool_loaded(&tool.name)); |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | pub(super) fn build_model_tool_catalog( |
| 102 | mut native_tools: Vec<Tool>, |
| 103 | mut mcp_tools: Vec<Tool>, |
| 104 | mode: AppMode, |
| 105 | ) -> Vec<Tool> { |
| 106 | apply_native_tool_deferral(&mut native_tools, mode); |
| 107 | apply_mcp_tool_deferral(&mut mcp_tools, mode); |
| 108 | // Sort each partition by name for prefix-cache stability (#263). The |
| 109 | // upstream `to_api_tools()` already sorts the registry's HashMap output; |
| 110 | // this catalog is built from caller-supplied Vecs which the test harness |
| 111 | // and (future) caller refactors may not pre-sort. Built-ins stay as a |
| 112 | // contiguous prefix ahead of MCP tools so adding/removing an MCP tool |
| 113 | // never shifts a built-in's position. |
| 114 | native_tools.sort_by(|a, b| a.name.cmp(&b.name)); |
| 115 | mcp_tools.sort_by(|a, b| a.name.cmp(&b.name)); |
| 116 | native_tools.extend(mcp_tools); |
| 117 | native_tools |
| 118 | } |
| 119 | |
| 120 | pub(super) fn ensure_advanced_tooling(catalog: &mut Vec<Tool>) { |
| 121 | if !catalog.iter().any(|t| t.name == CODE_EXECUTION_TOOL_NAME) { |
| 122 | catalog.push(Tool { |
| 123 | tool_type: Some(CODE_EXECUTION_TOOL_TYPE.to_string()), |
| 124 | name: CODE_EXECUTION_TOOL_NAME.to_string(), |
| 125 | description: "Execute Python code in a local sandboxed runtime and return stdout/stderr/return_code as JSON.".to_string(), |
| 126 | input_schema: json!({ |
| 127 | "type": "object", |
| 128 | "properties": { |
| 129 | "code": { "type": "string", "description": "Python source code to execute." } |
| 130 | }, |
| 131 | "required": ["code"] |
| 132 | }), |
| 133 | allowed_callers: Some(vec!["direct".to_string()]), |
| 134 | defer_loading: Some(false), |
| 135 | input_examples: None, |
| 136 | strict: None, |
| 137 | cache_control: None, |
| 138 | }); |
| 139 | } |
| 140 | |
| 141 | if !catalog.iter().any(|t| t.name == TOOL_SEARCH_REGEX_NAME) { |
| 142 | catalog.push(Tool { |
| 143 | tool_type: Some(TOOL_SEARCH_REGEX_TYPE.to_string()), |
| 144 | name: TOOL_SEARCH_REGEX_NAME.to_string(), |
| 145 | description: "Search deferred tool definitions using a regex query and return matching tool references.".to_string(), |
| 146 | input_schema: json!({ |
| 147 | "type": "object", |
| 148 | "properties": { |
| 149 | "query": { "type": "string", "description": "Regex pattern to search tool names/descriptions/schema." } |
| 150 | }, |
| 151 | "required": ["query"] |
| 152 | }), |
| 153 | allowed_callers: Some(vec!["direct".to_string()]), |
| 154 | defer_loading: Some(false), |
| 155 | input_examples: None, |
| 156 | strict: None, |
| 157 | cache_control: None, |
| 158 | }); |
| 159 | } |
| 160 | |
| 161 | if !catalog.iter().any(|t| t.name == TOOL_SEARCH_BM25_NAME) { |
| 162 | catalog.push(Tool { |
| 163 | tool_type: Some(TOOL_SEARCH_BM25_TYPE.to_string()), |
| 164 | name: TOOL_SEARCH_BM25_NAME.to_string(), |
| 165 | description: "Search deferred tool definitions using natural-language matching and return matching tool references.".to_string(), |
| 166 | input_schema: json!({ |
| 167 | "type": "object", |
| 168 | "properties": { |
| 169 | "query": { "type": "string", "description": "Natural language query for tool discovery." } |
| 170 | }, |
| 171 | "required": ["query"] |
| 172 | }), |
| 173 | allowed_callers: Some(vec!["direct".to_string()]), |
| 174 | defer_loading: Some(false), |
| 175 | input_examples: None, |
| 176 | strict: None, |
| 177 | cache_control: None, |
| 178 | }); |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | pub(super) fn initial_active_tools(catalog: &[Tool]) -> HashSet<String> { |
| 183 | let mut active = HashSet::new(); |
| 184 | for tool in catalog { |
| 185 | if !tool.defer_loading.unwrap_or(false) || is_tool_search_tool(&tool.name) { |
| 186 | active.insert(tool.name.clone()); |
| 187 | } |
| 188 | } |
| 189 | if active.is_empty() |
| 190 | && !catalog.is_empty() |
| 191 | && let Some(first) = catalog.first() |
| 192 | { |
| 193 | active.insert(first.name.clone()); |
| 194 | } |
| 195 | active |
| 196 | } |
| 197 | |
| 198 | fn active_tool_list_from_catalog(catalog: &[Tool], active: &HashSet<String>) -> Vec<Tool> { |
| 199 | // Two-pass for prefix-cache stability (#263). Always-loaded tools come |
| 200 | // first in their stable catalog order; tools that started life deferred |
| 201 | // and were activated mid-conversation by ToolSearch get appended at the |
| 202 | // tail. Otherwise activating a deferred tool shifts every later tool's |
| 203 | // byte offset and busts the cached prefix from that point onwards. |
| 204 | let mut head: Vec<Tool> = Vec::new(); |
| 205 | let mut tail: Vec<Tool> = Vec::new(); |
| 206 | for tool in catalog { |
| 207 | if !active.contains(&tool.name) { |
| 208 | continue; |
| 209 | } |
| 210 | if tool.defer_loading.unwrap_or(false) { |
| 211 | tail.push(tool.clone()); |
| 212 | } else { |
| 213 | head.push(tool.clone()); |
| 214 | } |
| 215 | } |
| 216 | head.extend(tail); |
| 217 | head |
| 218 | } |
| 219 | |
| 220 | pub(super) fn active_tools_for_step( |
| 221 | catalog: &[Tool], |
| 222 | active: &HashSet<String>, |
| 223 | force_update_plan: bool, |
| 224 | ) -> Vec<Tool> { |
| 225 | // DeepSeek reasoning models reject explicit named tool_choice forcing here, |
| 226 | // so for obvious quick-plan asks we narrow the first-step tool surface to |
| 227 | // update_plan instead. |
| 228 | if force_update_plan { |
| 229 | let forced: Vec<_> = catalog |
| 230 | .iter() |
| 231 | .filter(|tool| tool.name == "update_plan") |
| 232 | .cloned() |
| 233 | .collect(); |
| 234 | if !forced.is_empty() { |
| 235 | return forced; |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | active_tool_list_from_catalog(catalog, active) |
| 240 | } |
| 241 | |
| 242 | fn tool_search_haystack(tool: &Tool) -> String { |
| 243 | format!( |
| 244 | "{}\n{}\n{}", |
| 245 | tool.name.to_lowercase(), |
| 246 | tool.description.to_lowercase(), |
| 247 | tool.input_schema.to_string().to_lowercase() |
| 248 | ) |
| 249 | } |
| 250 | |
| 251 | fn discover_tools_with_regex(catalog: &[Tool], query: &str) -> Result<Vec<String>, ToolError> { |
| 252 | let regex = regex::Regex::new(query) |
| 253 | .map_err(|err| ToolError::invalid_input(format!("Invalid regex query: {err}")))?; |
| 254 | |
| 255 | let mut matches = Vec::new(); |
| 256 | for tool in catalog { |
| 257 | if is_tool_search_tool(&tool.name) { |
| 258 | continue; |
| 259 | } |
| 260 | let hay = tool_search_haystack(tool); |
| 261 | if regex.is_match(&hay) { |
| 262 | matches.push(tool.name.clone()); |
| 263 | } |
| 264 | if matches.len() >= 5 { |
| 265 | break; |
| 266 | } |
| 267 | } |
| 268 | Ok(matches) |
| 269 | } |
| 270 | |
| 271 | fn discover_tools_with_bm25_like(catalog: &[Tool], query: &str) -> Vec<String> { |
| 272 | let terms: Vec<String> = query |
| 273 | .split_whitespace() |
| 274 | .map(|term| term.trim().to_lowercase()) |
| 275 | .filter(|term| !term.is_empty()) |
| 276 | .collect(); |
| 277 | if terms.is_empty() { |
| 278 | return Vec::new(); |
| 279 | } |
| 280 | |
| 281 | let mut scored: Vec<(i64, String)> = Vec::new(); |
| 282 | for tool in catalog { |
| 283 | if is_tool_search_tool(&tool.name) { |
| 284 | continue; |
| 285 | } |
| 286 | let hay = tool_search_haystack(tool); |
| 287 | let mut score = 0i64; |
| 288 | for term in &terms { |
| 289 | if hay.contains(term) { |
| 290 | score += 1; |
| 291 | } |
| 292 | if tool.name.to_lowercase().contains(term) { |
| 293 | score += 2; |
| 294 | } |
| 295 | } |
| 296 | if score > 0 { |
| 297 | scored.push((score, tool.name.clone())); |
| 298 | } |
| 299 | } |
| 300 | scored.sort_by(|a, b| b.0.cmp(&a.0).then_with(|| a.1.cmp(&b.1))); |
| 301 | scored.into_iter().take(5).map(|(_, name)| name).collect() |
| 302 | } |
| 303 | |
| 304 | fn edit_distance(a: &str, b: &str) -> usize { |
| 305 | if a == b { |
| 306 | return 0; |
| 307 | } |
| 308 | if a.is_empty() { |
| 309 | return b.chars().count(); |
| 310 | } |
| 311 | if b.is_empty() { |
| 312 | return a.chars().count(); |
| 313 | } |
| 314 | |
| 315 | let b_chars: Vec<char> = b.chars().collect(); |
| 316 | let mut prev: Vec<usize> = (0..=b_chars.len()).collect(); |
| 317 | let mut curr = vec![0usize; b_chars.len() + 1]; |
| 318 | |
| 319 | for (i, a_ch) in a.chars().enumerate() { |
| 320 | curr[0] = i + 1; |
| 321 | for (j, b_ch) in b_chars.iter().enumerate() { |
| 322 | let cost = if a_ch == *b_ch { 0 } else { 1 }; |
| 323 | let delete = prev[j + 1] + 1; |
| 324 | let insert = curr[j] + 1; |
| 325 | let substitute = prev[j] + cost; |
| 326 | curr[j + 1] = delete.min(insert).min(substitute); |
| 327 | } |
| 328 | std::mem::swap(&mut prev, &mut curr); |
| 329 | } |
| 330 | |
| 331 | prev[b_chars.len()] |
| 332 | } |
| 333 | |
| 334 | fn suggest_tool_names(catalog: &[Tool], requested: &str, limit: usize) -> Vec<String> { |
| 335 | let requested = requested.trim().to_ascii_lowercase(); |
| 336 | if requested.is_empty() || limit == 0 { |
| 337 | return Vec::new(); |
| 338 | } |
| 339 | |
| 340 | let mut candidates: Vec<(u8, usize, String)> = Vec::new(); |
| 341 | for tool in catalog { |
| 342 | let candidate = tool.name.to_ascii_lowercase(); |
| 343 | let prefix_match = candidate.starts_with(&requested) || requested.starts_with(&candidate); |
| 344 | let contains_match = candidate.contains(&requested) || requested.contains(&candidate); |
| 345 | let distance = edit_distance(&candidate, &requested); |
| 346 | let close_typo = distance <= 3; |
| 347 | |
| 348 | if !(prefix_match || contains_match || close_typo) { |
| 349 | continue; |
| 350 | } |
| 351 | |
| 352 | let rank = if prefix_match { |
| 353 | 0 |
| 354 | } else if contains_match { |
| 355 | 1 |
| 356 | } else { |
| 357 | 2 |
| 358 | }; |
| 359 | candidates.push((rank, distance, tool.name.clone())); |
| 360 | } |
| 361 | |
| 362 | candidates.sort_by(|a, b| { |
| 363 | a.0.cmp(&b.0) |
| 364 | .then_with(|| a.1.cmp(&b.1)) |
| 365 | .then_with(|| a.2.cmp(&b.2)) |
| 366 | }); |
| 367 | candidates.dedup_by(|a, b| a.2 == b.2); |
| 368 | candidates |
| 369 | .into_iter() |
| 370 | .take(limit) |
| 371 | .map(|(_, _, name)| name) |
| 372 | .collect() |
| 373 | } |
| 374 | |
| 375 | pub(super) fn missing_tool_error_message(tool_name: &str, catalog: &[Tool]) -> String { |
| 376 | let suggestions = suggest_tool_names(catalog, tool_name, 3); |
| 377 | if suggestions.is_empty() { |
| 378 | return format!( |
| 379 | "Tool '{tool_name}' is not available in the current tool catalog. \ |
| 380 | Verify mode/feature flags, or use {TOOL_SEARCH_BM25_NAME} with a short query." |
| 381 | ); |
| 382 | } |
| 383 | |
| 384 | format!( |
| 385 | "Tool '{tool_name}' is not available in the current tool catalog. \ |
| 386 | Did you mean: {}? You can also use {TOOL_SEARCH_BM25_NAME} to discover tools.", |
| 387 | suggestions.join(", ") |
| 388 | ) |
| 389 | } |
| 390 | |
| 391 | pub(super) fn maybe_activate_requested_deferred_tool( |
| 392 | tool_name: &str, |
| 393 | catalog: &[Tool], |
| 394 | active_tools: &mut HashSet<String>, |
| 395 | ) -> bool { |
| 396 | let Some(def) = catalog.iter().find(|def| def.name == tool_name) else { |
| 397 | return false; |
| 398 | }; |
| 399 | |
| 400 | if !def.defer_loading.unwrap_or(false) || active_tools.contains(tool_name) { |
| 401 | return false; |
| 402 | } |
| 403 | |
| 404 | active_tools.insert(tool_name.to_string()) |
| 405 | } |
| 406 | |
| 407 | pub(super) fn execute_tool_search( |
| 408 | tool_name: &str, |
| 409 | input: &serde_json::Value, |
| 410 | catalog: &[Tool], |
| 411 | active_tools: &mut HashSet<String>, |
| 412 | ) -> Result<ToolResult, ToolError> { |
| 413 | let query = required_str(input, "query")?; |
| 414 | let discovered = if tool_name == TOOL_SEARCH_REGEX_NAME { |
| 415 | discover_tools_with_regex(catalog, query)? |
| 416 | } else { |
| 417 | discover_tools_with_bm25_like(catalog, query) |
| 418 | }; |
| 419 | |
| 420 | for name in &discovered { |
| 421 | active_tools.insert(name.clone()); |
| 422 | } |
| 423 | |
| 424 | let references = discovered |
| 425 | .iter() |
| 426 | .map(|name| json!({"type": "tool_reference", "tool_name": name})) |
| 427 | .collect::<Vec<_>>(); |
| 428 | |
| 429 | let payload = json!({ |
| 430 | "type": "tool_search_tool_search_result", |
| 431 | "tool_references": references, |
| 432 | }); |
| 433 | |
| 434 | Ok(ToolResult { |
| 435 | content: serde_json::to_string(&payload).unwrap_or_else(|_| payload.to_string()), |
| 436 | success: true, |
| 437 | metadata: Some(json!({ |
| 438 | "tool_references": discovered, |
| 439 | })), |
| 440 | }) |
| 441 | } |
| 442 | |
| 443 | pub(super) async fn execute_code_execution_tool( |
| 444 | input: &serde_json::Value, |
| 445 | workspace: &Path, |
| 446 | ) -> Result<ToolResult, ToolError> { |
| 447 | let code = required_str(input, "code")?; |
| 448 | let mut cmd = tokio::process::Command::new("python3"); |
| 449 | cmd.arg("-c"); |
| 450 | cmd.arg(code); |
| 451 | cmd.current_dir(workspace); |
| 452 | |
| 453 | let output = tokio::time::timeout(Duration::from_secs(120), cmd.output()) |
| 454 | .await |
| 455 | .map_err(|_| ToolError::Timeout { seconds: 120 }) |
| 456 | .and_then(|res| res.map_err(|e| ToolError::execution_failed(e.to_string())))?; |
| 457 | |
| 458 | let stdout = String::from_utf8_lossy(&output.stdout).to_string(); |
| 459 | let stderr = String::from_utf8_lossy(&output.stderr).to_string(); |
| 460 | let return_code = output.status.code().unwrap_or(-1); |
| 461 | let success = output.status.success(); |
| 462 | let payload = json!({ |
| 463 | "type": "code_execution_result", |
| 464 | "stdout": stdout, |
| 465 | "stderr": stderr, |
| 466 | "return_code": return_code, |
| 467 | "content": [], |
| 468 | }); |
| 469 | |
| 470 | Ok(ToolResult { |
| 471 | content: serde_json::to_string(&payload).unwrap_or_else(|_| payload.to_string()), |
| 472 | success, |
| 473 | metadata: Some(payload), |
| 474 | }) |
| 475 | } |
| 476 |