gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp): stdio/unix transports skip initialize handshake (#890) (#935)

Browse Source 
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
This commit is contained in:
Reid
2026-03-11 16:46:40 -07:00
committed by GitHub
parent 6321bb4688
commit c8cac0925d
4 changed files with 138 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cli/mcp.rs
+14 -4
View File
@@ -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
src/tools/mcp/client.rs
+88 -10
View File
@@ -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<String, String>,
/// 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<InitializeResult, ToolError> {
// 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<CallToolResult, ToolError> {
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!({
src/tools/mcp/stdio_transport.rs
+18 -4
View File
@@ -118,13 +118,27 @@ impl McpTransport for StdioMcpTransport {
request: &McpRequest,
_headers: &HashMap<String, String>,
) -> Result<McpResponse, ToolError> {
// 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
src/tools/mcp/unix_transport.rs
+18 -4
View File
@@ -91,13 +91,27 @@ impl McpTransport for UnixMcpTransport {
request: &McpRequest,
_headers: &HashMap<String, String>,
) -> Result<McpResponse, ToolError> {
// 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