diff --git a/crates/goose/src/agents/platform_extensions/code_execution.rs b/crates/goose/src/agents/platform_extensions/code_execution.rs index 03b812a23008..0a9ba0107ed6 100644 --- a/crates/goose/src/agents/platform_extensions/code_execution.rs +++ b/crates/goose/src/agents/platform_extensions/code_execution.rs @@ -271,6 +271,7 @@ impl CodeExecutionClient { ctx: &ToolCallContext, code_mode: &CodeMode, cancellation_token: CancellationToken, + rt: tokio::runtime::Handle, ) -> Result { let manager = self .context @@ -294,6 +295,7 @@ impl CodeExecutionClient { full_name, manager.clone(), cancellation_token.clone(), + rt.clone(), ); registry .add_callback(&cfg.id(), callback) @@ -377,7 +379,8 @@ impl CodeExecutionClient { let session_id = &ctx.session_id; let code_mode = self.get_code_mode(session_id).await?; let dispatch_token = cancellation_token.child_token(); - let registry = self.build_callback_registry(ctx, &code_mode, dispatch_token.clone())?; + let rt = tokio::runtime::Handle::current(); + let registry = self.build_callback_registry(ctx, &code_mode, dispatch_token.clone(), rt)?; let code = args.input.code.clone(); let disclosure = self.disclosure; @@ -474,12 +477,14 @@ fn create_tool_callback( full_name: String, manager: Arc, cancellation_token: CancellationToken, + rt: tokio::runtime::Handle, ) -> 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(); + let rt = rt.clone(); Box::pin(async move { let tool_call = { let mut params = CallToolRequestParams::new(full_name); @@ -488,39 +493,51 @@ fn create_tool_callback( } params }; - match manager - .dispatch_tool_call(&ctx, tool_call, cancellation_token) - .await - { - Ok(dispatch_result) => match dispatch_result.result.await { - Ok(result) => { - if let Some(sc) = &result.structured_content { - Ok(serde_json::to_value(sc).unwrap_or(Value::Null)) - } else { - // Filter to assistant-audience or no-audience content, - // skipping user-only content to avoid duplicated output - let text: String = result - .content - .iter() - .filter(|c| { - c.audience().is_none_or(|audiences| { - audiences.is_empty() || audiences.contains(&Role::Assistant) + + let handle = rt.spawn(async move { + match manager + .dispatch_tool_call(&ctx, tool_call, cancellation_token) + .await + { + Ok(dispatch_result) => match dispatch_result.result.await { + Ok(result) => { + if let Some(sc) = &result.structured_content { + Ok(serde_json::to_value(sc).unwrap_or(Value::Null)) + } else { + let text: String = result + .content + .iter() + .filter(|c| { + c.audience().is_none_or(|audiences| { + audiences.is_empty() + || audiences.contains(&Role::Assistant) + }) }) - }) - .filter_map(|c| match &c.raw { - RawContent::Text(t) => Some(t.text.clone()), - _ => None, - }) - .collect::>() - .join("\n"); - // Try to parse as JSON, otherwise return as string - Ok(serde_json::from_str(&text).unwrap_or(Value::String(text))) + .filter_map(|c| match &c.raw { + RawContent::Text(t) => Some(t.text.clone()), + _ => None, + }) + .collect::>() + .join("\n"); + Ok(serde_json::from_str(&text).unwrap_or(Value::String(text))) + } } - } - Err(e) => Err(format!("Tool error: {}", e.message)), - }, - Err(e) => Err(format!("Dispatch error: {e}")), + Err(e) => Err(format!("Tool error: {}", e.message)), + }, + Err(e) => Err(format!("Dispatch error: {e}")), + } + }); + + struct AbortOnDrop(tokio::task::AbortHandle); + impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } } + let _guard = AbortOnDrop(handle.abort_handle()); + handle + .await + .unwrap_or_else(|e| Err(format!("Callback task failed: {e}"))) }) as Pin> + Send>> }) } @@ -902,6 +919,66 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn callback_completing_from_main_runtime_does_not_hang() { + let rt = tokio::runtime::Handle::current(); + let (tx, rx) = tokio::sync::oneshot::channel::(); + let rx = Arc::new(std::sync::Mutex::new(Some(rx))); + + let callback: CallbackFn = Arc::new({ + let rt = rt.clone(); + move |_args: Option| { + let rt = rt.clone(); + let rx = rx.clone(); + Box::pin(async move { + let receiver = rx.lock().unwrap().take().expect("receiver taken once"); + let handle = rt.spawn(async move { + receiver.await.map_err(|_| "channel closed".to_string()) + }); + handle.await.unwrap_or_else(|e| Err(e.to_string())) + }) as Pin> + Send>> + } + }); + + rt.spawn(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = tx.send(serde_json::json!({"done": true})); + }); + + let cfg = CallbackConfig { + name: "ping".to_string(), + namespace: Some("Test".to_string()), + description: Some("ping".to_string()), + input_schema: None, + output_schema: None, + }; + let code_mode = CodeMode::default() + .with_callback(&cfg) + .expect("add callback"); + let registry = PctxRegistry::default(); + registry.add_callback(&cfg.id(), callback).unwrap(); + + let result = run_in_deno_runtime( + Duration::from_secs(5), + CancellationToken::new(), + CancellationToken::new(), + move || async move { + code_mode + .execute_typescript( + "async function run() { return await Test.ping(); }", + ToolDisclosure::default(), + Some(registry), + ) + .await + .map_err(|e| format!("execution error: {e}")) + }, + ) + .await + .expect("script should not time out"); + + assert!(result.success, "callback should succeed: {}", result.stderr); + } + #[test] fn catalog_moim_mentions_inspection_tools_without_function_names() { let moim = catalog_disclosure_moim(3);