From ed2c5a81263e4d8db7d2c03cea0d3942679e1f61 Mon Sep 17 00:00:00 2001 From: Samuel Yuan Date: Mon, 15 Jun 2026 17:59:33 -0700 Subject: [PATCH] [code-mode] wait for pending MCP elicitations --- .../tests/suite/v2/mcp_server_elicitation.rs | 58 +++++++------------ codex-rs/codex-mcp/src/connection_manager.rs | 9 +++ codex-rs/codex-mcp/src/elicitation.rs | 38 ++++++++++++ codex-rs/core/src/session/mcp.rs | 5 ++ codex-rs/core/src/tools/code_mode/mod.rs | 7 ++- 5 files changed, 80 insertions(+), 37 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs b/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs index b4bd1d377a..56d57eb79c 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs @@ -32,7 +32,6 @@ use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; use codex_config::types::AuthCredentialsStoreMode; -use core_test_support::assert_regex_match; use core_test_support::responses; use pretty_assertions::assert_eq; use rmcp::handler::server::ServerHandler; @@ -56,7 +55,6 @@ use rmcp::service::RoleServer; use rmcp::transport::StreamableHttpServerConfig; use rmcp::transport::StreamableHttpService; use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; -use serde_json::Value; use serde_json::json; use tempfile::TempDir; use tokio::net::TcpListener; @@ -73,9 +71,14 @@ const TOOL_CALL_ID: &str = "call-calendar-confirm"; const ELICITATION_MESSAGE: &str = "Allow this request?"; #[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn mcp_server_elicitation_round_trip() -> Result<()> { +async fn mcp_server_elicitation_pauses_after_code_mode_yield() -> Result<()> { let responses_server = responses::start_mock_server().await; - let tool_call_arguments = serde_json::to_string(&json!({}))?; + let code = format!( + r#"// @exec: {{"yield_time_ms": 1000}} +const pending = tools["{TOOL_NAMESPACE}{CALLABLE_TOOL_NAME}"]({{}}); +await pending; +"# + ); let response_mock = responses::mount_sse_sequence( &responses_server, vec![ @@ -86,12 +89,7 @@ async fn mcp_server_elicitation_round_trip() -> Result<()> { ]), responses::sse(vec![ responses::ev_response_created("resp-1"), - responses::ev_function_call_with_namespace( - TOOL_CALL_ID, - TOOL_NAMESPACE, - CALLABLE_TOOL_NAME, - &tool_call_arguments, - ), + responses::ev_custom_tool_call(TOOL_CALL_ID, "exec", &code), responses::ev_completed("resp-1"), ]), responses::sse(vec![ @@ -213,6 +211,13 @@ async fn mcp_server_elicitation_round_trip() -> Result<()> { } ); + tokio::time::sleep(std::time::Duration::from_millis(2000)).await; + assert_eq!( + response_mock.requests().len(), + 2, + "the model must not be sampled while the elicitation is unresolved" + ); + let resolved_request_id = request_id.clone(); mcp.send_response( request_id, @@ -266,32 +271,12 @@ async fn mcp_server_elicitation_round_trip() -> Result<()> { let requests = response_mock.requests(); assert_eq!(requests.len(), 3); - let function_call_output = requests[2].function_call_output(TOOL_CALL_ID); - assert_eq!( - function_call_output.get("type"), - Some(&Value::String("function_call_output".to_string())) - ); - assert_eq!( - function_call_output.get("call_id"), - Some(&Value::String(TOOL_CALL_ID.to_string())) - ); - let output = function_call_output - .get("output") - .and_then(Value::as_str) - .expect("function_call_output output should be a JSON string"); - let payload = assert_regex_match( - r#"(?s)^Wall time: [0-9]+(?:\.[0-9]+)? seconds\nOutput:\n(.*)$"#, - output, - ) - .get(1) - .expect("wall-time wrapped output should include payload") - .as_str(); - assert_eq!( - serde_json::from_str::(payload)?, - json!([{ - "type": "text", - "text": "accepted" - }]) + let (output, _) = requests[2] + .custom_tool_call_output_content_and_success(TOOL_CALL_ID) + .expect("follow-up request should contain the yielded exec output"); + assert!( + output.is_some_and(|text| text.starts_with("Script running with cell ID ")), + "expected the outer exec call to yield" ); apps_server_handle.abort(); @@ -476,6 +461,7 @@ mcp_oauth_credentials_store = "file" [features] apps = true +code_mode = true [model_providers.mock_provider] name = "Mock provider for test" diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 3596486568..7a8c33260d 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -421,6 +421,15 @@ impl McpConnectionManager { self.elicitation_requests.set_auto_deny(auto_deny); } + #[must_use = "dropping the guard marks the MCP elicitation inactive"] + pub fn begin_elicitation(&self) -> impl Drop + Send + 'static { + self.elicitation_requests.begin_activity() + } + + pub async fn wait_until_elicitations_complete(&self) { + self.elicitation_requests.wait_until_inactive().await; + } + pub async fn resolve_elicitation( &self, server_name: String, diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index a51cd7c623..ff0db943ca 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -31,6 +31,29 @@ use rmcp::model::ElicitationAction; use rmcp::model::RequestId; use tokio::sync::Mutex; use tokio::sync::oneshot; +use tokio::sync::watch; + +#[derive(Debug)] +#[must_use = "dropping the guard marks the MCP elicitation inactive"] +struct McpElicitationActivityGuard { + active_count: watch::Sender, +} + +fn begin_elicitation_activity(active_count: &watch::Sender) -> McpElicitationActivityGuard { + active_count.send_modify(|count| *count += 1); + McpElicitationActivityGuard { + active_count: active_count.clone(), + } +} + +impl Drop for McpElicitationActivityGuard { + fn drop(&mut self) { + self.active_count.send_modify(|count| { + debug_assert!(*count > 0); + *count -= 1; + }); + } +} #[derive(Debug, Clone)] pub struct ElicitationReviewRequest { @@ -51,6 +74,7 @@ pub type ElicitationReviewerHandle = Arc; #[derive(Clone)] pub(crate) struct ElicitationRequestManager { requests: Arc>, + active_elicitations: watch::Sender, pub(crate) approval_policy: Arc>, pub(crate) permission_profile: Arc>, auto_deny: Arc>, @@ -65,6 +89,7 @@ impl ElicitationRequestManager { ) -> Self { Self { requests: Arc::new(Mutex::new(HashMap::new())), + active_elicitations: watch::channel(0).0, approval_policy: Arc::new(StdMutex::new(approval_policy)), permission_profile: Arc::new(StdMutex::new(permission_profile)), auto_deny: Arc::new(StdMutex::new(false)), @@ -85,6 +110,15 @@ impl ElicitationRequestManager { } } + pub(crate) fn begin_activity(&self) -> impl Drop + Send + 'static { + begin_elicitation_activity(&self.active_elicitations) + } + + pub(crate) async fn wait_until_inactive(&self) { + let mut active_count = self.active_elicitations.subscribe(); + let _ = active_count.wait_for(|count| *count == 0).await; + } + pub(crate) async fn resolve( &self, server_name: String, @@ -106,12 +140,14 @@ impl ElicitationRequestManager { tx_event: Sender, ) -> SendElicitation { let elicitation_requests = self.requests.clone(); + let active_elicitations = self.active_elicitations.clone(); let approval_policy = self.approval_policy.clone(); let permission_profile = self.permission_profile.clone(); let auto_deny = self.auto_deny.clone(); let reviewer = self.reviewer.clone(); Box::new(move |id, elicitation| { let elicitation_requests = elicitation_requests.clone(); + let active_elicitations = active_elicitations.clone(); let tx_event = tx_event.clone(); let server_name = server_name.clone(); let approval_policy = approval_policy.clone(); @@ -160,6 +196,8 @@ impl ElicitationRequestManager { }); } + let _activity_guard = begin_elicitation_activity(&active_elicitations); + if let Some(reviewer) = reviewer.as_ref() { let request = ElicitationReviewRequest { server_name: server_name.clone(), diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index cc6e42f852..d28614c1c0 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -156,6 +156,11 @@ impl Session { }, }; + let _activity_guard = self + .services + .mcp_connection_manager + .load() + .begin_elicitation(); let (tx_response, rx_response) = oneshot::channel(); let prev_entry = { let mut active = self.active_turn.lock().await; diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ade23cf3b7..fa94de7061 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -149,10 +149,15 @@ pub(super) async fn handle_runtime_response( match response { RuntimeResponse::Yielded { content_items, .. } => { + let wall_time = started_at.elapsed(); + let mcp_connection_manager = exec.session.services.mcp_connection_manager.load_full(); + mcp_connection_manager + .wait_until_elicitations_complete() + .await; let mut content_items = into_function_call_output_content_items(content_items); sanitize_runtime_image_detail(exec.turn.as_ref(), &mut content_items); content_items = truncate_code_mode_result(content_items, max_output_tokens); - prepend_script_status(&mut content_items, &script_status, started_at.elapsed()); + prepend_script_status(&mut content_items, &script_status, wall_time); Ok(FunctionToolOutput::from_content(content_items, Some(true))) } RuntimeResponse::Terminated { content_items, .. } => {