From c8cac0925dbb2ee7d3eea573cce985136a17ce31 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Thu, 12 Mar 2026 07:46:40 +0800 Subject: [PATCH] fix(mcp): stdio/unix transports skip initialize handshake (#890) (#935) fixes #890 - Always call initialize() before list_tools()/call_tool(), removing the session_manager.is_some() guard that caused stdio/unix clients to skip the MCP protocol handshake entirely - Add local AtomicBool flag for idempotent initialization when no session manager is present - Fire-and-forget JSON-RPC notifications (id=None) in stdio/unix transports instead of registering a pending response that would block for 30s waiting on a reply that never comes - Fix mcp test panic on stdio/unix servers by using create_client_from_config() instead of new_with_config() which asserts HTTP-only transport --- src/cli/mcp.rs | 18 ++++-- src/tools/mcp/client.rs | 98 ++++++++++++++++++++++++++++---- src/tools/mcp/stdio_transport.rs | 22 +++++-- src/tools/mcp/unix_transport.rs | 22 +++++-- 4 files changed, 138 insertions(+), 22 deletions(-) diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index 5cc7b911..2293a6d6 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -12,9 +12,10 @@ use crate::config::Config; use crate::db::Database; use crate::secrets::SecretsStore; use crate::tools::mcp::{ - McpClient, McpServerConfig, McpSessionManager, OAuthConfig, + McpClient, McpProcessManager, McpServerConfig, McpSessionManager, OAuthConfig, auth::{authorize_mcp_server, is_authenticated}, config::{self, EffectiveTransport, McpServersFile}, + factory::create_client_from_config, }; /// Arguments for the `mcp add` subcommand. @@ -494,7 +495,7 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { let client = if has_tokens { // We have stored tokens, use authenticated client - McpClient::new_authenticated(server.clone(), session_manager, secrets, user_id) + McpClient::new_authenticated(server.clone(), session_manager.clone(), secrets, user_id) } else if server.requires_auth() { // OAuth configured but no tokens - need to authenticate println!(); @@ -505,8 +506,17 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { println!(); return Ok(()); } else { - // No OAuth and no tokens - try unauthenticated - McpClient::new_with_config(server.clone()) + // Use the factory to dispatch on transport type (HTTP, stdio, unix) + let process_manager = Arc::new(McpProcessManager::new()); + create_client_from_config( + server.clone(), + &session_manager, + &process_manager, + None, + "default", + ) + .await + .map_err(|e| anyhow::anyhow!("{}", e))? }; // Test connection diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index 61c9d5c7..377d0488 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use async_trait::async_trait; use tokio::sync::RwLock; @@ -57,6 +57,10 @@ pub struct McpClient { /// Custom headers to include in every request. custom_headers: HashMap, + + /// Whether the MCP initialize handshake has completed. + /// Used as a local idempotency guard when no session_manager is present. + initialized: AtomicBool, } impl McpClient { @@ -79,6 +83,7 @@ impl McpClient { user_id: "default".to_string(), server_config: None, custom_headers: HashMap::new(), + initialized: AtomicBool::new(false), } } @@ -101,6 +106,7 @@ impl McpClient { user_id: "default".to_string(), server_config: None, custom_headers: HashMap::new(), + initialized: AtomicBool::new(false), } } @@ -131,6 +137,7 @@ impl McpClient { secrets: None, user_id: "default".to_string(), custom_headers: config.headers.clone(), + initialized: AtomicBool::new(false), server_config: Some(config), } } @@ -162,6 +169,7 @@ impl McpClient { user_id: user_id.into(), server_config: Some(config), custom_headers, + initialized: AtomicBool::new(false), } } @@ -197,6 +205,7 @@ impl McpClient { user_id: user_id.into(), server_config, custom_headers, + initialized: AtomicBool::new(false), } } @@ -307,9 +316,14 @@ impl McpClient { /// Initialize the connection to the MCP server. pub async fn initialize(&self) -> Result { + // Fast path: already initialized (local flag or session manager) + if self.initialized.load(Ordering::Relaxed) { + return Ok(InitializeResult::default()); + } if let Some(ref session_manager) = self.session_manager && session_manager.is_initialized(&self.server_name).await { + self.initialized.store(true, Ordering::Relaxed); return Ok(InitializeResult::default()); } if let Some(ref session_manager) = self.session_manager { @@ -342,6 +356,7 @@ impl McpClient { if let Some(ref session_manager) = self.session_manager { session_manager.mark_initialized(&self.server_name).await; } + self.initialized.store(true, Ordering::Relaxed); let notification = McpRequest::initialized_notification(); let _ = self.send_request(notification).await; @@ -354,9 +369,7 @@ impl McpClient { if let Some(tools) = self.tools_cache.read().await.as_ref() { return Ok(tools.clone()); } - if self.session_manager.is_some() { - self.initialize().await?; - } + self.initialize().await?; let request = McpRequest::list_tools(self.next_request_id()); let response = self.send_request(request).await?; @@ -386,9 +399,7 @@ impl McpClient { name: &str, arguments: serde_json::Value, ) -> Result { - if self.session_manager.is_some() { - self.initialize().await?; - } + self.initialize().await?; let request = McpRequest::call_tool(self.next_request_id(), name, arguments); let response = self.send_request(request).await?; @@ -452,6 +463,7 @@ impl Clone for McpClient { user_id: self.user_id.clone(), server_config: self.server_config.clone(), custom_headers: self.custom_headers.clone(), + initialized: AtomicBool::new(self.initialized.load(Ordering::Relaxed)), } } } @@ -794,13 +806,34 @@ mod tests { #[tokio::test] async fn test_non_http_transport_skips_401_retry() { - let response = McpResponse { + // initialize response, then notification ack (consumed but ignored), + // then list_tools response + let init_response = McpResponse { jsonrpc: "2.0".to_string(), id: Some(1), + result: Some(serde_json::json!({ + "protocolVersion": "2024-11-05", + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1.0"} + })), + error: None, + }; + let notification_ack = McpResponse { + jsonrpc: "2.0".to_string(), + id: None, + result: None, + error: None, + }; + let list_response = McpResponse { + jsonrpc: "2.0".to_string(), + id: Some(2), result: Some(serde_json::json!({"tools": []})), error: None, }; - let transport = Arc::new(MockTransport::new(false, vec![response])); + let transport = Arc::new(MockTransport::new( + false, + vec![init_response, notification_ack, list_response], + )); let client = McpClient::new_with_transport( "test-stdio", transport.clone(), @@ -813,7 +846,8 @@ mod tests { assert!(result.is_ok()); assert_eq!(result.unwrap().len(), 0); let headers = transport.recorded_headers(); - assert_eq!(headers.len(), 1); + // 3 sends: initialize + notifications/initialized + list_tools + assert_eq!(headers.len(), 3); assert!(!headers[0].contains_key("Authorization")); assert!(!headers[0].contains_key("Mcp-Session-Id")); } @@ -826,6 +860,50 @@ mod tests { assert!(!mock_non_http.supports_http_features()); } + /// Regression test for issue #890: stdio clients must auto-initialize + /// even without a session manager, and the second call should be idempotent. + #[tokio::test] + async fn test_stdio_client_auto_initializes_without_session_manager() { + let init_response = McpResponse { + jsonrpc: "2.0".to_string(), + id: Some(1), + result: Some(serde_json::json!({ + "protocolVersion": "2024-11-05", + "capabilities": {}, + "serverInfo": {"name": "test", "version": "1.0"} + })), + error: None, + }; + let notification_ack = McpResponse { + jsonrpc: "2.0".to_string(), + id: None, + result: None, + error: None, + }; + let transport = Arc::new(MockTransport::new( + false, + vec![init_response, notification_ack], + )); + let client = McpClient::new_with_transport( + "test-stdio", + transport.clone(), + None, // no session manager + None, + "default", + None, + ); + + // First call should send initialize + notification + let result = client.initialize().await; + assert!(result.is_ok()); + assert_eq!(transport.recorded_headers().len(), 2); + + // Second call should be a no-op (idempotent via local flag) + let result2 = client.initialize().await; + assert!(result2.is_ok()); + assert_eq!(transport.recorded_headers().len(), 2); // no additional sends + } + #[test] fn test_strip_top_level_nulls_removes_null_fields() { let input = serde_json::json!({ diff --git a/src/tools/mcp/stdio_transport.rs b/src/tools/mcp/stdio_transport.rs index 95deb8b3..ed8c79aa 100644 --- a/src/tools/mcp/stdio_transport.rs +++ b/src/tools/mcp/stdio_transport.rs @@ -118,13 +118,27 @@ impl McpTransport for StdioMcpTransport { request: &McpRequest, _headers: &HashMap, ) -> Result { + // JSON-RPC notifications (no id) are fire-and-forget: the server + // will not send a response, so we must not wait for one. + if request.id.is_none() { + let mut stdin = self.stdin.lock().await; + write_jsonrpc_line(&mut *stdin, request).await?; + return Ok(McpResponse { + jsonrpc: "2.0".to_string(), + id: None, + result: None, + error: None, + }); + } + + let id = request.id.unwrap_or(0); let (tx, rx) = oneshot::channel(); // Register the pending response handler before writing the request, // so we don't miss a fast response from the child. { let mut pending = self.pending.lock().await; - pending.insert(request.id.unwrap_or(0), tx); + pending.insert(id, tx); } // Write the request to stdin. @@ -133,7 +147,7 @@ impl McpTransport for StdioMcpTransport { if let Err(e) = write_jsonrpc_line(&mut *stdin, request).await { // Remove the pending entry on write failure. let mut pending = self.pending.lock().await; - pending.remove(&request.id.unwrap_or(0)); + pending.remove(&id); return Err(e); } } @@ -145,7 +159,7 @@ impl McpTransport for StdioMcpTransport { Ok(Err(_)) => { // Sender was dropped (reader task ended). Clean up pending entry. let mut pending = self.pending.lock().await; - pending.remove(&request.id.unwrap_or(0)); + pending.remove(&id); Err(ToolError::ExternalService(format!( "[{}] MCP server closed connection before responding to request {:?}", self.server_name, request.id @@ -154,7 +168,7 @@ impl McpTransport for StdioMcpTransport { Err(_) => { // Timeout: remove the pending entry. let mut pending = self.pending.lock().await; - pending.remove(&request.id.unwrap_or(0)); + pending.remove(&id); Err(ToolError::ExternalService(format!( "[{}] Timeout waiting for response to request {:?} after {:?}", self.server_name, request.id, timeout diff --git a/src/tools/mcp/unix_transport.rs b/src/tools/mcp/unix_transport.rs index bf5865c6..e394d034 100644 --- a/src/tools/mcp/unix_transport.rs +++ b/src/tools/mcp/unix_transport.rs @@ -91,13 +91,27 @@ impl McpTransport for UnixMcpTransport { request: &McpRequest, _headers: &HashMap, ) -> Result { + // JSON-RPC notifications (no id) are fire-and-forget: the server + // will not send a response, so we must not wait for one. + if request.id.is_none() { + let mut writer = self.writer.lock().await; + write_jsonrpc_line(&mut *writer, request).await?; + return Ok(McpResponse { + jsonrpc: "2.0".to_string(), + id: None, + result: None, + error: None, + }); + } + + let id = request.id.unwrap_or(0); let (tx, rx) = oneshot::channel(); // Register the pending response handler before writing the request, // so we don't miss a fast response from the server. { let mut pending = self.pending.lock().await; - pending.insert(request.id.unwrap_or(0), tx); + pending.insert(id, tx); } // Write the request to the socket. @@ -106,7 +120,7 @@ impl McpTransport for UnixMcpTransport { if let Err(e) = write_jsonrpc_line(&mut *writer, request).await { // Remove the pending entry on write failure. let mut pending = self.pending.lock().await; - pending.remove(&request.id.unwrap_or(0)); + pending.remove(&id); return Err(e); } } @@ -118,7 +132,7 @@ impl McpTransport for UnixMcpTransport { Ok(Err(_)) => { // Sender was dropped (reader task ended). Clean up pending entry. let mut pending = self.pending.lock().await; - pending.remove(&request.id.unwrap_or(0)); + pending.remove(&id); Err(ToolError::ExternalService(format!( "[{}] MCP server closed connection before responding to request {:?}", self.server_name, request.id @@ -127,7 +141,7 @@ impl McpTransport for UnixMcpTransport { Err(_) => { // Timeout: remove the pending entry. let mut pending = self.pending.lock().await; - pending.remove(&request.id.unwrap_or(0)); + pending.remove(&id); Err(ToolError::ExternalService(format!( "[{}] Timeout waiting for response to request {:?} after {:?}", self.server_name, request.id, timeout