From d838c2346d05f4feb7928ea92e13e83dd479ba10 Mon Sep 17 00:00:00 2001 From: vkg-oai Date: Thu, 24 Sep 2026 00:25:57 +0000 Subject: [PATCH] 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 --- .../src/endpoint/responses_websocket.rs | 6 +- codex-rs/core/src/client.rs | 33 +++- codex-rs/core/src/codex_thread.rs | 18 ++ .../core/src/guardian/input_budget_tests.rs | 2 +- codex-rs/core/src/lib.rs | 1 - codex-rs/core/src/session/handlers.rs | 8 +- codex-rs/core/src/session/mod.rs | 4 +- codex-rs/core/src/session/session.rs | 3 +- .../startup_prewarm.rs} | 163 +++++++++++++----- codex-rs/core/src/session/tests.rs | 8 +- codex-rs/core/src/state/session.rs | 7 +- codex-rs/core/src/tasks/regular.rs | 2 +- .../core/tests/suite/client_websockets.rs | 86 +++++++++ codex-rs/core/tests/suite/rmcp_client.rs | 60 ++++--- 14 files changed, 310 insertions(+), 91 deletions(-) rename codex-rs/core/src/{session_startup_prewarm.rs => session/startup_prewarm.rs} (70%) diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 4890ef5cbf..56debdc6ad 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/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( diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 975aad120a..f69e28bd7b 100644 --- a/codex-rs/core/src/client.rs +++ b/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) -> 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 { 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(()); } diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index f7e10cfb33..b4fa389960 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/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> { + 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, diff --git a/codex-rs/core/src/guardian/input_budget_tests.rs b/codex-rs/core/src/guardian/input_budget_tests.rs index cbe095f4b7..b9dd1f1297 100644 --- a/codex-rs/core/src/guardian/input_budget_tests.rs +++ b/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, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 1ce7587fd6..2595ce5bb9 100644 --- a/codex-rs/core/src/lib.rs +++ b/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; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 8ab5964427..c6648e3df6 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -284,7 +284,13 @@ pub async fn set_thread_memory_mode(sess: &Arc, sub_id: String, mode: T } pub(super) async fn shutdown_session_runtime(sess: &Arc) { - 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; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 3a76a5b3c7..2aea6cfa22 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 22f3d42d9d..aa04cbb324 100644 --- a/codex-rs/core/src/session/session.rs +++ b/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 diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session/startup_prewarm.rs similarity index 70% rename from codex-rs/core/src/session_startup_prewarm.rs rename to codex-rs/core/src/session/startup_prewarm.rs index e6846aa243..7b8d5978b2 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/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>, + task: SessionStartupPrewarmTask, started_at: Instant, timeout: Duration, } +enum SessionStartupPrewarmTask { + Running(AbortOnDropHandle>), + Completed(SessionStartupPrewarmResolution), +} + pub(crate) enum SessionStartupPrewarmResolution { Cancelled, Ready(Box), @@ -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, base_instructions: String) { + pub(crate) async fn schedule_startup_prewarm(self: &Arc) { + 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, - base_instructions: String, + client_session: Option, ) -> CodexResult { 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 { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 14431708f0..8c8e1cf34e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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, diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index e3ddd2cc24..27e366fa0a 100644 --- a/codex-rs/core/src/state/session.rs +++ b/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, /// Retained after completion so later turns do not repeat speculative captures. pub(crate) shell_snapshot_prewarm: Option>, @@ -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(), diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 5068ebf5a8..bb8f00dc52 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/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; diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 2da2ff7116..d57c3b67fb 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/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!(); diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 11e6bdf92c..0fca7d0ea5 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/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;