Fix MCP timeouts to be recoverable

crates/goose-cli/src/commands/mcp.rs

@@ -5,6 +5,7 @@ use goose_mcp::{
 use mcp_server::router::RouterService;
 use mcp_server::{BoundedService, ByteTransport, Server};
 use tokio::io::{stdin, stdout};
+use tokio::time::{timeout, Duration};
 
 pub async fn run_server(name: &str) -> Result<()> {
     // Initialize logging
@@ -29,5 +30,15 @@ pub async fn run_server(name: &str) -> Result<()> {
     let transport = ByteTransport::new(stdin(), stdout());
 
     tracing::info!("Server initialized and ready to handle requests");
-    Ok(server.run(transport).await?)
+
+    // Add timeout and cancellation handling
+    let server_future = server.run(transport);
+    match timeout(Duration::from_secs(30), server_future).await {
+        Ok(result) => result,
+        Err(_) => {
+            tracing::warn!("Timeout occurred while running the server");
+            // Handle cancellation logic here if needed
+            Ok(())
+        }
+    }
 }

crates/goose-cli/src/commands/session.rs

@@ -11,6 +11,7 @@ use goose::providers::create;
 use std::path::Path;
 
 use mcp_client::transport::Error as McpClientError;
+use tokio::time::{timeout, Duration};
 
 pub async fn build_session(
     name: Option<String>,
@@ -176,3 +177,58 @@ fn display_session_info(resume: bool, provider: &str, model: &str, session_file:
         style(session_file.display()).dim().cyan(),
     );
 }
+
+// Add timeout mechanism using `tokio::time::timeout` in `agent_process_messages` function
+impl<'a> Session<'a> {
+    async fn agent_process_messages(&mut self) {
+        let mut stream = match self.agent.reply(&self.messages).await {
+            Ok(stream) => stream,
+            Err(e) => {
+                eprintln!("Error starting reply stream: {}", e);
+                return;
+            }
+        };
+
+        loop {
+            tokio::select! {
+                response = timeout(Duration::from_secs(30), stream.next()) => {
+                    match response {
+                        Ok(Some(Ok(message))) => {
+                            self.messages.push(message.clone());
+                            persist_messages(&self.session_file, &self.messages).unwrap_or_else(|e| eprintln!("Failed to persist messages: {}", e));
+                            self.prompt.hide_busy();
+                            self.prompt.render(Box::new(message.clone()));
+                            self.prompt.show_busy();
+                        }
+                        Ok(Some(Err(e))) => {
+                            eprintln!("Error: {}", e);
+                            drop(stream);
+                            self.rewind_messages();
+                            self.prompt.render(raw_message(r#"
+The error above was an exception we were not able to handle.\n\n
+These errors are often related to connection or authentication\n
+We've removed the conversation up to the most recent user message
+ - depending on the error you may be able to continue"#));
+                            break;
+                        }
+                        Ok(None) => break,
+                        Err(_) => {
+                            eprintln!("Timeout occurred while waiting for response");
+                            drop(stream);
+                            self.handle_interrupted_messages();
+                            break;
+                        }
+                    }
+                }
+                _ = tokio::signal::ctrl_c() => {
+                    // Kill any running processes when the client disconnects
+                    // TODO is this used? I suspect post MCP this is on the server instead
+                    // goose::process_store::kill_processes();
+                    drop(stream);
+                    self.handle_interrupted_messages();
+                    break;
+                }
+            }
+        }
+    }
+}

crates/goose-server/src/commands/mcp.rs

@@ -5,6 +5,7 @@ use goose_mcp::{
 use mcp_server::router::RouterService;
 use mcp_server::{BoundedService, ByteTransport, Server};
 use tokio::io::{stdin, stdout};
+use tokio::time::{timeout, Duration};
 
 pub async fn run(name: &str) -> Result<()> {
     // Initialize logging
@@ -28,5 +29,15 @@ pub async fn run(name: &str) -> Result<()> {
     let transport = ByteTransport::new(stdin(), stdout());
 
     tracing::info!("Server initialized and ready to handle requests");
-    Ok(server.run(transport).await?)
+
+    // Add timeout and cancellation handling
+    let server_future = server.run(transport);
+    match timeout(Duration::from_secs(30), server_future).await {
+        Ok(result) => result,
+        Err(_) => {
+            tracing::warn!("Timeout occurred while running the server");
+            // Handle cancellation logic here if needed
+            Ok(())
+        }
+    }
 }

crates/mcp-client/src/service.rs

@@ -18,6 +18,16 @@ impl<T: TransportHandle> McpService<T> {
             inner: Arc::new(transport),
         }
     }
+
+    pub async fn cancel(&self, request_id: &str) -> Result<(), Error> {
+        let cancel_message = JsonRpcMessage::Notification(JsonRpcNotification {
+            jsonrpc: "2.0".to_string(),
+            method: "cancel".to_string(),
+            params: Some(json!({ "request_id": request_id })),
+        });
+
+        self.inner.send(cancel_message).await.map(|_| ())
+    }
 }
 
 impl<T> Service<JsonRpcMessage> for McpService<T>