Interrupt running shell tool commands

crates/goose-server/src/routes/reply.rs

@@ -22,8 +22,10 @@ use std::{
     convert::Infallible,
     pin::Pin,
     task::{Context, Poll},
+    time::Duration,
 };
 use tokio::sync::mpsc;
+use tokio::time::timeout;
 use tokio_stream::wrappers::ReceiverStream;
 
 // Types matching the incoming JSON structure
@@ -139,7 +141,6 @@ fn convert_messages(incoming: Vec<IncomingMessage>) -> Vec<Message> {
 struct ProtocolFormatter;
 
 impl ProtocolFormatter {
-
     fn format_text(text: &str) -> String {
         let encoded_text = serde_json::to_string(text).unwrap_or_else(|_| String::new());
         format!("0:{}\n", encoded_text)
@@ -175,6 +176,10 @@ impl ProtocolFormatter {
         });
         format!("d:{}\n", finish)
     }
+
+    fn heartbeat() -> String {
+        "2:[]\n".to_string()
+    }
 }
 
 async fn stream_message(
@@ -294,18 +299,40 @@ async fn handler(
             }
         };
 
-        while let Some(response) = stream.next().await {
-            match response {
-                Ok(message) => {
-                    if let Err(e) = stream_message(message, &tx).await {
-                        tracing::error!("Error sending message through channel: {}", e);
-                        break;
+        loop {
+            tokio::select! {
+                response = timeout(Duration::from_millis(500), stream.next()) => {
+                    match response {
+                        Ok(Some(Ok(message))) => {
+                            if let Err(e) = stream_message(message, &tx).await {
+                                tracing::error!("Error sending message through channel: {}", e);
+                                break;
+                            }
+                        }
+                        Ok(Some(Err(e))) => {
+                            tracing::error!("Error processing message: {}", e);
+                            break;
+                        }
+                        Ok(None) => {
+                            break;
+                        }
+                        Err(_) => { // Heartbeat, used to detect disconnected clients and then end running tools.
+                            if let Err(e) = tx.try_send(ProtocolFormatter::heartbeat()) {
+                                match e {
+                                    mpsc::error::TrySendError::Closed(_) => {
+                                        // Client has disconnected, end the stream and close running tools (works by ending this process).
+                                        break;
+                                    }
+                                    mpsc::error::TrySendError::Full(_) => {
+                                        tracing::warn!("Error sending heartbeat message through channel: {}", e);
+                                        continue;
+                                    }
+                                }
+                            }
+                            continue;
+                        }
                     }
                 }
-                Err(e) => {
-                    tracing::error!("Error processing message: {}", e);
-                    break;
-                }
             }
         }
 

crates/goose/src/developer.rs

@@ -8,8 +8,8 @@ use serde_json::{json, Value};
 use std::collections::{HashMap, HashSet};
 use std::io::Cursor;
 use std::path::{Path, PathBuf};
-use std::process::Command;
 use std::sync::Mutex;
+use tokio::process::Command;
 use xcap::Monitor;
 
 use crate::errors::{AgentError, AgentResult};
@@ -192,9 +192,11 @@ impl DeveloperSystem {
 
         // Execute the command
         let output = Command::new("bash")
+            .kill_on_drop(true) // Critical so that the command is killed when the agent.reply stream is interrupted.
             .arg("-c")
             .arg(cmd_with_redirect)
             .output()
+            .await
             .map_err(|e| AgentError::ExecutionError(e.to_string()))?;
 
         let output_str = format!(