Detect client disconnects and cancel tool calls (#3782)

crates/goose-cli/src/scenario_tests/mock_client.rs

@@ -13,6 +13,7 @@ use rmcp::{
 use serde_json::Value;
 use std::collections::HashMap;
 use tokio::sync::mpsc::{self, Receiver};
+use tokio_util::sync::CancellationToken;
 
 pub struct MockClient {
     tools: HashMap<String, Tool>,
@@ -43,6 +44,7 @@ impl McpClientTrait for MockClient {
     async fn list_resources(
         &self,
         _next_cursor: Option<String>,
+        _cancel_token: CancellationToken,
     ) -> Result<ListResourcesResult, Error> {
         Ok(ListResourcesResult {
             resources: vec![],
@@ -54,11 +56,19 @@ impl McpClientTrait for MockClient {
         todo!()
     }
 
-    async fn read_resource(&self, _uri: &str) -> Result<ReadResourceResult, Error> {
+    async fn read_resource(
+        &self,
+        _uri: &str,
+        _cancel_token: CancellationToken,
+    ) -> Result<ReadResourceResult, Error> {
         Err(Error::UnexpectedResponse)
     }
 
-    async fn list_tools(&self, _: Option<String>) -> Result<ListToolsResult, Error> {
+    async fn list_tools(
+        &self,
+        _: Option<String>,
+        _cancel_token: CancellationToken,
+    ) -> Result<ListToolsResult, Error> {
         let rmcp_tools: Vec<rmcp::model::Tool> = self
             .tools
             .values()
@@ -77,7 +87,12 @@ impl McpClientTrait for MockClient {
         })
     }
 
-    async fn call_tool(&self, name: &str, arguments: Value) -> Result<CallToolResult, Error> {
+    async fn call_tool(
+        &self,
+        name: &str,
+        arguments: Value,
+        _cancel_token: CancellationToken,
+    ) -> Result<CallToolResult, Error> {
         if let Some(handler) = self.handlers.get(name) {
             match handler(&arguments) {
                 Ok(content) => Ok(CallToolResult {
@@ -91,14 +106,23 @@ impl McpClientTrait for MockClient {
         }
     }
 
-    async fn list_prompts(&self, _next_cursor: Option<String>) -> Result<ListPromptsResult, Error> {
+    async fn list_prompts(
+        &self,
+        _next_cursor: Option<String>,
+        _cancel_token: CancellationToken,
+    ) -> Result<ListPromptsResult, Error> {
         Ok(ListPromptsResult {
             prompts: vec![],
             next_cursor: None,
         })
     }
 
-    async fn get_prompt(&self, _name: &str, _arguments: Value) -> Result<GetPromptResult, Error> {
+    async fn get_prompt(
+        &self,
+        _name: &str,
+        _arguments: Value,
+        _cancel_token: CancellationToken,
+    ) -> Result<GetPromptResult, Error> {
         Err(Error::UnexpectedResponse)
     }
 

crates/goose-cli/src/scenario_tests/scenario_runner.rs

@@ -12,6 +12,7 @@ use goose::providers::{create, testprovider::TestProvider};
 use std::collections::{HashMap, HashSet};
 use std::path::Path;
 use std::sync::Arc;
+use tokio_util::sync::CancellationToken;
 
 pub const SCENARIO_TESTS_DIR: &str = "src/scenario_tests";
 
@@ -205,7 +206,10 @@ where
 
     let mut error = None;
     for message in &messages {
-        if let Err(e) = session.process_message(message.clone()).await {
+        if let Err(e) = session
+            .process_message(message.clone(), CancellationToken::default())
+            .await
+        {
             error = Some(e.to_string());
             break;
         }

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

@@ -364,7 +364,12 @@ impl Session {
     }
 
     /// Process a single message and get the response
-    pub(crate) async fn process_message(&mut self, message: Message) -> Result<()> {
+    pub(crate) async fn process_message(
+        &mut self,
+        message: Message,
+        cancel_token: CancellationToken,
+    ) -> Result<()> {
+        let cancel_token = cancel_token.clone();
         let message_text = message.as_concat_text();
 
         self.push_message(message);
@@ -405,7 +410,7 @@ impl Session {
             );
         }
 
-        self.process_agent_response(false).await?;
+        self.process_agent_response(false, cancel_token).await?;
         Ok(())
     }
 
@@ -414,7 +419,8 @@ impl Session {
         // Process initial message if provided
         if let Some(prompt) = prompt {
             let msg = Message::user().with_text(&prompt);
-            self.process_message(msg).await?;
+            self.process_message(msg, CancellationToken::default())
+                .await?;
         }
 
         // Initialize the completion cache
@@ -514,7 +520,8 @@ impl Session {
                             }
 
                             output::show_thinking();
-                            self.process_agent_response(true).await?;
+                            self.process_agent_response(true, CancellationToken::default())
+                                .await?;
                             output::hide_thinking();
                         }
                         RunMode::Plan => {
@@ -814,7 +821,8 @@ impl Session {
                     self.push_message(plan_message);
                     // act on the plan
                     output::show_thinking();
-                    self.process_agent_response(true).await?;
+                    self.process_agent_response(true, CancellationToken::default())
+                        .await?;
                     output::hide_thinking();
 
                     // Reset run & goose mode
@@ -842,12 +850,15 @@ impl Session {
     /// Process a single message and exit
     pub async fn headless(&mut self, prompt: String) -> Result<()> {
         let message = Message::user().with_text(&prompt);
-        self.process_message(message).await
+        self.process_message(message, CancellationToken::default())
+            .await
     }
 
-    async fn process_agent_response(&mut self, interactive: bool) -> Result<()> {
-        // Messages will be auto-compacted in agent.reply() if needed
-        let cancel_token = CancellationToken::new();
+    async fn process_agent_response(
+        &mut self,
+        interactive: bool,
+        cancel_token: CancellationToken,
+    ) -> Result<()> {
         let cancel_token_clone = cancel_token.clone();
 
         let session_config = self.session_file.as_ref().map(|s| {
@@ -1511,7 +1522,8 @@ impl Session {
 
                     if valid {
                         output::show_thinking();
-                        self.process_agent_response(true).await?;
+                        self.process_agent_response(true, CancellationToken::default())
+                            .await?;
                         output::hide_thinking();
                     }
                 }