Add timeout middleware for clients

crates/goose/src/agents/capabilities.rs

@@ -1,15 +1,17 @@
 use chrono::{DateTime, TimeZone, Utc};
+use mcp_client::McpService;
 use rust_decimal_macros::dec;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::LazyLock;
+use std::time::Duration;
 use tokio::sync::Mutex;
 use tracing::{debug, instrument};
 
 use super::system::{SystemConfig, SystemError, SystemInfo, SystemResult};
 use crate::prompt_template::load_prompt_file;
 use crate::providers::base::{Provider, ProviderUsage};
-use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient};
+use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
 use mcp_client::transport::{SseTransport, StdioTransport, Transport};
 use mcp_core::{Content, Tool, ToolCall, ToolError, ToolResult};
 
@@ -20,7 +22,7 @@ static DEFAULT_TIMESTAMP: LazyLock<DateTime<Utc>> =
 
 /// Manages MCP clients and their interactions
 pub struct Capabilities {
-    clients: HashMap<String, Arc<Mutex<McpClient>>>,
+    clients: HashMap<String, Arc<Mutex<Box<dyn McpClientTrait>>>>,
     instructions: HashMap<String, String>,
     provider: Box<dyn Provider>,
     provider_usage: Mutex<Vec<ProviderUsage>>,
@@ -87,14 +89,18 @@ impl Capabilities {
     /// Add a new MCP system based on the provided client type
     // TODO IMPORTANT need to ensure this times out if the system command is broken!
     pub async fn add_system(&mut self, config: SystemConfig) -> SystemResult<()> {
-        let mut client: McpClient = match config {
+        let mut client: Box<dyn McpClientTrait> = match config {
             SystemConfig::Sse { ref uri } => {
                 let transport = SseTransport::new(uri);
-                McpClient::new(transport.start().await?)
+                let handle = transport.start().await?;
+                let service = McpService::with_timeout(handle, Duration::from_secs(10));
+                Box::new(McpClient::new(service))
             }
             SystemConfig::Stdio { ref cmd, ref args } => {
                 let transport = StdioTransport::new(cmd, args.to_vec());
-                McpClient::new(transport.start().await?)
+                let handle = transport.start().await?;
+                let service = McpService::with_timeout(handle, Duration::from_secs(10));
+                Box::new(McpClient::new(service))
             }
         };
 
@@ -258,7 +264,10 @@ impl Capabilities {
     }
 
     /// Find and return a reference to the appropriate client for a tool call
-    fn get_client_for_tool(&self, prefixed_name: &str) -> Option<Arc<Mutex<McpClient>>> {
+    fn get_client_for_tool(
+        &self,
+        prefixed_name: &str,
+    ) -> Option<Arc<Mutex<Box<dyn McpClientTrait>>>> {
         prefixed_name
             .split_once("__")
             .and_then(|(client_name, _)| self.clients.get(client_name))

crates/mcp-client/examples/clients.rs

@@ -1,6 +1,7 @@
 use mcp_client::{
-    client::{ClientCapabilities, ClientInfo, McpClient},
+    client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait},
     transport::{SseTransport, StdioTransport, Transport},
+    McpService,
 };
 use rand::Rng;
 use rand::SeedableRng;
@@ -17,13 +18,24 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         )
         .init();
 
-    // Create two separate clients with stdio transport
-    let client1 = create_stdio_client("client1", "1.0.0").await?;
-    let client2 = create_stdio_client("client2", "1.0.0").await?;
-    let client3 = create_sse_client("client3", "1.0.0").await?;
+    let transport1 = StdioTransport::new("uvx", vec!["mcp-server-git".to_string()]);
+    let handle1 = transport1.start().await?;
+    let service1 = McpService::with_timeout(handle1, Duration::from_secs(30));
+    let client1 = McpClient::new(service1);
+
+    let transport2 = StdioTransport::new("uvx", vec!["mcp-server-git".to_string()]);
+    let handle2 = transport2.start().await?;
+    let service2 = McpService::with_timeout(handle2, Duration::from_secs(30));
+    let client2 = McpClient::new(service2);
+
+    let transport3 = SseTransport::new("http://localhost:8000/sse");
+    let handle3 = transport3.start().await?;
+    let service3 = McpService::with_timeout(handle3, Duration::from_secs(10));
+    let client3 = McpClient::new(service3);
 
     // Initialize both clients
-    let mut clients = vec![client1, client2, client3];
+    let mut clients: Vec<Box<dyn McpClientTrait>> =
+        vec![Box::new(client1), Box::new(client2), Box::new(client3)];
 
     // Initialize all clients
     for (i, client) in clients.iter_mut().enumerate() {
@@ -117,19 +129,3 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     Ok(())
 }
-
-async fn create_stdio_client(
-    _name: &str,
-    _version: &str,
-) -> Result<McpClient, Box<dyn std::error::Error>> {
-    let transport = StdioTransport::new("uvx", vec!["mcp-server-git".to_string()]);
-    Ok(McpClient::new(transport.start().await?))
-}
-
-async fn create_sse_client(
-    _name: &str,
-    _version: &str,
-) -> Result<McpClient, Box<dyn std::error::Error>> {
-    let transport = SseTransport::new("http://localhost:8000/sse");
-    Ok(McpClient::new(transport.start().await?))
-}

crates/mcp-client/examples/sse.rs

@@ -1,6 +1,7 @@
 use anyhow::Result;
-use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient};
+use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
 use mcp_client::transport::{SseTransport, Transport};
+use mcp_client::McpService;
 use std::time::Duration;
 use tracing_subscriber::EnvFilter;
 
@@ -21,8 +22,11 @@ async fn main() -> Result<()> {
     // Start transport
     let handle = transport.start().await?;
 
+    // Create the service with timeout middleware
+    let service = McpService::with_timeout(handle, Duration::from_secs(3));
+
     // Create client
-    let mut client = McpClient::new(handle);
+    let mut client = McpClient::new(service);
     println!("Client created\n");
 
     // Initialize

crates/mcp-client/examples/stdio.rs

@@ -1,6 +1,9 @@
 use anyhow::Result;
-use mcp_client::client::{ClientCapabilities, ClientInfo, Error as ClientError, McpClient};
-use mcp_client::transport::{StdioTransport, Transport};
+use mcp_client::{
+    ClientCapabilities, ClientInfo, Error as ClientError, McpClient, McpClientTrait, McpService,
+    StdioTransport, Transport,
+};
+use std::time::Duration;
 use tracing_subscriber::EnvFilter;
 
 #[tokio::main]
@@ -20,8 +23,11 @@ async fn main() -> Result<(), ClientError> {
     // 2) Start the transport to get a handle
     let transport_handle = transport.start().await?;
 
-    // 3) Create the client
-    let mut client = McpClient::new(transport_handle);
+    // 3) Create the service with timeout middleware
+    let service = McpService::with_timeout(transport_handle, Duration::from_secs(10));
+
+    // 4) Create the client with the middleware-wrapped service
+    let mut client = McpClient::new(service);
 
     // Initialize
     let server_info = client

crates/mcp-client/examples/stdio_integration.rs

@@ -1,8 +1,13 @@
+use std::time::Duration;
+
 // This example shows how to use the mcp-client crate to interact with a server that has a simple counter tool.
 // The server is started by running `cargo run -p mcp-server` in the root of the mcp-server crate.
 use anyhow::Result;
-use mcp_client::client::{ClientCapabilities, ClientInfo, Error as ClientError, McpClient};
+use mcp_client::client::{
+    ClientCapabilities, ClientInfo, Error as ClientError, McpClient, McpClientTrait,
+};
 use mcp_client::transport::{StdioTransport, Transport};
+use mcp_client::McpService;
 use tracing_subscriber::EnvFilter;
 
 #[tokio::main]
@@ -28,8 +33,11 @@ async fn main() -> Result<(), ClientError> {
     // Start the transport to get a handle
     let transport_handle = transport.start().await.unwrap();
 
+    // Create the service with timeout middleware
+    let service = McpService::with_timeout(transport_handle, Duration::from_secs(10));
+
     // Create client
-    let mut client = McpClient::new(transport_handle);
+    let mut client = McpClient::new(service);
 
     // Initialize
     let server_info = client