From 3b9d2faa172459df6adb207c9eebb53940b18cd8 Mon Sep 17 00:00:00 2001
From: Felipe Coury <felipe.coury@openai.com>
Date: Sat, 30 May 2026 15:36:48 -0300
Subject: [PATCH] test(core): make mcp readiness waits state aware

---
 codex-rs/core/src/codex_thread.rs   | 19 +++++++++++++++++++
 codex-rs/core/tests/common/lib.rs   | 28 +++-------------------------
 codex-rs/core/tests/suite/client.rs |  4 ++++
 3 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs
index dec1ce6895..b8c97dd55e 100644
--- a/codex-rs/core/src/codex_thread.rs
+++ b/codex-rs/core/src/codex_thread.rs
@@ -23,6 +23,7 @@ use codex_protocol::openai_models::ReasoningEffort;
 use codex_protocol::protocol::AdditionalContextEntry;
 use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::Event;
+use codex_protocol::protocol::McpStartupFailure;
 use codex_protocol::protocol::Op;
 use codex_protocol::protocol::SandboxPolicy;
 use codex_protocol::protocol::SessionConfiguredEvent;
@@ -142,6 +143,24 @@ impl CodexThread {
         self.codex.shutdown_and_wait().await
     }
 
+    #[doc(hidden)]
+    pub async fn wait_for_mcp_servers_ready(
+        &self,
+        server_names: &[String],
+    ) -> Vec<McpStartupFailure> {
+        let readiness = {
+            let manager = self
+                .codex
+                .session
+                .services
+                .mcp_connection_manager
+                .read()
+                .await;
+            manager.wait_for_servers_ready(server_names)
+        };
+        readiness.await
+    }
+
     /// Wait until the underlying session loop has terminated.
     pub async fn wait_until_terminated(&self) {
         self.codex.session_loop_termination.clone().await;
diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs
index 3d7fd21228..bfd10ce848 100644
--- a/codex-rs/core/tests/common/lib.rs
+++ b/codex-rs/core/tests/common/lib.rs
@@ -250,34 +250,12 @@ where
 
 /// Waits for a configured MCP server to finish startup and requires it to be ready.
 pub async fn wait_for_mcp_server(codex: &CodexThread, server_name: &str) -> anyhow::Result<()> {
-    use codex_protocol::protocol::EventMsg;
-
-    // Wait for the startup summary regardless of outcome, then interpret the
-    // requested server's ready, failed, or cancelled entry below.
-    let summary = loop {
-        let event = codex
-            .next_event()
-            .await
-            .expect("stream ended unexpectedly while waiting for MCP startup");
-        if let EventMsg::McpStartupComplete(summary) = event.msg {
-            break summary;
-        }
-    };
-    if let Some(failure) = summary
-        .failed
-        .iter()
-        .find(|failure| failure.server == server_name)
-    {
+    let server_names = vec![server_name.to_string()];
+    let failures = codex.wait_for_mcp_servers_ready(&server_names).await;
+    if let Some(failure) = failures.first() {
         let error = &failure.error;
         anyhow::bail!("MCP server {server_name} failed to start: {error}");
     }
-    if summary.cancelled.iter().any(|server| server == server_name) {
-        anyhow::bail!("MCP server {server_name} startup was cancelled");
-    }
-    assert!(
-        summary.ready.iter().any(|server| server == server_name),
-        "expected MCP server {server_name} to be ready; startup summary: {summary:?}"
-    );
     Ok(())
 }
 
diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs
index 7234c59c0f..07512f8777 100644
--- a/codex-rs/core/tests/suite/client.rs
+++ b/codex-rs/core/tests/suite/client.rs
@@ -67,6 +67,7 @@ use core_test_support::skip_if_no_network;
 use core_test_support::test_codex::TestCodex;
 use core_test_support::test_codex::test_codex;
 use core_test_support::wait_for_event;
+use core_test_support::wait_for_mcp_server;
 use dunce::canonicalize as normalize_path;
 use futures::StreamExt;
 use pretty_assertions::assert_eq;
@@ -1273,6 +1274,9 @@ async fn includes_apps_guidance_as_developer_message_for_chatgpt_auth() {
         .await
         .expect("create new conversation")
         .codex;
+    wait_for_mcp_server(&codex, "codex_apps")
+        .await
+        .expect("wait for apps MCP startup");
 
     codex
         .submit(Op::UserInput {