vtcode_mcp/
connection_pool.rs

1//! MCP connection pool for efficient provider management
2//!
3//! This module provides connection pooling and parallel initialization
4//! for MCP providers to eliminate sequential connection bottlenecks.
5
6use hashbrown::HashMap;
7use std::sync::Arc;
8use std::time::Duration;
9use tokio::sync::{RwLock, Semaphore};
10use tracing::{error, info, warn};
11
12use super::{McpElicitationHandler, McpProvider, McpSandboxContext, apply_ui_extension_gate};
13use rmcp::model::{ClientCapabilities, InitializeRequestParams};
14use vtcode_commons::MultiErrors;
15use vtcode_config::mcp::{McpAllowListConfig, McpProviderConfig};
16
17/// MCP connection pool for efficient provider management
18pub(crate) struct McpConnectionPool {
19    /// Active provider connections
20    providers: Arc<RwLock<HashMap<String, Arc<McpProvider>>>>,
21    /// Connection semaphore to limit concurrent connections
22    connection_semaphore: Arc<Semaphore>,
23    /// Maximum connections allowed concurrently
24    max_concurrent_connections: usize,
25    /// Connection timeout
26    connection_timeout: Duration,
27}
28
29impl McpConnectionPool {
30    pub(crate) fn new(max_concurrent_connections: usize, connection_timeout_seconds: u64) -> Self {
31        Self {
32            providers: Arc::new(RwLock::new(HashMap::new())),
33            connection_semaphore: Arc::new(Semaphore::new(max_concurrent_connections)),
34            max_concurrent_connections,
35            connection_timeout: Duration::from_secs(connection_timeout_seconds),
36        }
37    }
38
39    /// Initialize multiple providers in parallel with controlled concurrency
40    pub(crate) async fn initialize_providers_parallel(
41        &self,
42        provider_configs: Vec<McpProviderConfig>,
43        elicitation_handler: Option<Arc<dyn McpElicitationHandler>>,
44        sandbox_context: Option<McpSandboxContext>,
45        tool_timeout: Option<Duration>,
46        allowlist_snapshot: &McpAllowListConfig,
47        advertise_ui_apps: bool,
48    ) -> Result<Vec<(String, Arc<McpProvider>)>, McpPoolError> {
49        use futures::future::join_all;
50
51        // Create initialization tasks for each provider
52        let tasks: Vec<_> = provider_configs
53            .into_iter()
54            .map(|config| {
55                let elicitation_handler = elicitation_handler.clone();
56                let sandbox_context = sandbox_context.clone();
57                let allowlist_snapshot = allowlist_snapshot.clone();
58
59                async move {
60                    self.initialize_provider(
61                        config,
62                        elicitation_handler,
63                        sandbox_context,
64                        tool_timeout.unwrap_or(Duration::from_secs(30)),
65                        allowlist_snapshot,
66                        advertise_ui_apps,
67                    )
68                    .await
69                }
70            })
71            .collect();
72
73        // Execute all tasks in parallel
74        let results = join_all(tasks).await;
75
76        // Collect successful connections, accumulating errors
77        let mut successful_providers = Vec::new();
78        let mut errors: MultiErrors<McpPoolError> = MultiErrors::new();
79
80        for result in results {
81            if let Some(provider) = errors.collect_result(result) {
82                successful_providers.push(provider);
83            }
84        }
85
86        if !errors.is_empty() {
87            warn!("Some MCP provider connections failed: {errors}");
88        }
89
90        Ok(successful_providers)
91    }
92
93    /// Initialize a single provider with connection pooling
94    async fn initialize_provider(
95        &self,
96        config: McpProviderConfig,
97        elicitation_handler: Option<Arc<dyn McpElicitationHandler>>,
98        sandbox_context: Option<McpSandboxContext>,
99        tool_timeout: Duration,
100        allowlist_snapshot: McpAllowListConfig,
101        advertise_ui_apps: bool,
102    ) -> Result<(String, Arc<McpProvider>), McpPoolError> {
103        // Acquire semaphore permit to limit concurrent connections
104        let _permit = self
105            .connection_semaphore
106            .acquire()
107            .await
108            .map_err(|e| McpPoolError::SemaphoreError(e.to_string()))?;
109
110        info!("Initializing MCP provider '{}'", config.name);
111
112        // Connect to provider with timeout
113        let provider = tokio::time::timeout(
114            self.connection_timeout,
115            McpProvider::connect(config.clone(), elicitation_handler, sandbox_context, advertise_ui_apps),
116        )
117        .await
118        .map_err(|_e| McpPoolError::ConnectionTimeout(config.name.clone()))?
119        .map_err(|e| McpPoolError::ConnectionError(config.name.clone(), e.to_string()))?;
120
121        // Initialize the provider with proper parameters
122        let provider_startup_timeout = self.resolve_startup_timeout(&config);
123        let initialize_params = build_pool_initialize_params(advertise_ui_apps);
124        let tool_timeout_opt = Some(tool_timeout);
125
126        if let Err(err) = provider
127            .initialize(initialize_params, provider_startup_timeout, tool_timeout_opt, &allowlist_snapshot)
128            .await
129        {
130            return Err(McpPoolError::InitializationError(config.name.clone(), err.to_string()));
131        }
132
133        // Refresh tools
134        if let Err(err) = provider
135            .cached_tools_or_refresh_shared(&allowlist_snapshot, tool_timeout_opt)
136            .await
137        {
138            warn!("Failed to fetch tools for provider '{}': {}", config.name, err);
139        }
140
141        info!("Successfully initialized MCP provider '{}'", config.name);
142
143        Ok((config.name.clone(), Arc::new(provider)))
144    }
145
146    /// Get a provider by name
147    async fn get_provider(&self, name: &str) -> Option<Arc<McpProvider>> {
148        let providers = self.providers.read().await;
149        providers.get(name).cloned()
150    }
151
152    /// Get all active providers
153    async fn get_all_providers(&self) -> Vec<Arc<McpProvider>> {
154        let providers = self.providers.read().await;
155        providers.values().cloned().collect()
156    }
157
158    /// Remove a provider from the pool
159    pub async fn remove_provider(&self, name: &str) -> Option<Arc<McpProvider>> {
160        let mut providers = self.providers.write().await;
161        providers.remove(name)
162    }
163
164    /// Check if a provider exists in the pool
165    async fn has_provider(&self, name: &str) -> bool {
166        let providers = self.providers.read().await;
167        providers.contains_key(name)
168    }
169
170    /// Get connection pool statistics
171    async fn stats(&self) -> ConnectionPoolStats {
172        let providers = self.providers.read().await;
173        let semaphore = self.connection_semaphore.available_permits();
174
175        ConnectionPoolStats {
176            active_connections: providers.len(),
177            available_permits: semaphore,
178            max_connections: self.max_concurrent_connections,
179        }
180    }
181
182    /// Shutdown all providers gracefully
183    pub async fn shutdown_all(&self) {
184        let providers: Vec<_> = {
185            let mut providers = self.providers.write().await;
186            providers.drain().collect()
187        };
188
189        for (name, provider) in providers {
190            if let Err(err) = provider.shutdown().await {
191                error!("Failed to shutdown MCP provider '{}': {}", name, err);
192            }
193        }
194    }
195
196    /// Check the health of all active providers.
197    ///
198    /// Returns a map of provider name → `true` (healthy) / `false` (unhealthy).
199    pub async fn health_check(&self) -> HashMap<String, bool> {
200        let providers: Vec<_> = {
201            let providers = self.providers.read().await;
202            providers
203                .iter()
204                .map(|(name, provider)| (name.clone(), Arc::clone(provider)))
205                .collect()
206        };
207        let mut results = HashMap::with_capacity(providers.len());
208        for (name, provider) in providers {
209            let _previous = results.insert(name, provider.is_healthy().await);
210        }
211        results
212    }
213
214    /// Attempt to reconnect any unhealthy providers.
215    ///
216    /// Returns the names of providers that were successfully reconnected.
217    pub async fn reconnect_unhealthy(
218        &self,
219        startup_timeout: Option<Duration>,
220        tool_timeout: Option<Duration>,
221        allowlist: &McpAllowListConfig,
222    ) -> Vec<String> {
223        let providers: Vec<_> = {
224            let providers = self.providers.read().await;
225            providers
226                .iter()
227                .map(|(name, provider)| (name.clone(), Arc::clone(provider)))
228                .collect()
229        };
230        let mut reconnected = Vec::new();
231        for (name, provider) in providers {
232            if !provider.is_healthy().await {
233                info!("Provider '{}' is unhealthy, attempting reconnect", name);
234                match provider.reconnect(startup_timeout, tool_timeout, allowlist).await {
235                    Ok(()) => {
236                        info!("Successfully reconnected MCP provider '{}'", name);
237                        reconnected.push(name);
238                    }
239                    Err(err) => {
240                        error!("Failed to reconnect MCP provider '{}': {}", name, err);
241                    }
242                }
243            }
244        }
245        reconnected
246    }
247
248    /// Resolve startup timeout based on provider configuration
249    fn resolve_startup_timeout(&self, config: &McpProviderConfig) -> Option<Duration> {
250        config.startup_timeout_ms.map(Duration::from_millis)
251    }
252}
253
254/// Connection pool statistics
255#[derive(Debug, Clone)]
256pub struct ConnectionPoolStats {
257    active_connections: usize,
258    available_permits: usize,
259    max_connections: usize,
260}
261
262/// Enhanced MCP manager with connection pooling
263pub(crate) struct PooledMcpManager {
264    /// Connection pool for providers
265    pool: Arc<McpConnectionPool>,
266    /// Tool discovery cache
267    tool_cache: Arc<super::tool_discovery_cache::ToolDiscoveryCache>,
268}
269
270impl PooledMcpManager {
271    fn new(max_concurrent_connections: usize, connection_timeout_seconds: u64, tool_cache_capacity: usize) -> Self {
272        Self {
273            pool: Arc::new(McpConnectionPool::new(max_concurrent_connections, connection_timeout_seconds)),
274            tool_cache: Arc::new(super::tool_discovery_cache::ToolDiscoveryCache::new(tool_cache_capacity)),
275        }
276    }
277
278    /// Initialize providers with pooling and caching
279    pub async fn initialize_providers(
280        &self,
281        provider_configs: Vec<McpProviderConfig>,
282        elicitation_handler: Option<Arc<dyn McpElicitationHandler>>,
283        sandbox_context: Option<McpSandboxContext>,
284        tool_timeout: Option<Duration>,
285        allowlist_snapshot: &McpAllowListConfig,
286        advertise_ui_apps: bool,
287    ) -> Result<Vec<(String, Arc<McpProvider>)>, McpPoolError> {
288        // Initialize providers in parallel
289        let providers = self
290            .pool
291            .initialize_providers_parallel(
292                provider_configs,
293                elicitation_handler,
294                sandbox_context,
295                tool_timeout,
296                allowlist_snapshot,
297                advertise_ui_apps,
298            )
299            .await?;
300
301        // Add providers to the pool
302        let mut pool_providers = self.pool.providers.write().await;
303        for (name, provider) in &providers {
304            drop(pool_providers.insert(name.clone(), provider.clone()));
305        }
306
307        Ok(providers)
308    }
309
310    /// Execute a tool on a specific provider
311    pub async fn execute_tool(
312        &self,
313        provider_name: &str,
314        tool_name: &str,
315        arguments: serde_json::Value,
316        allowlist: &McpAllowListConfig,
317        tool_timeout: Option<Duration>,
318    ) -> Result<serde_json::Value, McpPoolError> {
319        let provider = self
320            .pool
321            .get_provider(provider_name)
322            .await
323            .ok_or_else(|| McpPoolError::ProviderNotFound(provider_name.to_string()))?;
324
325        // Convert arguments to proper format
326        let args_ref = &arguments;
327
328        // Execute the tool with correct signature
329        let result = provider
330            .call_tool(tool_name, args_ref, tool_timeout, allowlist)
331            .await
332            .map_err(|e| McpPoolError::ToolExecutionError(provider_name.to_string(), e.to_string()))?;
333
334        // Convert result to JSON value
335        Ok(serde_json::to_value(&result).unwrap_or(serde_json::Value::Null))
336    }
337
338    /// Check if a tool is read-only (safe to cache)
339    fn is_read_only_tool(&self, tool_name: &str) -> bool {
340        // This is a simple heuristic - in practice, you might want to
341        // check tool metadata or maintain a list of read-only tools
342        matches!(
343            tool_name,
344            "read_file"
345                | "list_directory"
346                | "search_files"
347                | "get_file_info"
348                | "read_environment"
349                | "get_system_info"
350                | "search_code"
351                | "analyze_code"
352        )
353    }
354
355    /// Get pool statistics
356    async fn stats(&self) -> PooledMcpStats {
357        let pool_stats = self.pool.stats().await;
358        let tool_cache_stats = self.tool_cache.stats();
359
360        PooledMcpStats {
361            connection_pool: pool_stats,
362            tool_cache: tool_cache_stats,
363        }
364    }
365
366    /// Shutdown all providers gracefully
367    pub async fn shutdown(&self) {
368        self.pool.shutdown_all().await;
369    }
370}
371
372/// Pooled MCP manager statistics
373#[derive(Debug, Clone)]
374pub struct PooledMcpStats {
375    connection_pool: ConnectionPoolStats,
376    tool_cache: super::tool_discovery_cache::ToolCacheStats,
377}
378
379/// Build initialize params for an MCP provider
380fn build_pool_initialize_params(advertise_ui_apps: bool) -> InitializeRequestParams {
381    let mut capabilities = ClientCapabilities::default();
382    apply_ui_extension_gate(&mut capabilities, advertise_ui_apps);
383    InitializeRequestParams::new(capabilities, super::utils::build_client_implementation())
384        .with_protocol_version(super::rmcp_client::latest_protocol_version())
385}
386
387/// MCP connection pool errors
388#[derive(Debug, thiserror::Error)]
389pub enum McpPoolError {
390    #[error("Connection timeout for provider '{0}'")]
391    ConnectionTimeout(String),
392
393    #[error("Connection error for provider '{0}': {1}")]
394    ConnectionError(String, String),
395
396    #[error("Initialization timeout for provider '{0}'")]
397    InitializationTimeout(String),
398
399    #[error("Initialization error for provider '{0}': {1}")]
400    InitializationError(String, String),
401
402    #[error("Provider not found: {0}")]
403    ProviderNotFound(String),
404
405    #[error("Tool execution error for provider '{0}': {1}")]
406    ToolExecutionError(String, String),
407
408    #[error("Semaphore error: {0}")]
409    SemaphoreError(String),
410}
411
412#[cfg(test)]
413pub(crate) mod tests {
414    use super::*;
415
416    #[tokio::test]
417    async fn test_connection_pool_creation() {
418        let pool = McpConnectionPool::new(5, 30);
419        let stats = pool.stats().await;
420
421        assert_eq!(stats.active_connections, 0);
422        assert_eq!(stats.max_connections, 5);
423        assert_eq!(stats.available_permits, 5);
424    }
425
426    #[tokio::test]
427    async fn test_connection_pool_semaphore_limits() {
428        let pool = McpConnectionPool::new(3, 30);
429
430        // Acquire 3 permits
431        let permit1 = pool.connection_semaphore.acquire().await.unwrap();
432        let _permit2 = pool.connection_semaphore.acquire().await.unwrap();
433        let _permit3 = pool.connection_semaphore.acquire().await.unwrap();
434
435        let stats = pool.stats().await;
436        assert_eq!(stats.available_permits, 0);
437
438        // Try to acquire another (would block if not in test)
439        drop(permit1);
440        let _permit4 = pool.connection_semaphore.acquire().await.unwrap();
441
442        let stats = pool.stats().await;
443        assert_eq!(stats.available_permits, 0);
444    }
445
446    #[tokio::test]
447    async fn test_pooled_manager_creation() {
448        let manager = PooledMcpManager::new(10, 30, 100);
449        let stats = manager.stats().await;
450
451        assert_eq!(stats.connection_pool.max_connections, 10);
452        assert_eq!(stats.connection_pool.active_connections, 0);
453    }
454
455    #[tokio::test]
456    async fn test_read_only_tool_detection() {
457        let manager = PooledMcpManager::new(5, 30, 50);
458
459        assert!(manager.is_read_only_tool("read_file"));
460        assert!(manager.is_read_only_tool("search_files"));
461        assert!(manager.is_read_only_tool("get_system_info"));
462        assert!(manager.is_read_only_tool("get_file_info"));
463
464        assert!(!manager.is_read_only_tool("write_file"));
465        assert!(!manager.is_read_only_tool("edit_file"));
466        assert!(!manager.is_read_only_tool("execute_command"));
467        assert!(!manager.is_read_only_tool("delete_file"));
468    }
469
470    #[test]
471    fn test_connection_pool_error_display() {
472        let error = McpPoolError::ConnectionTimeout("test_provider".to_string());
473        assert!(error.to_string().contains("test_provider"));
474
475        let error = McpPoolError::InitializationError("auth".to_string(), "invalid credentials".to_string());
476        assert!(error.to_string().contains("auth"));
477        assert!(error.to_string().contains("invalid credentials"));
478    }
479
480    #[tokio::test]
481    async fn test_pool_provider_not_found() {
482        let pool = McpConnectionPool::new(5, 30);
483        let provider = pool.get_provider("nonexistent").await;
484        assert!(provider.is_none());
485    }
486
487    #[tokio::test]
488    async fn test_pool_has_provider() {
489        let pool = McpConnectionPool::new(5, 30);
490        assert!(!pool.has_provider("test").await);
491    }
492
493    #[tokio::test]
494    async fn test_pool_get_all_providers_empty() {
495        let pool = McpConnectionPool::new(5, 30);
496        let providers = pool.get_all_providers().await;
497        assert_eq!(providers.len(), 0);
498    }
499
500    #[tokio::test]
501    async fn test_pool_stats() {
502        let pool = McpConnectionPool::new(7, 60);
503        let stats = pool.stats().await;
504
505        assert_eq!(stats.max_connections, 7);
506        assert_eq!(stats.available_permits, 7);
507        assert_eq!(stats.active_connections, 0);
508    }
509
510    #[test]
511    fn pool_initialize_params_gate_ui_extension_on_flag() {
512        let advertised = build_pool_initialize_params(true);
513        assert!(
514            advertised
515                .capabilities
516                .extensions
517                .as_ref()
518                .is_some_and(|extensions| extensions.contains_key("io.modelcontextprotocol/ui")),
519        );
520        let silent = build_pool_initialize_params(false);
521        assert_eq!(silent.capabilities.extensions, None);
522    }
523}
vtcode_mcp/connection_pool.rs

vtcode_mcp/
connection_pool.rs
