Bound code-mode execution with timeout and cancellation (#10214)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

crates/goose/src/agents/platform_extensions/code_execution.rs

@@ -23,6 +23,7 @@ use std::future::Future;
 use std::hash::{Hash, Hasher};
 use std::pin::Pin;
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::sync::RwLock;
 use tokio_util::sync::CancellationToken;
 
@@ -145,6 +146,7 @@ impl CodeExecutionClient {
         &self,
         ctx: &ToolCallContext,
         code_mode: &CodeMode,
+        cancellation_token: CancellationToken,
     ) -> Result<PctxRegistry, String> {
         let manager = self
             .context
@@ -163,7 +165,12 @@ impl CodeExecutionClient {
                     .unwrap_or_default(),
                 &cfg.name
             );
-            let callback = create_tool_callback(ctx.clone(), full_name, manager.clone());
+            let callback = create_tool_callback(
+                ctx.clone(),
+                full_name,
+                manager.clone(),
+                cancellation_token.clone(),
+            );
             registry
                 .add_callback(&cfg.id(), callback)
                 .map_err(|e| format!("Failed to register callback: {e}"))?;
@@ -203,6 +210,7 @@ impl CodeExecutionClient {
         &self,
         session_id: &str,
         arguments: Option<JsonObject>,
+        cancellation_token: CancellationToken,
     ) -> Result<Vec<Content>, String> {
         let input: ExecuteBashInput = arguments
             .map(|args| serde_json::from_value(Value::Object(args)))
@@ -212,23 +220,19 @@ impl CodeExecutionClient {
         let command = input.command;
         let code_mode = self.get_code_mode(session_id).await?;
 
-        // Deno runtime is not Send, so we need to run it in a blocking task
-        // with its own tokio runtime
-        let output = tokio::task::spawn_blocking(move || {
-            let rt = tokio::runtime::Builder::new_current_thread()
-                .enable_all()
-                .build()
-                .map_err(|e| format!("Failed to create runtime: {e}"))?;
-
-            rt.block_on(async move {
+        let dispatch_token = cancellation_token.child_token();
+        let output = run_in_deno_runtime(
+            execution_timeout(),
+            cancellation_token,
+            dispatch_token,
+            move || async move {
                 code_mode
                     .execute_bash(&command)
                     .await
-                    .map_err(|e| format!("Typescript execution error: {e}"))
-            })
-        })
-        .await
-        .map_err(|e| format!("Typescript execution task failed: {e}"))??;
+                    .map_err(|e| format!("Bash execution error: {e}"))
+            },
+        )
+        .await?;
 
         Ok(vec![Content::text(output.markdown())])
     }
@@ -238,6 +242,7 @@ impl CodeExecutionClient {
         &self,
         ctx: &ToolCallContext,
         arguments: Option<JsonObject>,
+        cancellation_token: CancellationToken,
     ) -> Result<Vec<Content>, String> {
         let args: ExecuteWithToolGraph = arguments
             .map(|args| serde_json::from_value(Value::Object(args)))
@@ -247,41 +252,110 @@ impl CodeExecutionClient {
 
         let session_id = &ctx.session_id;
         let code_mode = self.get_code_mode(session_id).await?;
-        let registry = self.build_callback_registry(ctx, &code_mode)?;
+        let dispatch_token = cancellation_token.child_token();
+        let registry = self.build_callback_registry(ctx, &code_mode, dispatch_token.clone())?;
         let code = args.input.code.clone();
         let disclosure = self.disclosure;
 
-        // Deno runtime is not Send, so we need to run it in a blocking task
-        // with its own tokio runtime
-        let output = tokio::task::spawn_blocking(move || {
-            let rt = tokio::runtime::Builder::new_current_thread()
-                .enable_all()
-                .build()
-                .map_err(|e| format!("Failed to create runtime: {e}"))?;
-
-            rt.block_on(async move {
+        let output = run_in_deno_runtime(
+            execution_timeout(),
+            cancellation_token,
+            dispatch_token,
+            move || async move {
                 code_mode
                     .execute_typescript(&code, disclosure, Some(registry))
                     .await
                     .map_err(|e| format!("Typescript execution error: {e}"))
-            })
-        })
-        .await
-        .map_err(|e| format!("Typescript execution task failed: {e}"))??;
+            },
+        )
+        .await?;
 
         Ok(vec![Content::text(output.markdown())])
     }
 }
 
+fn execution_timeout() -> Duration {
+    let secs = crate::config::Config::global()
+        .get_goose_default_extension_timeout()
+        .unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT);
+    Duration::from_secs(secs)
+}
+
+/// Deno runtime is not Send, so execution runs in a blocking task with its
+/// own tokio runtime. pctx serializes all executions behind a process-wide
+/// V8 mutex, so a hung script would wedge code execution for every session:
+/// bound the wait with the extension timeout and honor cancellation.
+///
+/// `dispatch_token` is the child token shared with the callback dispatches that
+/// a script makes back into Goose tools. When execution is abandoned (timeout
+/// or cancellation), the token is cancelled so an in-flight nested tool call
+/// (e.g. a long `developer.shell` command) is told to stop instead of running
+/// on in the background.
+/// Grace period for nested tool calls to observe a dispatched cancellation
+/// signal and clean up (e.g. kill child processes) before the task future is
+/// abandoned.
+const DISPATCH_DRAIN_TIMEOUT: Duration = Duration::from_millis(500);
+
+async fn run_in_deno_runtime<T, F, Fut>(
+    timeout: Duration,
+    cancellation_token: CancellationToken,
+    dispatch_token: CancellationToken,
+    task: F,
+) -> Result<T, String>
+where
+    F: FnOnce() -> Fut + Send + 'static,
+    Fut: Future<Output = Result<T, String>>,
+    T: Send + 'static,
+{
+    tokio::task::spawn_blocking(move || {
+        let rt = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .map_err(|e| format!("Failed to create runtime: {e}"))?;
+
+        rt.block_on(async move {
+            let task_future = task();
+            tokio::pin!(task_future);
+
+            tokio::select! {
+                _ = cancellation_token.cancelled() => {
+                    dispatch_token.cancel();
+                    let _ = tokio::time::timeout(
+                        DISPATCH_DRAIN_TIMEOUT,
+                        &mut task_future,
+                    ).await;
+                    Err("Execution cancelled".to_string())
+                }
+                _ = tokio::time::sleep(timeout) => {
+                    dispatch_token.cancel();
+                    let _ = tokio::time::timeout(
+                        DISPATCH_DRAIN_TIMEOUT,
+                        &mut task_future,
+                    ).await;
+                    Err(format!(
+                        "Execution timed out after {} seconds",
+                        timeout.as_secs()
+                    ))
+                }
+                result = &mut task_future => result,
+            }
+        })
+    })
+    .await
+    .map_err(|e| format!("Execution task failed: {e}"))?
+}
+
 fn create_tool_callback(
     ctx: ToolCallContext,
     full_name: String,
     manager: Arc<crate::agents::ExtensionManager>,
+    cancellation_token: CancellationToken,
 ) -> CallbackFn {
     Arc::new(move |args: Option<Value>| {
         let ctx = ctx.clone();
         let full_name = full_name.clone();
         let manager = manager.clone();
+        let cancellation_token = cancellation_token.clone();
         Box::pin(async move {
             let tool_call = {
                 let mut params = CallToolRequestParams::new(full_name);
@@ -291,7 +365,7 @@ fn create_tool_callback(
                 params
             };
             match manager
-                .dispatch_tool_call(&ctx, tool_call, CancellationToken::new())
+                .dispatch_tool_call(&ctx, tool_call, cancellation_token)
                 .await
             {
                 Ok(dispatch_result) => match dispatch_result.result.await {
@@ -447,7 +521,7 @@ impl McpClientTrait for CodeExecutionClient {
         ctx: &ToolCallContext,
         name: &str,
         arguments: Option<JsonObject>,
-        _cancellation_token: CancellationToken,
+        cancellation_token: CancellationToken,
     ) -> Result<CallToolResult, Error> {
         let session_id = &ctx.session_id;
         let result = match name {
@@ -456,8 +530,14 @@ impl McpClientTrait for CodeExecutionClient {
                 self.handle_get_function_details(session_id, arguments)
                     .await
             }
-            "execute_bash" => self.handle_execute_bash(session_id, arguments).await,
-            "execute_typescript" => self.handle_execute_typescript(ctx, arguments).await,
+            "execute_bash" => {
+                self.handle_execute_bash(session_id, arguments, cancellation_token)
+                    .await
+            }
+            "execute_typescript" => {
+                self.handle_execute_typescript(ctx, arguments, cancellation_token)
+                    .await
+            }
             _ => Err(format!("Unknown tool: {name}")),
         };
 
@@ -556,6 +636,148 @@ impl CodeModeState {
 mod tests {
     use super::*;
 
+    #[tokio::test]
+    async fn run_in_deno_runtime_times_out_on_hung_execution() {
+        let result: Result<(), String> = run_in_deno_runtime(
+            Duration::from_millis(50),
+            CancellationToken::new(),
+            CancellationToken::new(),
+            std::future::pending,
+        )
+        .await;
+
+        assert!(result.unwrap_err().contains("timed out"));
+    }
+
+    #[tokio::test]
+    async fn run_in_deno_runtime_honors_cancellation() {
+        let token = CancellationToken::new();
+        token.cancel();
+
+        let result: Result<(), String> = run_in_deno_runtime(
+            Duration::from_secs(60),
+            token,
+            CancellationToken::new(),
+            std::future::pending,
+        )
+        .await;
+
+        assert_eq!(result.unwrap_err(), "Execution cancelled");
+    }
+
+    #[tokio::test]
+    async fn run_in_deno_runtime_cancels_dispatch_token_when_abandoned() {
+        // On timeout, an in-flight nested tool call (via the dispatch token)
+        // must be told to stop rather than left running in the background.
+        let dispatch_token = CancellationToken::new();
+        let result: Result<(), String> = run_in_deno_runtime(
+            Duration::from_millis(50),
+            CancellationToken::new(),
+            dispatch_token.clone(),
+            std::future::pending,
+        )
+        .await;
+        assert!(result.unwrap_err().contains("timed out"));
+        assert!(dispatch_token.is_cancelled());
+
+        // On cancellation, the child dispatch token is likewise cancelled.
+        let outer = CancellationToken::new();
+        outer.cancel();
+        let dispatch_token = outer.child_token();
+        let result: Result<(), String> = run_in_deno_runtime(
+            Duration::from_secs(60),
+            outer,
+            dispatch_token.clone(),
+            std::future::pending,
+        )
+        .await;
+        assert_eq!(result.unwrap_err(), "Execution cancelled");
+        assert!(dispatch_token.is_cancelled());
+    }
+
+    #[tokio::test]
+    async fn run_in_deno_runtime_drains_task_on_timeout() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+
+        let dispatch_token = CancellationToken::new();
+        let observed = Arc::new(AtomicBool::new(false));
+        let task_token = dispatch_token.clone();
+        let task_observed = observed.clone();
+
+        let result: Result<(), String> = run_in_deno_runtime(
+            Duration::from_millis(50),
+            CancellationToken::new(),
+            dispatch_token.clone(),
+            move || async move {
+                task_token.cancelled().await;
+                task_observed.store(true, Ordering::SeqCst);
+                Ok(())
+            },
+        )
+        .await;
+
+        assert!(
+            result.unwrap_err().contains("timed out"),
+            "should report timeout"
+        );
+        assert!(
+            observed.load(Ordering::SeqCst),
+            "task should observe dispatch token cancellation before being dropped"
+        );
+    }
+
+    /// Exercises the real Deno/V8 stack: a script whose event loop never
+    /// resolves must time out instead of wedging forever, and a normal
+    /// script must run right after, proving pctx's process-wide V8 mutex
+    /// was released (i.e. one hung execution no longer blocks other sessions).
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn real_v8_hung_script_times_out_and_frees_the_runtime() {
+        let hung = CodeMode::default();
+        let hung_result = run_in_deno_runtime(
+            Duration::from_secs(2),
+            CancellationToken::new(),
+            CancellationToken::new(),
+            move || async move {
+                hung.execute_typescript(
+                    "async function run() { await new Promise(() => {}); }",
+                    ToolDisclosure::default(),
+                    None,
+                )
+                .await
+                .map_err(|e| format!("execution error: {e}"))
+            },
+        )
+        .await;
+        assert!(
+            hung_result.unwrap_err().contains("timed out"),
+            "hung script should time out"
+        );
+
+        let normal = CodeMode::default();
+        let normal_result = run_in_deno_runtime(
+            Duration::from_secs(60),
+            CancellationToken::new(),
+            CancellationToken::new(),
+            move || async move {
+                normal
+                    .execute_typescript(
+                        "async function run() { return 1 + 1; }",
+                        ToolDisclosure::default(),
+                        None,
+                    )
+                    .await
+                    .map_err(|e| format!("execution error: {e}"))
+            },
+        )
+        .await
+        .expect("normal script should run after a prior timeout");
+        assert!(
+            normal_result.success,
+            "normal script should succeed once the V8 mutex is released: {}",
+            normal_result.stderr
+        );
+    }
+
     #[test]
     fn catalog_moim_mentions_inspection_tools_without_function_names() {
         let moim = catalog_disclosure_moim(3);

crates/goose/src/agents/platform_extensions/developer/mod.rs

@@ -192,12 +192,15 @@ impl McpClientTrait for DeveloperClient {
         ctx: &ToolCallContext,
         name: &str,
         arguments: Option<JsonObject>,
-        _cancel_token: CancellationToken,
+        cancel_token: CancellationToken,
     ) -> Result<CallToolResult, Error> {
         let working_dir = ctx.working_dir.as_deref();
         match name {
             "shell" => match Self::parse_args::<ShellParams>(arguments) {
-                Ok(params) => Ok(self.shell_tool.shell_with_cwd(params, working_dir).await),
+                Ok(params) => Ok(self
+                    .shell_tool
+                    .shell_with_cwd(params, working_dir, cancel_token)
+                    .await),
                 Err(error) => Ok(ShellTool::error_result(&format!("Error: {error}"), None)),
             },
             "write" => match Self::parse_args::<FileWriteParams>(arguments) {

crates/goose/src/agents/platform_extensions/developer/shell.rs

@@ -18,6 +18,7 @@ use tokio::sync::OnceCell;
 #[cfg(not(windows))]
 use tokio::task::JoinHandle;
 use tokio_stream::{wrappers::SplitStream, StreamExt};
+use tokio_util::sync::CancellationToken;
 
 use crate::subprocess::SubprocessExt;
 
@@ -345,13 +346,15 @@ impl ShellTool {
     }
 
     pub async fn shell(&self, params: ShellParams) -> CallToolResult {
-        self.shell_with_cwd(params, None).await
+        self.shell_with_cwd(params, None, CancellationToken::new())
+            .await
     }
 
     pub async fn shell_with_cwd(
         &self,
         params: ShellParams,
         working_dir: Option<&std::path::Path>,
+        cancellation_token: CancellationToken,
     ) -> CallToolResult {
         if params.command.trim().is_empty() {
             return Self::error_result("Command cannot be empty.", None);
@@ -369,6 +372,7 @@ impl ShellTool {
             params.timeout_secs,
             working_dir,
             login_path_ref,
+            cancellation_token,
         )
         .await
         {
@@ -510,6 +514,7 @@ async fn run_command(
     timeout_secs: Option<u64>,
     working_dir: Option<&std::path::Path>,
     login_path: Option<&str>,
+    cancellation_token: CancellationToken,
 ) -> Result<ExecutionOutput, String> {
     let mut command = build_shell_command(command_line, working_dir, login_path);
 
@@ -536,23 +541,35 @@ async fn run_command(
 
     let mut timed_out = false;
     let exit_code = if let Some(timeout_secs) = timeout_secs.filter(|value| *value > 0) {
-        match tokio::time::timeout(Duration::from_secs(timeout_secs), child.wait()).await {
-            Ok(wait_result) => wait_result
-                .map_err(|error| format!("Failed waiting on shell command: {}", error))?
-                .code(),
-            Err(_) => {
-                timed_out = true;
+        tokio::select! {
+            result = tokio::time::timeout(Duration::from_secs(timeout_secs), child.wait()) => match result {
+                Ok(wait_result) => wait_result
+                    .map_err(|error| format!("Failed waiting on shell command: {}", error))?
+                    .code(),
+                Err(_) => {
+                    timed_out = true;
+                    let _ = child.start_kill();
+                    let _ = child.wait().await;
+                    None
+                }
+            },
+            _ = cancellation_token.cancelled() => {
                 let _ = child.start_kill();
                 let _ = child.wait().await;
                 None
             }
         }
     } else {
-        child
-            .wait()
-            .await
-            .map_err(|error| format!("Failed waiting on shell command: {}", error))?
-            .code()
+        tokio::select! {
+            result = child.wait() => result
+                .map_err(|error| format!("Failed waiting on shell command: {}", error))?
+                .code(),
+            _ = cancellation_token.cancelled() => {
+                let _ = child.start_kill();
+                let _ = child.wait().await;
+                None
+            }
+        }
     };
 
     const OUTPUT_DRAIN_TIMEOUT_MILLIS: u64 = 500;
@@ -847,6 +864,7 @@ mod tests {
                     timeout_secs: None,
                 },
                 Some(dir.path()),
+                CancellationToken::new(),
             )
             .await;
 
@@ -856,6 +874,41 @@ mod tests {
         assert_eq!(observed, expected);
     }
 
+    #[cfg(not(windows))]
+    #[tokio::test]
+    async fn shell_kills_child_on_cancellation() {
+        let tool = ShellTool::new_for_test().unwrap();
+        let token = CancellationToken::new();
+        let token_clone = token.clone();
+
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            token_clone.cancel();
+        });
+
+        let start = std::time::Instant::now();
+        let result = tool
+            .shell_with_cwd(
+                ShellParams {
+                    command: "sleep 30".to_string(),
+                    timeout_secs: None,
+                },
+                None,
+                token,
+            )
+            .await;
+
+        assert!(
+            start.elapsed().as_secs() < 5,
+            "shell should return quickly after cancellation, not wait for the command"
+        );
+        let shell_output = extract_shell_output(&result);
+        assert!(
+            shell_output.exit_code.is_none(),
+            "cancelled process should have no exit code"
+        );
+    }
+
     #[cfg(not(windows))]
     #[test]
     fn unix_shell_flavor_detects_nushell_names() {