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::{PersistentMemoryConfig, PtyConfig};
12use crate::core::memory_pool::MemoryPool;
13use crate::tool_policy::ToolPolicyManager;
14use crate::tools::handlers::PlanningWorkflowState;
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::circuit_breaker;
23use super::distributed::ToolConfigSnapshot;
24use super::execution_history::ToolExecutionHistory;
25use super::harness::HarnessContext;
26use super::inventory::ToolInventory;
27use super::pack::register_builtin_packs;
28use super::policy::ToolPolicyGateway;
29use super::pty;
30use super::resiliency::ResiliencyContext;
31use super::shell_policy::ShellPolicyChecker;
32use super::timeout::ToolTimeoutPolicy;
33
34fn spooler_config_from_dynamic_context(config: &DynamicContextConfig) -> SpoolerConfig {
35    SpoolerConfig {
36        enabled: config.enabled,
37        threshold_bytes: config.tool_output_threshold,
38        max_files: config.max_spooled_files,
39        max_age_secs: config.spool_max_age_secs,
40        include_file_reference: true,
41    }
42}
43
44#[derive(Debug, Default)]
45struct WorkspaceToolConfig {
46    tool: ToolConfigSnapshot,
47    spooler: SpoolerConfig,
48    persistent_memory: PersistentMemoryConfig,
49    persistent_memory_enabled: bool,
50    inspection_mode: vtcode_config::core::agent::InspectionMode,
51}
52
53/// Load the config needed by tool registration in one parse. Falls back to
54/// defaults on any load/parse error so a malformed config never blocks tool
55/// registration.
56fn load_workspace_tool_config(workspace_root: &Path) -> WorkspaceToolConfig {
57    match ConfigManager::load_from_workspace(workspace_root) {
58        Ok(manager) => workspace_tool_config_from_loaded(manager.config()),
59        Err(err) => {
60            tracing::warn!(
61                workspace = %workspace_root.display(),
62                error = %err,
63                "Failed to load workspace config for tool registry; using defaults"
64            );
65            WorkspaceToolConfig::default()
66        }
67    }
68}
69
70fn workspace_tool_config_from_loaded(config: &vtcode_config::loader::VTCodeConfig) -> WorkspaceToolConfig {
71    let persistent_memory_enabled = config.persistent_memory_enabled();
72    let mut persistent_memory = config.agent.persistent_memory.clone();
73    persistent_memory.enabled = persistent_memory_enabled;
74    WorkspaceToolConfig {
75        tool: ToolConfigSnapshot {
76            web_search: config.tools.web_search.clone(),
77            web_fetch: config.tools.web_fetch.clone(),
78        },
79        spooler: spooler_config_from_dynamic_context(&config.context.dynamic),
80        persistent_memory,
81        persistent_memory_enabled,
82        inspection_mode: config.agent.harness.inspection.mode,
83    }
84}
85
86impl ToolRegistry {
87    pub fn new(workspace_root: PathBuf) -> impl Future<Output = Self> {
88        Self::build(workspace_root, PtyConfig::default())
89    }
90
91    /// Build a registry reusing an already-loaded session config snapshot.
92    ///
93    /// Interactive startup already owns the merged `VTCodeConfig`; passing it
94    /// here avoids a second `ConfigManager::load_from_workspace` file read +
95    /// TOML parse on the first-paint path. Falls back to defaults for any
96    /// field the snapshot does not carry, matching `load_workspace_tool_config`
97    /// behavior on parse failure.
98    pub fn new_with_loaded_config(
99        workspace_root: PathBuf,
100        config: &vtcode_config::loader::VTCodeConfig,
101    ) -> impl Future<Output = Self> {
102        let snapshot = workspace_tool_config_from_loaded(config);
103        Self::build_with_snapshot(workspace_root, PtyConfig::default(), None, true, snapshot)
104    }
105
106    /// Build a first-paint registry: reuses the loaded session config snapshot
107    /// and skips workspace policy file I/O.
108    ///
109    /// Policy evaluation fails open to tool metadata defaults until hydration
110    /// attaches the manager via `ensure_workspace_policy_manager`, which runs
111    /// before any tool executes.
112    pub fn new_for_first_paint_with_loaded_config(
113        workspace_root: PathBuf,
114        config: &vtcode_config::loader::VTCodeConfig,
115    ) -> impl Future<Output = Self> {
116        let snapshot = workspace_tool_config_from_loaded(config);
117        Self::build_with_snapshot(workspace_root, PtyConfig::default(), None, false, snapshot)
118    }
119
120    pub fn new_with_config(workspace_root: PathBuf, pty_config: PtyConfig) -> impl Future<Output = Self> {
121        Self::build(workspace_root, pty_config)
122    }
123
124    /// Build a registry for offline metadata without creating a workspace policy file.
125    pub fn new_for_schema(workspace_root: PathBuf) -> impl Future<Output = Self> {
126        Self::build_with_policy(workspace_root, PtyConfig::default(), None, false)
127    }
128
129    pub fn new_with_custom_policy(
130        workspace_root: PathBuf,
131        policy_manager: ToolPolicyManager,
132    ) -> impl Future<Output = Self> {
133        Self::build_with_policy(workspace_root, PtyConfig::default(), Some(policy_manager), true)
134    }
135
136    pub fn new_with_custom_policy_and_config(
137        workspace_root: PathBuf,
138        pty_config: PtyConfig,
139        policy_manager: ToolPolicyManager,
140    ) -> impl Future<Output = Self> {
141        Self::build_with_policy(workspace_root, pty_config, Some(policy_manager), true)
142    }
143
144    async fn build(workspace_root: PathBuf, pty_config: PtyConfig) -> Self {
145        Self::build_with_policy(workspace_root, pty_config, None, true).await
146    }
147
148    async fn build_with_policy(
149        workspace_root: PathBuf,
150        pty_config: PtyConfig,
151        policy_manager: Option<ToolPolicyManager>,
152        initialize_policy: bool,
153    ) -> Self {
154        // Load the user-config snapshot *before* constructing the inventory so
155        // `WebFetchTool`/`WebSearchTool` are built with the user's allow/block
156        // lists, cooldown, and session cap rather than defaulting out. The
157        // snapshot is passed explicitly to active packs; parse the workspace
158        // config once because both the web tools and output spooler consume it.
159        let config_workspace = workspace_root.clone();
160        let workspace_config =
161            match tokio::task::spawn_blocking(move || load_workspace_tool_config(&config_workspace)).await {
162                Ok(config) => config,
163                Err(error) => {
164                    tracing::warn!(%error, "Workspace tool config task panicked; using defaults");
165                    WorkspaceToolConfig::default()
166                }
167            };
168        Self::build_with_snapshot(workspace_root, pty_config, policy_manager, initialize_policy, workspace_config).await
169    }
170
171    async fn build_with_snapshot(
172        workspace_root: PathBuf,
173        pty_config: PtyConfig,
174        policy_manager: Option<ToolPolicyManager>,
175        initialize_policy: bool,
176        workspace_config: WorkspaceToolConfig,
177    ) -> Self {
178        let WorkspaceToolConfig {
179            tool: tool_config,
180            spooler: spooler_config,
181            persistent_memory: persistent_memory_config,
182            persistent_memory_enabled,
183            inspection_mode,
184        } = workspace_config;
185        let edited_file_monitor = Arc::new(crate::tools::edited_file_monitor::EditedFileMonitor::new());
186        edited_file_monitor.source_inspection().set_mode(inspection_mode);
187        let inventory = ToolInventory::new(workspace_root.clone(), Arc::clone(&edited_file_monitor));
188        let planning_workflow_state = PlanningWorkflowState::new(workspace_root.clone());
189
190        register_builtin_packs(&inventory, &planning_workflow_state, &tool_config).await;
191
192        let pty_sessions = pty::PtySessionManager::new(workspace_root.clone(), pty_config);
193        let exec_sessions =
194            crate::tools::exec_session::ExecSessionManager::new(workspace_root.clone(), pty_sessions.clone());
195
196        let policy_gateway = match (policy_manager, initialize_policy) {
197            (Some(pm), _) => ToolPolicyGateway::with_policy_manager(pm),
198            (None, true) => ToolPolicyGateway::new(&workspace_root).await,
199            (None, false) => ToolPolicyGateway::without_persistence(),
200        };
201
202        let optimization_config = vtcode_config::OptimizationConfig::default();
203        let metrics = Arc::new(crate::metrics::MetricsCollector::new());
204        let hot_cache_size = std::num::NonZeroUsize::new(optimization_config.tool_registry.hot_cache_size)
205            .unwrap_or(std::num::NonZeroUsize::MIN);
206        let output_spooler = Arc::new(ToolOutputSpooler::with_config(&workspace_root, spooler_config));
207        // Age-expired leftovers from prior sessions are pruned immediately;
208        // periodic cleanup only fires every N spools inside one session, so
209        // short sessions never reached the threshold and stale spools piled up
210        // across runs. The count budget stays on the in-session periodic path:
211        // pinning is per-process, so a startup count prune in this session
212        // could delete young spools a concurrent session still reads.
213        if let Err(error) = output_spooler.cleanup_expired_files().await {
214            tracing::debug!(%error, "startup spool prune failed");
215        }
216
217        // Pre-allocate FxHashMaps with expected capacity for typical MCP tool sets.
218        // Most sessions register 10-50 MCP tools; start with room for 32 to
219        // avoid rehashing during initial discovery without wasting memory.
220        let mcp_tool_index = rustc_hash::FxHashMap::with_capacity_and_hasher(32, rustc_hash::FxBuildHasher);
221        let mcp_reverse_index = rustc_hash::FxHashMap::with_capacity_and_hasher(32, rustc_hash::FxBuildHasher);
222
223        let registry = Self {
224            inventory,
225            persistent_memory_config: Arc::new(persistent_memory_config),
226            persistent_memory_enabled,
227            edited_file_monitor,
228            policy_gateway: Arc::new(policy_gateway),
229            pty_sessions,
230            exec_sessions,
231            mcp_client: Arc::new(parking_lot::RwLock::new(None)),
232            mcp_tool_index: Arc::new(tokio::sync::RwLock::new(mcp_tool_index)),
233            mcp_reverse_index: Arc::new(tokio::sync::RwLock::new(mcp_reverse_index)),
234            timeout_policy: Arc::new(parking_lot::RwLock::new(ToolTimeoutPolicy::default())),
235            execution_history: ToolExecutionHistory::with_workspace_root(100, workspace_root.clone()),
236            patch_recovery_reads: Arc::new(Mutex::new(super::patch_recovery::PatchRecoveryReads::default())),
237            harness_context: HarnessContext::default(),
238            resiliency: Arc::new(Mutex::new(ResiliencyContext::default())),
239            mcp_circuit_breaker: Arc::new(circuit_breaker::McpCircuitBreaker::with_metrics(metrics.clone())),
240            shared_circuit_breaker: Arc::new(RwLock::new(None)),
241            initialized: Arc::new(std::sync::atomic::AtomicBool::new(false)),
242            tool_call_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
243            pty_poll_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
244            canonical_plans_dir: Arc::new(std::sync::OnceLock::new()),
245            metrics,
246            shell_policy: Arc::new(RwLock::new(ShellPolicyChecker::new())),
247            runtime_sandbox_config: Arc::new(RwLock::new(super::sandbox_facade::runtime_sandbox_config_default())),
248            agent_type: Arc::new(RwLock::new("unknown".to_owned())),
249            cached_available_tools: Arc::new(parking_lot::RwLock::new(None)),
250            active_tool_profile: Arc::new(RwLock::new(crate::config::ToolProfile::default())),
251            progress_callback: Arc::new(RwLock::new(None)),
252            active_pty_sessions: Arc::new(RwLock::new(None)),
253
254            memory_pool: Arc::new(MemoryPool::from_config(&optimization_config.memory_pool)),
255            hot_tool_cache: Arc::new(parking_lot::RwLock::new(lru::LruCache::new(hot_cache_size))),
256            optimization_config,
257            middleware: crate::tools::tool_middleware::MiddlewareChain::new(),
258
259            output_spooler,
260
261            planning_workflow_state,
262            planning_mode_policy_overrides: Arc::new(parking_lot::RwLock::new(None)),
263            safety_gateway: Arc::new(SafetyGateway::default()),
264            cgp_runtime_mode: Arc::new(RwLock::new(None)),
265            tool_assembly: Arc::new(RwLock::new(ToolAssembly::empty())),
266            tool_catalog_state: Arc::new(super::tool_catalog_facade::SessionToolCatalogState::new()),
267            subagent_controller: Arc::new(RwLock::new(None)),
268            matrix_worker: Arc::new(parking_lot::RwLock::new(None)),
269            matrix_coordinator: Arc::new(std::sync::atomic::AtomicBool::new(false)),
270            session_scheduler: Arc::new(tokio::sync::Mutex::new(crate::scheduler::SessionScheduler::new())),
271            session_model_tools: Arc::new(RwLock::new(None)),
272            self_ref: Arc::new(RwLock::new(None)),
273        };
274
275        registry.rebuild_tool_assembly().await;
276        registry.sync_policy_catalog().await;
277        registry.initialize_resiliency_trackers();
278        registry
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285    use tempfile::tempdir;
286
287    #[tokio::test]
288    async fn tool_registry_uses_workspace_dynamic_context_for_output_spooler() {
289        let temp = tempdir().unwrap();
290        std::fs::write(
291            temp.path().join("vtcode.toml"),
292            r#"[context.dynamic]
293enabled = true
294tool_output_threshold = 4096
295max_spooled_files = 7
296spool_max_age_secs = 12
297
298[workspace]
299use_root_config = true
300"#,
301        )
302        .unwrap();
303
304        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
305        let config = registry.output_spooler().config();
306
307        assert!(config.enabled);
308        assert_eq!(config.threshold_bytes, 4096);
309        assert_eq!(config.max_files, 7);
310        assert_eq!(config.max_age_secs, 12);
311    }
312
313    #[tokio::test]
314    async fn tool_registry_falls_back_to_default_spooler_config_on_invalid_workspace_config() {
315        let temp = tempdir().unwrap();
316        std::fs::write(
317            temp.path().join("vtcode.toml"),
318            r#"[context.dynamic]
319tool_output_threshold = "oops"
320"#,
321        )
322        .unwrap();
323
324        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
325        let config = registry.output_spooler().config();
326
327        assert!(config.enabled);
328        assert_eq!(config.threshold_bytes, DynamicContextConfig::default().tool_output_threshold);
329        assert_eq!(config.max_files, SpoolerConfig::default().max_files);
330        assert_eq!(config.max_age_secs, SpoolerConfig::default().max_age_secs);
331    }
332
333    #[tokio::test]
334    async fn tool_registry_captures_effective_memory_config() {
335        let temp = tempdir().expect("workspace");
336        std::fs::write(
337            temp.path().join("vtcode.toml"),
338            "[features]\nmemories = true\n\n[agent.persistent_memory]\nenabled = true\n\n[workspace]\nuse_root_config = true\n",
339        )
340        .expect("workspace config");
341
342        let snapshot = load_workspace_tool_config(temp.path());
343        assert!(snapshot.persistent_memory_enabled);
344        assert!(snapshot.persistent_memory.enabled);
345
346        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
347        assert!(registry.persistent_memory_enabled);
348        assert!(registry.persistent_memory_config.enabled);
349    }
350
351    #[tokio::test]
352    async fn schema_registry_does_not_create_workspace_policy() {
353        let temp = tempdir().expect("workspace");
354
355        let _registry = ToolRegistry::new_for_schema(temp.path().to_path_buf()).await;
356
357        assert!(!temp.path().join(".vtcode").exists());
358    }
359
360    #[tokio::test]
361    async fn loaded_config_registry_reuses_snapshot_without_workspace_parse() {
362        // Asymmetric: workspace file says 4096/7/12, loaded snapshot says
363        // 8192/3/99. The registry must reflect the snapshot, proving the
364        // first-paint path does not pay a second workspace TOML parse.
365        let temp = tempdir().expect("workspace");
366        std::fs::write(
367            temp.path().join("vtcode.toml"),
368            r#"[context.dynamic]
369enabled = true
370tool_output_threshold = 4096
371max_spooled_files = 7
372spool_max_age_secs = 12
373
374[workspace]
375use_root_config = true
376"#,
377        )
378        .expect("workspace config");
379
380        let mut loaded = vtcode_config::loader::VTCodeConfig::default();
381        loaded.context.dynamic.enabled = true;
382        loaded.context.dynamic.tool_output_threshold = 8192;
383        loaded.context.dynamic.max_spooled_files = 3;
384        loaded.context.dynamic.spool_max_age_secs = 99;
385
386        let registry = ToolRegistry::new_with_loaded_config(temp.path().to_path_buf(), &loaded).await;
387        let config = registry.output_spooler().config();
388
389        assert!(config.enabled);
390        assert_eq!(config.threshold_bytes, 8192);
391        assert_eq!(config.max_files, 3);
392        assert_eq!(config.max_age_secs, 99);
393    }
394
395    #[tokio::test]
396    async fn first_paint_registry_defers_policy_manager_until_ensure() {
397        let temp = tempdir().expect("workspace");
398        let loaded = vtcode_config::loader::VTCodeConfig::default();
399
400        let registry = ToolRegistry::new_for_first_paint_with_loaded_config(temp.path().to_path_buf(), &loaded).await;
401        assert!(!registry.has_policy_manager().await, "first-paint registry must skip policy file I/O");
402
403        registry.ensure_workspace_policy_manager(temp.path()).await;
404        assert!(
405            registry.has_policy_manager().await,
406            "hydration must attach the workspace policy manager before any tool runs"
407        );
408    }
409}