| 1 | //! Tool registry for managing and executing tools. |
| 2 | //! |
| 3 | //! The registry provides: |
| 4 | //! - Dynamic tool registration |
| 5 | //! - Tool lookup by name |
| 6 | //! - Conversion to API Tool format |
| 7 | //! - Filtering by capability |
| 8 | |
| 9 | use std::collections::HashMap; |
| 10 | use std::sync::{Arc, OnceLock}; |
| 11 | |
| 12 | use serde_json::Value; |
| 13 | |
| 14 | use crate::client::DeepSeekClient; |
| 15 | use crate::models::Tool; |
| 16 | |
| 17 | use super::schema_sanitize; |
| 18 | use super::spec::{ |
| 19 | ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, |
| 20 | }; |
| 21 | |
| 22 | // === Types === |
| 23 | |
| 24 | /// Registry that holds all available tools. |
| 25 | pub struct ToolRegistry { |
| 26 | tools: HashMap<String, Arc<dyn ToolSpec>>, |
| 27 | context: ToolContext, |
| 28 | /// Memoised serialised tool catalog. Rebuilt lazily on first |
| 29 | /// `to_api_tools` call after a mutation; pinned across reads so the |
| 30 | /// description and schema bytes stay byte-stable for DeepSeek's KV |
| 31 | /// prefix cache. Invalidated on `register` / `remove` / `clear`. |
| 32 | api_cache: OnceLock<Vec<Tool>>, |
| 33 | } |
| 34 | |
| 35 | impl ToolRegistry { |
| 36 | /// Create a new empty registry with the given context. |
| 37 | #[must_use] |
| 38 | pub fn new(context: ToolContext) -> Self { |
| 39 | Self { |
| 40 | tools: HashMap::new(), |
| 41 | context, |
| 42 | api_cache: OnceLock::new(), |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | /// Register a tool in the registry. |
| 47 | pub fn register(&mut self, tool: Arc<dyn ToolSpec>) { |
| 48 | let name = tool.name().to_string(); |
| 49 | if self.tools.insert(name.clone(), tool).is_some() { |
| 50 | tracing::warn!("Overwriting existing tool: {}", name); |
| 51 | } |
| 52 | self.invalidate_api_cache(); |
| 53 | } |
| 54 | |
| 55 | /// Register multiple tools at once. |
| 56 | pub fn register_all(&mut self, tools: Vec<Arc<dyn ToolSpec>>) { |
| 57 | for tool in tools { |
| 58 | self.register(tool); |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | /// Get a tool by name. |
| 63 | #[must_use] |
| 64 | pub fn get(&self, name: &str) -> Option<Arc<dyn ToolSpec>> { |
| 65 | self.tools.get(name).cloned() |
| 66 | } |
| 67 | |
| 68 | /// Check if a tool exists. |
| 69 | #[must_use] |
| 70 | pub fn contains(&self, name: &str) -> bool { |
| 71 | self.tools.contains_key(name) |
| 72 | } |
| 73 | |
| 74 | /// Get all registered tool names. |
| 75 | #[must_use] |
| 76 | #[allow(dead_code)] |
| 77 | pub fn names(&self) -> Vec<&str> { |
| 78 | self.tools.keys().map(std::string::String::as_str).collect() |
| 79 | } |
| 80 | |
| 81 | /// Get the number of registered tools. |
| 82 | #[must_use] |
| 83 | #[allow(dead_code)] |
| 84 | pub fn len(&self) -> usize { |
| 85 | self.tools.len() |
| 86 | } |
| 87 | |
| 88 | /// Check if the registry is empty. |
| 89 | #[must_use] |
| 90 | #[allow(dead_code)] |
| 91 | pub fn is_empty(&self) -> bool { |
| 92 | self.tools.is_empty() |
| 93 | } |
| 94 | |
| 95 | /// Get all registered tools. |
| 96 | #[must_use] |
| 97 | pub fn all(&self) -> Vec<Arc<dyn ToolSpec>> { |
| 98 | self.tools.values().cloned().collect() |
| 99 | } |
| 100 | |
| 101 | /// Execute a tool by name with the given input. |
| 102 | pub async fn execute(&self, name: &str, input: Value) -> Result<String, ToolError> { |
| 103 | let tool = self |
| 104 | .get(name) |
| 105 | .ok_or_else(|| ToolError::not_available(format!("tool '{name}' is not registered")))?; |
| 106 | |
| 107 | let result = tool.execute(input, &self.context).await?; |
| 108 | Ok(result.content) |
| 109 | } |
| 110 | |
| 111 | /// Execute a tool by name, returning the full `ToolResult`. |
| 112 | pub async fn execute_full(&self, name: &str, input: Value) -> Result<ToolResult, ToolError> { |
| 113 | let tool = self |
| 114 | .get(name) |
| 115 | .ok_or_else(|| ToolError::not_available(format!("tool '{name}' is not registered")))?; |
| 116 | |
| 117 | tool.execute(input, &self.context).await |
| 118 | } |
| 119 | |
| 120 | /// Execute a tool with an optional context override. |
| 121 | /// |
| 122 | /// This is used for retrying tools with elevated sandbox policies. |
| 123 | /// After execution, large results are routed through the workshop (#548). |
| 124 | pub async fn execute_full_with_context( |
| 125 | &self, |
| 126 | name: &str, |
| 127 | input: Value, |
| 128 | context_override: Option<&ToolContext>, |
| 129 | ) -> Result<ToolResult, ToolError> { |
| 130 | let tool = self |
| 131 | .get(name) |
| 132 | .ok_or_else(|| ToolError::not_available(format!("tool '{name}' is not registered")))?; |
| 133 | |
| 134 | let ctx = context_override.unwrap_or(&self.context); |
| 135 | let result = tool.execute(input.clone(), ctx).await?; |
| 136 | |
| 137 | // Large-output routing (#548): if the result exceeds the threshold and |
| 138 | // the caller did not request `raw=true`, synthesise via the workshop. |
| 139 | let raw_bypass = input.get("raw").and_then(|v| v.as_bool()).unwrap_or(false); |
| 140 | |
| 141 | if let Some(router) = ctx.large_output_router.as_ref() { |
| 142 | use crate::tools::large_output_router::{LargeOutputRouter, RouteDecision}; |
| 143 | match router.route(name, &result, raw_bypass) { |
| 144 | RouteDecision::PassThrough => {} |
| 145 | RouteDecision::Synthesise { |
| 146 | estimated_tokens, |
| 147 | threshold, |
| 148 | } => { |
| 149 | // Store the raw output in the workshop variable store. |
| 150 | if let Some(vars_arc) = ctx.workshop_vars.as_ref() { |
| 151 | let mut vars = vars_arc.lock().await; |
| 152 | vars.store_raw(name, &result.content); |
| 153 | } |
| 154 | |
| 155 | // Build a terse synthesis using the same model the registry |
| 156 | // was constructed for (workshop Flash model). For now we |
| 157 | // produce a structured header + truncated preview without |
| 158 | // a live API call so the engine stays dependency-free at |
| 159 | // the registry layer. A follow-up can wire in the Flash |
| 160 | // client when the async LLM call is safe here. |
| 161 | let preview_chars = 1_200usize; |
| 162 | let preview: String = result.content.chars().take(preview_chars).collect(); |
| 163 | let ellipsis = if result.content.chars().count() > preview_chars { |
| 164 | "\n… [output truncated — full text in workshop variable `last_tool_result`]" |
| 165 | } else { |
| 166 | "" |
| 167 | }; |
| 168 | let synthesis = format!("{preview}{ellipsis}"); |
| 169 | let wrapped = LargeOutputRouter::wrap_synthesis( |
| 170 | name, |
| 171 | &synthesis, |
| 172 | estimated_tokens, |
| 173 | threshold, |
| 174 | ); |
| 175 | tracing::debug!( |
| 176 | tool = name, |
| 177 | estimated_tokens, |
| 178 | threshold, |
| 179 | "large-output routed through workshop" |
| 180 | ); |
| 181 | return Ok(ToolResult::success(wrapped)); |
| 182 | } |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | Ok(result) |
| 187 | } |
| 188 | |
| 189 | /// Get the current tool context. |
| 190 | #[must_use] |
| 191 | pub fn context(&self) -> &ToolContext { |
| 192 | &self.context |
| 193 | } |
| 194 | |
| 195 | /// Convert all tools to API Tool format for sending to the model. |
| 196 | /// |
| 197 | /// Output is sorted by tool name for **prefix-cache stability** (#263). |
| 198 | /// Rust's `HashMap` uses a randomly-seeded hasher per process, so a raw |
| 199 | /// `self.tools.values()` iteration emits tools in a different order on |
| 200 | /// every `deepseek` launch, invalidating DeepSeek's KV prefix cache for |
| 201 | /// every cross-session resume. Sorting here matches the way Claude Code |
| 202 | /// stabilises its tool array (`assembleToolPool` in their reference). |
| 203 | /// |
| 204 | /// The serialised catalog is memoised on first call and pinned across |
| 205 | /// reads so each tool's `description()` and `input_schema()` are sampled |
| 206 | /// exactly once per registration. MCP adapters whose upstream description |
| 207 | /// drifts on reconnect would otherwise rewrite the catalog mid-session |
| 208 | /// and bust the prefix cache. The cache is invalidated on `register`, |
| 209 | /// `remove`, and `clear`. |
| 210 | #[must_use] |
| 211 | pub fn to_api_tools(&self) -> Vec<Tool> { |
| 212 | self.api_cache |
| 213 | .get_or_init(|| self.build_api_tools()) |
| 214 | .clone() |
| 215 | } |
| 216 | |
| 217 | fn build_api_tools(&self) -> Vec<Tool> { |
| 218 | let mut tools: Vec<&Arc<dyn ToolSpec>> = self.tools.values().collect(); |
| 219 | tools.sort_by(|a, b| a.name().cmp(b.name())); |
| 220 | tools |
| 221 | .into_iter() |
| 222 | .map(|tool| { |
| 223 | let mut schema = tool.input_schema(); |
| 224 | schema_sanitize::sanitize(&mut schema); |
| 225 | Tool { |
| 226 | tool_type: None, |
| 227 | name: tool.name().to_string(), |
| 228 | description: tool.description().to_string(), |
| 229 | input_schema: schema, |
| 230 | allowed_callers: Some(vec!["direct".to_string()]), |
| 231 | defer_loading: Some(tool.defer_loading()), |
| 232 | input_examples: None, |
| 233 | strict: None, |
| 234 | cache_control: None, |
| 235 | } |
| 236 | }) |
| 237 | .collect() |
| 238 | } |
| 239 | |
| 240 | fn invalidate_api_cache(&mut self) { |
| 241 | self.api_cache = OnceLock::new(); |
| 242 | } |
| 243 | |
| 244 | /// Convert tools to API Tool format with optional cache control on the last tool. |
| 245 | #[must_use] |
| 246 | pub fn to_api_tools_with_cache(&self, enable_cache: bool) -> Vec<Tool> { |
| 247 | let mut tools = self.to_api_tools(); |
| 248 | if enable_cache && let Some(last) = tools.last_mut() { |
| 249 | last.cache_control = Some(crate::models::CacheControl { |
| 250 | cache_type: "ephemeral".to_string(), |
| 251 | }); |
| 252 | } |
| 253 | tools |
| 254 | } |
| 255 | |
| 256 | /// Filter tools by capability. |
| 257 | #[must_use] |
| 258 | #[allow(dead_code)] |
| 259 | pub fn filter_by_capability(&self, capability: ToolCapability) -> Vec<Arc<dyn ToolSpec>> { |
| 260 | self.tools |
| 261 | .values() |
| 262 | .filter(|t| t.capabilities().contains(&capability)) |
| 263 | .cloned() |
| 264 | .collect() |
| 265 | } |
| 266 | |
| 267 | /// Get read-only tools. |
| 268 | #[must_use] |
| 269 | #[allow(dead_code)] |
| 270 | pub fn read_only_tools(&self) -> Vec<Arc<dyn ToolSpec>> { |
| 271 | self.tools |
| 272 | .values() |
| 273 | .filter(|t| t.is_read_only()) |
| 274 | .cloned() |
| 275 | .collect() |
| 276 | } |
| 277 | |
| 278 | /// Get tools that require approval. |
| 279 | #[must_use] |
| 280 | #[allow(dead_code)] |
| 281 | pub fn approval_required_tools(&self) -> Vec<Arc<dyn ToolSpec>> { |
| 282 | self.tools |
| 283 | .values() |
| 284 | .filter(|t| t.approval_requirement() == ApprovalRequirement::Required) |
| 285 | .cloned() |
| 286 | .collect() |
| 287 | } |
| 288 | |
| 289 | /// Get tools that suggest approval. |
| 290 | #[must_use] |
| 291 | #[allow(dead_code)] |
| 292 | pub fn approval_suggested_tools(&self) -> Vec<Arc<dyn ToolSpec>> { |
| 293 | self.tools |
| 294 | .values() |
| 295 | .filter(|t| { |
| 296 | matches!( |
| 297 | t.approval_requirement(), |
| 298 | ApprovalRequirement::Suggest | ApprovalRequirement::Required |
| 299 | ) |
| 300 | }) |
| 301 | .cloned() |
| 302 | .collect() |
| 303 | } |
| 304 | |
| 305 | /// Update the context (e.g., when workspace changes). |
| 306 | #[allow(dead_code)] |
| 307 | pub fn set_context(&mut self, context: ToolContext) { |
| 308 | self.context = context; |
| 309 | } |
| 310 | |
| 311 | /// Get a mutable reference to the current context. |
| 312 | #[must_use] |
| 313 | #[allow(dead_code)] |
| 314 | pub fn context_mut(&mut self) -> &mut ToolContext { |
| 315 | &mut self.context |
| 316 | } |
| 317 | |
| 318 | /// Remove a tool by name. |
| 319 | #[must_use] |
| 320 | #[allow(dead_code)] |
| 321 | pub fn remove(&mut self, name: &str) -> Option<Arc<dyn ToolSpec>> { |
| 322 | let removed = self.tools.remove(name); |
| 323 | if removed.is_some() { |
| 324 | self.invalidate_api_cache(); |
| 325 | } |
| 326 | removed |
| 327 | } |
| 328 | |
| 329 | /// Resolve a non-canonical tool name to a registered canonical name. |
| 330 | /// |
| 331 | /// Runs a deterministic ladder against the registered tool names: |
| 332 | /// 1. Lowercase exact match. |
| 333 | /// 2. Hyphens/spaces → underscores (read-file → read_file). |
| 334 | /// 3. CamelCase → snake_case (ReadFile → read_file). |
| 335 | /// 4. Strip trailing `_tool` / `-tool` suffix (twice). |
| 336 | /// 5. Fuzzy match via simple prefix/suffix similarity. |
| 337 | /// |
| 338 | /// Returns `None` when no resolution is found (let the caller surface |
| 339 | /// "Unknown tool"). |
| 340 | #[must_use] |
| 341 | pub fn resolve(&self, requested: &str) -> Option<&str> { |
| 342 | let names: Vec<&str> = self.tools.keys().map(String::as_str).collect(); |
| 343 | let lower = requested.to_lowercase(); |
| 344 | |
| 345 | // 1. lowercase exact |
| 346 | if let Some(n) = names.iter().find(|n| n.to_lowercase() == lower) { |
| 347 | return Some(n); |
| 348 | } |
| 349 | // 2. hyphen/space → underscore |
| 350 | let snaked = lower.replace(['-', ' '], "_"); |
| 351 | if let Some(n) = names.iter().find(|n| **n == snaked) { |
| 352 | return Some(n); |
| 353 | } |
| 354 | // 3. CamelCase → snake_case |
| 355 | let cc = to_snake_case(requested); |
| 356 | if let Some(n) = names.iter().find(|n| **n == cc) { |
| 357 | return Some(n); |
| 358 | } |
| 359 | // 4. strip _tool/-tool/tool suffix, twice |
| 360 | let mut stripped = cc.clone(); |
| 361 | for _ in 0..2 { |
| 362 | for suf in ["_tool", "-tool", "tool"] { |
| 363 | if let Some(s) = stripped.strip_suffix(suf) { |
| 364 | stripped = s.to_string(); |
| 365 | break; |
| 366 | } |
| 367 | } |
| 368 | } |
| 369 | if !stripped.is_empty() |
| 370 | && let Some(n) = names.iter().find(|n| **n == stripped) |
| 371 | { |
| 372 | return Some(n); |
| 373 | } |
| 374 | // 5. fuzzy: simple prefix match (at least 3 chars) |
| 375 | if lower.len() >= 3 { |
| 376 | for n in &names { |
| 377 | if n.len() >= 3 && (n.starts_with(&lower) || lower.starts_with(n)) { |
| 378 | return Some(n); |
| 379 | } |
| 380 | } |
| 381 | } |
| 382 | None |
| 383 | } |
| 384 | |
| 385 | /// Clear all tools from the registry. |
| 386 | #[allow(dead_code)] |
| 387 | pub fn clear(&mut self) { |
| 388 | self.tools.clear(); |
| 389 | self.invalidate_api_cache(); |
| 390 | } |
| 391 | } |
| 392 | |
| 393 | /// Builder for constructing a `ToolRegistry` with common tools. |
| 394 | pub struct ToolRegistryBuilder { |
| 395 | tools: Vec<Arc<dyn ToolSpec>>, |
| 396 | } |
| 397 | |
| 398 | impl ToolRegistryBuilder { |
| 399 | /// Create a new builder. |
| 400 | #[must_use] |
| 401 | pub fn new() -> Self { |
| 402 | Self { tools: Vec::new() } |
| 403 | } |
| 404 | |
| 405 | /// Add a custom tool. |
| 406 | #[must_use] |
| 407 | pub fn with_tool(mut self, tool: Arc<dyn ToolSpec>) -> Self { |
| 408 | self.tools.push(tool); |
| 409 | self |
| 410 | } |
| 411 | |
| 412 | /// Include file tools (read, write, edit, list). |
| 413 | #[must_use] |
| 414 | pub fn with_file_tools(self) -> Self { |
| 415 | use super::file::{EditFileTool, ListDirTool, ReadFileTool, WriteFileTool}; |
| 416 | self.with_tool(Arc::new(ReadFileTool)) |
| 417 | .with_tool(Arc::new(WriteFileTool)) |
| 418 | .with_tool(Arc::new(EditFileTool)) |
| 419 | .with_tool(Arc::new(ListDirTool)) |
| 420 | } |
| 421 | |
| 422 | /// Include only read-only file tools (read, list). |
| 423 | #[must_use] |
| 424 | pub fn with_read_only_file_tools(self) -> Self { |
| 425 | use super::file::{ListDirTool, ReadFileTool}; |
| 426 | self.with_tool(Arc::new(ReadFileTool)) |
| 427 | .with_tool(Arc::new(ListDirTool)) |
| 428 | } |
| 429 | |
| 430 | /// Include shell execution tool. |
| 431 | #[must_use] |
| 432 | pub fn with_shell_tools(self) -> Self { |
| 433 | use super::shell::{ExecShellTool, ShellCancelTool, ShellInteractTool, ShellWaitTool}; |
| 434 | self.with_tool(Arc::new(ExecShellTool)) |
| 435 | .with_tool(Arc::new(ShellWaitTool::new("exec_shell_wait"))) |
| 436 | .with_tool(Arc::new(ShellInteractTool::new("exec_shell_interact"))) |
| 437 | .with_tool(Arc::new(ShellCancelTool)) |
| 438 | .with_tool(Arc::new(ShellWaitTool::new("exec_wait"))) |
| 439 | .with_tool(Arc::new(ShellInteractTool::new("exec_interact"))) |
| 440 | } |
| 441 | |
| 442 | /// Include search tools (`grep_files`). |
| 443 | #[must_use] |
| 444 | pub fn with_search_tools(self) -> Self { |
| 445 | use super::file_search::FileSearchTool; |
| 446 | use super::search::GrepFilesTool; |
| 447 | self.with_tool(Arc::new(GrepFilesTool)) |
| 448 | .with_tool(Arc::new(FileSearchTool)) |
| 449 | } |
| 450 | |
| 451 | /// Include git inspection tools (`git_status`, `git_diff`). |
| 452 | #[must_use] |
| 453 | pub fn with_git_tools(self) -> Self { |
| 454 | use super::git::{GitDiffTool, GitStatusTool}; |
| 455 | self.with_tool(Arc::new(GitStatusTool)) |
| 456 | .with_tool(Arc::new(GitDiffTool)) |
| 457 | } |
| 458 | |
| 459 | /// Include git history tools (`git_log`, `git_show`, `git_blame`). |
| 460 | #[must_use] |
| 461 | pub fn with_git_history_tools(self) -> Self { |
| 462 | use super::git_history::{GitBlameTool, GitLogTool, GitShowTool}; |
| 463 | self.with_tool(Arc::new(GitLogTool)) |
| 464 | .with_tool(Arc::new(GitShowTool)) |
| 465 | .with_tool(Arc::new(GitBlameTool)) |
| 466 | } |
| 467 | |
| 468 | /// Include workspace diagnostics tool. |
| 469 | #[must_use] |
| 470 | pub fn with_diagnostics_tool(self) -> Self { |
| 471 | use super::diagnostics::DiagnosticsTool; |
| 472 | self.with_tool(Arc::new(DiagnosticsTool)) |
| 473 | } |
| 474 | |
| 475 | /// Include the `load_skill` tool (#434) so the model can pull a |
| 476 | /// SKILL.md body + companion file list into context with one |
| 477 | /// call instead of `read_file` + `list_dir` against the path |
| 478 | /// shown in the system prompt's `## Skills` section. |
| 479 | #[must_use] |
| 480 | pub fn with_skill_tools(self) -> Self { |
| 481 | use super::skill::LoadSkillTool; |
| 482 | self.with_tool(Arc::new(LoadSkillTool)) |
| 483 | } |
| 484 | |
| 485 | /// Include project mapping tools. |
| 486 | #[must_use] |
| 487 | pub fn with_project_tools(self) -> Self { |
| 488 | use super::project::ProjectMapTool; |
| 489 | self.with_tool(Arc::new(ProjectMapTool)) |
| 490 | } |
| 491 | |
| 492 | /// Include cargo test runner tool. |
| 493 | #[must_use] |
| 494 | pub fn with_test_runner_tool(self) -> Self { |
| 495 | use super::test_runner::RunTestsTool; |
| 496 | self.with_tool(Arc::new(RunTestsTool)) |
| 497 | } |
| 498 | |
| 499 | /// Include structured data validation tool (`validate_data`). |
| 500 | #[must_use] |
| 501 | pub fn with_validation_tools(self) -> Self { |
| 502 | use super::validate_data::ValidateDataTool; |
| 503 | self.with_tool(Arc::new(ValidateDataTool)) |
| 504 | } |
| 505 | |
| 506 | /// Include durable task, gate, PR-attempt, GitHub, and automation tools. |
| 507 | #[must_use] |
| 508 | pub fn with_runtime_task_tools(self) -> Self { |
| 509 | use super::automation::{ |
| 510 | AutomationCreateTool, AutomationDeleteTool, AutomationListTool, AutomationPauseTool, |
| 511 | AutomationReadTool, AutomationResumeTool, AutomationRunTool, AutomationUpdateTool, |
| 512 | }; |
| 513 | use super::github::{ |
| 514 | GithubCloseIssueTool, GithubCommentTool, GithubIssueContextTool, GithubPrContextTool, |
| 515 | }; |
| 516 | use super::tasks::{ |
| 517 | PrAttemptListTool, PrAttemptPreflightTool, PrAttemptReadTool, PrAttemptRecordTool, |
| 518 | TaskCancelTool, TaskCreateTool, TaskGateRunTool, TaskListTool, TaskReadTool, |
| 519 | TaskShellStartTool, TaskShellWaitTool, |
| 520 | }; |
| 521 | |
| 522 | self.with_tool(Arc::new(TaskCreateTool)) |
| 523 | .with_tool(Arc::new(TaskListTool)) |
| 524 | .with_tool(Arc::new(TaskReadTool)) |
| 525 | .with_tool(Arc::new(TaskCancelTool)) |
| 526 | .with_tool(Arc::new(TaskGateRunTool)) |
| 527 | .with_tool(Arc::new(TaskShellStartTool)) |
| 528 | .with_tool(Arc::new(TaskShellWaitTool)) |
| 529 | .with_tool(Arc::new(GithubIssueContextTool)) |
| 530 | .with_tool(Arc::new(GithubPrContextTool)) |
| 531 | .with_tool(Arc::new(PrAttemptRecordTool)) |
| 532 | .with_tool(Arc::new(PrAttemptListTool)) |
| 533 | .with_tool(Arc::new(PrAttemptReadTool)) |
| 534 | .with_tool(Arc::new(PrAttemptPreflightTool)) |
| 535 | .with_tool(Arc::new(AutomationCreateTool)) |
| 536 | .with_tool(Arc::new(AutomationListTool)) |
| 537 | .with_tool(Arc::new(AutomationReadTool)) |
| 538 | .with_tool(Arc::new(AutomationUpdateTool)) |
| 539 | .with_tool(Arc::new(AutomationPauseTool)) |
| 540 | .with_tool(Arc::new(AutomationResumeTool)) |
| 541 | .with_tool(Arc::new(AutomationDeleteTool)) |
| 542 | .with_tool(Arc::new(AutomationRunTool)) |
| 543 | .with_tool(Arc::new(GithubCommentTool)) |
| 544 | .with_tool(Arc::new(GithubCloseIssueTool)) |
| 545 | } |
| 546 | |
| 547 | /// Include web search tools. |
| 548 | #[must_use] |
| 549 | pub fn with_web_tools(self) -> Self { |
| 550 | use super::fetch_url::FetchUrlTool; |
| 551 | use super::finance::FinanceTool; |
| 552 | use super::web_run::WebRunTool; |
| 553 | use super::web_search::WebSearchTool; |
| 554 | self.with_tool(Arc::new(WebSearchTool)) |
| 555 | .with_tool(Arc::new(FetchUrlTool)) |
| 556 | .with_tool(Arc::new(FinanceTool::new())) |
| 557 | .with_tool(Arc::new(WebRunTool)) |
| 558 | } |
| 559 | |
| 560 | /// Previously registered the OpenAI-style `multi_tool_use.parallel` |
| 561 | /// meta-tool. DeepSeek-V4 has native parallel tool calls (multiple |
| 562 | /// `tool_calls` entries in one assistant turn) and the meta-tool name |
| 563 | /// triggered the model to hallucinate OpenAI-internal XML wrappers |
| 564 | /// (`<multi_tool_use.parallel><tool_name>…</tool_name>…`) instead of |
| 565 | /// emitting native calls. Kept as a no-op so existing callers compile; |
| 566 | /// the engine's compatibility dispatcher still handles legacy emissions. |
| 567 | #[must_use] |
| 568 | pub fn with_parallel_tool(self) -> Self { |
| 569 | self |
| 570 | } |
| 571 | |
| 572 | /// Include request_user_input tool. |
| 573 | #[must_use] |
| 574 | pub fn with_user_input_tool(self) -> Self { |
| 575 | use super::user_input::RequestUserInputTool; |
| 576 | self.with_tool(Arc::new(RequestUserInputTool)) |
| 577 | } |
| 578 | |
| 579 | /// Include patch tools (`apply_patch`). |
| 580 | #[must_use] |
| 581 | pub fn with_patch_tools(self) -> Self { |
| 582 | use super::apply_patch::ApplyPatchTool; |
| 583 | self.with_tool(Arc::new(ApplyPatchTool)) |
| 584 | } |
| 585 | |
| 586 | /// Include the `revert_turn` tool. Approval-gated since it mutates |
| 587 | /// the workspace; the model uses it when the user asks to "undo my |
| 588 | /// last edit". Backed by the per-workspace snapshot side-repo |
| 589 | /// (`crate::snapshot`). |
| 590 | #[must_use] |
| 591 | pub fn with_revert_turn_tool(self) -> Self { |
| 592 | use super::revert_turn::RevertTurnTool; |
| 593 | self.with_tool(Arc::new(RevertTurnTool)) |
| 594 | } |
| 595 | |
| 596 | /// Include the RLM tool (`rlm`). Runs the full recursive language-model |
| 597 | /// loop on a long input (file or inline content); the long input never |
| 598 | /// enters the calling model's context window. The Python REPL exposes |
| 599 | /// `llm_query` / `llm_query_batched` / `rlm_query` / `rlm_query_batched` |
| 600 | /// helpers for sub-LLM work — that's where parallel fan-out belongs. |
| 601 | #[must_use] |
| 602 | pub fn with_rlm_tool(self, client: Option<DeepSeekClient>, root_model: String) -> Self { |
| 603 | use super::rlm::RlmTool; |
| 604 | self.with_tool(Arc::new(RlmTool::new(client, root_model))) |
| 605 | } |
| 606 | |
| 607 | /// Include the review tool. |
| 608 | #[must_use] |
| 609 | pub fn with_review_tool(self, client: Option<DeepSeekClient>, model: String) -> Self { |
| 610 | use super::review::ReviewTool; |
| 611 | self.with_tool(Arc::new(ReviewTool::new(client, model))) |
| 612 | } |
| 613 | |
| 614 | /// Include the `recall_archive` tool — searches prior cycle archives |
| 615 | /// produced by the checkpoint-restart system (issue #127). |
| 616 | #[must_use] |
| 617 | pub fn with_recall_archive_tool(self) -> Self { |
| 618 | use super::recall_archive::RecallArchiveTool; |
| 619 | self.with_tool(Arc::new(RecallArchiveTool)) |
| 620 | } |
| 621 | |
| 622 | /// Include note tool. |
| 623 | #[must_use] |
| 624 | pub fn with_note_tool(self) -> Self { |
| 625 | use super::shell::NoteTool; |
| 626 | self.with_tool(Arc::new(NoteTool)) |
| 627 | } |
| 628 | |
| 629 | /// Include the FIM (Fill-in-the-Middle) edit tool. |
| 630 | #[must_use] |
| 631 | pub fn with_fim_tool(self, client: Option<DeepSeekClient>, model: String) -> Self { |
| 632 | use super::fim::FimEditTool; |
| 633 | self.with_tool(Arc::new(FimEditTool::new(client, model))) |
| 634 | } |
| 635 | |
| 636 | /// Include the `remember` tool — model-callable bullet-add into the |
| 637 | /// user memory file (#489). Only register when the user has opted |
| 638 | /// in to the memory feature; without that, the tool would surface |
| 639 | /// in the model's catalog but always fail with "memory disabled". |
| 640 | #[must_use] |
| 641 | pub fn with_remember_tool(self) -> Self { |
| 642 | use super::remember::RememberTool; |
| 643 | self.with_tool(Arc::new(RememberTool)) |
| 644 | } |
| 645 | |
| 646 | /// Include MCP tools from a connected pool as first-class registry |
| 647 | /// citizens. Each MCP tool is wrapped in a lightweight adapter that |
| 648 | /// implements `ToolSpec`, so the unified `ToolRegistryBuilder` flow |
| 649 | /// handles them alongside native tools. |
| 650 | /// |
| 651 | /// MCP tools are marked `defer_loading` by default (except discovery |
| 652 | /// helpers) to keep the model-visible catalog compact. |
| 653 | #[must_use] |
| 654 | #[allow(dead_code)] |
| 655 | pub fn with_mcp_tools( |
| 656 | mut self, |
| 657 | mcp_pool: std::sync::Arc<tokio::sync::Mutex<crate::mcp::McpPool>>, |
| 658 | ) -> Self { |
| 659 | // Snapshot the current tool list from the pool (non-blocking). |
| 660 | // The adapter lazily resolves at execution time via the pool. |
| 661 | if let Ok(pool) = mcp_pool.try_lock() { |
| 662 | for (name, tool) in pool.all_tools() { |
| 663 | let adapter = Arc::new(McpToolAdapter { |
| 664 | name: name.clone(), |
| 665 | tool: tool.clone(), |
| 666 | pool: mcp_pool.clone(), |
| 667 | }); |
| 668 | self.tools.push(adapter); |
| 669 | } |
| 670 | } |
| 671 | self |
| 672 | } |
| 673 | |
| 674 | /// Include all agent tools (file tools + shell + note + search + patch). |
| 675 | #[must_use] |
| 676 | pub fn with_agent_tools(self, allow_shell: bool) -> Self { |
| 677 | let builder = self |
| 678 | .with_file_tools() |
| 679 | .with_note_tool() |
| 680 | .with_search_tools() |
| 681 | .with_web_tools() |
| 682 | .with_user_input_tool() |
| 683 | .with_parallel_tool() |
| 684 | .with_patch_tools() |
| 685 | .with_git_tools() |
| 686 | .with_git_history_tools() |
| 687 | .with_diagnostics_tool() |
| 688 | .with_project_tools() |
| 689 | .with_skill_tools() |
| 690 | .with_test_runner_tool() |
| 691 | .with_validation_tools() |
| 692 | .with_runtime_task_tools() |
| 693 | .with_revert_turn_tool(); |
| 694 | |
| 695 | if allow_shell { |
| 696 | builder.with_shell_tools() |
| 697 | } else { |
| 698 | builder |
| 699 | } |
| 700 | } |
| 701 | |
| 702 | /// Include the full agent tool surface: every tool family the parent gets |
| 703 | /// in Agent mode, including review, RLM, and the sub-agent management |
| 704 | /// family (so children can recurse). Used by both the parent's Agent-mode |
| 705 | /// registry build (`core/engine.rs`) and by every sub-agent |
| 706 | /// (`subagent::SubAgentToolRegistry`) — keeping them in lockstep. |
| 707 | /// |
| 708 | /// `allow_shell` mirrors the session's shell permission. `manager` and |
| 709 | /// `runtime` are the sub-agent runtime — children pass through their own |
| 710 | /// runtime so grandchildren can spawn within the same depth/cancellation |
| 711 | /// envelope. |
| 712 | #[must_use] |
| 713 | #[allow(clippy::too_many_arguments)] |
| 714 | pub fn with_full_agent_surface( |
| 715 | self, |
| 716 | client: Option<DeepSeekClient>, |
| 717 | model: String, |
| 718 | manager: super::subagent::SharedSubAgentManager, |
| 719 | runtime: super::subagent::SubAgentRuntime, |
| 720 | allow_shell: bool, |
| 721 | todo_list: super::todo::SharedTodoList, |
| 722 | plan_state: super::plan::SharedPlanState, |
| 723 | ) -> Self { |
| 724 | self.with_agent_tools(allow_shell) |
| 725 | .with_todo_tool(todo_list) |
| 726 | .with_plan_tool(plan_state) |
| 727 | .with_review_tool(client.clone(), model.clone()) |
| 728 | .with_rlm_tool(client, model) |
| 729 | .with_recall_archive_tool() |
| 730 | .with_subagent_tools(manager, runtime) |
| 731 | } |
| 732 | |
| 733 | /// Include the todo tool with a shared `TodoList`. |
| 734 | #[must_use] |
| 735 | pub fn with_todo_tool(self, todo_list: super::todo::SharedTodoList) -> Self { |
| 736 | use super::todo::{TodoAddTool, TodoListTool, TodoUpdateTool, TodoWriteTool}; |
| 737 | self.with_tool(Arc::new(TodoWriteTool::checklist(todo_list.clone()))) |
| 738 | .with_tool(Arc::new(TodoAddTool::checklist(todo_list.clone()))) |
| 739 | .with_tool(Arc::new(TodoUpdateTool::checklist(todo_list.clone()))) |
| 740 | .with_tool(Arc::new(TodoListTool::checklist(todo_list.clone()))) |
| 741 | .with_tool(Arc::new(TodoWriteTool::new(todo_list.clone()))) |
| 742 | .with_tool(Arc::new(TodoAddTool::new(todo_list.clone()))) |
| 743 | .with_tool(Arc::new(TodoUpdateTool::new(todo_list.clone()))) |
| 744 | .with_tool(Arc::new(TodoListTool::new(todo_list))) |
| 745 | } |
| 746 | |
| 747 | /// Include the plan tool with a shared `PlanState`. |
| 748 | #[must_use] |
| 749 | pub fn with_plan_tool(self, plan_state: super::plan::SharedPlanState) -> Self { |
| 750 | use super::plan::UpdatePlanTool; |
| 751 | self.with_tool(Arc::new(UpdatePlanTool::new(plan_state))) |
| 752 | } |
| 753 | |
| 754 | /// Include sub-agent management tools. |
| 755 | #[must_use] |
| 756 | pub fn with_subagent_tools( |
| 757 | self, |
| 758 | manager: super::subagent::SharedSubAgentManager, |
| 759 | runtime: super::subagent::SubAgentRuntime, |
| 760 | ) -> Self { |
| 761 | use super::subagent::{ |
| 762 | AgentAssignTool, AgentCancelTool, AgentCloseTool, AgentListTool, AgentResultTool, |
| 763 | AgentResumeTool, AgentSendInputTool, AgentSpawnTool, AgentWaitTool, |
| 764 | DelegateToAgentTool, |
| 765 | }; |
| 766 | |
| 767 | self.with_tool(Arc::new(AgentSpawnTool::new( |
| 768 | manager.clone(), |
| 769 | runtime.clone(), |
| 770 | ))) |
| 771 | .with_tool(Arc::new(AgentSpawnTool::with_name( |
| 772 | manager.clone(), |
| 773 | runtime.clone(), |
| 774 | "spawn_agent", |
| 775 | ))) |
| 776 | .with_tool(Arc::new(DelegateToAgentTool::new( |
| 777 | manager.clone(), |
| 778 | runtime.clone(), |
| 779 | ))) |
| 780 | .with_tool(Arc::new(AgentResultTool::new(manager.clone()))) |
| 781 | .with_tool(Arc::new(AgentSendInputTool::new( |
| 782 | manager.clone(), |
| 783 | "send_input", |
| 784 | ))) |
| 785 | .with_tool(Arc::new(AgentAssignTool::new( |
| 786 | manager.clone(), |
| 787 | "agent_assign", |
| 788 | ))) |
| 789 | .with_tool(Arc::new(AgentAssignTool::new( |
| 790 | manager.clone(), |
| 791 | "assign_agent", |
| 792 | ))) |
| 793 | .with_tool(Arc::new(AgentWaitTool::new(manager.clone(), "wait"))) |
| 794 | .with_tool(Arc::new(AgentSendInputTool::new( |
| 795 | manager.clone(), |
| 796 | "agent_send_input", |
| 797 | ))) |
| 798 | .with_tool(Arc::new(AgentWaitTool::new(manager.clone(), "agent_wait"))) |
| 799 | .with_tool(Arc::new(AgentResumeTool::new( |
| 800 | manager.clone(), |
| 801 | runtime.clone(), |
| 802 | ))) |
| 803 | .with_tool(Arc::new(AgentCloseTool::new(manager.clone()))) |
| 804 | .with_tool(Arc::new(AgentCancelTool::new(manager.clone()))) |
| 805 | .with_tool(Arc::new(AgentListTool::new(manager))) |
| 806 | } |
| 807 | |
| 808 | /// Build the registry with the given context. |
| 809 | #[must_use] |
| 810 | pub fn build(self, context: ToolContext) -> ToolRegistry { |
| 811 | let mut registry = ToolRegistry::new(context); |
| 812 | registry.register_all(self.tools); |
| 813 | registry |
| 814 | } |
| 815 | } |
| 816 | |
| 817 | impl Default for ToolRegistryBuilder { |
| 818 | fn default() -> Self { |
| 819 | Self::new() |
| 820 | } |
| 821 | } |
| 822 | |
| 823 | /// Convert CamelCase to snake_case. |
| 824 | fn to_snake_case(s: &str) -> String { |
| 825 | let mut out = String::with_capacity(s.len() + 4); |
| 826 | for (i, ch) in s.chars().enumerate() { |
| 827 | if ch.is_uppercase() { |
| 828 | if i > 0 { |
| 829 | out.push('_'); |
| 830 | } |
| 831 | out.push(ch.to_ascii_lowercase()); |
| 832 | } else { |
| 833 | out.push(ch); |
| 834 | } |
| 835 | } |
| 836 | out |
| 837 | } |
| 838 | |
| 839 | /// Adapter that wraps an MCP tool definition so it can live in the |
| 840 | /// unified `ToolRegistry` alongside native tools (§5.B). |
| 841 | #[allow(dead_code)] |
| 842 | struct McpToolAdapter { |
| 843 | name: String, |
| 844 | tool: crate::mcp::McpTool, |
| 845 | pool: std::sync::Arc<tokio::sync::Mutex<crate::mcp::McpPool>>, |
| 846 | } |
| 847 | |
| 848 | #[async_trait::async_trait] |
| 849 | impl ToolSpec for McpToolAdapter { |
| 850 | fn name(&self) -> &str { |
| 851 | &self.name |
| 852 | } |
| 853 | |
| 854 | fn description(&self) -> &str { |
| 855 | // McpTool.description is Option<String>; fall back to the |
| 856 | // prefixed name when absent. |
| 857 | self.tool.description.as_deref().unwrap_or(&self.name) |
| 858 | } |
| 859 | |
| 860 | fn input_schema(&self) -> Value { |
| 861 | self.tool.input_schema.clone() |
| 862 | } |
| 863 | |
| 864 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 865 | // Conservatively treat MCP tools as requiring approval and |
| 866 | // network access unless they're known discovery helpers. |
| 867 | let name_lower = self.name.to_lowercase(); |
| 868 | if name_lower.contains("list_mcp") |
| 869 | || name_lower.contains("read_mcp") |
| 870 | || name_lower.contains("mcp_read") |
| 871 | || name_lower.contains("mcp_get_prompt") |
| 872 | { |
| 873 | vec![ToolCapability::ReadOnly] |
| 874 | } else { |
| 875 | vec![ToolCapability::Network, ToolCapability::RequiresApproval] |
| 876 | } |
| 877 | } |
| 878 | |
| 879 | fn defer_loading(&self) -> bool { |
| 880 | // Discovery helpers stay loaded; everything else is deferred. |
| 881 | let keep_loaded = matches!( |
| 882 | self.name.as_str(), |
| 883 | "list_mcp_resources" |
| 884 | | "list_mcp_resource_templates" |
| 885 | | "mcp_read_resource" |
| 886 | | "read_mcp_resource" |
| 887 | | "mcp_get_prompt" |
| 888 | ); |
| 889 | !keep_loaded |
| 890 | } |
| 891 | |
| 892 | async fn execute(&self, input: Value, _context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 893 | let mut pool = self.pool.lock().await; |
| 894 | let result = pool |
| 895 | .call_tool(&self.name, input) |
| 896 | .await |
| 897 | .map_err(|e| ToolError::execution_failed(format!("MCP tool failed: {e}")))?; |
| 898 | let content = serde_json::to_string_pretty(&result).unwrap_or_else(|_| result.to_string()); |
| 899 | Ok(ToolResult::success(content)) |
| 900 | } |
| 901 | } |
| 902 | |
| 903 | // === Unit Tests === |
| 904 | |
| 905 | #[cfg(test)] |
| 906 | mod tests { |
| 907 | use std::sync::Arc; |
| 908 | |
| 909 | use serde_json::{Value, json}; |
| 910 | use tempfile::tempdir; |
| 911 | |
| 912 | use crate::tools::ToolRegistryBuilder; |
| 913 | use crate::tools::spec::{ |
| 914 | ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, required_str, |
| 915 | }; |
| 916 | |
| 917 | use super::ToolRegistry; |
| 918 | |
| 919 | /// A simple test tool for unit testing |
| 920 | struct TestTool { |
| 921 | name: String, |
| 922 | description: String, |
| 923 | } |
| 924 | |
| 925 | #[async_trait::async_trait] |
| 926 | impl ToolSpec for TestTool { |
| 927 | fn name(&self) -> &str { |
| 928 | &self.name |
| 929 | } |
| 930 | |
| 931 | fn description(&self) -> &str { |
| 932 | &self.description |
| 933 | } |
| 934 | |
| 935 | fn input_schema(&self) -> Value { |
| 936 | json!({ |
| 937 | "type": "object", |
| 938 | "properties": { |
| 939 | "message": { "type": "string" } |
| 940 | }, |
| 941 | "required": ["message"] |
| 942 | }) |
| 943 | } |
| 944 | |
| 945 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 946 | vec![ToolCapability::ReadOnly] |
| 947 | } |
| 948 | |
| 949 | async fn execute( |
| 950 | &self, |
| 951 | input: Value, |
| 952 | _context: &ToolContext, |
| 953 | ) -> Result<ToolResult, ToolError> { |
| 954 | let message = required_str(&input, "message")?; |
| 955 | Ok(ToolResult::success(format!("Echo: {message}"))) |
| 956 | } |
| 957 | } |
| 958 | |
| 959 | fn make_test_tool(name: &str) -> Arc<TestTool> { |
| 960 | Arc::new(TestTool { |
| 961 | name: name.to_string(), |
| 962 | description: "A test tool".to_string(), |
| 963 | }) |
| 964 | } |
| 965 | |
| 966 | #[test] |
| 967 | fn test_registry_register_and_get() { |
| 968 | let tmp = tempdir().expect("tempdir"); |
| 969 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 970 | let mut registry = ToolRegistry::new(ctx); |
| 971 | |
| 972 | let tool = make_test_tool("test_tool"); |
| 973 | registry.register(tool); |
| 974 | |
| 975 | assert!(registry.contains("test_tool")); |
| 976 | assert!(!registry.contains("nonexistent")); |
| 977 | assert_eq!(registry.len(), 1); |
| 978 | } |
| 979 | |
| 980 | #[test] |
| 981 | fn test_registry_names() { |
| 982 | let tmp = tempdir().expect("tempdir"); |
| 983 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 984 | let mut registry = ToolRegistry::new(ctx); |
| 985 | |
| 986 | registry.register(make_test_tool("tool_a")); |
| 987 | registry.register(make_test_tool("tool_b")); |
| 988 | |
| 989 | let names = registry.names(); |
| 990 | assert_eq!(names.len(), 2); |
| 991 | assert!(names.contains(&"tool_a")); |
| 992 | assert!(names.contains(&"tool_b")); |
| 993 | } |
| 994 | |
| 995 | #[test] |
| 996 | fn test_registry_to_api_tools() { |
| 997 | let tmp = tempdir().expect("tempdir"); |
| 998 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 999 | let mut registry = ToolRegistry::new(ctx); |
| 1000 | |
| 1001 | registry.register(make_test_tool("my_tool")); |
| 1002 | |
| 1003 | let api_tools = registry.to_api_tools(); |
| 1004 | assert_eq!(api_tools.len(), 1); |
| 1005 | assert_eq!(api_tools[0].name, "my_tool"); |
| 1006 | assert_eq!(api_tools[0].description, "A test tool"); |
| 1007 | } |
| 1008 | |
| 1009 | #[test] |
| 1010 | fn api_tools_with_cache_marks_last_tool_ephemeral() { |
| 1011 | let tmp = tempdir().expect("tempdir"); |
| 1012 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1013 | let mut registry = ToolRegistry::new(ctx); |
| 1014 | |
| 1015 | registry.register(make_test_tool("tool_a")); |
| 1016 | registry.register(make_test_tool("tool_b")); |
| 1017 | |
| 1018 | let api_tools = registry.to_api_tools_with_cache(true); |
| 1019 | assert_eq!(api_tools.len(), 2); |
| 1020 | assert!(api_tools[0].cache_control.is_none()); |
| 1021 | assert_eq!( |
| 1022 | api_tools[1] |
| 1023 | .cache_control |
| 1024 | .as_ref() |
| 1025 | .map(|c| c.cache_type.as_str()), |
| 1026 | Some("ephemeral") |
| 1027 | ); |
| 1028 | } |
| 1029 | |
| 1030 | /// Tool whose `description()` advances through a script of pre-built |
| 1031 | /// strings, one per call. Used to demonstrate that the api-tools cache |
| 1032 | /// pins the description bytes on first read instead of re-sampling them |
| 1033 | /// each turn (#263 follow-up; mirrors reference-cc's `getToolSchemaCache`). |
| 1034 | struct VaryingDescriptionTool { |
| 1035 | name: String, |
| 1036 | descriptions: Vec<String>, |
| 1037 | next: std::sync::atomic::AtomicUsize, |
| 1038 | } |
| 1039 | |
| 1040 | impl VaryingDescriptionTool { |
| 1041 | fn new(name: &str, descriptions: &[&str]) -> Self { |
| 1042 | Self { |
| 1043 | name: name.to_string(), |
| 1044 | descriptions: descriptions.iter().map(|s| (*s).to_string()).collect(), |
| 1045 | next: std::sync::atomic::AtomicUsize::new(0), |
| 1046 | } |
| 1047 | } |
| 1048 | } |
| 1049 | |
| 1050 | #[async_trait::async_trait] |
| 1051 | impl ToolSpec for VaryingDescriptionTool { |
| 1052 | fn name(&self) -> &str { |
| 1053 | &self.name |
| 1054 | } |
| 1055 | |
| 1056 | fn description(&self) -> &str { |
| 1057 | let idx = self |
| 1058 | .next |
| 1059 | .fetch_add(1, std::sync::atomic::Ordering::SeqCst) |
| 1060 | .min(self.descriptions.len() - 1); |
| 1061 | &self.descriptions[idx] |
| 1062 | } |
| 1063 | |
| 1064 | fn input_schema(&self) -> Value { |
| 1065 | json!({"type": "object", "properties": {}, "required": []}) |
| 1066 | } |
| 1067 | |
| 1068 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 1069 | vec![ToolCapability::ReadOnly] |
| 1070 | } |
| 1071 | |
| 1072 | async fn execute( |
| 1073 | &self, |
| 1074 | _input: Value, |
| 1075 | _context: &ToolContext, |
| 1076 | ) -> Result<ToolResult, ToolError> { |
| 1077 | Ok(ToolResult::success("ok".to_string())) |
| 1078 | } |
| 1079 | } |
| 1080 | |
| 1081 | #[test] |
| 1082 | fn to_api_tools_pins_description_bytes_across_calls() { |
| 1083 | // Regression for the cache-stability follow-up: an MCP adapter that |
| 1084 | // returns a different `description()` on reconnect (or any other |
| 1085 | // tool whose description isn't a `&'static str`) would otherwise |
| 1086 | // rewrite the catalog bytes mid-session and miss the prefix cache. |
| 1087 | // The registry pins the first call's value until it's mutated. |
| 1088 | let tmp = tempdir().expect("tempdir"); |
| 1089 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1090 | let mut registry = ToolRegistry::new(ctx); |
| 1091 | registry.register(Arc::new(VaryingDescriptionTool::new( |
| 1092 | "varying", |
| 1093 | &["first description", "second description"], |
| 1094 | ))); |
| 1095 | |
| 1096 | let first = registry.to_api_tools(); |
| 1097 | let second = registry.to_api_tools(); |
| 1098 | |
| 1099 | assert_eq!(first.len(), 1); |
| 1100 | assert_eq!(first[0].description, "first description"); |
| 1101 | assert_eq!( |
| 1102 | first, second, |
| 1103 | "api-tools catalog must be byte-identical across reads with no mutation in between" |
| 1104 | ); |
| 1105 | } |
| 1106 | |
| 1107 | #[test] |
| 1108 | fn register_invalidates_api_tools_cache() { |
| 1109 | // Counter-test: when a real change happens (a new tool registers, |
| 1110 | // an existing one is removed, or `clear` is called), the cache must |
| 1111 | // be discarded so the next read reflects the live registry. |
| 1112 | let tmp = tempdir().expect("tempdir"); |
| 1113 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1114 | let mut registry = ToolRegistry::new(ctx); |
| 1115 | registry.register(Arc::new(VaryingDescriptionTool::new( |
| 1116 | "varying", |
| 1117 | &["first description", "second description"], |
| 1118 | ))); |
| 1119 | |
| 1120 | let before = registry.to_api_tools(); |
| 1121 | assert_eq!(before.len(), 1); |
| 1122 | |
| 1123 | registry.register(make_test_tool("late_arrival")); |
| 1124 | |
| 1125 | let after = registry.to_api_tools(); |
| 1126 | assert_eq!(after.len(), 2, "cache must rebuild after register"); |
| 1127 | assert!(after.iter().any(|t| t.name == "varying")); |
| 1128 | assert!(after.iter().any(|t| t.name == "late_arrival")); |
| 1129 | // The varying tool's description advances on cache rebuild — the |
| 1130 | // first read above sampled `first description`; this rebuild samples |
| 1131 | // `second description`. The point is just that the bytes *can* |
| 1132 | // change after a real mutation, not that they always do. |
| 1133 | let varying_after = after |
| 1134 | .iter() |
| 1135 | .find(|t| t.name == "varying") |
| 1136 | .expect("varying tool present"); |
| 1137 | assert_eq!(varying_after.description, "second description"); |
| 1138 | } |
| 1139 | |
| 1140 | #[test] |
| 1141 | fn remove_and_clear_invalidate_api_tools_cache() { |
| 1142 | let tmp = tempdir().expect("tempdir"); |
| 1143 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1144 | let mut registry = ToolRegistry::new(ctx); |
| 1145 | registry.register(make_test_tool("alpha")); |
| 1146 | registry.register(make_test_tool("beta")); |
| 1147 | |
| 1148 | let before = registry.to_api_tools(); |
| 1149 | assert_eq!(before.len(), 2); |
| 1150 | |
| 1151 | let _ = registry.remove("alpha"); |
| 1152 | let after_remove = registry.to_api_tools(); |
| 1153 | assert_eq!(after_remove.len(), 1); |
| 1154 | assert_eq!(after_remove[0].name, "beta"); |
| 1155 | |
| 1156 | registry.clear(); |
| 1157 | let after_clear = registry.to_api_tools(); |
| 1158 | assert!(after_clear.is_empty(), "cache must clear with the registry"); |
| 1159 | } |
| 1160 | |
| 1161 | #[test] |
| 1162 | fn to_api_tools_emits_alphabetical_order_regardless_of_registration_order() { |
| 1163 | // Regression for #263: HashMap iteration is non-deterministic across |
| 1164 | // process launches, which busts DeepSeek's KV prefix cache for every |
| 1165 | // cross-session resume. `to_api_tools` must emit by name regardless |
| 1166 | // of registration order so two consecutive calls (and two distinct |
| 1167 | // launches) produce byte-identical output. |
| 1168 | let tmp = tempdir().expect("tempdir"); |
| 1169 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1170 | |
| 1171 | let order_a = { |
| 1172 | let mut registry = ToolRegistry::new(ctx.clone()); |
| 1173 | registry.register(make_test_tool("zebra")); |
| 1174 | registry.register(make_test_tool("alpha")); |
| 1175 | registry.register(make_test_tool("mango")); |
| 1176 | registry |
| 1177 | .to_api_tools() |
| 1178 | .iter() |
| 1179 | .map(|t| t.name.clone()) |
| 1180 | .collect::<Vec<_>>() |
| 1181 | }; |
| 1182 | |
| 1183 | let order_b = { |
| 1184 | let mut registry = ToolRegistry::new(ctx.clone()); |
| 1185 | registry.register(make_test_tool("alpha")); |
| 1186 | registry.register(make_test_tool("mango")); |
| 1187 | registry.register(make_test_tool("zebra")); |
| 1188 | registry |
| 1189 | .to_api_tools() |
| 1190 | .iter() |
| 1191 | .map(|t| t.name.clone()) |
| 1192 | .collect::<Vec<_>>() |
| 1193 | }; |
| 1194 | |
| 1195 | assert_eq!(order_a, vec!["alpha", "mango", "zebra"]); |
| 1196 | assert_eq!(order_a, order_b); |
| 1197 | } |
| 1198 | |
| 1199 | #[test] |
| 1200 | fn test_registry_remove() { |
| 1201 | let tmp = tempdir().expect("tempdir"); |
| 1202 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1203 | let mut registry = ToolRegistry::new(ctx); |
| 1204 | |
| 1205 | registry.register(make_test_tool("removable")); |
| 1206 | assert!(registry.contains("removable")); |
| 1207 | |
| 1208 | let _ = registry.remove("removable"); |
| 1209 | assert!(!registry.contains("removable")); |
| 1210 | } |
| 1211 | |
| 1212 | #[test] |
| 1213 | fn test_registry_clear() { |
| 1214 | let tmp = tempdir().expect("tempdir"); |
| 1215 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1216 | let mut registry = ToolRegistry::new(ctx); |
| 1217 | |
| 1218 | registry.register(make_test_tool("tool1")); |
| 1219 | registry.register(make_test_tool("tool2")); |
| 1220 | assert_eq!(registry.len(), 2); |
| 1221 | |
| 1222 | registry.clear(); |
| 1223 | assert!(registry.is_empty()); |
| 1224 | } |
| 1225 | |
| 1226 | #[tokio::test] |
| 1227 | async fn test_registry_execute() { |
| 1228 | let tmp = tempdir().expect("tempdir"); |
| 1229 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1230 | let mut registry = ToolRegistry::new(ctx); |
| 1231 | |
| 1232 | registry.register(make_test_tool("echo")); |
| 1233 | |
| 1234 | let result = registry |
| 1235 | .execute("echo", json!({"message": "hello"})) |
| 1236 | .await |
| 1237 | .expect("execute"); |
| 1238 | |
| 1239 | assert_eq!(result, "Echo: hello"); |
| 1240 | } |
| 1241 | |
| 1242 | #[tokio::test] |
| 1243 | async fn test_registry_execute_unknown_tool() { |
| 1244 | let tmp = tempdir().expect("tempdir"); |
| 1245 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1246 | let registry = ToolRegistry::new(ctx); |
| 1247 | |
| 1248 | let result = registry.execute("nonexistent", json!({})).await; |
| 1249 | assert!(result.is_err()); |
| 1250 | } |
| 1251 | |
| 1252 | #[test] |
| 1253 | fn test_builder_basic() { |
| 1254 | let tmp = tempdir().expect("tempdir"); |
| 1255 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1256 | |
| 1257 | let registry = ToolRegistryBuilder::new() |
| 1258 | .with_tool(make_test_tool("custom")) |
| 1259 | .build(ctx); |
| 1260 | |
| 1261 | assert!(registry.contains("custom")); |
| 1262 | } |
| 1263 | |
| 1264 | #[test] |
| 1265 | fn test_filter_by_capability() { |
| 1266 | let tmp = tempdir().expect("tempdir"); |
| 1267 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1268 | let mut registry = ToolRegistry::new(ctx); |
| 1269 | |
| 1270 | registry.register(make_test_tool("readonly_tool")); |
| 1271 | |
| 1272 | let readonly = registry.filter_by_capability(ToolCapability::ReadOnly); |
| 1273 | assert_eq!(readonly.len(), 1); |
| 1274 | |
| 1275 | let writes = registry.filter_by_capability(ToolCapability::WritesFiles); |
| 1276 | assert_eq!(writes.len(), 0); |
| 1277 | } |
| 1278 | |
| 1279 | #[test] |
| 1280 | fn test_read_only_tools() { |
| 1281 | let tmp = tempdir().expect("tempdir"); |
| 1282 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1283 | let mut registry = ToolRegistry::new(ctx); |
| 1284 | |
| 1285 | registry.register(make_test_tool("reader")); |
| 1286 | |
| 1287 | let readonly = registry.read_only_tools(); |
| 1288 | assert_eq!(readonly.len(), 1); |
| 1289 | assert_eq!(readonly[0].name(), "reader"); |
| 1290 | } |
| 1291 | |
| 1292 | #[test] |
| 1293 | fn test_builder_with_web_tools_includes_finance() { |
| 1294 | let tmp = tempdir().expect("tempdir"); |
| 1295 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1296 | |
| 1297 | let registry = ToolRegistryBuilder::new().with_web_tools().build(ctx); |
| 1298 | |
| 1299 | assert!(registry.contains("finance")); |
| 1300 | } |
| 1301 | |
| 1302 | #[test] |
| 1303 | fn test_builder_with_agent_tools_includes_finance() { |
| 1304 | let tmp = tempdir().expect("tempdir"); |
| 1305 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 1306 | |
| 1307 | let registry = ToolRegistryBuilder::new() |
| 1308 | .with_agent_tools(false) |
| 1309 | .build(ctx); |
| 1310 | |
| 1311 | assert!(registry.contains("finance")); |
| 1312 | } |
| 1313 | } |
| 1314 |