vtcode_core/tools/registry/
builder.rs1use 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}