From e9f0d43288d9aa17f472ed3239139df1dad6656c Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Mon, 27 Apr 2026 01:05:59 +0300 Subject: [PATCH] Extract MCP elicitation handling Move ElicitationRequestManager and elicitation policy handling into codex-mcp/src/elicitation.rs. Update the manager, RMCP client layer, and tests to import elicitation behavior from the new module. Co-authored-by: Codex --- codex-rs/codex-mcp/src/client.rs | 2 +- codex-rs/codex-mcp/src/elicitation.rs | 180 ++++++++++++++++++ codex-rs/codex-mcp/src/lib.rs | 1 + codex-rs/codex-mcp/src/manager.rs | 2 +- codex-rs/codex-mcp/src/mcp_connection.rs | 171 ----------------- .../src/mcp_connection_manager_tests.rs | 4 +- 6 files changed, 185 insertions(+), 175 deletions(-) create mode 100644 codex-rs/codex-mcp/src/elicitation.rs diff --git a/codex-rs/codex-mcp/src/client.rs b/codex-rs/codex-mcp/src/client.rs index 54fb664585..8ea17a5a45 100644 --- a/codex-rs/codex-mcp/src/client.rs +++ b/codex-rs/codex-mcp/src/client.rs @@ -8,13 +8,13 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use crate::elicitation::ElicitationRequestManager; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; use crate::mcp_connection::CachedCodexAppsToolsLoad; use crate::mcp_connection::CodexAppsToolsCacheContext; use crate::mcp_connection::DEFAULT_STARTUP_TIMEOUT; use crate::mcp_connection::DEFAULT_TOOL_TIMEOUT; -use crate::mcp_connection::ElicitationRequestManager; use crate::mcp_connection::MCP_SANDBOX_STATE_META_CAPABILITY; use crate::mcp_connection::McpRuntimeEnvironment; use crate::mcp_connection::ToolFilter; diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs new file mode 100644 index 0000000000..ee3125e922 --- /dev/null +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -0,0 +1,180 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; + +use crate::mcp::mcp_permission_prompt_is_auto_approved; +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use async_channel::Sender; +use codex_protocol::approvals::ElicitationRequest; +use codex_protocol::approvals::ElicitationRequestEvent; +use codex_protocol::mcp::RequestId as ProtocolRequestId; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::SandboxPolicy; +use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::SendElicitation; +use futures::future::FutureExt; +use rmcp::model::CreateElicitationRequestParams; +use rmcp::model::ElicitationAction; +use rmcp::model::RequestId; +use tokio::sync::Mutex; +use tokio::sync::oneshot; + +type ResponderMap = HashMap<(String, RequestId), oneshot::Sender>; + +pub(crate) fn elicitation_is_rejected_by_policy(approval_policy: AskForApproval) -> bool { + match approval_policy { + AskForApproval::Never => true, + AskForApproval::OnFailure => false, + AskForApproval::OnRequest => false, + AskForApproval::UnlessTrusted => false, + AskForApproval::Granular(granular_config) => !granular_config.allows_mcp_elicitations(), + } +} + +fn can_auto_accept_elicitation(elicitation: &CreateElicitationRequestParams) -> bool { + match elicitation { + CreateElicitationRequestParams::FormElicitationParams { + requested_schema, .. + } => { + // Auto-accept confirm/approval elicitations without schema requirements. + requested_schema.properties.is_empty() + } + CreateElicitationRequestParams::UrlElicitationParams { .. } => false, + } +} + +#[derive(Clone)] +pub(crate) struct ElicitationRequestManager { + requests: Arc>, + pub(crate) approval_policy: Arc>, + pub(crate) sandbox_policy: Arc>, +} + +impl ElicitationRequestManager { + pub(crate) fn new(approval_policy: AskForApproval, sandbox_policy: SandboxPolicy) -> Self { + Self { + requests: Arc::new(Mutex::new(HashMap::new())), + approval_policy: Arc::new(StdMutex::new(approval_policy)), + sandbox_policy: Arc::new(StdMutex::new(sandbox_policy)), + } + } + + pub(crate) async fn resolve( + &self, + server_name: String, + id: RequestId, + response: ElicitationResponse, + ) -> Result<()> { + self.requests + .lock() + .await + .remove(&(server_name, id)) + .ok_or_else(|| anyhow!("elicitation request not found"))? + .send(response) + .map_err(|e| anyhow!("failed to send elicitation response: {e:?}")) + } + + pub(crate) fn make_sender( + &self, + server_name: String, + tx_event: Sender, + ) -> SendElicitation { + let elicitation_requests = self.requests.clone(); + let approval_policy = self.approval_policy.clone(); + let sandbox_policy = self.sandbox_policy.clone(); + Box::new(move |id, elicitation| { + let elicitation_requests = elicitation_requests.clone(); + let tx_event = tx_event.clone(); + let server_name = server_name.clone(); + let approval_policy = approval_policy.clone(); + let sandbox_policy = sandbox_policy.clone(); + async move { + let approval_policy = approval_policy + .lock() + .map(|policy| *policy) + .unwrap_or(AskForApproval::Never); + let sandbox_policy = sandbox_policy + .lock() + .map(|policy| policy.clone()) + .unwrap_or_else(|_| SandboxPolicy::new_read_only_policy()); + if mcp_permission_prompt_is_auto_approved(approval_policy, &sandbox_policy) + && can_auto_accept_elicitation(&elicitation) + { + return Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(serde_json::json!({})), + meta: None, + }); + } + + if elicitation_is_rejected_by_policy(approval_policy) { + return Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + }); + } + + let request = match elicitation { + CreateElicitationRequestParams::FormElicitationParams { + meta, + message, + requested_schema, + } => ElicitationRequest::Form { + meta: meta + .map(serde_json::to_value) + .transpose() + .context("failed to serialize MCP elicitation metadata")?, + message, + requested_schema: serde_json::to_value(requested_schema) + .context("failed to serialize MCP elicitation schema")?, + }, + CreateElicitationRequestParams::UrlElicitationParams { + meta, + message, + url, + elicitation_id, + } => ElicitationRequest::Url { + meta: meta + .map(serde_json::to_value) + .transpose() + .context("failed to serialize MCP elicitation metadata")?, + message, + url, + elicitation_id, + }, + }; + let (tx, rx) = oneshot::channel(); + { + let mut lock = elicitation_requests.lock().await; + lock.insert((server_name.clone(), id.clone()), tx); + } + let _ = tx_event + .send(Event { + id: "mcp_elicitation_request".to_string(), + msg: EventMsg::ElicitationRequest(ElicitationRequestEvent { + turn_id: None, + server_name, + id: match id.clone() { + rmcp::model::NumberOrString::String(value) => { + ProtocolRequestId::String(value.to_string()) + } + rmcp::model::NumberOrString::Number(value) => { + ProtocolRequestId::Integer(value) + } + }, + request, + }), + }) + .await; + rx.await + .context("elicitation request channel closed unexpectedly") + } + .boxed() + }) + } +} diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 1291b75fd4..b603571ef3 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -39,6 +39,7 @@ pub use mcp_connection::declared_openai_file_input_param_names; pub use mcp_connection::filter_non_codex_apps_mcp_tools_only; pub(crate) mod client; +pub(crate) mod elicitation; pub(crate) mod manager; pub(crate) mod mcp; pub(crate) mod mcp_connection; diff --git a/codex-rs/codex-mcp/src/manager.rs b/codex-rs/codex-mcp/src/manager.rs index 211e394b71..f8189e492c 100644 --- a/codex-rs/codex-mcp/src/manager.rs +++ b/codex-rs/codex-mcp/src/manager.rs @@ -11,11 +11,11 @@ use crate::client::MCP_TOOLS_LIST_DURATION_METRIC; use crate::client::ManagedClient; use crate::client::StartupOutcomeError; use crate::client::list_tools_for_client_uncached; +use crate::elicitation::ElicitationRequestManager; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; use crate::mcp_connection::CodexAppsToolsCacheContext; use crate::mcp_connection::CodexAppsToolsCacheKey; -use crate::mcp_connection::ElicitationRequestManager; use crate::mcp_connection::McpRuntimeEnvironment; use crate::mcp_connection::ToolInfo; use crate::mcp_connection::emit_duration; diff --git a/codex-rs/codex-mcp/src/mcp_connection.rs b/codex-rs/codex-mcp/src/mcp_connection.rs index 68ce48b2d3..c694f53a5c 100644 --- a/codex-rs/codex-mcp/src/mcp_connection.rs +++ b/codex-rs/codex-mcp/src/mcp_connection.rs @@ -7,36 +7,23 @@ use std::collections::HashSet; use std::env; use std::path::PathBuf; use std::sync::Arc; -use std::sync::Mutex as StdMutex; use std::time::Duration; use std::time::Instant; use crate::McpAuthStatusEntry; use crate::client::StartupOutcomeError; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; -use crate::mcp::mcp_permission_prompt_is_auto_approved; pub(crate) use crate::mcp_tool_names::qualify_tools; -use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use async_channel::Sender; use codex_exec_server::Environment; use codex_protocol::ToolName; -use codex_protocol::approvals::ElicitationRequest; -use codex_protocol::approvals::ElicitationRequestEvent; -use codex_protocol::mcp::RequestId as ProtocolRequestId; use codex_protocol::models::PermissionProfile; -use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpStartupUpdateEvent; use codex_protocol::protocol::SandboxPolicy; -use codex_rmcp_client::ElicitationResponse; -use codex_rmcp_client::SendElicitation; -use futures::future::FutureExt; -use rmcp::model::CreateElicitationRequestParams; -use rmcp::model::ElicitationAction; -use rmcp::model::RequestId; use rmcp::model::Tool; use serde::Deserialize; @@ -45,8 +32,6 @@ use serde_json::Map; use serde_json::Value as JsonValue; use sha1::Digest; use sha1::Sha1; -use tokio::sync::Mutex; -use tokio::sync::oneshot; use url::Url; use codex_config::McpServerConfig; @@ -255,162 +240,6 @@ pub(crate) enum CachedCodexAppsToolsLoad { Invalid, } -type ResponderMap = HashMap<(String, RequestId), oneshot::Sender>; - -pub(crate) fn elicitation_is_rejected_by_policy(approval_policy: AskForApproval) -> bool { - match approval_policy { - AskForApproval::Never => true, - AskForApproval::OnFailure => false, - AskForApproval::OnRequest => false, - AskForApproval::UnlessTrusted => false, - AskForApproval::Granular(granular_config) => !granular_config.allows_mcp_elicitations(), - } -} - -fn can_auto_accept_elicitation(elicitation: &CreateElicitationRequestParams) -> bool { - match elicitation { - CreateElicitationRequestParams::FormElicitationParams { - requested_schema, .. - } => { - // Auto-accept confirm/approval elicitations without schema requirements. - requested_schema.properties.is_empty() - } - CreateElicitationRequestParams::UrlElicitationParams { .. } => false, - } -} - -#[derive(Clone)] -pub(crate) struct ElicitationRequestManager { - requests: Arc>, - pub(crate) approval_policy: Arc>, - pub(crate) sandbox_policy: Arc>, -} - -impl ElicitationRequestManager { - pub(crate) fn new(approval_policy: AskForApproval, sandbox_policy: SandboxPolicy) -> Self { - Self { - requests: Arc::new(Mutex::new(HashMap::new())), - approval_policy: Arc::new(StdMutex::new(approval_policy)), - sandbox_policy: Arc::new(StdMutex::new(sandbox_policy)), - } - } - - pub(crate) async fn resolve( - &self, - server_name: String, - id: RequestId, - response: ElicitationResponse, - ) -> Result<()> { - self.requests - .lock() - .await - .remove(&(server_name, id)) - .ok_or_else(|| anyhow!("elicitation request not found"))? - .send(response) - .map_err(|e| anyhow!("failed to send elicitation response: {e:?}")) - } - - pub(crate) fn make_sender( - &self, - server_name: String, - tx_event: Sender, - ) -> SendElicitation { - let elicitation_requests = self.requests.clone(); - let approval_policy = self.approval_policy.clone(); - let sandbox_policy = self.sandbox_policy.clone(); - Box::new(move |id, elicitation| { - let elicitation_requests = elicitation_requests.clone(); - let tx_event = tx_event.clone(); - let server_name = server_name.clone(); - let approval_policy = approval_policy.clone(); - let sandbox_policy = sandbox_policy.clone(); - async move { - let approval_policy = approval_policy - .lock() - .map(|policy| *policy) - .unwrap_or(AskForApproval::Never); - let sandbox_policy = sandbox_policy - .lock() - .map(|policy| policy.clone()) - .unwrap_or_else(|_| SandboxPolicy::new_read_only_policy()); - if mcp_permission_prompt_is_auto_approved(approval_policy, &sandbox_policy) - && can_auto_accept_elicitation(&elicitation) - { - return Ok(ElicitationResponse { - action: ElicitationAction::Accept, - content: Some(serde_json::json!({})), - meta: None, - }); - } - - if elicitation_is_rejected_by_policy(approval_policy) { - return Ok(ElicitationResponse { - action: ElicitationAction::Decline, - content: None, - meta: None, - }); - } - - let request = match elicitation { - CreateElicitationRequestParams::FormElicitationParams { - meta, - message, - requested_schema, - } => ElicitationRequest::Form { - meta: meta - .map(serde_json::to_value) - .transpose() - .context("failed to serialize MCP elicitation metadata")?, - message, - requested_schema: serde_json::to_value(requested_schema) - .context("failed to serialize MCP elicitation schema")?, - }, - CreateElicitationRequestParams::UrlElicitationParams { - meta, - message, - url, - elicitation_id, - } => ElicitationRequest::Url { - meta: meta - .map(serde_json::to_value) - .transpose() - .context("failed to serialize MCP elicitation metadata")?, - message, - url, - elicitation_id, - }, - }; - let (tx, rx) = oneshot::channel(); - { - let mut lock = elicitation_requests.lock().await; - lock.insert((server_name.clone(), id.clone()), tx); - } - let _ = tx_event - .send(Event { - id: "mcp_elicitation_request".to_string(), - msg: EventMsg::ElicitationRequest(ElicitationRequestEvent { - turn_id: None, - server_name, - id: match id.clone() { - rmcp::model::NumberOrString::String(value) => { - ProtocolRequestId::String(value.to_string()) - } - rmcp::model::NumberOrString::Number(value) => { - ProtocolRequestId::Integer(value) - } - }, - request, - }), - }) - .await; - rx.await - .context("elicitation request channel closed unexpectedly") - } - .boxed() - }) - } -} - const META_OPENAI_FILE_PARAMS: &str = "openai/fileParams"; /// Returns the model-visible view of a tool while preserving the raw metadata diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs b/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs index 5e99b36c7e..d3d3bf7c22 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs @@ -4,12 +4,12 @@ use crate::client::ManagedClient; use crate::client::StartupOutcomeError; use crate::client::elicitation_capability_for_server; use crate::declared_openai_file_input_param_names; +use crate::elicitation::ElicitationRequestManager; +use crate::elicitation::elicitation_is_rejected_by_policy; use crate::mcp_connection::CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION; use crate::mcp_connection::CodexAppsToolsCacheContext; -use crate::mcp_connection::ElicitationRequestManager; use crate::mcp_connection::ToolFilter; use crate::mcp_connection::ToolInfo; -use crate::mcp_connection::elicitation_is_rejected_by_policy; use crate::mcp_connection::filter_tools; use crate::mcp_connection::load_startup_cached_codex_apps_tools_snapshot; use crate::mcp_connection::read_cached_codex_apps_tools;