diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index d3bf17b8..7780ff80 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -209,6 +209,12 @@ impl McpClient { } } + /// Attach a session manager for Streamable HTTP session tracking. + pub fn with_session_manager(mut self, session_manager: Arc) -> Self { + self.session_manager = Some(session_manager); + self + } + /// Get the server name. pub fn server_name(&self) -> &str { &self.server_name @@ -219,6 +225,11 @@ impl McpClient { &self.server_url } + /// Whether this client has a session manager attached. + pub fn has_session_manager(&self) -> bool { + self.session_manager.is_some() + } + /// Get the next request ID. fn next_request_id(&self) -> u64 { self.next_id.fetch_add(1, Ordering::SeqCst) @@ -716,6 +727,17 @@ mod tests { assert!(client.session_manager.is_none()); } + #[test] + fn test_with_session_manager() { + let client = McpClient::new("http://localhost:8080"); + assert!(!client.has_session_manager()); + + let session_manager = Arc::new(McpSessionManager::new()); + let client = client.with_session_manager(session_manager); + + assert!(client.has_session_manager()); + } + #[test] fn test_next_request_id_monotonically_increasing() { let client = McpClient::new("http://localhost:1234"); diff --git a/src/tools/mcp/factory.rs b/src/tools/mcp/factory.rs index b5acb3f9..1cc714bc 100644 --- a/src/tools/mcp/factory.rs +++ b/src/tools/mcp/factory.rs @@ -88,11 +88,40 @@ pub async fn create_client_from_config( user_id, )) } else { - Ok(McpClient::new_with_config(server)) + Ok(McpClient::new_with_config(server) + .with_session_manager(Arc::clone(session_manager))) } } else { - Ok(McpClient::new_with_config(server)) + Ok(McpClient::new_with_config(server) + .with_session_manager(Arc::clone(session_manager))) } } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_factory_non_oauth_http_has_session_manager() { + let server = McpServerConfig::new("test-server", "http://localhost:9999"); + let session_manager = Arc::new(McpSessionManager::new()); + let process_manager = Arc::new(McpProcessManager::new()); + + let client = create_client_from_config( + server, + &session_manager, + &process_manager, + None, + "test-user", + ) + .await + .expect("factory should succeed for HTTP config"); + + assert!( + client.has_session_manager(), + "non-OAuth HTTP clients must carry a session manager" + ); + } +}