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

Fix MCP elicitation deadlock and improve UX (#6650)

Browse Source
This commit is contained in:
Rizel Scarlett
2026-01-23 14:58:57 -05:00
committed by GitHub
parent 8233f0afac
commit c1a916efd9
2 changed files with 121 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+43 -31
View File
@@ -1288,7 +1288,7 @@ impl Agent {
let mut combined = stream::select_all(with_id);
let mut all_install_successful = true;
while let Some((request_id, item)) = combined.next().await {
loop {
if is_token_cancelled(&cancel_token) {
break;
}
@@ -1297,43 +1297,55 @@ impl Agent {
yield AgentEvent::Message(msg);
}
match item {
ToolStreamItem::Result(output) => {
let output = call_tool_result::validate(output);
tokio::select! {
biased;
// Platform extensions use meta as a way to publish notifications. Ideally we'd
// send the notifications directly, but the current plumbing doesn't support that
// well:
if let Ok(ref call_result) = output {
if let Some(ref meta) = call_result.meta {
if let Some(notification_data) = meta.0.get("platform_notification") {
if let Some(method) = notification_data.get("method").and_then(|v| v.as_str()) {
let params = notification_data.get("params").cloned();
let custom_notification = rmcp::model::CustomNotification::new(
method.to_string(),
params,
);
tool_item = combined.next() => {
match tool_item {
Some((request_id, item)) => {
match item {
ToolStreamItem::Result(output) => {
let output = call_tool_result::validate(output);
let server_notification = rmcp::model::ServerNotification::CustomNotification(custom_notification);
yield AgentEvent::McpNotification((request_id.clone(), server_notification));
if let Ok(ref call_result) = output {
if let Some(ref meta) = call_result.meta {
if let Some(notification_data) = meta.0.get("platform_notification") {
if let Some(method) = notification_data.get("method").and_then(|v| v.as_str()) {
let params = notification_data.get("params").cloned();
let custom_notification = rmcp::model::CustomNotification::new(
method.to_string(),
params,
);
let server_notification = rmcp::model::ServerNotification::CustomNotification(custom_notification);
yield AgentEvent::McpNotification((request_id.clone(), server_notification));
}
}
}
}
if enable_extension_request_ids.contains(&request_id)
&& output.is_err()
{
all_install_successful = false;
}
if let Some(response_msg) = request_to_response_map.get(&request_id) {
let metadata = request_metadata.get(&request_id).and_then(|m| m.as_ref());
let mut response = response_msg.lock().await;
*response = response.clone().with_tool_response_with_metadata(request_id, output, metadata);
}
}
ToolStreamItem::Message(msg) => {
yield AgentEvent::McpNotification((request_id, msg));
}
}
}
}
if enable_extension_request_ids.contains(&request_id)
&& output.is_err()
{
all_install_successful = false;
}
if let Some(response_msg) = request_to_response_map.get(&request_id) {
let metadata = request_metadata.get(&request_id).and_then(|m| m.as_ref());
let mut response = response_msg.lock().await;
*response = response.clone().with_tool_response_with_metadata(request_id, output, metadata);
None => break,
}
}
ToolStreamItem::Message(msg) => {
yield AgentEvent::McpNotification((request_id, msg));
_ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
// Continue loop to drain elicitation messages
}
}
}