codex: address PR review feedback (#16041)

codex-rs/tui/src/chatwidget.rs

@@ -26,6 +26,7 @@
 //! progress indicators; once commentary completes and stream queues drain, we
 //! re-show it so users still see turn-in-progress state between output bursts.
 use std::collections::BTreeMap;
+use std::collections::BTreeSet;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::collections::VecDeque;
@@ -2963,26 +2964,17 @@ impl ChatWidget {
         let mut failed = Vec::new();
         let mut cancelled = Vec::new();
 
-        if let Some(expected_servers) = &self.mcp_startup_expected_servers
-            && !expected_servers.is_empty()
-        {
-            for name in expected_servers {
-                match current.get(name) {
-                    Some(McpStartupStatus::Ready) => {}
-                    Some(McpStartupStatus::Failed { .. }) => failed.push(name.clone()),
-                    Some(McpStartupStatus::Cancelled | McpStartupStatus::Starting) | None => {
-                        cancelled.push(name.clone());
-                    }
-                }
-            }
-        } else {
-            for (name, state) in current {
-                match state {
-                    McpStartupStatus::Ready => {}
-                    McpStartupStatus::Failed { .. } => failed.push(name.clone()),
-                    McpStartupStatus::Cancelled | McpStartupStatus::Starting => {
-                        cancelled.push(name.clone());
-                    }
+        let mut server_names: BTreeSet<String> = current.keys().cloned().collect();
+        if let Some(expected_servers) = &self.mcp_startup_expected_servers {
+            server_names.extend(expected_servers.iter().cloned());
+        }
+
+        for name in server_names {
+            match current.get(&name) {
+                Some(McpStartupStatus::Ready) => {}
+                Some(McpStartupStatus::Failed { .. }) => failed.push(name),
+                Some(McpStartupStatus::Cancelled | McpStartupStatus::Starting) | None => {
+                    cancelled.push(name);
                 }
             }
         }

codex-rs/tui/src/chatwidget/tests.rs

@@ -12161,6 +12161,40 @@ async fn app_server_mcp_startup_after_lag_with_empty_expected_servers_preserves_
     assert!(!chat.bottom_pane.is_task_running());
 }
 
+#[tokio::test]
+async fn app_server_mcp_startup_after_lag_includes_runtime_servers_with_expected_set() {
+    let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await;
+    chat.show_welcome_banner = false;
+    chat.set_mcp_startup_expected_servers(["alpha".to_string()]);
+
+    chat.on_mcp_startup_update(McpStartupUpdateEvent {
+        server: "alpha".to_string(),
+        status: McpStartupStatus::Ready,
+    });
+    chat.on_mcp_startup_update(McpStartupUpdateEvent {
+        server: "runtime".to_string(),
+        status: McpStartupStatus::Failed {
+            error: "MCP client for `runtime` failed to start: handshake failed".to_string(),
+        },
+    });
+
+    let warning_text = drain_insert_history(&mut rx)
+        .iter()
+        .map(|lines| lines_to_single_string(lines))
+        .collect::<String>();
+    assert!(warning_text.contains("MCP client for `runtime` failed to start: handshake failed"));
+    assert!(chat.bottom_pane.is_task_running());
+
+    chat.finish_mcp_startup_after_lag();
+
+    let summary_text = drain_insert_history(&mut rx)
+        .iter()
+        .map(|lines| lines_to_single_string(lines))
+        .collect::<String>();
+    assert!(summary_text.contains("MCP startup incomplete (failed: runtime)"));
+    assert!(!chat.bottom_pane.is_task_running());
+}
+
 #[tokio::test]
 async fn app_server_mcp_startup_next_round_after_lag_can_settle_without_starting_updates() {
     let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await;