Allow idle threads to prewarm and repair WebSocket connections (#47701)

## Why

A previously used WebSocket can close while a thread is idle. Hosts need a way to prepare the connection before the next turn without repeating warmup on a healthy connection.

## What changed

- Expose `CodexThread::prewarm()` using the existing startup warmup and next-turn handoff.
- Reuse healthy connections and pending warmups, retry failed warmups, and repair closed connections. Treat a finished WebSocket reader task as a closed connection.
- Skip scheduling during an active turn or shutdown, and stop warmup admission before shutdown takes the pending task.
- Expose `active_turn_environment_selections()` so hosts can inspect the environments captured by the active turn, including those still starting or failed, when authorizing steering.

## Testing

Add regression coverage for repeated idle prewarm calls, healthy connection reuse with preserved response continuation, reconnection with `generate: false`, and warmup rejection after shutdown. Update delayed MCP discovery coverage to verify prewarm repairs an eagerly closed socket and the first turn reuses it.

GitOrigin-RevId: aa164b91a399de2375eda3bf28397725439fce93

codex-rs/codex-api/src/endpoint/responses_websocket.rs

@@ -223,7 +223,11 @@ impl ResponsesWebsocketConnection {
     }
 
     pub async fn is_closed(&self) -> bool {
-        self.stream.lock().await.is_none()
+        self.stream
+            .lock()
+            .await
+            .as_ref()
+            .is_none_or(|stream| stream.pump_task.is_finished())
     }
 
     #[instrument(

codex-rs/core/src/client.rs

@@ -100,6 +100,7 @@ use codex_tools::create_tools_json_for_responses_lite;
 use codex_tools::create_tools_raw_json_for_responses_api;
 use eventsource_stream::Event;
 use eventsource_stream::EventStreamError;
+use futures::FutureExt;
 use futures::StreamExt;
 use http::HeaderMap as ApiHeaderMap;
 use http::HeaderValue;
@@ -417,6 +418,17 @@ fn is_internal_metadata_destination(provider: &ApiProvider) -> bool {
 }
 
 impl WebsocketSession {
+    /// A nonblocking resume check. A busy stream is already in use, so resume
+    /// leaves it alone. The next request still validates provider, auth, and headers.
+    fn is_prewarmed(&self, auth_owner_generation: Option<u64>) -> bool {
+        self.auth_owner_generation == auth_owner_generation
+            && self.last_request.is_some()
+            && self
+                .connection
+                .as_ref()
+                .is_some_and(|connection| connection.is_closed().now_or_never() != Some(true))
+    }
+
     fn reset(&mut self, reason: Option<&'static str>) {
         // Per-socket backend metrics call a resend after reconnect "initial".
         // Retain the loss reason across reconnects/turns until the next send.
@@ -618,6 +630,15 @@ impl ModelClient {
         self.state.provider.auth_manager()
     }
 
+    pub(crate) fn is_websocket_prewarmed(&self) -> bool {
+        let auth_owner_generation = self.auth_owner_generation();
+        self.state
+            .cached_websocket_session
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .is_prewarmed(auth_owner_generation)
+    }
+
     fn auth_owner_generation(&self) -> Option<u64> {
         self.auth_manager().map(|manager| {
             manager
@@ -1340,6 +1361,11 @@ impl Drop for ModelClientSession {
 }
 
 impl ModelClientSession {
+    pub(crate) fn is_websocket_prewarmed(&self) -> bool {
+        self.websocket_session
+            .is_prewarmed(self.client.auth_owner_generation())
+    }
+
     #[allow(clippy::too_many_arguments)]
     /// Builds shared Responses API transport options and request-body options.
     ///
@@ -2126,7 +2152,12 @@ impl ModelClientSession {
         if !self.client.responses_websocket_enabled() {
             return Ok(());
         }
-        if self.websocket_session.last_request.is_some() {
+        // A previous request only makes prewarm unnecessary while its socket is
+        // still alive. Reconnecting below resets the old continuation normally.
+        if self.websocket_session.last_request.is_some()
+            && let Some(connection) = self.websocket_session.connection.as_ref()
+            && !connection.is_closed().await
+        {
             return Ok(());
         }
 

codex-rs/core/src/codex_thread.rs

@@ -239,6 +239,13 @@ impl CodexThread {
         self.session.services.session_telemetry.clone()
     }
 
+    /// Schedule the same background model warmup used at startup for an idle thread.
+    /// The next turn consumes the warmup through the existing startup handoff.
+    /// Call after installing host services such as the thread's attestation routing.
+    pub async fn prewarm(&self) {
+        self.session.schedule_startup_prewarm().await;
+    }
+
     /// Whether analytics is enabled for this thread after configuration and host overrides.
     pub fn analytics_enabled(&self) -> bool {
         self.session.services.analytics_events_client.is_enabled()
@@ -525,6 +532,17 @@ impl CodexThread {
         self.session.inject_if_running(items).await
     }
 
+    /// Environment selections captured by the active turn, before later settings updates.
+    /// Includes environments that are still starting or have failed. Hosts use this snapshot
+    /// to authorize steering against every executor that the active turn selected.
+    pub async fn active_turn_environment_selections(
+        &self,
+    ) -> Option<Vec<TurnEnvironmentSelection>> {
+        let active = self.session.active_turn.lock().await;
+        let task = active.as_ref()?.task.as_ref()?;
+        Some(task.turn_context.initial_environments.all_selections())
+    }
+
     /// Captures a regular turn only after its input is recorded. The caller must flush the rollout.
     pub async fn interrupted_turn(
         &self,

codex-rs/core/src/guardian/input_budget_tests.rs

@@ -49,7 +49,7 @@ async fn cancelled_startup_does_not_record_unselected_review_evidence() {
         .insert(PendingReviewContext(context));
     session
         .set_session_startup_prewarm(
-            crate::session_startup_prewarm::SessionStartupPrewarmHandle::new(
+            crate::session::startup_prewarm::SessionStartupPrewarmHandle::new(
                 tokio::spawn(std::future::pending()),
                 std::time::Instant::now(),
                 crate::client::WEBSOCKET_CONNECT_TIMEOUT,

codex-rs/core/src/lib.rs

@@ -128,7 +128,6 @@ pub(crate) mod mentions {
 mod sandbox_tags;
 pub mod sandboxing;
 mod session_prefix;
-mod session_startup_prewarm;
 mod skills;
 pub(crate) use skills::maybe_emit_implicit_skill_invocation;
 pub(crate) use skills::skills_load_input_from_config;

codex-rs/core/src/session/handlers.rs

@@ -284,7 +284,13 @@ pub async fn set_thread_memory_mode(sess: &Arc<Session>, sub_id: String, mode: T
 }
 
 pub(super) async fn shutdown_session_runtime(sess: &Arc<Session>) {
-    if let Some(startup_prewarm) = sess.take_session_startup_prewarm().await {
+    let startup_prewarm = {
+        let mut state = sess.state.lock().await;
+        // Stop admission and take the current warmup together so resume cannot replace it.
+        state.shutting_down = true;
+        state.take_session_startup_prewarm()
+    };
+    if let Some(startup_prewarm) = startup_prewarm {
         startup_prewarm.abort().await;
     }
     let _ = sess.conversation.shutdown().await;

codex-rs/core/src/session/mod.rs

@@ -253,6 +253,7 @@ mod rollout_budget;
 mod rollout_reconstruction;
 #[allow(clippy::module_inception)]
 pub(crate) mod session;
+pub(crate) mod startup_prewarm;
 mod step_activation;
 pub(crate) mod step_context;
 pub(crate) mod step_settings;
@@ -296,7 +297,7 @@ use crate::mcp::McpManager;
 use crate::mcp::McpThreadIdentity;
 use crate::network_policy_decision::execpolicy_network_rule_amendment;
 use crate::rollout::map_session_init_error;
-use crate::session_startup_prewarm::SessionStartupPrewarmHandle;
+use crate::session::startup_prewarm::SessionStartupPrewarmHandle;
 use crate::shell;
 use crate::state::AcceptedUserInputResponse;
 use crate::state::AutoCompactWindowIds;
@@ -1994,6 +1995,7 @@ impl Session {
             .map_or_else(Vec::new, |instructions| instructions.sources().collect())
     }
 
+    #[cfg(test)]
     pub(crate) async fn set_session_startup_prewarm(
         &self,
         startup_prewarm: SessionStartupPrewarmHandle,

codex-rs/core/src/session/session.rs

@@ -1872,8 +1872,7 @@ impl Session {
             )
             .await?;
             sess.start_mcp_prewarm_worker(mcp_prewarm_rx, mcp_auth_changes);
-            sess.schedule_startup_prewarm(sess.get_prompt_base_instructions().await.text)
-                .await;
+            sess.schedule_startup_prewarm().await;
             let session_start_source = match &initial_history {
                 InitialHistory::Forked(_) if forked_from_id.is_some() => {
                     codex_hooks::SessionStartSource::Fork

codex-rs/core/src/session/startup_prewarm.rs

@@ -1,7 +1,12 @@
+//! Background model warmup shared by startup and idle-thread resume.
+//! One scheduled task owns the prepared client session until the next regular
+//! turn consumes it; shutdown and turn cancellation use the same handoff.
+
 use std::sync::Arc;
 use std::time::Duration;
 use std::time::Instant;
 
+use futures::FutureExt;
 use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 use tokio_util::task::AbortOnDropHandle;
@@ -27,11 +32,16 @@ use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::SubAgentSource;
 
 pub(crate) struct SessionStartupPrewarmHandle {
-    task: AbortOnDropHandle<CodexResult<ModelClientSession>>,
+    task: SessionStartupPrewarmTask,
     started_at: Instant,
     timeout: Duration,
 }
 
+enum SessionStartupPrewarmTask {
+    Running(AbortOnDropHandle<CodexResult<ModelClientSession>>),
+    Completed(SessionStartupPrewarmResolution),
+}
+
 pub(crate) enum SessionStartupPrewarmResolution {
     Cancelled,
     Ready(Box<ModelClientSession>),
@@ -48,15 +58,17 @@ impl SessionStartupPrewarmHandle {
         timeout: Duration,
     ) -> Self {
         Self {
-            task: AbortOnDropHandle::new(task),
+            task: SessionStartupPrewarmTask::Running(AbortOnDropHandle::new(task)),
             started_at,
             timeout,
         }
     }
 
     pub(crate) async fn abort(self) {
-        self.task.abort();
-        let _ = self.task.await;
+        if let SessionStartupPrewarmTask::Running(task) = self.task {
+            task.abort();
+            let _ = task.await;
+        }
     }
 
     #[instrument(name = "startup_prewarm.resolve", level = "trace", skip_all)]
@@ -67,48 +79,55 @@ impl SessionStartupPrewarmHandle {
     ) -> SessionStartupPrewarmResolution {
         let resolve_started_at = Instant::now();
         let Self {
-            mut task,
+            task,
             started_at,
             timeout,
         } = self;
         let age_at_first_turn = started_at.elapsed();
         let remaining = timeout.saturating_sub(age_at_first_turn);
 
-        let resolution = if task.is_finished() {
-            Self::resolution_from_join_result(task.await, started_at)
-        } else {
-            match tokio::select! {
-                _ = cancellation_token.cancelled() => None,
-                result = tokio::time::timeout(remaining, &mut task) => Some(result),
-            } {
-                Some(Ok(result)) => Self::resolution_from_join_result(result, started_at),
-                Some(Err(_elapsed)) => {
-                    task.abort();
-                    info!("startup websocket prewarm timed out before the first turn could use it");
-                    SessionStartupPrewarmResolution::Unavailable {
-                        status: "timed_out",
-                        prewarm_duration: Some(started_at.elapsed()),
+        let resolution = match task {
+            SessionStartupPrewarmTask::Completed(resolution) => resolution,
+            SessionStartupPrewarmTask::Running(mut task) => {
+                if task.is_finished() {
+                    Self::resolution_from_join_result(task.await, started_at)
+                } else {
+                    match tokio::select! {
+                        _ = cancellation_token.cancelled() => None,
+                        result = tokio::time::timeout(remaining, &mut task) => Some(result),
+                    } {
+                        Some(Ok(result)) => Self::resolution_from_join_result(result, started_at),
+                        Some(Err(_elapsed)) => {
+                            task.abort();
+                            info!(
+                                "startup websocket prewarm timed out before the first turn could use it"
+                            );
+                            SessionStartupPrewarmResolution::Unavailable {
+                                status: "timed_out",
+                                prewarm_duration: Some(started_at.elapsed()),
+                            }
+                        }
+                        None => {
+                            task.abort();
+                            session_telemetry.record_startup_phase(
+                                "startup_prewarm_resolve",
+                                resolve_started_at.elapsed(),
+                                Some("cancelled"),
+                            );
+                            session_telemetry.record_duration(
+                                STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC,
+                                age_at_first_turn,
+                                &[("status", "cancelled")],
+                            );
+                            session_telemetry.record_duration(
+                                STARTUP_PREWARM_DURATION_METRIC,
+                                started_at.elapsed(),
+                                &[("status", "cancelled")],
+                            );
+                            return SessionStartupPrewarmResolution::Cancelled;
+                        }
                     }
                 }
-                None => {
-                    task.abort();
-                    session_telemetry.record_startup_phase(
-                        "startup_prewarm_resolve",
-                        resolve_started_at.elapsed(),
-                        Some("cancelled"),
-                    );
-                    session_telemetry.record_duration(
-                        STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC,
-                        age_at_first_turn,
-                        &[("status", "cancelled")],
-                    );
-                    session_telemetry.record_duration(
-                        STARTUP_PREWARM_DURATION_METRIC,
-                        started_at.elapsed(),
-                        &[("status", "cancelled")],
-                    );
-                    return SessionStartupPrewarmResolution::Cancelled;
-                }
             }
         };
         let status = match &resolution {
@@ -185,7 +204,58 @@ impl SessionStartupPrewarmHandle {
 }
 
 impl Session {
-    pub(crate) async fn schedule_startup_prewarm(self: &Arc<Self>, base_instructions: String) {
+    pub(crate) async fn schedule_startup_prewarm(self: &Arc<Self>) {
+        let websocket_connect_timeout = self.provider().await.websocket_connect_timeout();
+        let mut state = self.state.lock().await;
+        if state.shutting_down {
+            return;
+        }
+        // Publish the warmup before another turn can be admitted. No network work
+        // or awaits occur while these guards are held.
+        let Ok(active_turn) = self.active_turn.try_lock() else {
+            return;
+        };
+        if active_turn.is_some() {
+            return;
+        }
+        if let Some(prewarm) = state.startup_prewarm.as_mut() {
+            if let SessionStartupPrewarmTask::Running(task) = &mut prewarm.task {
+                let Some(completed) = task.now_or_never() else {
+                    return;
+                };
+                prewarm.task = SessionStartupPrewarmTask::Completed(
+                    SessionStartupPrewarmHandle::resolution_from_join_result(
+                        completed,
+                        prewarm.started_at,
+                    ),
+                );
+            }
+            if let SessionStartupPrewarmTask::Completed(SessionStartupPrewarmResolution::Ready(
+                client_session,
+            )) = &prewarm.task
+                && client_session.is_websocket_prewarmed()
+            {
+                return;
+            }
+        } else if self.services.model_client.is_websocket_prewarmed() {
+            return;
+        }
+
+        // Keep ownership explicit when repairing a completed warmup's connection.
+        let client_session = match state.startup_prewarm.take() {
+            Some(prewarm) => match prewarm.task {
+                SessionStartupPrewarmTask::Completed(SessionStartupPrewarmResolution::Ready(
+                    client_session,
+                )) => Some(*client_session),
+                SessionStartupPrewarmTask::Completed(_) => None,
+                SessionStartupPrewarmTask::Running(_) => {
+                    state.startup_prewarm = Some(prewarm);
+                    return;
+                }
+            },
+            None => None,
+        };
+
         if self.features().enabled(Feature::CodeModePrewarm)
             && self.services.code_mode_service.is_available()
         {
@@ -210,14 +280,12 @@ impl Session {
         }
 
         let session_telemetry = self.services.session_telemetry.clone();
-        let websocket_connect_timeout = self.provider().await.websocket_connect_timeout();
         let started_at = Instant::now();
         let startup_prewarm_session = Arc::clone(self);
         let startup_prewarm = tokio::spawn(
             async move {
                 let result =
-                    schedule_startup_prewarm_inner(startup_prewarm_session, base_instructions)
-                        .await;
+                    schedule_startup_prewarm_inner(startup_prewarm_session, client_session).await;
                 let status = if result.is_ok() { "ready" } else { "failed" };
                 session_telemetry.record_startup_phase(
                     "startup_prewarm_total",
@@ -237,12 +305,11 @@ impl Session {
                 thread.id = %self.thread_id(),
             )),
         );
-        self.set_session_startup_prewarm(SessionStartupPrewarmHandle::new(
+        state.set_session_startup_prewarm(SessionStartupPrewarmHandle::new(
             startup_prewarm,
             started_at,
             websocket_connect_timeout,
-        ))
-        .await;
+        ));
     }
 
     pub(crate) async fn consume_startup_prewarm_for_regular_turn(
@@ -263,9 +330,10 @@ impl Session {
 
 async fn schedule_startup_prewarm_inner(
     session: Arc<Session>,
-    base_instructions: String,
+    client_session: Option<ModelClientSession>,
 ) -> CodexResult<ModelClientSession> {
     let prewarm_started_at = Instant::now();
+    let base_instructions = session.get_prompt_base_instructions().await.text;
     let startup_turn_context = session
         .new_startup_prewarm_turn_with_sub_id(INITIAL_SUBMIT_ID.to_owned())
         .await;
@@ -307,7 +375,8 @@ async fn schedule_startup_prewarm_inner(
         window_number,
         context_window_id,
     );
-    let mut client_session = session.services.model_client.new_session();
+    let mut client_session =
+        client_session.unwrap_or_else(|| session.services.model_client.new_session());
     // Start the handshake with the expected route while capturing tools for generate=false.
     let (step_context, ()) = tokio::try_join!(
         async {

codex-rs/core/src/session/tests.rs

@@ -616,7 +616,7 @@ async fn regular_turn_emits_turn_started_with_trace_id_without_waiting_for_start
     });
 
     sess.set_session_startup_prewarm(
-        crate::session_startup_prewarm::SessionStartupPrewarmHandle::new(
+        crate::session::startup_prewarm::SessionStartupPrewarmHandle::new(
             handle,
             std::time::Instant::now(),
             crate::client::WEBSOCKET_CONNECT_TIMEOUT,
@@ -772,7 +772,7 @@ async fn interrupting_regular_turn_waiting_on_startup_prewarm_emits_turn_aborted
     });
 
     sess.set_session_startup_prewarm(
-        crate::session_startup_prewarm::SessionStartupPrewarmHandle::new(
+        crate::session::startup_prewarm::SessionStartupPrewarmHandle::new(
             handle,
             std::time::Instant::now(),
             crate::client::WEBSOCKET_CONNECT_TIMEOUT,
@@ -8291,6 +8291,10 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() {
     assert!(session.async_hook_results.is_empty());
     assert!(result_sender.is_closed());
 
+    assert!(session.services.model_client.responses_websocket_enabled());
+    session.schedule_startup_prewarm().await;
+    assert!(session.state.lock().await.startup_prewarm.is_none());
+
     assert_eq!(
         codex_thread_store::InMemoryThreadStoreCalls {
             create_thread: 1,

codex-rs/core/src/state/session.rs

@@ -17,8 +17,8 @@ use crate::context_manager::ContextManager;
 use crate::context_manager::HistoryReplacement;
 use crate::session::PreviousTurnSettings;
 use crate::session::session::SessionConfiguration;
+use crate::session::startup_prewarm::SessionStartupPrewarmHandle;
 use crate::session::time_reminder::CurrentTimeReminderState;
-use crate::session_startup_prewarm::SessionStartupPrewarmHandle;
 use codex_history::ResponseItemEnvelope;
 use codex_protocol::SessionId;
 use codex_protocol::ThreadId;
@@ -91,7 +91,9 @@ pub(crate) struct SessionState {
     auto_compact_window: AutoCompactWindow,
     /// Original request effort for the current model while configuration updates remain active.
     pub(crate) reasoning_effort_pin: ReasoningEffortPin,
-    /// Startup prewarmed session prepared during session initialization.
+    /// Set under the state lock before shutdown takes the last warmup handle.
+    pub(crate) shutting_down: bool,
+    /// Background model warmup scheduled at startup or while resuming an idle thread.
     pub(crate) startup_prewarm: Option<SessionStartupPrewarmHandle>,
     /// Retained after completion so later turns do not repeat speculative captures.
     pub(crate) shell_snapshot_prewarm: Option<AbortOnDropHandle<()>>,
@@ -133,6 +135,7 @@ impl SessionState {
             last_started_turn_id: None,
             auto_compact_window: AutoCompactWindow::new_with_ids(auto_compact_window_ids),
             reasoning_effort_pin: ReasoningEffortPin::Unset,
+            shutting_down: false,
             startup_prewarm: None,
             shell_snapshot_prewarm: None,
             current_time_reminder: CurrentTimeReminderState::default(),

codex-rs/core/src/tasks/regular.rs

@@ -6,11 +6,11 @@ use tokio_util::sync::CancellationToken;
 
 use crate::session::TurnInput;
 use crate::session::session::Session;
+use crate::session::startup_prewarm::SessionStartupPrewarmResolution;
 use crate::session::turn::McpStartupRequirements;
 use crate::session::turn::run_hooks_and_record_inputs;
 use crate::session::turn::run_turn;
 use crate::session::turn_context::TurnContext;
-use crate::session_startup_prewarm::SessionStartupPrewarmResolution;
 use crate::state::TaskKind;
 use codex_thread_store::PersistContext;
 use tracing::Instrument;

codex-rs/core/tests/suite/client_websockets.rs

@@ -9,7 +9,9 @@ use codex_core::ResponseEvent;
 use codex_core::TurnInputRequest;
 use codex_core::X_CODEX_ROUTING_HINT_HEADER;
 use codex_core::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER;
+use codex_core::test_support::EmptyUserInstructionsProvider;
 use codex_core::test_support::with_parent_turn;
+use codex_extension_api::ExtensionRegistryBuilder;
 use codex_features::Feature;
 use codex_http_client::OutboundProxyPolicy;
 use codex_login::AuthManager;
@@ -50,6 +52,7 @@ use codex_rollout_trace::RawTraceEventPayload;
 use codex_rollout_trace::TraceWriter;
 use codex_rollout_trace::replay_bundle;
 use core_test_support::TestCodexResponsesRequestKind;
+use core_test_support::ThreadIdle;
 use core_test_support::load_default_config_for_test;
 use core_test_support::responses::WebSocketConnectionConfig;
 use core_test_support::responses::WebSocketTestServer;
@@ -60,6 +63,7 @@ use core_test_support::responses::start_websocket_server;
 use core_test_support::responses::start_websocket_server_with_headers;
 use core_test_support::responses_metadata as test_responses_metadata;
 use core_test_support::skip_if_no_network;
+use core_test_support::test_codex::RecordingUserInstructionsProvider;
 use core_test_support::test_codex::test_codex;
 use core_test_support::tracing::install_test_tracing;
 use core_test_support::wait_for_event;
@@ -605,6 +609,88 @@ async fn responses_websocket_preconnect_reuses_connection_without_replacing_turn
     server.shutdown().await;
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn responses_websocket_resume_prewarm_reuses_and_repairs_connection() -> anyhow::Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_websocket_server_with_headers(vec![
+        WebSocketConnectionConfig {
+            requests: vec![
+                vec![ev_response_created("warm-1"), ev_completed("warm-1")],
+                vec![ev_response_created("resp-1"), ev_completed("resp-1")],
+                vec![ev_response_created("resp-2"), ev_completed("resp-2")],
+            ],
+            response_headers: Vec::new(),
+            accept_delay: None,
+            close_after_requests: true,
+        },
+        WebSocketConnectionConfig {
+            requests: vec![
+                vec![ev_response_created("warm-2"), ev_completed("warm-2")],
+                vec![ev_response_created("resp-3"), ev_completed("resp-3")],
+            ],
+            response_headers: Vec::new(),
+            accept_delay: Some(Duration::from_millis(100)),
+            close_after_requests: false,
+        },
+    ])
+    .await;
+    let mut extensions = ExtensionRegistryBuilder::new();
+    extensions.thread_lifecycle_contributor(Arc::new(ThreadIdle));
+    let instructions = Arc::new(RecordingUserInstructionsProvider::new(Arc::new(
+        EmptyUserInstructionsProvider,
+    )));
+    let test = test_codex()
+        .with_extensions(Arc::new(extensions.build()))
+        .with_user_instructions_provider(instructions.clone())
+        .build_with_websocket_server(&server)
+        .await?;
+    test.submit_text_turn("hello").await?;
+    ThreadIdle::wait(&test.codex).await;
+
+    // A healthy warm resume preserves the response baseline and skips preparation.
+    let instruction_loads = instructions.load_count();
+    test.codex.prewarm().await;
+    test.codex.prewarm().await;
+    test.submit_text_turn("continue").await?;
+
+    assert_eq!(server.handshakes().len(), 1);
+    assert_eq!(instructions.load_count(), instruction_loads + 1);
+    let connection = server.single_connection();
+    assert_eq!(connection.len(), 3);
+    assert_eq!(connection[2].body_json()["previous_response_id"], "resp-1");
+
+    // Turn idle does not synchronize with the reader observing the server's close.
+    // Retry resume until it sees the close; pending attempts must still share one socket.
+    ThreadIdle::wait(&test.codex).await;
+    let warmup = tokio::time::timeout(Duration::from_secs(10), async {
+        loop {
+            test.codex.prewarm().await;
+            test.codex.prewarm().await;
+            tokio::select! {
+                request = server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0) => break request,
+                _ = tokio::time::sleep(Duration::from_millis(10)) => {}
+            }
+        }
+    })
+    .await?;
+    assert_eq!(warmup.body_json()["generate"], false);
+    assert!(warmup.body_json().get("previous_response_id").is_none());
+
+    test.submit_text_turn("continue after reconnect").await?;
+    assert_eq!(server.handshakes().len(), 2);
+    let connections = server.connections();
+    assert_eq!(connections[1].len(), 2);
+    assert_eq!(
+        connections[1][1].body_json()["previous_response_id"],
+        "warm-2"
+    );
+
+    test.codex.shutdown_and_wait().await?;
+    server.shutdown().await;
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn responses_websocket_request_prewarm_reuses_connection() {
     skip_if_no_network!();

codex-rs/core/tests/suite/rmcp_client.rs

@@ -1626,7 +1626,7 @@ enum PrewarmMcpScenario {
 
 #[test_case(PrewarmMcpScenario::Responses; "responses")]
 #[test_case(PrewarmMcpScenario::ResponsesLite; "responses lite")]
-#[test_case(PrewarmMcpScenario::EagerSocketCloses; "first turn reconnects when the eager socket closes")]
+#[test_case(PrewarmMcpScenario::EagerSocketCloses; "prewarm reconnects when the eager socket closes")]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn startup_prewarm_connects_before_delayed_mcp_tools_then_reuses_socket(
     scenario: PrewarmMcpScenario,
@@ -1648,11 +1648,12 @@ async fn startup_prewarm_connects_before_delayed_mcp_tools_then_reuses_socket(
         responses::ev_assistant_message("message-1", "done"),
         responses::ev_completed("turn-1"),
     ];
+    let requests = vec![warmup_events, turn_events];
     let connections = if close_eager_socket {
         // A connection without scripted requests closes as soon as it is accepted.
-        vec![Vec::new(), vec![turn_events]]
+        vec![Vec::new(), requests]
     } else {
-        vec![vec![warmup_events, turn_events]]
+        vec![requests]
     };
     let server = responses::start_websocket_server(connections).await;
     let command = stdio_server_bin()?;
@@ -1729,37 +1730,24 @@ async fn startup_prewarm_connects_before_delayed_mcp_tools_then_reuses_socket(
     wait_for_mcp_server(&fixture.codex, "delayed_prewarm").await?;
 
     let namespace = "mcp__delayed_prewarm";
-    if close_eager_socket {
-        fixture.submit_text_turn("hello").await?;
-        let first_turn = server
-            .wait_for_request(/*connection_index*/ 1, /*request_index*/ 0)
-            .await
-            .body_json();
-        assert_eq!(first_turn["type"], "response.create");
-        assert_eq!(first_turn.get("generate"), None);
-        assert_eq!(first_turn.get("previous_response_id"), None);
-        assert!(responses::namespace_child_tool(&first_turn, namespace, "echo").is_some());
-        assert!(server.connections()[0].is_empty());
-        let handshakes = server.handshakes();
-        assert_eq!(handshakes.len(), 2);
-        assert_eq!(
-            handshakes[1].header(X_CODEX_ROUTING_HINT_HEADER),
-            handshakes[0].header(X_CODEX_ROUTING_HINT_HEADER)
-        );
-        fixture.codex.shutdown_and_wait().await?;
-        server.shutdown().await;
-        return Ok(());
-    }
-
-    let prewarm = tokio::time::timeout(
-        Duration::from_secs(5),
-        server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0),
-    )
+    let connection_index = usize::from(close_eager_socket);
+    let prewarm = tokio::time::timeout(Duration::from_secs(5), async {
+        loop {
+            tokio::select! {
+                request = server.wait_for_request(connection_index, /*request_index*/ 0) => break request,
+                // The reader may observe the close after the initial warmup tries to send.
+                _ = tokio::time::sleep(Duration::from_millis(10)), if close_eager_socket => {
+                    fixture.codex.prewarm().await;
+                }
+            }
+        }
+    })
     .await
     .context("prewarm should send its request after MCP discovery")?
     .body_json();
     assert_eq!(prewarm["type"].as_str(), Some("response.create"));
     assert_eq!(prewarm["generate"].as_bool(), Some(false));
+    assert!(prewarm.get("previous_response_id").is_none());
     assert_eq!(prewarm["service_tier"].as_str(), Some(service_tier));
     let tool_body = if use_responses_lite {
         &prewarm["input"][0]
@@ -1785,13 +1773,23 @@ async fn startup_prewarm_connects_before_delayed_mcp_tools_then_reuses_socket(
     fixture.submit_text_turn("hello").await?;
     let first_turn = tokio::time::timeout(
         Duration::from_secs(5),
-        server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 1),
+        server.wait_for_request(connection_index, /*request_index*/ 1),
     )
     .await
     .context("first turn should reuse the prewarmed socket")?
     .body_json();
+    assert_eq!(first_turn["type"], "response.create");
+    assert!(first_turn.get("generate").is_none());
     assert_eq!(first_turn["previous_response_id"], "warm-1");
-    assert_eq!(server.handshakes().len(), 1);
+    let handshakes = server.handshakes();
+    assert_eq!(handshakes.len(), connection_index + 1);
+    if close_eager_socket {
+        assert!(server.connections()[0].is_empty());
+        assert_eq!(
+            handshakes[1].header(X_CODEX_ROUTING_HINT_HEADER),
+            handshakes[0].header(X_CODEX_ROUTING_HINT_HEADER)
+        );
+    }
 
     fixture.codex.shutdown_and_wait().await?;
     server.shutdown().await;