feat(mcp): elicitation support (#5965)

2025-12-05 21:29:59 -05:00
parent 1db40709bb
commit 06155ce0ca
24 changed files with 992 additions and 22 deletions
@@ -0,0 +1,99 @@
+use anyhow::Result;
+use serde_json::Value;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::{mpsc, Mutex, RwLock};
+use tokio::time::timeout;
+use tracing::warn;
+use uuid::Uuid;
+
+use crate::conversation::message::{Message, MessageContent};
+
+struct PendingRequest {
+    response_tx: Option<tokio::sync::oneshot::Sender<Value>>,
+}
+
+pub struct ActionRequiredManager {
+    pending: Arc<RwLock<HashMap<String, Arc<Mutex<PendingRequest>>>>>,
+    request_tx: mpsc::UnboundedSender<Message>,
+    pub request_rx: Mutex<mpsc::UnboundedReceiver<Message>>,
+}
+
+impl ActionRequiredManager {
+    fn new() -> Self {
+        let (request_tx, request_rx) = mpsc::unbounded_channel();
+        Self {
+            pending: Arc::new(RwLock::new(HashMap::new())),
+            request_tx,
+            request_rx: Mutex::new(request_rx),
+        }
+    }
+
+    pub fn global() -> &'static Self {
+        static INSTANCE: once_cell::sync::Lazy<ActionRequiredManager> =
+            once_cell::sync::Lazy::new(ActionRequiredManager::new);
+        &INSTANCE
+    }
+
+    pub async fn request_and_wait(
+        &self,
+        message: String,
+        schema: Value,
+        timeout_duration: Duration,
+    ) -> Result<Value> {
+        let id = Uuid::new_v4().to_string();
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        let pending_request = PendingRequest {
+            response_tx: Some(tx),
+        };
+
+        self.pending
+            .write()
+            .await
+            .insert(id.clone(), Arc::new(Mutex::new(pending_request)));
+
+        let action_required_message = Message::assistant().with_content(
+            MessageContent::action_required_elicitation(id.clone(), message, schema),
+        );
+
+        if let Err(e) = self.request_tx.send(action_required_message) {
+            warn!("Failed to send action required message: {}", e);
+        }
+
+        let result = match timeout(timeout_duration, rx).await {
+            Ok(Ok(user_data)) => Ok(user_data),
+            Ok(Err(_)) => {
+                warn!("Response channel closed for request: {}", id);
+                Err(anyhow::anyhow!("Response channel closed"))
+            }
+            Err(_) => {
+                warn!("Timeout waiting for response: {}", id);
+                Err(anyhow::anyhow!("Timeout waiting for user response"))
+            }
+        };
+
+        self.pending.write().await.remove(&id);
+
+        result
+    }
+
+    pub async fn submit_response(&self, request_id: String, user_data: Value) -> Result<()> {
+        let pending_arc = {
+            let pending = self.pending.read().await;
+            pending
+                .get(&request_id)
+                .cloned()
+                .ok_or_else(|| anyhow::anyhow!("Request not found: {}", request_id))?
+        };
+
+        let mut pending = pending_arc.lock().await;
+        if let Some(tx) = pending.response_tx.take() {
+            if tx.send(user_data).is_err() {
+                warn!("Failed to send response through oneshot channel");
+            }
+        }
+
+        Ok(())
+    }
+}
@@ -8,6 +8,10 @@ use futures::stream::BoxStream;
 use futures::{stream, FutureExt, Stream, StreamExt, TryStreamExt};
 use uuid::Uuid;

+use super::final_output_tool::FinalOutputTool;
+use super::platform_tools;
+use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
+use crate::action_required_manager::ActionRequiredManager;
 use crate::agents::extension::{ExtensionConfig, ExtensionError, ExtensionResult, ToolInfo};
 use crate::agents::extension_manager::{get_parameter_names, ExtensionManager};
 use crate::agents::extension_manager_extension::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE;
@@ -25,6 +29,7 @@ use crate::agents::subagent_execution_tool::subagent_execute_task_tool::{
    self, SUBAGENT_EXECUTE_TASK_TOOL_NAME,
 };
 use crate::agents::subagent_execution_tool::tasks_manager::TasksManager;
+use crate::agents::subagent_task_config::TaskConfig;
 use crate::agents::tool_route_manager::ToolRouteManager;
 use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
 use crate::agents::types::SessionConfig;
@@ -33,6 +38,9 @@ use crate::config::{get_enabled_extensions, Config, GooseMode};
 use crate::context_mgmt::{
    check_if_compaction_needed, compact_messages, DEFAULT_COMPACTION_THRESHOLD,
 };
+use crate::conversation::message::{
+    ActionRequiredData, Message, MessageContent, SystemNotificationType, ToolRequest,
+};
 use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation};
 use crate::mcp_utils::ToolResult;
 use crate::permission::permission_inspector::PermissionInspector;
@@ -41,7 +49,10 @@ use crate::permission::PermissionConfirmation;
 use crate::providers::base::Provider;
 use crate::providers::errors::ProviderError;
 use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
+use crate::scheduler_trait::SchedulerTrait;
 use crate::security::security_inspector::SecurityInspector;
+use crate::session::extension_data::{EnabledExtensionsState, ExtensionState};
+use crate::session::{Session, SessionManager};
 use crate::tool_inspection::ToolInspectionManager;
 use crate::tool_monitor::RepetitionInspector;
 use crate::utils::is_token_cancelled;
@@ -55,15 +66,6 @@ use tokio::sync::{mpsc, Mutex};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info, instrument, warn};

-use super::final_output_tool::FinalOutputTool;
-use super::platform_tools;
-use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
-use crate::agents::subagent_task_config::TaskConfig;
-use crate::conversation::message::{Message, SystemNotificationType, ToolRequest};
-use crate::scheduler_trait::SchedulerTrait;
-use crate::session::extension_data::{EnabledExtensionsState, ExtensionState};
-use crate::session::{Session, SessionManager};
-
 const DEFAULT_MAX_TURNS: u32 = 1000;
 const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation...";
 pub const MANUAL_COMPACT_TRIGGERS: &[&str] =
@@ -241,6 +243,17 @@ impl Agent {
            | RetryResult::SuccessChecksPassed => Ok(false),
        }
    }
+    async fn drain_elicitation_messages(session_id: &str) -> Vec<Message> {
+        let mut messages = Vec::new();
+        let mut elicitation_rx = ActionRequiredManager::global().request_rx.lock().await;
+        while let Ok(elicitation_message) = elicitation_rx.try_recv() {
+            if let Err(e) = SessionManager::add_message(session_id, &elicitation_message).await {
+                warn!("Failed to save elicitation message to session: {}", e);
+            }
+            messages.push(elicitation_message);
+        }
+        messages
+    }

    async fn prepare_reply_context(
        &self,
@@ -784,6 +797,29 @@ impl Agent {
        session_config: SessionConfig,
        cancel_token: Option<CancellationToken>,
    ) -> Result<BoxStream<'_, Result<AgentEvent>>> {
+        for content in &user_message.content {
+            if let MessageContent::ActionRequired(action_required) = content {
+                if let ActionRequiredData::ElicitationResponse { id, user_data } =
+                    &action_required.data
+                {
+                    if let Err(e) = ActionRequiredManager::global()
+                        .submit_response(id.clone(), user_data.clone())
+                        .await
+                    {
+                        let error_text = format!("Failed to submit elicitation response: {}", e);
+                        error!(error_text);
+                        return Ok(Box::pin(stream::once(async {
+                            Ok(AgentEvent::Message(
+                                Message::assistant().with_text(error_text),
+                            ))
+                        })));
+                    }
+                    SessionManager::add_message(&session_config.id, &user_message).await?;
+                    return Ok(Box::pin(futures::stream::empty()));
+                }
+            }
+        }
+
        let message_text = user_message.as_concat_text();
        let is_manual_compact = MANUAL_COMPACT_TRIGGERS.contains(&message_text.trim());

@@ -1133,6 +1169,11 @@ impl Agent {
                                        if is_token_cancelled(&cancel_token) {
                                            break;
                                        }
+
+                                        for msg in Self::drain_elicitation_messages(&session_config.id).await {
+                                            yield AgentEvent::Message(msg);
+                                        }
+
                                        match item {
                                            ToolStreamItem::Result(output) => {
                                                if enable_extension_request_ids.contains(&request_id)
@@ -1151,6 +1192,11 @@ impl Agent {
                                        }
                                    }

+                                    // check for remaining elicitation messages after all tools complete
+                                    for msg in Self::drain_elicitation_messages(&session_config.id).await {
+                                        yield AgentEvent::Message(msg);
+                                    }
+
                                    if all_install_successful && !enable_extension_request_ids.is_empty() {
                                        if let Err(e) = self.save_extension_state(&session_config).await {
                                            warn!("Failed to save extension state after runtime changes: {}", e);
@@ -1,6 +1,10 @@
+use crate::action_required_manager::ActionRequiredManager;
 use crate::agents::types::SharedProvider;
 use crate::session_context::SESSION_ID_HEADER;
-use rmcp::model::{Content, ErrorCode, JsonObject};
+use rmcp::model::{
+    Content, CreateElicitationRequestParam, CreateElicitationResult, ElicitationAction, ErrorCode,
+    JsonObject,
+};
 /// MCP client implementation for Goose
 use rmcp::{
    model::{
@@ -218,10 +222,46 @@ impl ClientHandler for GooseClient {
        })
    }

+    async fn create_elicitation(
+        &self,
+        request: CreateElicitationRequestParam,
+        _context: RequestContext<RoleClient>,
+    ) -> Result<CreateElicitationResult, ErrorData> {
+        let schema_value = serde_json::to_value(&request.requested_schema).map_err(|e| {
+            ErrorData::new(
+                ErrorCode::INTERNAL_ERROR,
+                format!("Failed to serialize elicitation schema: {}", e),
+                None,
+            )
+        })?;
+
+        ActionRequiredManager::global()
+            .request_and_wait(
+                request.message.clone(),
+                schema_value,
+                Duration::from_secs(300),
+            )
+            .await
+            .map(|user_data| CreateElicitationResult {
+                action: ElicitationAction::Accept,
+                content: Some(user_data),
+            })
+            .map_err(|e| {
+                ErrorData::new(
+                    ErrorCode::INTERNAL_ERROR,
+                    format!("Elicitation request timed out or failed: {}", e),
+                    None,
+                )
+            })
+    }
+
    fn get_info(&self) -> ClientInfo {
        ClientInfo {
            protocol_version: ProtocolVersion::V_2025_03_26,
-            capabilities: ClientCapabilities::builder().enable_sampling().build(),
+            capabilities: ClientCapabilities::builder()
+                .enable_sampling()
+                .enable_elicitation()
+                .build(),
            client_info: Implementation {
                name: "goose".to_string(),
                version: std::env::var("GOOSE_MCP_CLIENT_VERSION")
@@ -1,4 +1,4 @@
-use crate::conversation::message::MessageMetadata;
+use crate::conversation::message::{ActionRequiredData, MessageMetadata};
 use crate::conversation::message::{Message, MessageContent};
 use crate::conversation::{merge_consecutive_messages, Conversation};
 use crate::prompt_template::render_global_file;
@@ -376,12 +376,15 @@ fn format_message_for_compacting(msg: &Message) -> String {
                format!("tool_confirmation_request: {}", req.tool_name)
            }
            MessageContent::ActionRequired(action) => match &action.data {
-                crate::conversation::message::ActionRequiredData::ToolConfirmation {
-                    tool_name,
-                    ..
-                } => {
+                ActionRequiredData::ToolConfirmation { tool_name, .. } => {
                    format!("action_required(tool_confirmation): {}", tool_name)
                }
+                ActionRequiredData::Elicitation { message, .. } => {
+                    format!("action_required(elicitation): {}", message)
+                }
+                ActionRequiredData::ElicitationResponse { id, .. } => {
+                    format!("action_required(elicitation_response): {}", id)
+                }
            },
            MessageContent::FrontendToolRequest(req) => {
                if let Ok(call) = &req.tool_call {
@@ -111,6 +111,15 @@ pub enum ActionRequiredData {
        arguments: JsonObject,
        prompt: Option<String>,
    },
+    Elicitation {
+        id: String,
+        message: String,
+        requested_schema: serde_json::Value,
+    },
+    ElicitationResponse {
+        id: String,
+        user_data: serde_json::Value,
+    },
 }

 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
@@ -192,6 +201,12 @@ impl fmt::Display for MessageContent {
                ActionRequiredData::ToolConfirmation { tool_name, .. } => {
                    write!(f, "[ActionRequired: ToolConfirmation for {}]", tool_name)
                }
+                ActionRequiredData::Elicitation { message, .. } => {
+                    write!(f, "[ActionRequired: Elicitation - {}]", message)
+                }
+                ActionRequiredData::ElicitationResponse { id, .. } => {
+                    write!(f, "[ActionRequired: ElicitationResponse for {}]", id)
+                }
            },
            MessageContent::FrontendToolRequest(r) => match &r.tool_call {
                Ok(tool_call) => write!(f, "[FrontendToolRequest: {}]", tool_call.name),
@@ -274,6 +289,32 @@ impl MessageContent {
        })
    }

+    pub fn action_required_elicitation<S: Into<String>>(
+        id: S,
+        message: String,
+        requested_schema: serde_json::Value,
+    ) -> Self {
+        MessageContent::ActionRequired(ActionRequired {
+            data: ActionRequiredData::Elicitation {
+                id: id.into(),
+                message,
+                requested_schema,
+            },
+        })
+    }
+
+    pub fn action_required_elicitation_response<S: Into<String>>(
+        id: S,
+        user_data: serde_json::Value,
+    ) -> Self {
+        MessageContent::ActionRequired(ActionRequired {
+            data: ActionRequiredData::ElicitationResponse {
+                id: id.into(),
+                user_data,
+            },
+        })
+    }
+
    pub fn thinking<S1: Into<String>, S2: Into<String>>(thinking: S1, signature: S2) -> Self {
        MessageContent::Thinking(ThinkingContent {
            thinking: thinking.into(),
@@ -1,3 +1,4 @@
+pub mod action_required_manager;
 pub mod agents;
 pub mod config;
 pub mod context_mgmt;
@@ -1,4 +1,4 @@
-STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
+STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"elicitation": {}, "sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
 STDERR:   [2m2025-10-30T14:41:09.117156Z[0m [32m INFO[0m [1;32mgoose_mcp::mcp_server_runner[0m[32m: [32mStarting MCP server[0m
 STDERR:     [2;3mat[0m crates/goose-mcp/src/mcp_server_runner.rs:18
 STDERR: 
@@ -1,4 +1,4 @@
-STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
+STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"elicitation": {}, "sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
 STDERR: GitHub MCP Server running on stdio
 STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabilities":{"logging":{},"prompts":{},"resources":{"subscribe":true,"listChanged":true},"tools":{"listChanged":true}},"serverInfo":{"name":"github-mcp-server","version":"version"}}}
 STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
@@ -1,4 +1,4 @@
-STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
+STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"elicitation": {}, "sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
 STDERR: Starting default (STDIO) server...
 STDERR: Starting logs update interval
 STDOUT: {"result":{"protocolVersion":"2025-03-26","capabilities":{"prompts":{},"resources":{"subscribe":true},"tools":{},"logging":{},"completions":{}},"serverInfo":{"name":"example-servers/everything","title":"Everything Example Server","version":"1.0.0"},"instructions":"Testing and demonstration server for MCP protocol features.\n\n## Resources\n\nResources 1-100 follow pattern: even IDs contain text, odd IDs contain binary data. Resources are paginated at 10 items per page with cursor-based navigation.\n\n## Key dependencies\n\nProgress notifications require `_meta.progressToken` in tool calls. Resource subscriptions generate updates every 10 seconds.\n\n## Performance characteristics\n\nServer generates automatic log messages every 20 seconds (filtered by current log level), stderr notifications every 30 seconds, and resource update notifications every 10 seconds for active subscriptions. Log level changes affect message filtering in real-time.\n\n## Multi-modal testing\n\n`complex_prompt` includes both text arguments and image content for testing client multi-modal handling. `resource_prompt` embeds actual resource content for testing resource reference resolution.\n\nArgument completion is available for prompt parameters and resource IDs. Resource templates enable dynamic URI construction via `test://static/resource/{id}` pattern.\n\n## Easter egg\n\nIf asked about server instructions, respond with \"🎉 Server instructions are working! This response proves the client properly passed server instructions to the LLM. This demonstrates MCP's instructions feature in action.\"\n"},"jsonrpc":"2.0","id":0}
@@ -1,4 +1,4 @@
-STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
+STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"elicitation": {}, "sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
 STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabilities":{"experimental":{},"prompts":{"listChanged":false},"tools":{"listChanged":false}},"serverInfo":{"name":"mcp-fetch","version":"1.19.0"}}}
 STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
 STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{"progressToken":0},"name":"fetch","arguments":{"url":"https://example.com"}}}