Skip to main content

vtcode_core/tools/registry/
builder.rs

1//! ToolRegistry construction helpers.
2
3use std::future::Future;
4use std::path::Path;
5use std::path::PathBuf;
6use std::sync::Arc;
7use std::sync::RwLock;
8
9use parking_lot::Mutex;
10
11use crate::config::PtyConfig;
12use crate::core::memory_pool::MemoryPool;
13use crate::tool_policy::ToolPolicyManager;
14use crate::tools::handlers::PlanModeState;
15use crate::tools::output_spooler::{SpoolerConfig, ToolOutputSpooler};
16use crate::tools::safety_gateway::SafetyGateway;
17use vtcode_config::DynamicContextConfig;
18use vtcode_config::loader::ConfigManager;
19
20use super::ToolRegistry;
21use super::assembly::ToolAssembly;
22use super::builtins::register_builtin_tools;
23use super::circuit_breaker;
24use super::execution_history::ToolExecutionHistory;
25use super::harness::HarnessContext;
26use super::inventory::ToolInventory;
27use super::policy::ToolPolicyGateway;
28use super::pty;
29use super::resiliency::ResiliencyContext;
30use super::shell_policy::ShellPolicyChecker;
31use super::timeout::ToolTimeoutPolicy;
32
33fn spooler_config_from_dynamic_context(config: &DynamicContextConfig) -> SpoolerConfig {
34    SpoolerConfig {
35        enabled: config.enabled,
36        threshold_bytes: config.tool_output_threshold,
37        max_files: config.max_spooled_files,
38        max_age_secs: config.spool_max_age_secs,
39        include_file_reference: true,
40    }
41}
42
43fn load_workspace_spooler_config(workspace_root: &Path) -> SpoolerConfig {
44    match ConfigManager::load_from_workspace(workspace_root) {
45        Ok(manager) => spooler_config_from_dynamic_context(&manager.config().context.dynamic),
46        Err(err) => {
47            tracing::warn!(
48                workspace = %workspace_root.display(),
49                error = %err,
50                "Failed to load workspace config for output spooler; using defaults"
51            );
52            SpoolerConfig::default()
53        }
54    }
55}
56
57impl ToolRegistry {
58    pub fn new(workspace_root: PathBuf) -> impl Future<Output = Self> {
59        Self::build(workspace_root, PtyConfig::default())
60    }
61
62    pub fn new_with_config(
63        workspace_root: PathBuf,
64        pty_config: PtyConfig,
65    ) -> impl Future<Output = Self> {
66        Self::build(workspace_root, pty_config)
67    }
68
69    pub fn new_with_custom_policy(
70        workspace_root: PathBuf,
71        policy_manager: ToolPolicyManager,
72    ) -> impl Future<Output = Self> {
73        Self::build_with_policy(workspace_root, PtyConfig::default(), Some(policy_manager))
74    }
75
76    pub fn new_with_custom_policy_and_config(
77        workspace_root: PathBuf,
78        pty_config: PtyConfig,
79        policy_manager: ToolPolicyManager,
80    ) -> impl Future<Output = Self> {
81        Self::build_with_policy(workspace_root, pty_config, Some(policy_manager))
82    }
83
84    async fn build(workspace_root: PathBuf, pty_config: PtyConfig) -> Self {
85        Self::build_with_policy(workspace_root, pty_config, None).await
86    }
87
88    async fn build_with_policy(
89        workspace_root: PathBuf,
90        pty_config: PtyConfig,
91        policy_manager: Option<ToolPolicyManager>,
92    ) -> Self {
93        let edited_file_monitor =
94            Arc::new(crate::tools::edited_file_monitor::EditedFileMonitor::new());
95        let inventory =
96            ToolInventory::new(workspace_root.clone(), Arc::clone(&edited_file_monitor));
97        let plan_mode_state = PlanModeState::new(workspace_root.clone());
98
99        register_builtin_tools(&inventory, &plan_mode_state);
100
101        let pty_sessions = pty::PtySessionManager::new(workspace_root.clone(), pty_config);
102        let exec_sessions = crate::tools::exec_session::ExecSessionManager::new(
103            workspace_root.clone(),
104            pty_sessions.clone(),
105        );
106
107        let policy_gateway = match policy_manager {
108            Some(pm) => ToolPolicyGateway::with_policy_manager(pm),
109            None => ToolPolicyGateway::new(&workspace_root).await,
110        };
111
112        let optimization_config = vtcode_config::OptimizationConfig::default();
113        let metrics = Arc::new(crate::metrics::MetricsCollector::new());
114        let hot_cache_size =
115            std::num::NonZeroUsize::new(optimization_config.tool_registry.hot_cache_size)
116                .unwrap_or(std::num::NonZeroUsize::MIN);
117        let output_spooler = Arc::new(ToolOutputSpooler::with_config(
118            &workspace_root,
119            load_workspace_spooler_config(&workspace_root),
120        ));
121
122        let registry = Self {
123            inventory,
124            edited_file_monitor,
125            policy_gateway: Arc::new(tokio::sync::Mutex::new(policy_gateway)),
126            pty_sessions,
127            exec_sessions,
128            mcp_client: Arc::new(RwLock::new(None)),
129            mcp_tool_index: Arc::new(tokio::sync::RwLock::new(rustc_hash::FxHashMap::default())),
130            mcp_reverse_index: Arc::new(tokio::sync::RwLock::new(rustc_hash::FxHashMap::default())),
131            timeout_policy: Arc::new(RwLock::new(ToolTimeoutPolicy::default())),
132            execution_history: ToolExecutionHistory::new(100),
133            harness_context: HarnessContext::default(),
134            resiliency: Arc::new(Mutex::new(ResiliencyContext::default())),
135            mcp_circuit_breaker: Arc::new(circuit_breaker::McpCircuitBreaker::with_metrics(
136                metrics.clone(),
137            )),
138            shared_circuit_breaker: Arc::new(RwLock::new(None)),
139            initialized: Arc::new(std::sync::atomic::AtomicBool::new(false)),
140            tool_call_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
141            pty_poll_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
142            metrics,
143            shell_policy: Arc::new(RwLock::new(ShellPolicyChecker::new())),
144            runtime_sandbox_config: Arc::new(RwLock::new(
145                super::sandbox_facade::runtime_sandbox_config_default(),
146            )),
147            agent_type: Arc::new(RwLock::new("unknown".to_owned())),
148            cached_available_tools: Arc::new(RwLock::new(None)),
149            progress_callback: Arc::new(RwLock::new(None)),
150            active_pty_sessions: Arc::new(RwLock::new(None)),
151
152            memory_pool: Arc::new(MemoryPool::from_config(&optimization_config.memory_pool)),
153            hot_tool_cache: Arc::new(parking_lot::RwLock::new(lru::LruCache::new(hot_cache_size))),
154            optimization_config,
155
156            output_spooler,
157
158            plan_mode_state,
159            safety_gateway: Arc::new(SafetyGateway::default()),
160            cgp_runtime_mode: Arc::new(RwLock::new(None)),
161            tool_assembly: Arc::new(RwLock::new(ToolAssembly::empty())),
162            tool_catalog_state: Arc::new(super::tool_catalog_facade::SessionToolCatalogState::new()),
163            subagent_controller: Arc::new(RwLock::new(None)),
164            session_scheduler: Arc::new(tokio::sync::Mutex::new(
165                crate::scheduler::SessionScheduler::new(),
166            )),
167        };
168
169        registry.rebuild_tool_assembly().await;
170        registry.sync_policy_catalog().await;
171        registry.initialize_resiliency_trackers();
172        registry
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use tempfile::tempdir;
180
181    #[tokio::test]
182    async fn tool_registry_uses_workspace_dynamic_context_for_output_spooler() {
183        let temp = tempdir().unwrap();
184        std::fs::write(
185            temp.path().join("vtcode.toml"),
186            r#"[context.dynamic]
187enabled = true
188tool_output_threshold = 4096
189max_spooled_files = 7
190spool_max_age_secs = 12
191"#,
192        )
193        .unwrap();
194
195        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
196        let config = registry.output_spooler().config();
197
198        assert!(config.enabled);
199        assert_eq!(config.threshold_bytes, 4096);
200        assert_eq!(config.max_files, 7);
201        assert_eq!(config.max_age_secs, 12);
202    }
203
204    #[tokio::test]
205    async fn tool_registry_falls_back_to_default_spooler_config_on_invalid_workspace_config() {
206        let temp = tempdir().unwrap();
207        std::fs::write(
208            temp.path().join("vtcode.toml"),
209            r#"[context.dynamic]
210tool_output_threshold = "oops"
211"#,
212        )
213        .unwrap();
214
215        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
216        let config = registry.output_spooler().config();
217
218        assert!(config.enabled);
219        assert_eq!(
220            config.threshold_bytes,
221            DynamicContextConfig::default().tool_output_threshold
222        );
223        assert_eq!(config.max_files, SpoolerConfig::default().max_files);
224        assert_eq!(config.max_age_secs, SpoolerConfig::default().max_age_secs);
225    }
226}