diff --git a/src/crates/assembly/core/src/service/mcp/server/manager/external_lifecycle.rs b/src/crates/assembly/core/src/service/mcp/server/manager/external_lifecycle.rs new file mode 100644 index 0000000000..21fdd978db --- /dev/null +++ b/src/crates/assembly/core/src/service/mcp/server/manager/external_lifecycle.rs @@ -0,0 +1,356 @@ +use super::*; + +impl MCPServerManager { + /// Adds a runtime-only MCP server without saving it to user or project config. + pub async fn add_ephemeral_server(&self, config: MCPServerConfig) -> BitFunResult<()> { + config.validate()?; + + let server_id = config.id.clone(); + if self.runtime.contains(&server_id).await { + return Err(BitFunError::Configuration(format!( + "MCP server already exists: {}", + server_id + ))); + } + + self.runtime.insert_runtime_config(config.clone()).await?; + self.runtime.register(&config).await?; + + if config.enabled && config.auto_start { + if let Err(error) = self.start_server(&server_id).await { + let _ = self.remove_ephemeral_server(&server_id).await; + return Err(error); + } + } + + Ok(()) + } + + async fn external_start_token_matches(&self, server_id: &str, expected: &Arc<()>) -> bool { + let start_tokens = self.ephemeral_start_tokens.read().await; + external_start_token_is_current(start_tokens.get(server_id), expected) + } + + async fn remove_ephemeral_server_for_start(&self, server_id: &str, expected: &Arc<()>) -> bool { + let _lifecycle_guard = self.ephemeral_lifecycle.lock().await; + if !self.external_start_token_matches(server_id, expected).await { + return false; + } + if let Err(error) = self.remove_ephemeral_server(server_id).await { + warn!( + "Could not clean up failed external MCP startup: id={} error={}", + server_id, error + ); + } + true + } + + /// Installs a product-approved runtime-only server. A matching retirement + /// can be cancelled without restarting the process, which keeps rapid + /// disable/enable actions from interrupting unrelated session work. + pub async fn install_external_ephemeral_server( + &self, + config: MCPServerConfig, + workspace_key: String, + ) -> BitFunResult<()> { + config.validate()?; + let _lifecycle_guard = self.ephemeral_lifecycle.lock().await; + let server_id = config.id.clone(); + let start_token = Arc::new(()); + self.ephemeral_start_tokens + .write() + .await + .insert(server_id.clone(), Arc::clone(&start_token)); + self.ephemeral_workspace_scopes + .write() + .await + .insert(server_id.clone(), workspace_key); + self.ephemeral_ready_servers + .write() + .await + .remove(&server_id); + let cancelled_retirement = self + .ephemeral_retirements + .write() + .await + .remove(&server_id) + .map(|cancelled| { + cancelled.store(true, Ordering::Release); + true + }) + .unwrap_or(false); + + if cancelled_retirement && self.runtime.contains(&server_id).await { + if let Err(error) = self.runtime.insert_runtime_config(config.clone()).await { + let _ = self.remove_ephemeral_server(&server_id).await; + return Err(error.into()); + } + let connection = self.runtime.process_connection(&server_id).await; + if let Some(connection) = connection { + self.runtime + .add_connection(server_id.clone(), connection.clone()) + .await; + if let Err(error) = self + .refresh_mcp_tools(&server_id, &config.name, connection.clone()) + .await + { + let _ = self.remove_ephemeral_server(&server_id).await; + return Err(error); + } + self.start_connection_event_listener(&server_id, &config.name, connection.clone()) + .await; + self.warm_catalog_caches(&server_id, connection).await; + self.ephemeral_ready_servers + .write() + .await + .insert(server_id.clone()); + } else { + let _ = self.remove_ephemeral_server(&server_id).await; + return Err(BitFunError::MCPError( + "External MCP server did not retain its connection".to_string(), + )); + } + return Ok(()); + } + if self.runtime.contains(&server_id).await { + self.ephemeral_workspace_scopes + .write() + .await + .remove(&server_id); + self.ephemeral_start_tokens.write().await.remove(&server_id); + return Err(BitFunError::Configuration(format!( + "MCP server already exists: {}", + server_id + ))); + } + + if let Err(error) = self.runtime.insert_runtime_config(config.clone()).await { + self.ephemeral_workspace_scopes + .write() + .await + .remove(&server_id); + self.ephemeral_start_tokens.write().await.remove(&server_id); + return Err(error.into()); + } + if let Err(error) = self.runtime.register(&config).await { + self.runtime.remove_runtime_config(&server_id).await; + self.ephemeral_workspace_scopes + .write() + .await + .remove(&server_id); + self.ephemeral_start_tokens.write().await.remove(&server_id); + return Err(error.into()); + } + if config.enabled && config.auto_start { + // External source refresh and product-surface reads must not wait + // for a third-party process or network handshake. Registration is + // synchronous so status reads immediately see Loading; startup is + // bounded in the background and cleans up only this runtime item. + const EXTERNAL_START_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + let manager = self.clone(); + tokio::spawn(async move { + let startup = tokio::time::timeout( + EXTERNAL_START_TIMEOUT, + manager.start_server_with_external_token( + &server_id, + Some(Arc::clone(&start_token)), + ), + ) + .await; + match startup { + Ok(Ok(())) => { + if manager + .external_start_token_matches(&server_id, &start_token) + .await + { + crate::external_sources::notify_external_tool_registry_changed(); + } + } + Ok(Err(error)) => { + warn!( + "External ephemeral MCP server failed to start: id={} error={}", + server_id, error + ); + if manager + .remove_ephemeral_server_for_start(&server_id, &start_token) + .await + { + crate::external_sources::notify_external_tool_registry_changed(); + } + } + Err(_) => { + warn!( + "External ephemeral MCP server startup timed out: id={}", + server_id + ); + if manager + .remove_ephemeral_server_for_start(&server_id, &start_token) + .await + { + crate::external_sources::notify_external_tool_registry_changed(); + } + } + } + }); + } + Ok(()) + } + + /// Withdraws new tool/resource access immediately, then lets already-held + /// connection users finish before the process is reclaimed. The grace is + /// bounded so a deleted or malicious server cannot remain indefinitely. + pub async fn retire_external_ephemeral_server(&self, server_id: &str) -> BitFunResult<()> { + const RETIREMENT_GRACE: std::time::Duration = std::time::Duration::from_secs(30); + const RETIREMENT_RECLAIM_ATTEMPTS: usize = 3; + const RETIREMENT_RETRY_DELAY: std::time::Duration = std::time::Duration::from_millis(250); + let _lifecycle_guard = self.ephemeral_lifecycle.lock().await; + self.ephemeral_start_tokens.write().await.remove(server_id); + if !self.runtime.contains(server_id).await { + self.runtime.remove_runtime_config(server_id).await; + self.ephemeral_ready_servers.write().await.remove(server_id); + self.ephemeral_workspace_scopes + .write() + .await + .remove(server_id); + return Ok(()); + } + + if let Some(previous) = self + .ephemeral_retirements + .write() + .await + .insert(server_id.to_string(), Arc::new(AtomicBool::new(false))) + { + previous.store(true, Ordering::Release); + } + let cancelled = self + .ephemeral_retirements + .read() + .await + .get(server_id) + .cloned() + .expect("retirement marker was just inserted"); + let connection = self.runtime.get_connection(server_id).await; + + self.ephemeral_ready_servers.write().await.remove(server_id); + Self::unregister_mcp_tools(server_id).await; + self.stop_connection_event_listener(server_id).await; + self.runtime.remove_connection(server_id).await; + self.runtime.remove_catalog(server_id).await; + self.runtime.remove_runtime_config(server_id).await; + self.clear_reconnect_state(server_id).await; + + let manager = self.clone(); + let server_id = server_id.to_string(); + tokio::spawn(async move { + let started = std::time::Instant::now(); + loop { + if cancelled.load(Ordering::Acquire) { + return; + } + let references = connection.as_ref().map_or(0, Arc::strong_count); + if should_finish_ephemeral_retirement( + references, + started.elapsed(), + RETIREMENT_GRACE, + ) { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + + for attempt in 1..=RETIREMENT_RECLAIM_ATTEMPTS { + let lifecycle_guard = manager.ephemeral_lifecycle.lock().await; + if cancelled.load(Ordering::Acquire) { + return; + } + let should_remove = manager + .ephemeral_retirements + .read() + .await + .get(&server_id) + .is_some_and(|current| Arc::ptr_eq(current, &cancelled)); + if !should_remove { + return; + } + match manager.runtime.unregister(&server_id).await { + Ok(()) => { + manager + .ephemeral_retirements + .write() + .await + .remove(&server_id); + Self::unregister_mcp_tools(&server_id).await; + manager.stop_connection_event_listener(&server_id).await; + manager.runtime.remove_connection(&server_id).await; + manager.runtime.remove_catalog(&server_id).await; + manager + .ephemeral_ready_servers + .write() + .await + .remove(&server_id); + manager + .ephemeral_workspace_scopes + .write() + .await + .remove(&server_id); + return; + } + Err(error) if attempt < RETIREMENT_RECLAIM_ATTEMPTS => { + warn!( + "Could not reclaim retired ephemeral MCP server; retrying: id={} attempt={} error={}", + server_id, attempt, error + ); + } + Err(error) => { + warn!( + "Could not reclaim retired ephemeral MCP server; retaining ownership for a later retry: id={} attempts={} error={}", + server_id, RETIREMENT_RECLAIM_ATTEMPTS, error + ); + return; + } + } + drop(lifecycle_guard); + tokio::time::sleep(RETIREMENT_RETRY_DELAY).await; + } + }); + Ok(()) + } + + /// Removes a runtime-only MCP server and its registered tools without touching persisted config. + pub async fn remove_ephemeral_server(&self, server_id: &str) -> BitFunResult<()> { + info!("Removing ephemeral MCP server: id={}", server_id); + + if !self.runtime.contains(server_id).await { + self.runtime.remove_runtime_config(server_id).await; + self.clear_reconnect_state(server_id).await; + self.runtime.remove_catalog(server_id).await; + Self::unregister_mcp_tools(server_id).await; + return Ok(()); + } + + let stop_result = self.stop_server(server_id).await; + self.stop_connection_event_listener(server_id).await; + self.clear_reconnect_state(server_id).await; + self.runtime.remove_catalog(server_id).await; + self.ephemeral_ready_servers.write().await.remove(server_id); + self.ephemeral_start_tokens.write().await.remove(server_id); + self.ephemeral_workspace_scopes + .write() + .await + .remove(server_id); + + if let Err(error) = stop_result { + warn!( + "Failed to stop ephemeral MCP server; retaining runtime ownership for retry: id={} error={}", + server_id, error + ); + return Err(error); + } + + self.runtime.unregister(server_id).await?; + self.runtime.remove_runtime_config(server_id).await; + info!("Unregistered ephemeral MCP server: id={}", server_id); + Ok(()) + } +} diff --git a/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs b/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs index e5ae3162a3..965a47608b 100644 --- a/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs +++ b/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs @@ -1,6 +1,7 @@ use super::*; use bitfun_services_integrations::mcp::server::{ - mcp_server_is_running, mcp_should_start_after_config_update, resolve_mcp_local_command, + mcp_server_is_running, mcp_should_start_after_config_update, MCPProcessStartContext, + MCPProcessStartOutcome, }; impl MCPServerManager { @@ -17,12 +18,6 @@ impl MCPServerManager { }) } - fn resolve_local_command(command: &str) -> BitFunResult<(String, &'static str)> { - let runtime_root = crate::infrastructure::get_path_manager_arc().managed_runtimes_dir(); - let resolved = resolve_mcp_local_command(command, runtime_root)?; - Ok((resolved.command, resolved.source_label)) - } - /// Initializes all servers. pub async fn initialize_all(&self) -> BitFunResult<()> { info!("Initializing all MCP servers"); @@ -182,7 +177,7 @@ impl MCPServerManager { self.start_server_with_external_token(server_id, None).await } - async fn start_server_with_external_token( + pub(super) async fn start_server_with_external_token( &self, server_id: &str, expected_external_start_token: Option>, @@ -206,88 +201,36 @@ impl MCPServerManager { } self.runtime.ensure_registered(&config).await?; - - let process = self.runtime.get_process(server_id).await.ok_or_else(|| { - error!("MCP server not registered: id={}", server_id); - BitFunError::NotFound(format!("MCP server not registered: {}", server_id)) - })?; - - let mut proc = process.write().await; - - let status = proc.status().await; - if mcp_server_is_running(status) { + if mcp_server_is_running(self.runtime.process_status(server_id).await?) { warn!("MCP server already running: id={}", server_id); return Ok(()); } - match config.server_type { - super::super::MCPServerType::Local => { - let command = config.command.as_ref().ok_or_else(|| { - error!("Missing command for local MCP server: id={}", server_id); - BitFunError::Configuration("Missing command for local MCP server".to_string()) - })?; - - let (resolved_command, source_label) = Self::resolve_local_command(command)?; - - info!( - "Starting local MCP server: command={} source={} id={}", - resolved_command, source_label, server_id - ); - - proc.start_with_environment_policy( - &resolved_command, - &config.args, - &config.env, - config.working_directory.as_deref().map(Path::new), - config.inherits_parent_environment(), - ) - .await - .map_err(|e| { - error!( - "Failed to start local MCP server process: id={} command={} source={} error={}", - server_id, resolved_command, source_label, e - ); - e - })?; - } - super::super::MCPServerType::Remote => { - let transport = config.resolved_transport(); - if transport != crate::service::mcp::server::MCPServerTransport::StreamableHttp { - error!( - "Remote MCP transport not supported yet: id={} transport={}", - server_id, - transport.as_str() - ); - return Err(BitFunError::NotImplemented(format!( - "Remote MCP transport '{}' is not yet supported", - transport.as_str() - ))); - } - - config.url.as_ref().ok_or_else(|| { - error!("Missing URL for remote MCP server: id={}", server_id); - BitFunError::Configuration("Missing URL for remote MCP server".to_string()) - })?; - - info!( - "Connecting to remote MCP server: transport={} id={}", - transport.as_str(), - server_id + let start_context = match config.server_type { + super::super::MCPServerType::Local => MCPProcessStartContext::Local { + managed_runtimes_dir: crate::infrastructure::get_path_manager_arc() + .managed_runtimes_dir(), + }, + super::super::MCPServerType::Remote => MCPProcessStartContext::Remote { + data_dir: crate::infrastructure::try_get_path_manager_arc()?.user_data_dir(), + }, + }; + let connection = match self + .runtime + .start_process(&config, start_context) + .await + .inspect_err(|error| { + error!( + "Failed to start MCP server runtime: id={} error={}", + server_id, error ); - - let data_dir = crate::infrastructure::try_get_path_manager_arc()?.user_data_dir(); - proc.start_remote(data_dir, &config).await.map_err(|e| { - error!( - "Failed to connect to remote MCP server: id={} error={}", - server_id, e - ); - e - })?; + })? { + MCPProcessStartOutcome::AlreadyRunning => { + warn!("MCP server already running: id={}", server_id); + return Ok(()); } - } - - let connection = proc.connection(); - drop(proc); + MCPProcessStartOutcome::Started { connection } => connection, + }; let external_workspace_scope = self .ephemeral_workspace_scopes .read() @@ -321,52 +264,40 @@ impl MCPServerManager { } } - if let Some(connection) = connection { - self.runtime - .add_connection(server_id.to_string(), connection.clone()) - .await; + self.runtime + .add_connection(server_id.to_string(), connection.clone()) + .await; - match self - .register_mcp_tools(server_id, &config.name, connection.clone()) - .await - { - Ok(count) => { - info!( - "Registered {} MCP tools: server_name={} server_id={}", - count, config.name, server_id - ); - } - Err(e) => { - warn!( - "Failed to register MCP tools: server_name={} server_id={} error={}", - config.name, server_id, e - ); - if external_workspace_scope.is_some() { - self.runtime.remove_connection(server_id).await; - return Err(e); - } + match self + .register_mcp_tools(server_id, &config.name, connection.clone()) + .await + { + Ok(count) => { + info!( + "Registered {} MCP tools: server_name={} server_id={}", + count, config.name, server_id + ); + } + Err(e) => { + warn!( + "Failed to register MCP tools: server_name={} server_id={} error={}", + config.name, server_id, e + ); + if external_workspace_scope.is_some() { + self.runtime.remove_connection(server_id).await; + return Err(e); } } + } - self.start_connection_event_listener(server_id, &config.name, connection.clone()) - .await; - self.warm_catalog_caches(server_id, connection).await; - if external_workspace_scope.is_some() { - self.ephemeral_ready_servers - .write() - .await - .insert(server_id.to_string()); - } - } else { - warn!( - "Connection not available, server may not have started correctly: id={}", - server_id - ); - if external_workspace_scope.is_some() { - return Err(BitFunError::MCPError( - "External MCP server did not establish a connection".to_string(), - )); - } + self.start_connection_event_listener(server_id, &config.name, connection.clone()) + .await; + self.warm_catalog_caches(server_id, connection).await; + if external_workspace_scope.is_some() { + self.ephemeral_ready_servers + .write() + .await + .insert(server_id.to_string()); } info!("MCP server started successfully: id={}", server_id); @@ -380,16 +311,7 @@ impl MCPServerManager { self.stop_connection_event_listener(server_id).await; - let process = - self.runtime.get_process(server_id).await.ok_or_else(|| { - BitFunError::NotFound(format!("MCP server not found: {}", server_id)) - })?; - - let mut proc = process.write().await; - let stop_result = proc.stop().await; - - self.runtime.remove_connection(server_id).await; - self.runtime.remove_catalog(server_id).await; + let stop_result = self.runtime.stop_process(server_id).await; Self::unregister_mcp_tools(server_id).await; @@ -399,39 +321,10 @@ impl MCPServerManager { /// Restarts a server. pub async fn restart_server(&self, server_id: &str) -> BitFunResult<()> { info!("Restarting MCP server: id={}", server_id); - - let config = self.runtime_server_config(server_id).await?; - - match config.server_type { - super::super::MCPServerType::Local => { - self.ensure_registered(server_id).await?; - - let process = self.runtime.get_process(server_id).await.ok_or_else(|| { - BitFunError::NotFound(format!("MCP server not found: {}", server_id)) - })?; - let mut proc = process.write().await; - - let command = config - .command - .as_ref() - .ok_or_else(|| BitFunError::Configuration("Missing command".to_string()))?; - proc.restart_with_environment_policy( - command, - &config.args, - &config.env, - config.working_directory.as_deref().map(Path::new), - config.inherits_parent_environment(), - ) - .await?; - } - super::super::MCPServerType::Remote => { - self.ensure_registered(server_id).await?; - let _ = self.stop_server(server_id).await; - self.start_server(server_id).await?; - } - } - - Ok(()) + self.runtime_server_config(server_id).await?; + self.ensure_registered(server_id).await?; + self.stop_server(server_id).await?; + self.start_server(server_id).await } /// Returns server status. @@ -440,13 +333,10 @@ impl MCPServerManager { let _ = self.ensure_registered(server_id).await; } - let process = - self.runtime.get_process(server_id).await.ok_or_else(|| { - BitFunError::NotFound(format!("MCP server not found: {}", server_id)) - })?; - - let proc = process.read().await; - Ok(proc.status().await) + self.runtime + .process_status(server_id) + .await + .map_err(Into::into) } /// Returns the current status detail/message for one server. @@ -455,13 +345,10 @@ impl MCPServerManager { let _ = self.ensure_registered(server_id).await; } - let process = - self.runtime.get_process(server_id).await.ok_or_else(|| { - BitFunError::NotFound(format!("MCP server not found: {}", server_id)) - })?; - - let proc = process.read().await; - Ok(proc.status_message().await) + self.runtime + .process_status_message(server_id) + .await + .map_err(Into::into) } /// Returns statuses of all servers. @@ -508,363 +395,6 @@ impl MCPServerManager { Ok(()) } - /// Adds a runtime-only MCP server without saving it to user or project config. - pub async fn add_ephemeral_server(&self, config: MCPServerConfig) -> BitFunResult<()> { - config.validate()?; - - let server_id = config.id.clone(); - if self.runtime.contains(&server_id).await { - return Err(BitFunError::Configuration(format!( - "MCP server already exists: {}", - server_id - ))); - } - - self.runtime.insert_runtime_config(config.clone()).await?; - self.runtime.register(&config).await?; - - if config.enabled && config.auto_start { - if let Err(error) = self.start_server(&server_id).await { - let _ = self.remove_ephemeral_server(&server_id).await; - return Err(error); - } - } - - Ok(()) - } - - async fn external_start_token_matches(&self, server_id: &str, expected: &Arc<()>) -> bool { - let start_tokens = self.ephemeral_start_tokens.read().await; - external_start_token_is_current(start_tokens.get(server_id), expected) - } - - async fn remove_ephemeral_server_for_start(&self, server_id: &str, expected: &Arc<()>) -> bool { - let _lifecycle_guard = self.ephemeral_lifecycle.lock().await; - if !self.external_start_token_matches(server_id, expected).await { - return false; - } - if let Err(error) = self.remove_ephemeral_server(server_id).await { - warn!( - "Could not clean up failed external MCP startup: id={} error={}", - server_id, error - ); - } - true - } - - /// Installs a product-approved runtime-only server. A matching retirement - /// can be cancelled without restarting the process, which keeps rapid - /// disable/enable actions from interrupting unrelated session work. - pub async fn install_external_ephemeral_server( - &self, - config: MCPServerConfig, - workspace_key: String, - ) -> BitFunResult<()> { - config.validate()?; - let _lifecycle_guard = self.ephemeral_lifecycle.lock().await; - let server_id = config.id.clone(); - let start_token = Arc::new(()); - self.ephemeral_start_tokens - .write() - .await - .insert(server_id.clone(), Arc::clone(&start_token)); - self.ephemeral_workspace_scopes - .write() - .await - .insert(server_id.clone(), workspace_key); - self.ephemeral_ready_servers - .write() - .await - .remove(&server_id); - let cancelled_retirement = self - .ephemeral_retirements - .write() - .await - .remove(&server_id) - .map(|cancelled| { - cancelled.store(true, Ordering::Release); - true - }) - .unwrap_or(false); - - if cancelled_retirement && self.runtime.contains(&server_id).await { - if let Err(error) = self.runtime.insert_runtime_config(config.clone()).await { - let _ = self.remove_ephemeral_server(&server_id).await; - return Err(error.into()); - } - let connection = if let Some(process) = self.runtime.get_process(&server_id).await { - process.read().await.connection() - } else { - None - }; - if let Some(connection) = connection { - self.runtime - .add_connection(server_id.clone(), connection.clone()) - .await; - if let Err(error) = self - .refresh_mcp_tools(&server_id, &config.name, connection.clone()) - .await - { - let _ = self.remove_ephemeral_server(&server_id).await; - return Err(error); - } - self.start_connection_event_listener(&server_id, &config.name, connection.clone()) - .await; - self.warm_catalog_caches(&server_id, connection).await; - self.ephemeral_ready_servers - .write() - .await - .insert(server_id.clone()); - } else { - let _ = self.remove_ephemeral_server(&server_id).await; - return Err(BitFunError::MCPError( - "External MCP server did not retain its connection".to_string(), - )); - } - return Ok(()); - } - if self.runtime.contains(&server_id).await { - self.ephemeral_workspace_scopes - .write() - .await - .remove(&server_id); - self.ephemeral_start_tokens.write().await.remove(&server_id); - return Err(BitFunError::Configuration(format!( - "MCP server already exists: {}", - server_id - ))); - } - - if let Err(error) = self.runtime.insert_runtime_config(config.clone()).await { - self.ephemeral_workspace_scopes - .write() - .await - .remove(&server_id); - self.ephemeral_start_tokens.write().await.remove(&server_id); - return Err(error.into()); - } - if let Err(error) = self.runtime.register(&config).await { - self.runtime.remove_runtime_config(&server_id).await; - self.ephemeral_workspace_scopes - .write() - .await - .remove(&server_id); - self.ephemeral_start_tokens.write().await.remove(&server_id); - return Err(error.into()); - } - if config.enabled && config.auto_start { - // External source refresh and product-surface reads must not wait - // for a third-party process or network handshake. Registration is - // synchronous so status reads immediately see Loading; startup is - // bounded in the background and cleans up only this runtime item. - const EXTERNAL_START_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); - let manager = self.clone(); - tokio::spawn(async move { - let startup = tokio::time::timeout( - EXTERNAL_START_TIMEOUT, - manager.start_server_with_external_token( - &server_id, - Some(Arc::clone(&start_token)), - ), - ) - .await; - match startup { - Ok(Ok(())) => { - if manager - .external_start_token_matches(&server_id, &start_token) - .await - { - crate::external_sources::notify_external_tool_registry_changed(); - } - } - Ok(Err(error)) => { - warn!( - "External ephemeral MCP server failed to start: id={} error={}", - server_id, error - ); - if manager - .remove_ephemeral_server_for_start(&server_id, &start_token) - .await - { - crate::external_sources::notify_external_tool_registry_changed(); - } - } - Err(_) => { - warn!( - "External ephemeral MCP server startup timed out: id={}", - server_id - ); - if manager - .remove_ephemeral_server_for_start(&server_id, &start_token) - .await - { - crate::external_sources::notify_external_tool_registry_changed(); - } - } - } - }); - } - Ok(()) - } - - /// Withdraws new tool/resource access immediately, then lets already-held - /// connection users finish before the process is reclaimed. The grace is - /// bounded so a deleted or malicious server cannot remain indefinitely. - pub async fn retire_external_ephemeral_server(&self, server_id: &str) -> BitFunResult<()> { - const RETIREMENT_GRACE: std::time::Duration = std::time::Duration::from_secs(30); - const RETIREMENT_RECLAIM_ATTEMPTS: usize = 3; - const RETIREMENT_RETRY_DELAY: std::time::Duration = std::time::Duration::from_millis(250); - let _lifecycle_guard = self.ephemeral_lifecycle.lock().await; - self.ephemeral_start_tokens.write().await.remove(server_id); - if !self.runtime.contains(server_id).await { - self.runtime.remove_runtime_config(server_id).await; - self.ephemeral_ready_servers.write().await.remove(server_id); - self.ephemeral_workspace_scopes - .write() - .await - .remove(server_id); - return Ok(()); - } - - if let Some(previous) = self - .ephemeral_retirements - .write() - .await - .insert(server_id.to_string(), Arc::new(AtomicBool::new(false))) - { - previous.store(true, Ordering::Release); - } - let cancelled = self - .ephemeral_retirements - .read() - .await - .get(server_id) - .cloned() - .expect("retirement marker was just inserted"); - let connection = self.runtime.get_connection(server_id).await; - - self.ephemeral_ready_servers.write().await.remove(server_id); - Self::unregister_mcp_tools(server_id).await; - self.stop_connection_event_listener(server_id).await; - self.runtime.remove_connection(server_id).await; - self.runtime.remove_catalog(server_id).await; - self.runtime.remove_runtime_config(server_id).await; - self.clear_reconnect_state(server_id).await; - - let manager = self.clone(); - let server_id = server_id.to_string(); - tokio::spawn(async move { - let started = std::time::Instant::now(); - loop { - if cancelled.load(Ordering::Acquire) { - return; - } - let references = connection.as_ref().map_or(0, Arc::strong_count); - if should_finish_ephemeral_retirement( - references, - started.elapsed(), - RETIREMENT_GRACE, - ) { - break; - } - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - } - - for attempt in 1..=RETIREMENT_RECLAIM_ATTEMPTS { - let lifecycle_guard = manager.ephemeral_lifecycle.lock().await; - if cancelled.load(Ordering::Acquire) { - return; - } - let should_remove = manager - .ephemeral_retirements - .read() - .await - .get(&server_id) - .is_some_and(|current| Arc::ptr_eq(current, &cancelled)); - if !should_remove { - return; - } - match manager.runtime.unregister(&server_id).await { - Ok(()) => { - manager - .ephemeral_retirements - .write() - .await - .remove(&server_id); - Self::unregister_mcp_tools(&server_id).await; - manager.stop_connection_event_listener(&server_id).await; - manager.runtime.remove_connection(&server_id).await; - manager.runtime.remove_catalog(&server_id).await; - manager - .ephemeral_ready_servers - .write() - .await - .remove(&server_id); - manager - .ephemeral_workspace_scopes - .write() - .await - .remove(&server_id); - return; - } - Err(error) if attempt < RETIREMENT_RECLAIM_ATTEMPTS => { - warn!( - "Could not reclaim retired ephemeral MCP server; retrying: id={} attempt={} error={}", - server_id, attempt, error - ); - } - Err(error) => { - warn!( - "Could not reclaim retired ephemeral MCP server; retaining ownership for a later retry: id={} attempts={} error={}", - server_id, RETIREMENT_RECLAIM_ATTEMPTS, error - ); - return; - } - } - drop(lifecycle_guard); - tokio::time::sleep(RETIREMENT_RETRY_DELAY).await; - } - }); - Ok(()) - } - - /// Removes a runtime-only MCP server and its registered tools without touching persisted config. - pub async fn remove_ephemeral_server(&self, server_id: &str) -> BitFunResult<()> { - info!("Removing ephemeral MCP server: id={}", server_id); - - if !self.runtime.contains(server_id).await { - self.runtime.remove_runtime_config(server_id).await; - self.clear_reconnect_state(server_id).await; - self.runtime.remove_catalog(server_id).await; - Self::unregister_mcp_tools(server_id).await; - return Ok(()); - } - - let stop_result = self.stop_server(server_id).await; - self.stop_connection_event_listener(server_id).await; - self.clear_reconnect_state(server_id).await; - self.runtime.remove_catalog(server_id).await; - self.ephemeral_ready_servers.write().await.remove(server_id); - self.ephemeral_start_tokens.write().await.remove(server_id); - self.ephemeral_workspace_scopes - .write() - .await - .remove(server_id); - - if let Err(error) = stop_result { - warn!( - "Failed to stop ephemeral MCP server; retaining runtime ownership for retry: id={} error={}", - server_id, error - ); - return Err(error); - } - - self.runtime.unregister(server_id).await?; - self.runtime.remove_runtime_config(server_id).await; - info!("Unregistered ephemeral MCP server: id={}", server_id); - Ok(()) - } - /// Removes a server. pub async fn remove_server(&self, server_id: &str) -> BitFunResult<()> { info!("Removing MCP server: id={}", server_id); diff --git a/src/crates/assembly/core/src/service/mcp/server/manager/mod.rs b/src/crates/assembly/core/src/service/mcp/server/manager/mod.rs index 6bf428c5f5..e96b03b683 100644 --- a/src/crates/assembly/core/src/service/mcp/server/manager/mod.rs +++ b/src/crates/assembly/core/src/service/mcp/server/manager/mod.rs @@ -5,6 +5,7 @@ mod auth; mod catalog; +mod external_lifecycle; mod interaction; mod lifecycle; mod reconnect; diff --git a/src/crates/assembly/core/src/service/mcp/server/manager/tests.rs b/src/crates/assembly/core/src/service/mcp/server/manager/tests.rs index 35797fb258..ad1fa48715 100644 --- a/src/crates/assembly/core/src/service/mcp/server/manager/tests.rs +++ b/src/crates/assembly/core/src/service/mcp/server/manager/tests.rs @@ -1,48 +1,5 @@ -use bitfun_services_integrations::mcp::server::{ - compute_mcp_backoff_delay, detect_mcp_list_changed_kind, MCPListChangedKind, -}; use std::time::Duration; -#[test] -fn backoff_delay_grows_exponentially_and_caps() { - let base = Duration::from_secs(2); - let max = Duration::from_secs(60); - - assert_eq!( - compute_mcp_backoff_delay(base, max, 1), - Duration::from_secs(2) - ); - assert_eq!( - compute_mcp_backoff_delay(base, max, 2), - Duration::from_secs(4) - ); - assert_eq!( - compute_mcp_backoff_delay(base, max, 5), - Duration::from_secs(32) - ); - assert_eq!( - compute_mcp_backoff_delay(base, max, 10), - Duration::from_secs(60) - ); -} - -#[test] -fn detect_list_changed_kind_supports_three_catalogs() { - assert_eq!( - detect_mcp_list_changed_kind("notifications/tools/list_changed"), - Some(MCPListChangedKind::Tools) - ); - assert_eq!( - detect_mcp_list_changed_kind("notifications/prompts/list_changed"), - Some(MCPListChangedKind::Prompts) - ); - assert_eq!( - detect_mcp_list_changed_kind("notifications/resources/list_changed"), - Some(MCPListChangedKind::Resources) - ); - assert_eq!(detect_mcp_list_changed_kind("notifications/unknown"), None); -} - #[test] fn ephemeral_retirement_waits_for_in_flight_connection_users_but_is_bounded() { let grace = Duration::from_secs(30); diff --git a/src/crates/services/services-integrations/src/mcp/server/mod.rs b/src/crates/services/services-integrations/src/mcp/server/mod.rs index d3c5bc5cc7..4faa26c45f 100644 --- a/src/crates/services/services-integrations/src/mcp/server/mod.rs +++ b/src/crates/services/services-integrations/src/mcp/server/mod.rs @@ -30,7 +30,7 @@ pub use runtime_policy::{ mcp_server_is_running, mcp_server_is_starting_or_running, mcp_should_start_after_config_update, MCPListChangedKind, MCPReconnectRuntimeDecision, }; -pub use runtime_state::MCPServerRuntimeState; +pub use runtime_state::{MCPProcessStartContext, MCPProcessStartOutcome, MCPServerRuntimeState}; /// MCP server type. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] diff --git a/src/crates/services/services-integrations/src/mcp/server/runtime_state.rs b/src/crates/services/services-integrations/src/mcp/server/runtime_state.rs index f58d090a28..ef5b0ada1a 100644 --- a/src/crates/services/services-integrations/src/mcp/server/runtime_state.rs +++ b/src/crates/services/services-integrations/src/mcp/server/runtime_state.rs @@ -5,14 +5,56 @@ //! OAuth callback UI. use super::{ - MCPCatalogCache, MCPConnection, MCPConnectionPool, MCPReconnectTracker, MCPRuntimeResult, - MCPServerConfig, MCPServerProcess, MCPServerRegistry, MCPServerStatus, + mcp_server_is_running, resolve_mcp_local_command, MCPCatalogCache, MCPConnection, + MCPConnectionPool, MCPReconnectTracker, MCPRuntimeError, MCPRuntimeResult, MCPServerConfig, + MCPServerProcess, MCPServerRegistry, MCPServerStatus, MCPServerTransport, MCPServerType, }; use crate::mcp::protocol::{MCPPrompt, MCPResource}; +use log::info; +use std::fmt; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; use tokio::sync::RwLock; +/// Result of starting one registered MCP process. +/// +/// A newly started connection is returned without publishing it to the shared +/// connection pool. Product assembly can therefore complete its own +/// generation/retirement checks before making the connection visible. +pub enum MCPProcessStartOutcome { + AlreadyRunning, + Started { connection: Arc }, +} + +/// Host-owned paths required to start one MCP process. +/// +/// Keeping the server kind in the context prevents callers from passing an +/// unrelated empty or placeholder path across the service boundary. +#[derive(Debug)] +pub enum MCPProcessStartContext { + Local { managed_runtimes_dir: PathBuf }, + Remote { data_dir: PathBuf }, +} + +impl MCPProcessStartContext { + fn server_type(&self) -> MCPServerType { + match self { + Self::Local { .. } => MCPServerType::Local, + Self::Remote { .. } => MCPServerType::Remote, + } + } +} + +impl fmt::Debug for MCPProcessStartOutcome { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::AlreadyRunning => formatter.write_str("AlreadyRunning"), + Self::Started { .. } => formatter.write_str("Started"), + } + } +} + pub struct MCPServerRuntimeState { registry: MCPServerRegistry, connection_pool: MCPConnectionPool, @@ -54,10 +96,134 @@ impl MCPServerRuntimeState { self.registry.clear().await } - pub async fn get_process(&self, server_id: &str) -> Option>> { + async fn get_process(&self, server_id: &str) -> Option>> { self.registry.get_process(server_id).await } + pub async fn start_process( + &self, + config: &MCPServerConfig, + context: MCPProcessStartContext, + ) -> MCPRuntimeResult { + if !config.enabled { + return Err(MCPRuntimeError::configuration(format!( + "MCP server is disabled: {}", + config.id + ))); + } + if config.server_type != context.server_type() { + let server_type = match config.server_type { + MCPServerType::Local => "local", + MCPServerType::Remote => "remote", + }; + return Err(MCPRuntimeError::configuration(format!( + "MCP process start context does not match server type '{server_type}'" + ))); + } + + self.ensure_registered(config).await?; + let process = self.get_process(&config.id).await.ok_or_else(|| { + MCPRuntimeError::not_found(format!("MCP server not registered: {}", config.id)) + })?; + let mut process = process.write().await; + + if mcp_server_is_running(process.status().await) { + return Ok(MCPProcessStartOutcome::AlreadyRunning); + } + + match (config.server_type, context) { + ( + MCPServerType::Local, + MCPProcessStartContext::Local { + managed_runtimes_dir, + }, + ) => { + let command = config.command.as_ref().ok_or_else(|| { + MCPRuntimeError::configuration("Missing command for local MCP server") + })?; + let resolved = resolve_mcp_local_command(command, managed_runtimes_dir)?; + info!( + "Starting local MCP server: command={} source={} id={}", + resolved.command, resolved.source_label, config.id + ); + process + .start_with_environment_policy( + &resolved.command, + &config.args, + &config.env, + config.working_directory.as_deref().map(Path::new), + config.inherits_parent_environment(), + ) + .await?; + } + (MCPServerType::Remote, MCPProcessStartContext::Remote { data_dir }) => { + let transport = config.resolved_transport(); + if transport != MCPServerTransport::StreamableHttp { + return Err(MCPRuntimeError::not_implemented(format!( + "Remote MCP transport '{}' is not yet supported", + transport.as_str() + ))); + } + if config.url.is_none() { + return Err(MCPRuntimeError::configuration( + "Missing URL for remote MCP server", + )); + } + info!( + "Connecting to remote MCP server: transport={} id={}", + transport.as_str(), + config.id + ); + process.start_remote(data_dir, config).await?; + } + _ => unreachable!("MCP start context was validated before process registration"), + } + + let connection = process.connection().ok_or_else(|| { + MCPRuntimeError::mcp(format!( + "MCP server '{}' started without a connection", + config.id + )) + })?; + Ok(MCPProcessStartOutcome::Started { connection }) + } + + pub async fn stop_process(&self, server_id: &str) -> MCPRuntimeResult<()> { + let process = self.get_process(server_id).await.ok_or_else(|| { + MCPRuntimeError::not_found(format!("MCP server not found: {server_id}")) + })?; + let stop_result = process.write().await.stop().await; + + self.remove_connection(server_id).await; + self.remove_catalog(server_id).await; + stop_result + } + + pub async fn process_connection(&self, server_id: &str) -> Option> { + let process = self.get_process(server_id).await?; + let process = process.read().await; + process.connection() + } + + pub async fn process_status(&self, server_id: &str) -> MCPRuntimeResult { + let process = self.get_process(server_id).await.ok_or_else(|| { + MCPRuntimeError::not_found(format!("MCP server not found: {server_id}")) + })?; + let process = process.read().await; + Ok(process.status().await) + } + + pub async fn process_status_message( + &self, + server_id: &str, + ) -> MCPRuntimeResult> { + let process = self.get_process(server_id).await.ok_or_else(|| { + MCPRuntimeError::not_found(format!("MCP server not found: {server_id}")) + })?; + let process = process.read().await; + Ok(process.status_message().await) + } + pub async fn get_all_server_ids(&self) -> Vec { self.registry.get_all_server_ids().await } diff --git a/src/crates/services/services-integrations/tests/mcp_contracts.rs b/src/crates/services/services-integrations/tests/mcp_contracts.rs index 18e132fa1d..dbfe4659a4 100644 --- a/src/crates/services/services-integrations/tests/mcp_contracts.rs +++ b/src/crates/services/services-integrations/tests/mcp_contracts.rs @@ -25,8 +25,8 @@ use bitfun_services_integrations::mcp::server::{ compute_mcp_backoff_delay, detect_mcp_list_changed_kind, is_mcp_auth_error_message, mcp_reconnect_runtime_decision, mcp_server_is_running, mcp_should_start_after_config_update, merge_mcp_remote_headers, MCPCatalogCache, MCPConnectionPool, MCPListChangedKind, - MCPReconnectRuntimeDecision, MCPRuntimeErrorKind, MCPRuntimeResult, MCPServerConfig, - MCPServerProcess, MCPServerRuntimeState, MCPServerStatus, MCPServerTransport, MCPServerType, + MCPProcessStartContext, MCPReconnectRuntimeDecision, MCPRuntimeErrorKind, MCPRuntimeResult, + MCPServerConfig, MCPServerRuntimeState, MCPServerStatus, MCPServerTransport, MCPServerType, }; use bitfun_services_integrations::mcp::{ build_mcp_tool_descriptor, build_mcp_tool_name, normalize_name_for_mcp, @@ -1283,7 +1283,7 @@ async fn mcp_dynamic_tool_provider_preserves_manifest_order_and_metadata_snapsho } #[tokio::test] -async fn mcp_server_process_owner_preserves_unsupported_remote_transport_contract() { +async fn mcp_runtime_state_owner_preserves_unsupported_remote_transport_contract() { let mut config = make_mcp_config( "remote-sse", ConfigLocation::User, @@ -1293,23 +1293,37 @@ async fn mcp_server_process_owner_preserves_unsupported_remote_transport_contrac ); config.transport = Some(MCPServerTransport::Sse); - let mut process = MCPServerProcess::new( - "remote-sse".to_string(), - "Remote SSE".to_string(), - MCPServerType::Remote, + let runtime = MCPServerRuntimeState::new(); + runtime.register(&config).await.expect("register process"); + assert_eq!( + runtime + .process_status("remote-sse") + .await + .expect("registered process status"), + MCPServerStatus::Uninitialized ); - assert_eq!(process.status().await, MCPServerStatus::Uninitialized); - assert_eq!(process.server_type(), MCPServerType::Remote); - let error = process - .start_remote(std::env::temp_dir(), &config) + let error = runtime + .start_process( + &config, + MCPProcessStartContext::Remote { + data_dir: std::env::temp_dir(), + }, + ) .await .unwrap_err(); assert_eq!(error.kind(), MCPRuntimeErrorKind::NotImplemented); assert!(error .to_string() .contains("Remote MCP transport 'sse' is not yet supported")); - assert_eq!(process.status().await, MCPServerStatus::Uninitialized); + assert_eq!( + runtime + .process_status("remote-sse") + .await + .expect("registered process status"), + MCPServerStatus::Uninitialized + ); + assert!(runtime.process_connection("remote-sse").await.is_none()); let pool = MCPConnectionPool::new(); assert!(pool.get_all_server_ids().await.is_empty()); @@ -1546,6 +1560,20 @@ async fn mcp_runtime_state_owns_registry_runtime_config_and_reconnect_state() { config.auto_start = false; assert!(runtime.is_empty().await); + let error = runtime + .start_process( + &config, + MCPProcessStartContext::Remote { + data_dir: std::env::temp_dir(), + }, + ) + .await + .expect_err("local config must reject remote start context"); + assert_eq!(error.kind(), MCPRuntimeErrorKind::Configuration); + assert!(error + .to_string() + .contains("does not match server type 'local'")); + assert!(runtime.is_empty().await); runtime .insert_runtime_config(config.clone()) @@ -1558,7 +1586,14 @@ async fn mcp_runtime_state_owns_registry_runtime_config_and_reconnect_state() { assert!(runtime.contains("runtime-only").await); assert_eq!(runtime.get_all_server_ids().await, vec!["runtime-only"]); - assert!(runtime.get_process("runtime-only").await.is_some()); + assert_eq!( + runtime + .process_status("runtime-only") + .await + .expect("registered process status"), + MCPServerStatus::Uninitialized + ); + assert!(runtime.process_connection("runtime-only").await.is_none()); assert_eq!( runtime.get_all_statuses().await, vec![("runtime-only".to_string(), MCPServerStatus::Uninitialized)]