[code-mode] wait for pending MCP elicitations

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::<Value>(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"

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,

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<usize>,
+}
+
+fn begin_elicitation_activity(active_count: &watch::Sender<usize>) -> 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<dyn ElicitationReviewer>;
 #[derive(Clone)]
 pub(crate) struct ElicitationRequestManager {
     requests: Arc<Mutex<ResponderMap>>,
+    active_elicitations: watch::Sender<usize>,
     pub(crate) approval_policy: Arc<StdMutex<AskForApproval>>,
     pub(crate) permission_profile: Arc<StdMutex<PermissionProfile>>,
     auto_deny: Arc<StdMutex<bool>>,
@@ -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<Event>,
     ) -> 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(),

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;

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, .. } => {