retores use of oneshot in server launch

Author: dingfeli

crates/agent/src/agent/mcp/mod.rs

@@ -182,14 +182,22 @@ impl McpManagerHandle {
         &mut self,
         name: String,
         config: McpServerConfig,
-    ) -> Result<McpManagerResponse, McpManagerError> {
-        self.request_tx
+    ) -> Result<oneshot::Receiver<LaunchServerResult>, McpManagerError> {
+        match self
+            .request_tx
             .send_recv(McpManagerRequest::LaunchServer {
                 server_name: name,
                 config,
             })
             .await
-            .unwrap_or(Err(McpManagerError::Channel))
+            .unwrap_or(Err(McpManagerError::Channel))?
+        {
+            McpManagerResponse::LaunchServer(rx) => Ok(rx),
+            other => Err(McpManagerError::Custom(format!(
+                "received unexpected response: {:?}",
+                other
+            ))),
+        }
     }
 
     pub async fn get_tool_specs(&self, server_name: String) -> Result<Vec<ToolSpec>, McpManagerError> {
@@ -257,7 +265,7 @@ pub struct McpManager {
 
     cred_path: PathBuf,
 
-    initializing_servers: HashMap<String, McpServerActorHandle>,
+    initializing_servers: HashMap<String, (McpServerActorHandle, oneshot::Sender<LaunchServerResult>)>,
     servers: HashMap<String, McpServerActorHandle>,
     event_buf: Vec<McpServerActorEvent>,
 }
@@ -332,8 +340,10 @@ impl McpManager {
                 }
                 let event_tx = self.server_event_tx.clone();
                 let handle = McpServerActor::spawn(name.clone(), config, self.cred_path.clone(), event_tx);
-                self.initializing_servers.insert(name, handle);
-                Ok(McpManagerResponse::LaunchServer)
+                let (tx, rx) = oneshot::channel();
+
+                self.initializing_servers.insert(name, (handle, tx));
+                Ok(McpManagerResponse::LaunchServer(rx))
             },
             McpManagerRequest::GetToolSpecs { server_name } => match self.servers.get(&server_name) {
                 Some(handle) => Ok(McpManagerResponse::ToolSpecs(handle.get_tool_specs().await?)),
@@ -365,17 +375,25 @@ impl McpManager {
                 list_tools_duration: _,
                 list_prompts_duration: _,
             } => {
-                let Some(handle) = self.initializing_servers.remove(server_name) else {
+                let Some((handle, result_tx)) = self.initializing_servers.remove(server_name) else {
                     warn!(?server_name, ?evt, "event was not from an initializing MCP server");
                     return;
                 };
 
+                if let Err(e) = result_tx.send(Ok(())) {
+                    error!(?server_name, ?e, "failed to send server initialized message");
+                }
+
                 if self.servers.insert(server_name.clone(), handle).is_some() {
                     warn!(?server_name, "duplicated server. old server dropped");
                 }
             },
-            McpServerActorEvent::InitializeError { server_name, error: _ } => {
-                self.initializing_servers.remove(server_name);
+            McpServerActorEvent::InitializeError { server_name, error } => {
+                if let Some((_, result_tx)) = self.initializing_servers.remove(server_name) {
+                    if let Err(e) = result_tx.send(Err(McpManagerError::Custom(error.clone()))) {
+                        error!(?server_name, ?e, "failed to send server initialized message");
+                    }
+                }
             },
             McpServerActorEvent::OauthRequest { server_name, oauth_url } => {
                 info!(?server_name, ?oauth_url, "received oauth request");
@@ -418,14 +436,16 @@ pub enum McpManagerRequest {
 
 #[derive(Debug)]
 pub enum McpManagerResponse {
-    LaunchServer,
+    LaunchServer(oneshot::Receiver<LaunchServerResult>),
     ToolSpecs(Vec<ToolSpec>),
     Prompts(Vec<Prompt>),
     ExecuteTool(oneshot::Receiver<ExecuteToolResult>),
 }
 
 pub type ExecuteToolResult = Result<CallToolResult, McpServerActorError>;
 
+type LaunchServerResult = Result<(), McpManagerError>;
+
 #[derive(Debug, Clone, Serialize, Deserialize, thiserror::Error)]
 pub enum McpManagerError {
     #[error("Server with the name {} is not initialized", .name)]

crates/agent/src/agent/mod.rs

@@ -57,6 +57,7 @@ use agent_loop::{
 };
 use chrono::Utc;
 use consts::MAX_RESOURCE_FILE_LENGTH;
+use futures::stream::FuturesUnordered;
 use mcp::actor::McpServerActorEvent;
 use permissions::evaluate_tool_permission;
 use protocol::{
@@ -94,9 +95,11 @@ use task_executor::{
 };
 use tokio::sync::{
     broadcast,
+    mpsc,
     oneshot,
 };
 use tokio::time::Instant;
+use tokio_stream::StreamExt as _;
 use tokio_util::sync::CancellationToken;
 use tool_utils::{
     SanitizedToolSpecs,
@@ -344,54 +347,78 @@ impl Agent {
                 warn!(?self.cached_mcp_configs.overridden_configs, "ignoring overridden configs");
             }
 
-            let mut total_servers_to_be_loaded = 0_usize;
+            let mut results = FuturesUnordered::new();
 
             for config in self
                 .cached_mcp_configs
                 .configs
                 .iter()
                 .filter(|config| config.is_enabled())
             {
-                if let Err(e) = self
+                let Ok(rx) = self
                     .mcp_manager_handle
                     .launch_server(config.server_name.clone(), config.config.clone())
                     .await
-                {
-                    warn!(?config.server_name, ?e, "failed to launch MCP config, skipping");
-                } else {
-                    total_servers_to_be_loaded += 1;
-                }
+                else {
+                    warn!(?config.server_name, "failed to launch MCP config, skipping");
+                    continue;
+                };
+                let name = config.server_name.clone();
+                results.push(async move { (name, rx.await) });
             }
 
+            // Continually loop through the receivers until all have completed.
+            let mut launched_servers = Vec::new();
+            let (success_tx, mut success_rx) = mpsc::channel(8);
+            let mut failed_servers = Vec::new();
+            let (failed_tx, mut failed_rx) = mpsc::channel(8);
+            let init_results_handle = tokio::spawn(async move {
+                while let Some((name, res)) = results.next().await {
+                    debug!(?name, ?res, "received result from LaunchServer request");
+                    let Ok(res) = res else {
+                        warn!(?name, "channel unexpectedly dropped during MCP initialization");
+                        let _ = failed_tx.send(name).await;
+                        continue;
+                    };
+                    match res {
+                        Ok(_) => {
+                            let _ = success_tx.send(name).await;
+                        },
+                        Err(err) => {
+                            error!(?name, ?err, "failed to launch MCP server");
+                            let _ = failed_tx.send(name).await;
+                        },
+                    }
+                }
+            });
+
             let timeout_at = Instant::now() + self.settings.mcp_init_timeout;
             loop {
                 tokio::select! {
-                    evt = self.mcp_manager_handle.recv() => {
-                        let evt = match evt {
-                            Ok(evt) => evt,
-                            Err(e) => {
-                                error!(?e, "mcp manager handle channel closed");
-                                break;
-                            }
+                    name = success_rx.recv() => {
+                        let Some(name) = name else {
+                            // If None is returned in either success/failed receivers, then the
+                            // senders have dropped, meaning initialization has completed.
+                            break;
                         };
-
-                        if matches!(evt, McpServerActorEvent::Initialized{ .. } | McpServerActorEvent::InitializeError { .. }) {
-                             total_servers_to_be_loaded = total_servers_to_be_loaded.saturating_sub(1);
-                        }
-                        self.handle_mcp_server_actor_events(evt).await;
-
-                        if total_servers_to_be_loaded == 0 {
-                            info!("all mcp servers loaded before timeout");
+                        debug!(?name, "MCP server successfully initialized");
+                        launched_servers.push(name);
+                    },
+                    name = failed_rx.recv() => {
+                        let Some(name) = name else {
                             break;
-                        }
+                        };
+                        warn!(?name, "MCP server failed initialization");
+                        failed_servers.push(name);
                     },
-
                     _ = tokio::time::sleep_until(timeout_at) => {
                         warn!("timed out before all MCP servers could be initialized");
                         break;
                     },
                 }
             }
+            info!(?launched_servers, ?failed_servers, "MCP server initialization finished");
+            init_results_handle.abort();
         }
 
         // Next, run agent spawn hooks.