vtcode_core/tools/registry/
mod.rs

1//! Tool registry and function declarations
2
3mod astgrep;
4mod builtins;
5mod cache;
6mod declarations;
7mod error;
8mod executors;
9mod legacy;
10mod policy;
11mod pty;
12mod registration;
13mod utils;
14
15pub use declarations::{
16    build_function_declarations, build_function_declarations_for_level,
17    build_function_declarations_with_mode,
18};
19pub use error::{ToolErrorType, ToolExecutionError, classify_error};
20pub use registration::{ToolExecutorFn, ToolHandler, ToolRegistration};
21
22use builtins::register_builtin_tools;
23use utils::normalize_tool_output;
24
25use crate::config::PtyConfig;
26use crate::config::ToolsConfig;
27use crate::config::constants::tools;
28use crate::tool_policy::{ToolPolicy, ToolPolicyManager};
29use crate::tools::ast_grep::AstGrepEngine;
30use crate::tools::grep_search::GrepSearchManager;
31use anyhow::{Result, anyhow};
32use serde_json::Value;
33use std::collections::{HashMap, HashSet};
34use std::path::PathBuf;
35use std::sync::Arc;
36use std::sync::atomic::AtomicUsize;
37use tracing::{debug, warn};
38
39use super::bash_tool::BashTool;
40use super::command::CommandTool;
41use super::curl_tool::CurlTool;
42use super::file_ops::FileOpsTool;
43use super::plan::PlanManager;
44use super::pty::PtyManager;
45use super::search::SearchTool;
46use super::simple_search::SimpleSearchTool;
47use super::srgn::SrgnTool;
48use crate::mcp_client::{McpClient, McpToolExecutor, McpToolInfo};
49
50#[cfg(test)]
51use super::traits::Tool;
52#[cfg(test)]
53use crate::config::types::CapabilityLevel;
54
55#[derive(Clone)]
56pub struct ToolRegistry {
57    workspace_root: PathBuf,
58    search_tool: SearchTool,
59    simple_search_tool: SimpleSearchTool,
60    bash_tool: BashTool,
61    file_ops_tool: FileOpsTool,
62    command_tool: CommandTool,
63    curl_tool: CurlTool,
64    grep_search: Arc<GrepSearchManager>,
65    ast_grep_engine: Option<Arc<AstGrepEngine>>,
66    tool_policy: Option<ToolPolicyManager>,
67    pty_manager: PtyManager,
68    pty_config: PtyConfig,
69    active_pty_sessions: Arc<AtomicUsize>,
70    srgn_tool: SrgnTool,
71    plan_manager: PlanManager,
72    mcp_client: Option<Arc<McpClient>>,
73    mcp_tool_index: HashMap<String, Vec<String>>,
74    tool_registrations: Vec<ToolRegistration>,
75    tool_lookup: HashMap<&'static str, usize>,
76    preapproved_tools: HashSet<String>,
77    full_auto_allowlist: Option<HashSet<String>>,
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
81pub enum ToolPermissionDecision {
82    Allow,
83    Deny,
84    Prompt,
85}
86
87impl ToolRegistry {
88    pub fn new(workspace_root: PathBuf) -> Self {
89        Self::build(workspace_root, PtyConfig::default(), true)
90    }
91
92    pub fn new_with_config(workspace_root: PathBuf, pty_config: PtyConfig) -> Self {
93        Self::build(workspace_root, pty_config, true)
94    }
95
96    pub fn new_with_features(workspace_root: PathBuf, todo_planning_enabled: bool) -> Self {
97        Self::build(workspace_root, PtyConfig::default(), todo_planning_enabled)
98    }
99
100    pub fn new_with_config_and_features(
101        workspace_root: PathBuf,
102        pty_config: PtyConfig,
103        todo_planning_enabled: bool,
104    ) -> Self {
105        Self::build(workspace_root, pty_config, todo_planning_enabled)
106    }
107
108    fn build(workspace_root: PathBuf, pty_config: PtyConfig, todo_planning_enabled: bool) -> Self {
109        let grep_search = Arc::new(GrepSearchManager::new(workspace_root.clone()));
110
111        let search_tool = SearchTool::new(workspace_root.clone(), grep_search.clone());
112        let simple_search_tool = SimpleSearchTool::new(workspace_root.clone());
113        let bash_tool = BashTool::new(workspace_root.clone());
114        let file_ops_tool = FileOpsTool::new(workspace_root.clone(), grep_search.clone());
115        let command_tool = CommandTool::new(workspace_root.clone());
116        let curl_tool = CurlTool::new();
117        let srgn_tool = SrgnTool::new(workspace_root.clone());
118        let plan_manager = PlanManager::new();
119        let pty_manager = PtyManager::new(workspace_root.clone(), pty_config.clone());
120
121        let ast_grep_engine = match AstGrepEngine::new() {
122            Ok(engine) => Some(Arc::new(engine)),
123            Err(err) => {
124                eprintln!("Warning: Failed to initialize AST-grep engine: {}", err);
125                None
126            }
127        };
128
129        let policy_manager = match ToolPolicyManager::new_with_workspace(&workspace_root) {
130            Ok(manager) => Some(manager),
131            Err(err) => {
132                eprintln!("Warning: Failed to initialize tool policy manager: {}", err);
133                None
134            }
135        };
136
137        let mut registry = Self {
138            workspace_root,
139            search_tool,
140            simple_search_tool,
141            bash_tool,
142            file_ops_tool,
143            command_tool,
144            curl_tool,
145            grep_search,
146            ast_grep_engine,
147            tool_policy: policy_manager,
148            pty_manager,
149            pty_config,
150            active_pty_sessions: Arc::new(AtomicUsize::new(0)),
151            srgn_tool,
152            plan_manager,
153            mcp_client: None,
154            mcp_tool_index: HashMap::new(),
155            tool_registrations: Vec::new(),
156            tool_lookup: HashMap::new(),
157            preapproved_tools: HashSet::new(),
158            full_auto_allowlist: None,
159        };
160
161        register_builtin_tools(&mut registry, todo_planning_enabled);
162        registry
163    }
164
165    pub fn register_tool(&mut self, registration: ToolRegistration) -> Result<()> {
166        if self.tool_lookup.contains_key(registration.name()) {
167            return Err(anyhow!(format!(
168                "Tool '{}' is already registered",
169                registration.name()
170            )));
171        }
172
173        let index = self.tool_registrations.len();
174        self.tool_lookup.insert(registration.name(), index);
175        self.tool_registrations.push(registration);
176        Ok(())
177    }
178
179    pub fn available_tools(&self) -> Vec<String> {
180        self.tool_registrations
181            .iter()
182            .map(|registration| registration.name().to_string())
183            .collect()
184    }
185
186    fn mcp_policy_keys(&self) -> Vec<String> {
187        let mut keys = Vec::new();
188        for (provider, tools) in &self.mcp_tool_index {
189            for tool in tools {
190                keys.push(format!("mcp::{}::{}", provider, tool));
191            }
192        }
193        keys
194    }
195
196    fn find_mcp_provider(&self, tool_name: &str) -> Option<String> {
197        for (provider, tools) in &self.mcp_tool_index {
198            if tools.iter().any(|candidate| candidate == tool_name) {
199                return Some(provider.clone());
200            }
201        }
202        None
203    }
204
205    pub fn enable_full_auto_mode(&mut self, allowed_tools: &[String]) {
206        let mut normalized: HashSet<String> = HashSet::new();
207        if allowed_tools
208            .iter()
209            .any(|tool| tool.trim() == tools::WILDCARD_ALL)
210        {
211            for tool in self.available_tools() {
212                normalized.insert(tool);
213            }
214        } else {
215            for tool in allowed_tools {
216                let trimmed = tool.trim();
217                if !trimmed.is_empty() {
218                    normalized.insert(trimmed.to_string());
219                }
220            }
221        }
222
223        self.full_auto_allowlist = Some(normalized);
224    }
225
226    pub fn current_full_auto_allowlist(&self) -> Option<Vec<String>> {
227        self.full_auto_allowlist.as_ref().map(|set| {
228            let mut items: Vec<String> = set.iter().cloned().collect();
229            items.sort();
230            items
231        })
232    }
233
234    pub fn has_tool(&self, name: &str) -> bool {
235        self.tool_lookup.contains_key(name)
236    }
237
238    pub fn with_ast_grep(mut self, engine: Arc<AstGrepEngine>) -> Self {
239        self.ast_grep_engine = Some(engine);
240        self
241    }
242
243    pub fn workspace_root(&self) -> &PathBuf {
244        &self.workspace_root
245    }
246
247    pub fn pty_manager(&self) -> &PtyManager {
248        &self.pty_manager
249    }
250
251    pub fn plan_manager(&self) -> PlanManager {
252        self.plan_manager.clone()
253    }
254
255    pub fn current_plan(&self) -> crate::tools::TaskPlan {
256        self.plan_manager.snapshot()
257    }
258
259    pub async fn initialize_async(&mut self) -> Result<()> {
260        Ok(())
261    }
262
263    pub fn apply_config_policies(&mut self, tools_config: &ToolsConfig) -> Result<()> {
264        if let Ok(policy_manager) = self.policy_manager_mut() {
265            policy_manager.apply_tools_config(tools_config)?;
266        }
267
268        Ok(())
269    }
270
271    pub async fn execute_tool(&mut self, name: &str, args: Value) -> Result<Value> {
272        if let Some(allowlist) = &self.full_auto_allowlist
273            && !allowlist.contains(name)
274        {
275            let error = ToolExecutionError::new(
276                name.to_string(),
277                ToolErrorType::PolicyViolation,
278                format!(
279                    "Tool '{}' is not permitted while full-auto mode is active",
280                    name
281                ),
282            );
283            return Ok(error.to_json_value());
284        }
285
286        let skip_policy_prompt = self.preapproved_tools.remove(name);
287
288        if !skip_policy_prompt
289            && let Ok(policy_manager) = self.policy_manager_mut()
290            && !policy_manager.should_execute_tool(name)?
291        {
292            let error = ToolExecutionError::new(
293                name.to_string(),
294                ToolErrorType::PolicyViolation,
295                format!("Tool '{}' execution denied by policy", name),
296            );
297            return Ok(error.to_json_value());
298        }
299
300        let args = match self.apply_policy_constraints(name, args) {
301            Ok(args) => args,
302            Err(err) => {
303                let error = ToolExecutionError::with_original_error(
304                    name.to_string(),
305                    ToolErrorType::InvalidParameters,
306                    "Failed to apply policy constraints".to_string(),
307                    err.to_string(),
308                );
309                return Ok(error.to_json_value());
310            }
311        };
312
313        let registration = match self
314            .tool_lookup
315            .get(name)
316            .and_then(|index| self.tool_registrations.get(*index))
317        {
318            Some(registration) => registration,
319            None => {
320                // If not found in standard registry, check if it's an MCP tool
321                if let Some(mcp_client) = &self.mcp_client {
322                    // Check if it's an MCP tool (prefixed with "mcp_")
323                    if name.starts_with("mcp_") {
324                        let actual_tool_name = &name[4..]; // Remove "mcp_" prefix
325                        match mcp_client.has_mcp_tool(actual_tool_name).await {
326                            Ok(true) => {
327                                debug!(
328                                    "MCP tool '{}' found, executing via MCP client",
329                                    actual_tool_name
330                                );
331                                return self.execute_mcp_tool(actual_tool_name, args).await;
332                            }
333                            Ok(false) => {
334                                if let Some(resolved_name) =
335                                    self.resolve_mcp_tool_alias(actual_tool_name).await
336                                {
337                                    if resolved_name != actual_tool_name {
338                                        debug!(
339                                            "Resolved MCP tool alias '{}' to '{}'",
340                                            actual_tool_name, resolved_name
341                                        );
342                                        return self.execute_mcp_tool(&resolved_name, args).await;
343                                    }
344                                }
345
346                                // MCP client doesn't have this tool either
347                                let error = ToolExecutionError::new(
348                                    name.to_string(),
349                                    ToolErrorType::ToolNotFound,
350                                    format!("Unknown MCP tool: {}", actual_tool_name),
351                                );
352                                return Ok(error.to_json_value());
353                            }
354                            Err(e) => {
355                                warn!(
356                                    "Error checking MCP tool availability for '{}': {}",
357                                    actual_tool_name, e
358                                );
359                                let error = ToolExecutionError::with_original_error(
360                                    name.to_string(),
361                                    ToolErrorType::ExecutionError,
362                                    format!(
363                                        "Failed to verify MCP tool '{}' due to provider errors",
364                                        actual_tool_name
365                                    ),
366                                    e.to_string(),
367                                );
368                                return Ok(error.to_json_value());
369                            }
370                        }
371                    } else {
372                        // Check if MCP client has a tool with this exact name
373                        match mcp_client.has_mcp_tool(name).await {
374                            Ok(true) => {
375                                debug!(
376                                    "Tool '{}' not found in registry, delegating to MCP client",
377                                    name
378                                );
379                                return self.execute_mcp_tool(name, args).await;
380                            }
381                            Ok(false) => {
382                                // MCP client doesn't have this tool either
383                                let error = ToolExecutionError::new(
384                                    name.to_string(),
385                                    ToolErrorType::ToolNotFound,
386                                    format!("Unknown tool: {}", name),
387                                );
388                                return Ok(error.to_json_value());
389                            }
390                            Err(e) => {
391                                warn!("Error checking MCP tool availability for '{}': {}", name, e);
392                                let error = ToolExecutionError::with_original_error(
393                                    name.to_string(),
394                                    ToolErrorType::ExecutionError,
395                                    format!(
396                                        "Failed to verify MCP tool '{}' due to provider errors",
397                                        name
398                                    ),
399                                    e.to_string(),
400                                );
401                                return Ok(error.to_json_value());
402                            }
403                        }
404                    }
405                } else {
406                    // No MCP client available
407                    let error = ToolExecutionError::new(
408                        name.to_string(),
409                        ToolErrorType::ToolNotFound,
410                        format!("Unknown tool: {}", name),
411                    );
412                    return Ok(error.to_json_value());
413                }
414            }
415        };
416
417        let uses_pty = registration.uses_pty();
418        if uses_pty && let Err(err) = self.start_pty_session() {
419            let error = ToolExecutionError::with_original_error(
420                name.to_string(),
421                ToolErrorType::ExecutionError,
422                "Failed to start PTY session".to_string(),
423                err.to_string(),
424            );
425            return Ok(error.to_json_value());
426        }
427
428        let handler = registration.handler();
429        let result = match handler {
430            ToolHandler::RegistryFn(executor) => executor(self, args).await,
431            ToolHandler::TraitObject(tool) => tool.execute(args).await,
432        };
433
434        if uses_pty {
435            self.end_pty_session();
436        }
437
438        match result {
439            Ok(value) => Ok(normalize_tool_output(value)),
440            Err(err) => {
441                let error_type = classify_error(&err);
442                let error = ToolExecutionError::with_original_error(
443                    name.to_string(),
444                    error_type,
445                    format!("Tool execution failed: {}", err),
446                    err.to_string(),
447                );
448                Ok(error.to_json_value())
449            }
450        }
451    }
452
453    /// Set the MCP client for this registry
454    pub fn with_mcp_client(mut self, mcp_client: Arc<McpClient>) -> Self {
455        self.mcp_client = Some(mcp_client);
456        self
457    }
458
459    /// Attach an MCP client without consuming the registry
460    pub fn set_mcp_client(&mut self, mcp_client: Arc<McpClient>) {
461        self.mcp_client = Some(mcp_client);
462        self.mcp_tool_index.clear();
463    }
464
465    /// Get the MCP client if available
466    pub fn mcp_client(&self) -> Option<&Arc<McpClient>> {
467        self.mcp_client.as_ref()
468    }
469
470    /// List all MCP tools
471    pub async fn list_mcp_tools(&self) -> Result<Vec<McpToolInfo>> {
472        if let Some(mcp_client) = &self.mcp_client {
473            mcp_client.list_mcp_tools().await
474        } else {
475            Ok(Vec::new())
476        }
477    }
478
479    /// Check if an MCP tool exists
480    pub async fn has_mcp_tool(&self, tool_name: &str) -> bool {
481        if let Some(mcp_client) = &self.mcp_client {
482            match mcp_client.has_mcp_tool(tool_name).await {
483                Ok(true) => true,
484                Ok(false) => false,
485                Err(_) => {
486                    // Log error but return false to continue operation
487                    false
488                }
489            }
490        } else {
491            false
492        }
493    }
494
495    /// Execute an MCP tool
496    pub async fn execute_mcp_tool(&self, tool_name: &str, args: Value) -> Result<Value> {
497        if let Some(mcp_client) = &self.mcp_client {
498            mcp_client.execute_mcp_tool(tool_name, args).await
499        } else {
500            Err(anyhow::anyhow!("MCP client not available"))
501        }
502    }
503
504    async fn resolve_mcp_tool_alias(&self, tool_name: &str) -> Option<String> {
505        let Some(mcp_client) = &self.mcp_client else {
506            return None;
507        };
508
509        let normalized = normalize_mcp_tool_identifier(tool_name);
510        if normalized.is_empty() {
511            return None;
512        }
513
514        let tools = match mcp_client.list_mcp_tools().await {
515            Ok(list) => list,
516            Err(err) => {
517                warn!(
518                    "Failed to list MCP tools while resolving alias '{}': {}",
519                    tool_name, err
520                );
521                return None;
522            }
523        };
524
525        for tool in tools {
526            if normalize_mcp_tool_identifier(&tool.name) == normalized {
527                return Some(tool.name);
528            }
529        }
530
531        None
532    }
533
534    /// Refresh MCP tools (reconnect to providers and update tool lists)
535    pub async fn refresh_mcp_tools(&mut self) -> Result<()> {
536        if let Some(mcp_client) = &self.mcp_client {
537            debug!(
538                "Refreshing MCP tools for {} providers",
539                mcp_client.get_status().provider_count
540            );
541
542            let tools = mcp_client.list_mcp_tools().await?;
543            let mut provider_map: HashMap<String, Vec<String>> = HashMap::new();
544
545            for tool in tools {
546                provider_map
547                    .entry(tool.provider.clone())
548                    .or_default()
549                    .push(tool.name.clone());
550            }
551
552            for tools in provider_map.values_mut() {
553                tools.sort();
554                tools.dedup();
555            }
556
557            self.mcp_tool_index = provider_map;
558
559            if let Some(policy_manager) = self.tool_policy.as_mut() {
560                policy_manager.update_mcp_tools(&self.mcp_tool_index)?;
561                let allowlist = policy_manager.mcp_allowlist().clone();
562                mcp_client.update_allowlist(allowlist);
563            }
564
565            self.sync_policy_available_tools();
566            Ok(())
567        } else {
568            debug!("No MCP client configured, nothing to refresh");
569            Ok(())
570        }
571    }
572}
573
574impl ToolRegistry {
575    /// Prompt for permission before starting long-running tool executions to avoid spinner conflicts
576    pub fn preflight_tool_permission(&mut self, name: &str) -> Result<bool> {
577        match self.evaluate_tool_policy(name)? {
578            ToolPermissionDecision::Allow => Ok(true),
579            ToolPermissionDecision::Deny => Ok(false),
580            ToolPermissionDecision::Prompt => Ok(true),
581        }
582    }
583
584    pub fn evaluate_tool_policy(&mut self, name: &str) -> Result<ToolPermissionDecision> {
585        if let Some(tool_name) = name.strip_prefix("mcp_") {
586            return self.evaluate_mcp_tool_policy(name, tool_name);
587        }
588
589        if let Some(allowlist) = self.full_auto_allowlist.as_ref() {
590            if !allowlist.contains(name) {
591                return Ok(ToolPermissionDecision::Deny);
592            }
593
594            if let Some(policy_manager) = self.tool_policy.as_mut() {
595                match policy_manager.get_policy(name) {
596                    ToolPolicy::Deny => return Ok(ToolPermissionDecision::Deny),
597                    ToolPolicy::Allow | ToolPolicy::Prompt => {
598                        self.preapproved_tools.insert(name.to_string());
599                        return Ok(ToolPermissionDecision::Allow);
600                    }
601                }
602            }
603
604            self.preapproved_tools.insert(name.to_string());
605            return Ok(ToolPermissionDecision::Allow);
606        }
607
608        if let Some(policy_manager) = self.tool_policy.as_mut() {
609            match policy_manager.get_policy(name) {
610                ToolPolicy::Allow => {
611                    self.preapproved_tools.insert(name.to_string());
612                    Ok(ToolPermissionDecision::Allow)
613                }
614                ToolPolicy::Deny => Ok(ToolPermissionDecision::Deny),
615                ToolPolicy::Prompt => {
616                    if ToolPolicyManager::is_auto_allow_tool(name) {
617                        policy_manager.set_policy(name, ToolPolicy::Allow)?;
618                        self.preapproved_tools.insert(name.to_string());
619                        Ok(ToolPermissionDecision::Allow)
620                    } else {
621                        Ok(ToolPermissionDecision::Prompt)
622                    }
623                }
624            }
625        } else {
626            self.preapproved_tools.insert(name.to_string());
627            Ok(ToolPermissionDecision::Allow)
628        }
629    }
630
631    fn evaluate_mcp_tool_policy(
632        &mut self,
633        full_name: &str,
634        tool_name: &str,
635    ) -> Result<ToolPermissionDecision> {
636        let provider = match self.find_mcp_provider(tool_name) {
637            Some(provider) => provider,
638            None => {
639                // Unknown provider for this tool; default to prompt for safety
640                return Ok(ToolPermissionDecision::Prompt);
641            }
642        };
643
644        if let Some(allowlist) = self.full_auto_allowlist.as_ref() {
645            if !allowlist.contains(full_name) {
646                return Ok(ToolPermissionDecision::Deny);
647            }
648
649            if let Some(policy_manager) = self.tool_policy.as_mut() {
650                match policy_manager.get_mcp_tool_policy(&provider, tool_name) {
651                    ToolPolicy::Deny => return Ok(ToolPermissionDecision::Deny),
652                    ToolPolicy::Allow | ToolPolicy::Prompt => {
653                        self.preapproved_tools.insert(full_name.to_string());
654                        return Ok(ToolPermissionDecision::Allow);
655                    }
656                }
657            }
658
659            self.preapproved_tools.insert(full_name.to_string());
660            return Ok(ToolPermissionDecision::Allow);
661        }
662
663        if let Some(policy_manager) = self.tool_policy.as_mut() {
664            match policy_manager.get_mcp_tool_policy(&provider, tool_name) {
665                ToolPolicy::Allow => {
666                    self.preapproved_tools.insert(full_name.to_string());
667                    Ok(ToolPermissionDecision::Allow)
668                }
669                ToolPolicy::Deny => Ok(ToolPermissionDecision::Deny),
670                ToolPolicy::Prompt => Ok(ToolPermissionDecision::Prompt),
671            }
672        } else {
673            self.preapproved_tools.insert(full_name.to_string());
674            Ok(ToolPermissionDecision::Allow)
675        }
676    }
677
678    pub fn mark_tool_preapproved(&mut self, name: &str) {
679        self.preapproved_tools.insert(name.to_string());
680    }
681
682    pub fn persist_mcp_tool_policy(&mut self, name: &str, policy: ToolPolicy) -> Result<()> {
683        if !name.starts_with("mcp_") {
684            return Ok(());
685        }
686
687        let Some(tool_name) = name.strip_prefix("mcp_") else {
688            return Ok(());
689        };
690
691        let Some(provider) = self.find_mcp_provider(tool_name) else {
692            return Ok(());
693        };
694
695        if let Some(manager) = self.tool_policy.as_mut() {
696            manager.set_mcp_tool_policy(&provider, tool_name, policy)?;
697        }
698
699        Ok(())
700    }
701}
702
703fn normalize_mcp_tool_identifier(value: &str) -> String {
704    let mut normalized = String::new();
705    for ch in value.chars() {
706        if ch.is_ascii_alphanumeric() {
707            normalized.push(ch.to_ascii_lowercase());
708        }
709    }
710    normalized
711}
712
713#[cfg(test)]
714mod tests {
715    use super::*;
716    use async_trait::async_trait;
717    use serde_json::json;
718    use tempfile::TempDir;
719
720    const CUSTOM_TOOL_NAME: &str = "custom_test_tool";
721
722    struct CustomEchoTool;
723
724    #[async_trait]
725    impl Tool for CustomEchoTool {
726        async fn execute(&self, args: Value) -> Result<Value> {
727            Ok(json!({
728                "success": true,
729                "args": args,
730            }))
731        }
732
733        fn name(&self) -> &'static str {
734            CUSTOM_TOOL_NAME
735        }
736
737        fn description(&self) -> &'static str {
738            "Custom echo tool for testing"
739        }
740    }
741
742    #[tokio::test]
743    async fn registers_builtin_tools() -> Result<()> {
744        let temp_dir = TempDir::new()?;
745        let registry = ToolRegistry::new(temp_dir.path().to_path_buf());
746        let available = registry.available_tools();
747
748        assert!(available.contains(&tools::READ_FILE.to_string()));
749        assert!(available.contains(&tools::RUN_TERMINAL_CMD.to_string()));
750        assert!(available.contains(&tools::CURL.to_string()));
751        Ok(())
752    }
753
754    #[tokio::test]
755    async fn allows_registering_custom_tools() -> Result<()> {
756        let temp_dir = TempDir::new()?;
757        let mut registry = ToolRegistry::new(temp_dir.path().to_path_buf());
758
759        registry.register_tool(ToolRegistration::from_tool_instance(
760            CUSTOM_TOOL_NAME,
761            CapabilityLevel::CodeSearch,
762            CustomEchoTool,
763        ))?;
764
765        registry.sync_policy_available_tools();
766
767        registry.allow_all_tools().ok();
768
769        let available = registry.available_tools();
770        assert!(available.contains(&CUSTOM_TOOL_NAME.to_string()));
771
772        let response = registry
773            .execute_tool(CUSTOM_TOOL_NAME, json!({"input": "value"}))
774            .await?;
775        assert!(response["success"].as_bool().unwrap_or(false));
776        Ok(())
777    }
778
779    #[tokio::test]
780    async fn full_auto_allowlist_enforced() -> Result<()> {
781        let temp_dir = TempDir::new()?;
782        let mut registry = ToolRegistry::new(temp_dir.path().to_path_buf());
783
784        registry.enable_full_auto_mode(&vec![tools::READ_FILE.to_string()]);
785
786        assert!(registry.preflight_tool_permission(tools::READ_FILE)?);
787        assert!(!registry.preflight_tool_permission(tools::RUN_TERMINAL_CMD)?);
788
789        Ok(())
790    }
791
792    #[test]
793    fn normalizes_mcp_tool_identifiers() {
794        assert_eq!(
795            normalize_mcp_tool_identifier("sequential-thinking"),
796            "sequentialthinking"
797        );
798        assert_eq!(
799            normalize_mcp_tool_identifier("Context7.Lookup"),
800            "context7lookup"
801        );
802        assert_eq!(normalize_mcp_tool_identifier("alpha_beta"), "alphabeta");
803    }
804}