diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 715f982a5132..48a3382526f6 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -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); - - // 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, - ); - - let server_notification = rmcp::model::ServerNotification::CustomNotification(custom_notification); - yield AgentEvent::McpNotification((request_id.clone(), server_notification)); + tokio::select! { + biased; + + tool_item = combined.next() => { + match tool_item { + Some((request_id, item)) => { + match item { + ToolStreamItem::Result(output) => { + let output = call_tool_result::validate(output); + + 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 } } } diff --git a/ui/desktop/src/components/ElicitationRequest.tsx b/ui/desktop/src/components/ElicitationRequest.tsx index 9dcc95fef767..98deb045a4ca 100644 --- a/ui/desktop/src/components/ElicitationRequest.tsx +++ b/ui/desktop/src/components/ElicitationRequest.tsx @@ -1,8 +1,10 @@ -import { useState } from 'react'; +import { useState, useEffect, useRef } from 'react'; import { ActionRequired } from '../api'; import JsonSchemaForm from './ui/JsonSchemaForm'; import type { JsonSchema } from './ui/JsonSchemaForm'; +const ELICITATION_TIMEOUT_SECONDS = 300; + interface ElicitationRequestProps { isCancelledMessage: boolean; isClicked: boolean; @@ -10,6 +12,12 @@ interface ElicitationRequestProps { onSubmit: (elicitationId: string, userData: Record) => void; } +function formatTime(seconds: number): string { + const mins = Math.floor(seconds / 60); + const secs = seconds % 60; + return `${mins}:${secs.toString().padStart(2, '0')}`; +} + export default function ElicitationRequest({ isCancelledMessage, isClicked, @@ -17,6 +25,24 @@ export default function ElicitationRequest({ onSubmit, }: ElicitationRequestProps) { const [submitted, setSubmitted] = useState(isClicked); + const [timeRemaining, setTimeRemaining] = useState(ELICITATION_TIMEOUT_SECONDS); + const startTimeRef = useRef(Date.now()); + + useEffect(() => { + if (submitted || isCancelledMessage || isClicked) return; + + const interval = setInterval(() => { + const elapsed = Math.floor((Date.now() - startTimeRef.current) / 1000); + const remaining = Math.max(0, ELICITATION_TIMEOUT_SECONDS - elapsed); + setTimeRemaining(remaining); + + if (remaining === 0) { + clearInterval(interval); + } + }, 1000); + + return () => clearInterval(interval); + }, [submitted, isCancelledMessage, isClicked]); if (actionRequiredContent.data.actionType !== 'elicitation') { return null; @@ -57,10 +83,39 @@ export default function ElicitationRequest({ ); } + const isUrgent = timeRemaining <= 60; + const isExpired = timeRemaining === 0; + + if (isExpired) { + return ( +
+
+ + + + This request has expired. The extension will need to ask again. +
+
+ ); + } + return (
- {message || 'Goose needs some information from you.'} +
+ {message || 'Goose needs some information from you.'} +
+
+ + + + + Waiting for your response ({formatTime(timeRemaining)} remaining) + +
);