Skip to content

fix(acp): dispatch Code Mode callbacks on main runtime to prevent hangs #10484

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Abhijay007 merged 2 commits into from
Jul 17, 2026
Merged

fix(acp): dispatch Code Mode callbacks on main runtime to prevent hangs #10484

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6ce8e4e
fix(acp): dispatch Code Mode callbacks on main runtime to prevent hangs
Abhijay007 Jul 15, 2026
d7539ea
fix(acp): abort detached callback task when Code Mode future is dropped
Abhijay007 Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 108 additions & 31 deletions crates/goose/src/agents/platform_extensions/code_execution.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,7 @@ impl CodeExecutionClient {
ctx: &ToolCallContext,
code_mode: &CodeMode,
cancellation_token: CancellationToken,
rt: tokio::runtime::Handle,
) -> Result<PctxRegistry, String> {
let manager = self
.context
Expand All @@ -294,6 +295,7 @@ impl CodeExecutionClient {
full_name,
manager.clone(),
cancellation_token.clone(),
rt.clone(),
);
registry
.add_callback(&cfg.id(), callback)
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -474,12 +477,14 @@ fn create_tool_callback(
full_name: String,
manager: Arc<crate::agents::ExtensionManager>,
cancellation_token: CancellationToken,
rt: tokio::runtime::Handle,
) -> 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();
let rt = rt.clone();
Box::pin(async move {
let tool_call = {
let mut params = CallToolRequestParams::new(full_name);
Expand All @@ -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 {
Comment thread
Abhijay007 marked this conversation as resolved.
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::<Vec<_>>()
.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::<Vec<_>>()
.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<Box<dyn Future<Output = Result<Value, String>> + Send>>
})
}
Expand Down Expand Up @@ -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::<serde_json::Value>();
let rx = Arc::new(std::sync::Mutex::new(Some(rx)));

let callback: CallbackFn = Arc::new({
let rt = rt.clone();
move |_args: Option<Value>| {
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<Box<dyn Future<Output = Result<Value, String>> + 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);
Expand Down