tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Abhijay Jain
2026-07-17 15:51:34 +05:30
committed by GitHub
parent 01ceb16b2b
commit f13ae10c65
1 changed files with 108 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/platform_extensions/code_execution.rs
+108 -31
View File
@@ -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
@@ -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<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);
@@ -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::<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>>
})
}
@@ -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);