diff --git a/crates/goose/src/agents/platform_extensions/code_execution.rs b/crates/goose/src/agents/platform_extensions/code_execution.rs index 6fcf612cab..42f415b003 100644 --- a/crates/goose/src/agents/platform_extensions/code_execution.rs +++ b/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 { 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, + cancellation_token: CancellationToken, ) -> Result, 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, + cancellation_token: CancellationToken, ) -> Result, 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( + timeout: Duration, + cancellation_token: CancellationToken, + dispatch_token: CancellationToken, + task: F, +) -> Result +where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future>, + 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, + cancellation_token: CancellationToken, ) -> CallbackFn { Arc::new(move |args: Option| { 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, - _cancellation_token: CancellationToken, + cancellation_token: CancellationToken, ) -> Result { 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); diff --git a/crates/goose/src/agents/platform_extensions/developer/mod.rs b/crates/goose/src/agents/platform_extensions/developer/mod.rs index 71607aa385..e8eb1bdff7 100644 --- a/crates/goose/src/agents/platform_extensions/developer/mod.rs +++ b/crates/goose/src/agents/platform_extensions/developer/mod.rs @@ -192,12 +192,15 @@ impl McpClientTrait for DeveloperClient { ctx: &ToolCallContext, name: &str, arguments: Option, - _cancel_token: CancellationToken, + cancel_token: CancellationToken, ) -> Result { let working_dir = ctx.working_dir.as_deref(); match name { "shell" => match Self::parse_args::(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::(arguments) { diff --git a/crates/goose/src/agents/platform_extensions/developer/shell.rs b/crates/goose/src/agents/platform_extensions/developer/shell.rs index 27201809d3..c7bbfddd93 100644 --- a/crates/goose/src/agents/platform_extensions/developer/shell.rs +++ b/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, working_dir: Option<&std::path::Path>, login_path: Option<&str>, + cancellation_token: CancellationToken, ) -> Result { 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() {