Skip to content

Commit c1a916e

Browse files
blackgirlbytesblackgirlbytes
authored
Fix MCP elicitation deadlock and improve UX (#6650)
1 parent 8233f0a commit c1a916e

2 files changed

Lines changed: 123 additions & 35 deletions

File tree

‎crates/goose/src/agents/agent.rs‎

Lines changed: 45 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1288,7 +1288,7 @@ impl Agent {
12881288
let mut combined = stream::select_all(with_id);
12891289
let mut all_install_successful = true;
12901290

1291-
while let Some((request_id, item)) = combined.next().await {
1291+
loop {
12921292
if is_token_cancelled(&cancel_token) {
12931293
break;
12941294
}
@@ -1297,43 +1297,55 @@ impl Agent {
12971297
yield AgentEvent::Message(msg);
12981298
}
12991299

1300-
match item {
1301-
ToolStreamItem::Result(output) => {
1302-
let output = call_tool_result::validate(output);
1303-
1304-
// Platform extensions use meta as a way to publish notifications. Ideally we'd
1305-
// send the notifications directly, but the current plumbing doesn't support that
1306-
// well:
1307-
if let Ok(ref call_result) = output {
1308-
if let Some(ref meta) = call_result.meta {
1309-
if let Some(notification_data) = meta.0.get("platform_notification") {
1310-
if let Some(method) = notification_data.get("method").and_then(|v| v.as_str()) {
1311-
let params = notification_data.get("params").cloned();
1312-
let custom_notification = rmcp::model::CustomNotification::new(
1313-
method.to_string(),
1314-
params,
1315-
);
1316-
1317-
let server_notification = rmcp::model::ServerNotification::CustomNotification(custom_notification);
1318-
yield AgentEvent::McpNotification((request_id.clone(), server_notification));
1300+
tokio::select! {
1301+
biased;
1302+
1303+
tool_item = combined.next() => {
1304+
match tool_item {
1305+
Some((request_id, item)) => {
1306+
match item {
1307+
ToolStreamItem::Result(output) => {
1308+
let output = call_tool_result::validate(output);
1309+
1310+
if let Ok(ref call_result) = output {
1311+
if let Some(ref meta) = call_result.meta {
1312+
if let Some(notification_data) = meta.0.get("platform_notification") {
1313+
if let Some(method) = notification_data.get("method").and_then(|v| v.as_str()) {
1314+
let params = notification_data.get("params").cloned();
1315+
let custom_notification = rmcp::model::CustomNotification::new(
1316+
method.to_string(),
1317+
params,
1318+
);
1319+
1320+
let server_notification = rmcp::model::ServerNotification::CustomNotification(custom_notification);
1321+
yield AgentEvent::McpNotification((request_id.clone(), server_notification));
1322+
}
1323+
}
1324+
}
1325+
}
1326+
1327+
if enable_extension_request_ids.contains(&request_id)
1328+
&& output.is_err()
1329+
{
1330+
all_install_successful = false;
1331+
}
1332+
if let Some(response_msg) = request_to_response_map.get(&request_id) {
1333+
let metadata = request_metadata.get(&request_id).and_then(|m| m.as_ref());
1334+
let mut response = response_msg.lock().await;
1335+
*response = response.clone().with_tool_response_with_metadata(request_id, output, metadata);
1336+
}
1337+
}
1338+
ToolStreamItem::Message(msg) => {
1339+
yield AgentEvent::McpNotification((request_id, msg));
13191340
}
13201341
}
13211342
}
1322-
}
1323-
1324-
if enable_extension_request_ids.contains(&request_id)
1325-
&& output.is_err()
1326-
{
1327-
all_install_successful = false;
1328-
}
1329-
if let Some(response_msg) = request_to_response_map.get(&request_id) {
1330-
let metadata = request_metadata.get(&request_id).and_then(|m| m.as_ref());
1331-
let mut response = response_msg.lock().await;
1332-
*response = response.clone().with_tool_response_with_metadata(request_id, output, metadata);
1343+
None => break,
13331344
}
13341345
}
1335-
ToolStreamItem::Message(msg) => {
1336-
yield AgentEvent::McpNotification((request_id, msg));
1346+
1347+
_ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
1348+
// Continue loop to drain elicitation messages
13371349
}
13381350
}
13391351
}

‎ui/desktop/src/components/ElicitationRequest.tsx‎

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,48 @@
1-
import { useState } from 'react';
1+
import { useState, useEffect, useRef } from 'react';
22
import { ActionRequired } from '../api';
33
import JsonSchemaForm from './ui/JsonSchemaForm';
44
import type { JsonSchema } from './ui/JsonSchemaForm';
55

6+
const ELICITATION_TIMEOUT_SECONDS = 300;
7+
68
interface ElicitationRequestProps {
79
isCancelledMessage: boolean;
810
isClicked: boolean;
911
actionRequiredContent: ActionRequired & { type: 'actionRequired' };
1012
onSubmit: (elicitationId: string, userData: Record<string, unknown>) => void;
1113
}
1214

15+
function formatTime(seconds: number): string {
16+
const mins = Math.floor(seconds / 60);
17+
const secs = seconds % 60;
18+
return `${mins}:${secs.toString().padStart(2, '0')}`;
19+
}
20+
1321
export default function ElicitationRequest({
1422
isCancelledMessage,
1523
isClicked,
1624
actionRequiredContent,
1725
onSubmit,
1826
}: ElicitationRequestProps) {
1927
const [submitted, setSubmitted] = useState(isClicked);
28+
const [timeRemaining, setTimeRemaining] = useState(ELICITATION_TIMEOUT_SECONDS);
29+
const startTimeRef = useRef(Date.now());
30+
31+
useEffect(() => {
32+
if (submitted || isCancelledMessage || isClicked) return;
33+
34+
const interval = setInterval(() => {
35+
const elapsed = Math.floor((Date.now() - startTimeRef.current) / 1000);
36+
const remaining = Math.max(0, ELICITATION_TIMEOUT_SECONDS - elapsed);
37+
setTimeRemaining(remaining);
38+
39+
if (remaining === 0) {
40+
clearInterval(interval);
41+
}
42+
}, 1000);
43+
44+
return () => clearInterval(interval);
45+
}, [submitted, isCancelledMessage, isClicked]);
2046

2147
if (actionRequiredContent.data.actionType !== 'elicitation') {
2248
return null;
@@ -57,17 +83,67 @@ export default function ElicitationRequest({
5783
);
5884
}
5985

86+
const isUrgent = timeRemaining <= 60;
87+
const isExpired = timeRemaining === 0;
88+
89+
if (isExpired) {
90+
return (
91+
<div className="goose-message-content bg-background-muted rounded-2xl px-4 py-2 text-textStandard">
92+
<div className="flex items-center gap-2 text-textSubtle">
93+
<svg
94+
className="w-5 h-5"
95+
xmlns="http://www.w3.org/2000/svg"
96+
fill="none"
97+
viewBox="0 0 24 24"
98+
stroke="currentColor"
99+
strokeWidth={2}
100+
>
101+
<path
102+
strokeLinecap="round"
103+
strokeLinejoin="round"
104+
d="M12 8v4l3 3m6-3a9 9 0 11-18 0 9 9 0 0118 0z"
105+
/>
106+
</svg>
107+
<span>This request has expired. The extension will need to ask again.</span>
108+
</div>
109+
</div>
110+
);
111+
}
112+
60113
return (
61114
<div className="flex flex-col">
62115
<div className="goose-message-content bg-background-muted rounded-2xl rounded-b-none px-4 py-2 text-textStandard">
63-
{message || 'Goose needs some information from you.'}
116+
<div className="flex justify-between items-start gap-4">
117+
<span>{message || 'Goose needs some information from you.'}</span>
118+
</div>
64119
</div>
65120
<div className="goose-message-content bg-background-default border border-borderSubtle dark:border-gray-700 rounded-b-2xl px-4 py-3">
66121
<JsonSchemaForm
67122
schema={requested_schema as JsonSchema}
68123
onSubmit={handleSubmit}
69124
submitLabel="Submit"
70125
/>
126+
<div
127+
className={`mt-3 pt-3 border-t border-borderSubtle flex items-center gap-2 text-sm ${isUrgent ? 'text-red-500' : 'text-textSubtle'}`}
128+
>
129+
<svg
130+
className="w-4 h-4 animate-pulse"
131+
xmlns="http://www.w3.org/2000/svg"
132+
fill="none"
133+
viewBox="0 0 24 24"
134+
stroke="currentColor"
135+
strokeWidth={2}
136+
>
137+
<path
138+
strokeLinecap="round"
139+
strokeLinejoin="round"
140+
d="M12 8v4l3 3m6-3a9 9 0 11-18 0 9 9 0 0118 0z"
141+
/>
142+
</svg>
143+
<span>
144+
Waiting for your response ({formatTime(timeRemaining)} remaining)
145+
</span>
146+
</div>
71147
</div>
72148
</div>
73149
);

0 commit comments

Comments
 (0)