diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index 716270d11e..588756dcb9 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -1855,12 +1855,8 @@ mod tests { use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; - // Windows Bazel CI can take longer than a few seconds for the websocket - // client connection attempt to reach the local test listener. - #[cfg(windows)] + // State initialization can delay the first local connection on loaded Bazel workers. pub(super) const TEST_HTTP_ACCEPT_TIMEOUT: Duration = Duration::from_secs(30); - #[cfg(not(windows))] - pub(super) const TEST_HTTP_ACCEPT_TIMEOUT: Duration = Duration::from_secs(5); pub(super) const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; pub(super) const TEST_REMOTE_CONTROL_SERVER_TOKEN: &str = "Remote Control Token"; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 68900517ab..f182909905 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1129,9 +1129,7 @@ impl ModelClient { metadata.remove("guardian_credits_requested"); metadata.remove("parent_response_id"); } - let guardian_reviewer = responses_headers - .get("x-codex-guardian") - .is_some_and(|value| value == "reviewer"); + let guardian_reviewer = is_guardian_reviewer(responses_headers); if guardian_reviewer && let Some(parent_response_id) = parent_response_id { metadata.get_or_insert_with(HashMap::new).insert( "parent_response_id".to_owned(), @@ -1156,10 +1154,11 @@ impl ModelClient { fn build_routing_hint_header( &self, auth: Option<&CodexAuth>, + responses_headers: &ApiHeaderMap, model: &str, service_tier: Option<&str>, ) -> Option { - if !self.uses_codex_backend(auth) { + if !self.uses_codex_backend(auth) || is_guardian_reviewer(responses_headers) { return None; } @@ -1448,52 +1447,85 @@ impl ModelClientSession { pub async fn preconnect_websocket( &mut self, model_info: &ModelInfo, + service_tier: Option, session_telemetry: &SessionTelemetry, responses_metadata: &CodexResponsesMetadata, - ) -> std::result::Result<(), ApiError> { + ) -> Result<()> { if !self.client.responses_websocket_enabled() { return Ok(()); } - let client_setup = self - .client - .current_client_setup(ClientRouting::Workspace) - .await - .map_err(|err| { - ApiError::Stream(format!( - "failed to build websocket prewarm client setup: {err}" - )) - })?; - let connection_key = - ResponsesConnectionKey::new(&client_setup.api_provider, client_setup.auth_revision); - if self.websocket_session.connection.is_some() - && self.websocket_session.connection_key.as_ref() == Some(&connection_key) - { - return Ok(()); + let provider = Arc::clone(&self.client.state.provider); + let auth_manager = provider.auth_manager(); + let mut auth_recovery = auth_manager + .as_ref() + .map(AuthManager::unauthorized_recovery); + let mut provider_auth_recovery_attempted = false; + let mut pending_retry = PendingUnauthorizedRetry::default(); + loop { + let client_setup = self + .client + .current_client_setup(ClientRouting::Workspace) + .await?; + let auth_context = AuthRequestTelemetryContext::new( + client_setup.auth.as_ref().map(CodexAuth::auth_mode), + client_setup.api_auth.as_ref(), + client_setup.agent_identity_telemetry.clone(), + pending_retry, + ); + let responses_headers = self + .client + .responses_headers(client_setup.auth.as_ref(), &model_info.slug); + let mut websocket_metadata = responses_metadata.clone(); + websocket_metadata.routing_hint = self.client.build_routing_hint_header( + client_setup.auth.as_ref(), + &responses_headers, + &model_info.slug, + model_info + .service_tier_for_request(service_tier.clone()) + .as_deref(), + ); + match self + .websocket_connection(WebsocketConnectParams { + session_telemetry, + api_provider: client_setup.api_provider, + auth_revision: client_setup.auth_revision, + api_auth: client_setup.api_auth, + auth_owner_generation: client_setup.auth_owner_generation, + responses_metadata: &websocket_metadata, + auth_context, + request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"), + responses_headers: &responses_headers, + }) + .await + { + Ok(_) => return Ok(()), + Err(ApiError::Transport(TransportError::Http { status, .. })) + if status == StatusCode::UPGRADE_REQUIRED => + { + self.try_switch_fallback_transport(session_telemetry, model_info); + return Ok(()); + } + Err(ApiError::Transport(unauthorized_transport)) + if provider.is_recoverable_auth_error(&unauthorized_transport) => + { + pending_retry = PendingUnauthorizedRetry::from_recovery( + handle_unauthorized( + unauthorized_transport, + &mut auth_recovery, + &mut provider_auth_recovery_attempted, + session_telemetry, + &provider, + self.client.event_sender.as_ref(), + responses_metadata.turn_id.as_deref(), + ) + .await?, + ); + } + Err(err) => return Err(provider.map_api_error(err)), + } } - self.websocket_session.reset(Some("other")); - let auth_context = AuthRequestTelemetryContext::new( - client_setup.auth.as_ref().map(CodexAuth::auth_mode), - client_setup.api_auth.as_ref(), - client_setup.agent_identity_telemetry.clone(), - PendingUnauthorizedRetry::default(), - ); - let responses_headers = self - .client - .responses_headers(client_setup.auth.as_ref(), &model_info.slug); - self.websocket_connection(WebsocketConnectParams { - session_telemetry, - api_provider: client_setup.api_provider, - auth_revision: client_setup.auth_revision, - api_auth: client_setup.api_auth, - auth_owner_generation: client_setup.auth_owner_generation, - responses_metadata, - auth_context, - request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"), - responses_headers: &responses_headers, - }) - .await?; - Ok(()) } + /// Returns a websocket connection for this turn. #[instrument( name = "model_client.websocket_connection", @@ -1674,19 +1706,15 @@ impl ModelClientSession { client_setup.auth.as_ref(), &responses_headers, ); - let guardian_reviewer = responses_headers - .get("x-codex-guardian") - .is_some_and(|value| value == "reviewer"); - if guardian_reviewer { + if is_guardian_reviewer(&responses_headers) { request.service_tier = None; } - if !guardian_reviewer - && let Some(header_value) = self.client.build_routing_hint_header( - client_setup.auth.as_ref(), - &request.model, - request.service_tier.as_deref(), - ) - { + if let Some(header_value) = self.client.build_routing_hint_header( + client_setup.auth.as_ref(), + &responses_headers, + &request.model, + request.service_tier.as_deref(), + ) { options .extra_headers .insert(X_CODEX_ROUTING_HINT_HEADER, header_value); @@ -1828,10 +1856,7 @@ impl ModelClientSession { responses_metadata, include_internal, )?; - let guardian_reviewer = responses_headers - .get("x-codex-guardian") - .is_some_and(|value| value == "reviewer"); - if guardian_reviewer { + if is_guardian_reviewer(&responses_headers) { request.service_tier = None; } request.access_programs = cyber_access_program::for_auth( @@ -1839,15 +1864,12 @@ impl ModelClientSession { prompt.cyber_access_program, ); let mut websocket_metadata = responses_metadata.clone(); - websocket_metadata.routing_hint = if !guardian_reviewer { - self.client.build_routing_hint_header( - client_setup.auth.as_ref(), - &request.model, - request.service_tier.as_deref(), - ) - } else { - None - }; + websocket_metadata.routing_hint = self.client.build_routing_hint_header( + client_setup.auth.as_ref(), + &responses_headers, + &request.model, + request.service_tier.as_deref(), + ); let request_session_telemetry = if warmup { // `generate=false` prewarm is connection setup, not an inference request. session_telemetry.clone() @@ -2216,6 +2238,12 @@ fn stamp_ws_stream_request_start_ms(request: &mut ResponsesWsRequest<'_>) { ); } +fn is_guardian_reviewer(responses_headers: &ApiHeaderMap) -> bool { + responses_headers + .get("x-codex-guardian") + .is_some_and(|value| value == "reviewer") +} + /// Builds the extra headers attached to Responses API requests. /// /// These headers implement Codex-specific conventions: diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5f5a279c7a..815c350519 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1071,7 +1071,7 @@ pub(crate) fn new_submission_id() -> String { Uuid::now_v7().to_string() } -fn get_service_tier( +pub(crate) fn get_service_tier( configured_service_tier: Option, fast_mode_enabled: bool, model_info: &ModelInfo, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 5f9e6256db..051f6b0a11 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -709,7 +709,7 @@ impl Session { ) } - fn with_window_and_fork_metadata( + pub(crate) fn with_window_and_fork_metadata( &self, turn_context: &TurnContext, responses_metadata: CodexResponsesMetadata, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 43b66cb97c..e6846aa243 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -23,6 +23,8 @@ use codex_otel::STARTUP_PREWARM_DURATION_METRIC; use codex_otel::SessionTelemetry; use codex_protocol::error::Result as CodexResult; use codex_protocol::models::BaseInstructions; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; pub(crate) struct SessionStartupPrewarmHandle { task: AbortOnDropHandle>, @@ -273,19 +275,63 @@ async fn schedule_startup_prewarm_inner( /*status*/ None, ); let startup_cancellation_token = CancellationToken::new(); - let built_tools_started_at = Instant::now(); - // Startup prewarm runs before run_turn and needs its own tool-building snapshot. - let step_context = session - .capture_step_context( - Arc::clone(&startup_turn_context), - &startup_cancellation_token, + let preconnect_model_info = Arc::clone(startup_turn_context.model_info()); + // Spawned subagents inherit the root's selection, with the same feature and model filtering + // that capture applies to the actual request. + let preconnect_service_tier = if matches!( + startup_turn_context.session_source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ) { + crate::session::get_service_tier( + session.services.agent_control.service_tier(), + session.features().enabled(Feature::FastMode), + &preconnect_model_info, ) - .await?; - startup_turn_context.session_telemetry.record_startup_phase( - "startup_prewarm_build_tools", - built_tools_started_at.elapsed(), - /*status*/ None, + } else { + startup_turn_context.initial_settings.service_tier.clone() + }; + // The handshake precedes tool capture; attach the finalized step metadata to the warmup. + let (window_id, window_number, context_window_id) = session.current_window().await; + let mut handshake_metadata = startup_turn_context + .turn_metadata_state + .to_responses_metadata( + session.installation_id.clone(), + window_id, + CodexResponsesRequestKind::Prewarm, + ); + crate::turn_metadata::ExecutionMetadata::from_settings(&startup_turn_context.initial_settings) + .apply_to(&mut handshake_metadata); + let handshake_metadata = session.with_window_and_fork_metadata( + &startup_turn_context, + handshake_metadata, + window_number, + context_window_id, ); + let mut client_session = 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 { + let built_tools_started_at = Instant::now(); + let step_context = session + .capture_step_context( + Arc::clone(&startup_turn_context), + &startup_cancellation_token, + ) + .await?; + startup_turn_context.session_telemetry.record_startup_phase( + "startup_prewarm_build_tools", + built_tools_started_at.elapsed(), + /*status*/ None, + ); + Ok(step_context) + }, + client_session.preconnect_websocket( + &preconnect_model_info, + preconnect_service_tier, + &startup_turn_context.session_telemetry, + &handshake_metadata, + ), + )?; let build_prompt_started_at = Instant::now(); let startup_prompt = build_prompt( Vec::new(), @@ -300,10 +346,10 @@ async fn schedule_startup_prewarm_inner( build_prompt_started_at.elapsed(), /*status*/ None, ); + // Tool discovery may have updated Responses Lite metadata since the eager handshake. let responses_metadata = session .responses_metadata(step_context.as_ref(), CodexResponsesRequestKind::Prewarm) .await; - let mut client_session = session.services.model_client.new_session(); let websocket_warmup_started_at = Instant::now(); // Prewarm establishes the request baseline before the first turn can change effort. client_session diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 408a6234b3..5c65a36ade 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -16,6 +16,7 @@ use serde_json::Value; use tokio::net::TcpListener; use tokio::sync::Notify; use tokio::sync::oneshot; +use tokio::sync::watch; use tokio_tungstenite::accept_hdr_async_with_config; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::extensions::ExtensionsConfig; @@ -577,6 +578,7 @@ pub struct WebSocketTestServer { connections: Arc>>>, handshakes: Arc>>, request_log_updated: Arc, + closed_connections: watch::Receiver, shutdown: oneshot::Sender<()>, task: tokio::task::JoinHandle<()>, } @@ -598,6 +600,16 @@ impl WebSocketTestServer { connections.first().cloned().unwrap_or_default() } + pub async fn wait_for_connections(&self, expected: usize, timeout: Duration) -> bool { + tokio::time::timeout(timeout, async { + while self.connections.lock().unwrap().len() < expected { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .is_ok() + } + pub async fn wait_for_request( &self, connection_index: usize, @@ -619,6 +631,17 @@ impl WebSocketTestServer { } } + /// Waits for the server to finish reading any frames preceding the socket close. + pub async fn wait_for_closed_connections(&self, expected: usize, timeout: Duration) -> bool { + let mut closed_connections = self.closed_connections.clone(); + tokio::time::timeout( + timeout, + closed_connections.wait_for(|count| *count >= expected), + ) + .await + .is_ok_and(|result| result.is_ok()) + } + pub fn handshakes(&self) -> Vec { self.handshakes.lock().unwrap().clone() } @@ -1206,6 +1229,7 @@ pub async fn start_websocket_server_with_headers( let connections_log = Arc::new(Mutex::new(Vec::new())); let handshakes_log = Arc::new(Mutex::new(Vec::new())); let request_log_updated = Arc::new(Notify::new()); + let (closed_tx, closed_connections) = watch::channel(0); let requests = Arc::clone(&connections_log); let handshakes = Arc::clone(&handshakes_log); let request_log = Arc::clone(&request_log_updated); @@ -1351,6 +1375,7 @@ pub async fn start_websocket_server_with_headers( if close_after_requests { let _ = ws_stream.close(None).await; + closed_tx.send_modify(|count| *count += 1); } else { let _ = shutdown_rx.await; return; @@ -1367,6 +1392,7 @@ pub async fn start_websocket_server_with_headers( connections: connections_log, handshakes: handshakes_log, request_log_updated, + closed_connections, shutdown: shutdown_tx, task, } diff --git a/codex-rs/core/tests/suite/agent_websocket.rs b/codex-rs/core/tests/suite/agent_websocket.rs index 466c20bf4b..75f8c6db10 100644 --- a/codex-rs/core/tests/suite/agent_websocket.rs +++ b/codex-rs/core/tests/suite/agent_websocket.rs @@ -1,7 +1,10 @@ use anyhow::Result; use codex_core::TurnInputRequest; +use codex_core::X_CODEX_ROUTING_HINT_HEADER; use codex_features::Feature; +use codex_login::CodexAuth; use codex_protocol::config_types::ServiceTier; +use codex_protocol::openai_models::ModelServiceTier; use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ThreadSettingsOverrides; @@ -10,6 +13,7 @@ use core_test_support::responses::WebSocketConnectionConfig; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_exec_command_call; +use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; use core_test_support::responses::start_websocket_server; use core_test_support::responses::start_websocket_server_with_headers; @@ -18,10 +22,133 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::Value; +use serde_json::json; use std::time::Duration; const WS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06"; +#[test_case::test_case(true; "fast mode enabled")] +#[test_case::test_case(false; "fast mode disabled")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn spawned_agent_prewarm_handshake_inherits_effective_root_service_tier( + fast_mode_enabled: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_websocket_server(vec![ + vec![ + vec![ev_response_created("root-warm"), ev_completed("root-warm")], + vec![ + ev_response_created("root-spawn"), + ev_function_call_with_namespace( + "spawn-worker", + "collaboration", + "spawn_agent", + &json!({ "message": "hello", "task_name": "worker", "fork_turns": "none" }) + .to_string(), + ), + ev_completed("root-spawn"), + ], + vec![ev_response_created("root-done"), ev_completed("root-done")], + ], + vec![ + vec![ + ev_response_created("child-warm"), + ev_completed("child-warm"), + ], + vec![ + ev_response_created("child-done"), + ev_completed("child-done"), + ], + ], + ]) + .await; + let tier = ServiceTier::Fast.request_value(); + let test = test_codex() + .with_model("gpt-5.4") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_info_override("gpt-5.4", move |model| { + model.service_tiers.push(ModelServiceTier { + id: tier.to_string(), + name: "Fast".to_string(), + description: "Priority processing".to_string(), + }); + }) + .with_config(move |config| { + for feature in [Feature::Collab, Feature::MultiAgentV2] { + config.features.enable(feature).expect("spawn feature"); + } + if fast_mode_enabled { + config + .features + .enable(Feature::FastMode) + .expect("fast mode"); + } else { + config + .features + .disable(Feature::FastMode) + .expect("fast mode"); + } + config.service_tier = Some(ServiceTier::Fast.request_value().to_string()); + }) + .build_with_websocket_server(&server) + .await?; + let mut created_threads = test.thread_manager.subscribe_thread_created(); + tokio::time::timeout( + Duration::from_secs(5), + server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0), + ) + .await?; + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "spawn a worker".to_string(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + service_tier: Some(Some(tier.to_string())), + ..Default::default() + }), + ) + .await?; + let child_thread_id = + tokio::time::timeout(Duration::from_secs(5), created_threads.recv()).await??; + let child = test.thread_manager.get_thread(child_thread_id).await?; + + assert!( + server + .wait_for_handshakes(/*expected*/ 2, Duration::from_secs(5)) + .await, + "child startup should connect before its first turn" + ); + let handshakes = server.handshakes(); + let expected_hint = if fast_mode_enabled { + format!("model=gpt-5.4;tier={tier}") + } else { + "model=gpt-5.4".to_string() + }; + assert_eq!( + handshakes[1].header(X_CODEX_ROUTING_HINT_HEADER), + Some(expected_hint) + ); + let child_warmup = tokio::time::timeout( + Duration::from_secs(5), + server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0), + ) + .await? + .body_json(); + assert_eq!(child_warmup["generate"], false); + assert_eq!( + child_warmup["service_tier"].as_str(), + fast_mode_enabled.then_some(tier) + ); + assert_eq!(server.handshakes().len(), 2); + child.shutdown_and_wait().await?; + test.codex.shutdown_and_wait().await?; + server.shutdown().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn websocket_model_switch_to_responses_lite_omits_top_level_tools() -> Result<()> { skip_if_no_network!(Ok(())); @@ -388,26 +515,57 @@ async fn websocket_v2_test_codex_shell_chain() -> Result<()> { Ok(()) } +#[test_case::test_case(None, Some("ultrafast"); "standard to ultrafast")] +#[test_case::test_case(Some("ultrafast"), None; "ultrafast to standard")] +#[test_case::test_case(Some("priority"), Some("ultrafast"); "priority to ultrafast")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn websocket_v2_first_turn_uses_updated_fast_tier_after_startup_prewarm() -> Result<()> { +async fn websocket_v2_first_turn_sends_full_create_when_tier_changes_after_startup_prewarm( + startup_tier: Option<&'static str>, + turn_tier: Option<&'static str>, +) -> Result<()> { skip_if_no_network!(Ok(())); + let warmup_turn_state = "startup-turn-state"; let server = start_websocket_server(vec![vec![ - vec![ev_response_created("warm-1"), ev_completed("warm-1")], + vec![ + ev_response_created("warm-1"), + json!({ + "type": "response.metadata", + "headers": {"x-codex-turn-state": warmup_turn_state}, + }), + ev_completed("warm-1"), + ], vec![ ev_response_created("resp-1"), - ev_assistant_message("msg_1", "fast"), + ev_assistant_message("msg_1", "done"), ev_completed("resp-1"), ], ]]) .await; - let mut builder = test_codex().with_config(|config| { - config - .features - .enable(Feature::ResponsesWebsocketsV2) - .expect("test config should allow feature update"); - }); + let model = "gpt-5.6-sol"; + let mut builder = test_codex() + .with_model_info_override(model, |model_info| { + // Exercise tier transitions independently of the bundled model catalog. + model_info.service_tiers = ["priority", "ultrafast"] + .into_iter() + .map(|tier| ModelServiceTier { + id: tier.to_string(), + name: tier.to_string(), + description: String::new(), + }) + .collect(); + }) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + for feature in [Feature::ResponsesWebsocketsV2, Feature::FastMode] { + config + .features + .enable(feature) + .expect("test config should allow feature update"); + } + config.service_tier = startup_tier.map(str::to_string); + }); let test = builder.build_with_websocket_server(&server).await?; let warmup = server @@ -416,68 +574,21 @@ async fn websocket_v2_first_turn_uses_updated_fast_tier_after_startup_prewarm() .body_json(); assert_eq!(warmup["type"].as_str(), Some("response.create")); assert_eq!(warmup["generate"].as_bool(), Some(false)); - assert_eq!(warmup.get("service_tier"), None); + assert_eq!(warmup["service_tier"].as_str(), startup_tier); - test.submit_turn_with_service_tier("hello", Some(ServiceTier::Fast.request_value())) + test.submit_turn_with_service_tier("hello", turn_tier) .await?; assert_eq!(server.handshakes().len(), 1); - let connection = server.single_connection(); - assert_eq!(connection.len(), 2); - let first_turn = connection - .get(1) - .expect("missing first turn request") - .body_json(); - - assert_eq!(first_turn["type"].as_str(), Some("response.create")); - assert_eq!(first_turn["service_tier"].as_str(), Some("priority")); - assert_eq!(first_turn.get("previous_response_id"), None); - assert!( - first_turn - .get("input") - .and_then(Value::as_array) - .is_some_and(|items| !items.is_empty()) + assert_eq!( + server + .single_handshake() + .header(X_CODEX_ROUTING_HINT_HEADER), + Some(match startup_tier { + Some(tier) => format!("model={model};tier={tier}"), + None => format!("model={model}"), + }) ); - - server.shutdown().await; - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn websocket_v2_first_turn_drops_fast_tier_after_startup_prewarm() -> Result<()> { - skip_if_no_network!(Ok(())); - - let server = start_websocket_server(vec![vec![ - vec![ev_response_created("warm-1"), ev_completed("warm-1")], - vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg_1", "standard"), - ev_completed("resp-1"), - ], - ]]) - .await; - - let mut builder = test_codex().with_config(|config| { - config - .features - .enable(Feature::ResponsesWebsocketsV2) - .expect("test config should allow feature update"); - config.service_tier = Some(ServiceTier::Fast.request_value().to_string()); - }); - let test = builder.build_with_websocket_server(&server).await?; - - let warmup = server - .wait_for_request(/*connection_index*/ 0, /*request_index*/ 0) - .await - .body_json(); - assert_eq!(warmup["type"].as_str(), Some("response.create")); - assert_eq!(warmup["generate"].as_bool(), Some(false)); - assert_eq!(warmup["service_tier"].as_str(), Some("priority")); - - test.submit_turn_with_service_tier("hello", /*service_tier*/ None) - .await?; - - assert_eq!(server.handshakes().len(), 1); let connection = server.single_connection(); assert_eq!(connection.len(), 2); let first_turn = connection @@ -486,7 +597,13 @@ async fn websocket_v2_first_turn_drops_fast_tier_after_startup_prewarm() -> Resu .body_json(); assert_eq!(first_turn["type"].as_str(), Some("response.create")); - assert_eq!(first_turn.get("service_tier"), None); + assert_eq!(first_turn["model"].as_str(), Some(model)); + assert_eq!(first_turn["service_tier"].as_str(), turn_tier); + assert_ne!(first_turn["generate"].as_bool(), Some(false)); + assert_eq!( + first_turn["client_metadata"]["x-codex-turn-state"].as_str(), + Some(warmup_turn_state) + ); assert_eq!(first_turn.get("previous_response_id"), None); assert!( first_turn diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index cae43a74ee..48cc71457e 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1520,6 +1520,42 @@ async fn provider_auth_command_refreshes_after_401() { send_provider_auth_request(&server, auth_fixture.auth()).await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn provider_auth_command_refreshes_during_websocket_preconnect() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let auth_fixture = ProviderAuthCommandFixture::new(&["first-token", "second-token"]).unwrap(); + Mock::given(method("GET")) + .and(path("/v1/responses")) + .and(header("authorization", "Bearer first-token")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + // Auth must recover during preconnect, before the ordinary HTTP fallback turn. + Mock::given(method("GET")) + .and(path("/v1/responses")) + .and(header("authorization", "Bearer second-token")) + .respond_with(ResponseTemplate::new(426)) + .expect(1) + .mount(&server) + .await; + mount_sse_once_match( + &server, + header("authorization", "Bearer second-token"), + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + ) + .await; + + let mut provider = + ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))); + provider.requires_openai_auth = false; + provider.auth = Some(auth_fixture.auth()); + provider.supports_websockets = true; + send_request_with_provider(provider).await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn provider_auth_command_recovers_after_initial_resolution_failure() { skip_if_no_network!(); @@ -1629,6 +1665,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth #[expect(clippy::unwrap_used)] async fn send_request_with_provider(provider: ModelProviderInfo) { + let preconnect = provider.supports_websockets; let codex_home = TempDir::new().unwrap(); let mut config = load_default_config_for_test(&codex_home).await; config.model_provider_id = provider.name.clone(); @@ -1678,6 +1715,17 @@ async fn send_request_with_provider(provider: ModelProviderInfo) { ); let responses_metadata = test_turn_responses_metadata(&client, thread_id); let mut client_session = client.new_session(); + if preconnect { + client_session + .preconnect_websocket( + &model_info, + /*service_tier*/ None, + &session_telemetry, + &responses_metadata, + ) + .await + .expect("preconnect should recover authentication before the first turn"); + } let mut prompt = Prompt::default(); prompt.input.push(ResponseItem::Message { id: None, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index b9a0492e1a..a652b37558 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -539,7 +539,7 @@ async fn responses_websocket_reuses_connection_with_per_turn_trace_payloads() { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn responses_websocket_preconnect_does_not_replace_turn_trace_payload() { +async fn responses_websocket_preconnect_reuses_connection_without_replacing_turn_trace() { skip_if_no_network!(); let _trace_test_context = install_test_tracing("client-websocket-test"); @@ -556,11 +556,28 @@ async fn responses_websocket_preconnect_does_not_replace_turn_trace_payload() { client_session .preconnect_websocket( &harness.model_info, + /*service_tier*/ None, &harness.session_telemetry, &responses_metadata, ) .await .expect("websocket preconnect failed"); + assert!(server.single_connection().is_empty()); + let handshake = server.single_handshake(); + assert_eq!( + handshake.header(USER_AGENT_HEADER), + Some(codex_login::default_client::get_codex_user_agent()) + ); + assert_eq!( + handshake.header("x-codex-window-id"), + Some(TEST_WINDOW_ID.to_string()) + ); + assert_eq!(handshake.header("x-codex-turn-metadata"), None); + assert_eq!(handshake.header(X_CODEX_ROUTING_HINT_HEADER), None); + assert_eq!( + handshake.header(OPENAI_BETA_HEADER), + Some(WS_V2_BETA_HEADER_VALUE.to_string()) + ); let prompt = prompt_with_input(vec![message_item("hello")]); let expected_trace = async { @@ -577,45 +594,13 @@ async fn responses_websocket_preconnect_does_not_replace_turn_trace_payload() { assert_eq!(connection.len(), 1); let request = connection.first().expect("missing request").body_json(); assert_request_trace_matches(&request, &expected_trace); - - server.shutdown().await; -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn responses_websocket_preconnect_reuses_connection() { - skip_if_no_network!(); - - let server = start_websocket_server(vec![vec![vec![ - ev_response_created("resp-1"), - ev_completed("resp-1"), - ]]]) - .await; - - let harness = websocket_harness(&server).await; - let mut client_session = harness.client.new_session(); - let responses_metadata = websocket_connection_metadata(&harness); - client_session - .preconnect_websocket( - &harness.model_info, - &harness.session_telemetry, - &responses_metadata, - ) - .await - .expect("websocket preconnect failed"); - let prompt = prompt_with_input(vec![message_item("hello")]); - stream_until_complete(&mut client_session, &harness, &prompt).await; - - assert_eq!(server.handshakes().len(), 1); - assert_eq!( - server.single_handshake().header(USER_AGENT_HEADER), - Some(codex_login::default_client::get_codex_user_agent()) - ); - assert_eq!( - server.single_handshake().header("x-codex-window-id"), - Some(TEST_WINDOW_ID.to_string()) - ); - let connection = server.single_connection(); - assert_eq!(connection.len(), 1); + let turn_metadata: serde_json::Value = serde_json::from_str( + request["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .expect("turn metadata"), + ) + .expect("valid turn metadata"); + assert_eq!(turn_metadata["request_kind"], "turn"); server.shutdown().await; } @@ -657,10 +642,15 @@ async fn responses_websocket_request_prewarm_reuses_connection() { stream_until_complete(&mut client_session, &harness, &prompt).await; assert_eq!(server.handshakes().len(), 1); + let handshake = server.single_handshake(); assert_eq!( - server.single_handshake().header(USER_AGENT_HEADER), + handshake.header(USER_AGENT_HEADER), Some(codex_login::default_client::get_codex_user_agent()) ); + assert_eq!( + handshake.header(OPENAI_BETA_HEADER), + Some(WS_V2_BETA_HEADER_VALUE.to_string()) + ); let connection = server.single_connection(); assert_eq!(connection.len(), 2); let warmup = connection @@ -674,6 +664,10 @@ async fn responses_websocket_request_prewarm_reuses_connection() { assert_eq!(warmup["type"].as_str(), Some("response.create")); assert_eq!(warmup["generate"].as_bool(), Some(false)); + assert_eq!( + warmup["input"], + serde_json::to_value(&prompt.input).unwrap() + ); assert_eq!( warmup["stream_options"]["reasoning_summary_delivery"].as_str(), Some("sequential_cutoff") @@ -692,6 +686,13 @@ async fn responses_websocket_request_prewarm_reuses_connection() { assert_eq!(follow_up["type"].as_str(), Some("response.create")); assert_eq!(follow_up["previous_response_id"].as_str(), Some("warm-1")); assert_eq!(follow_up["input"], serde_json::json!([])); + let follow_up_turn_metadata: serde_json::Value = serde_json::from_str( + follow_up["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .expect("follow-up turn metadata"), + ) + .expect("valid follow-up turn metadata"); + assert_eq!(follow_up_turn_metadata["request_kind"], "turn"); assert_eq!( follow_up["stream_options"]["reasoning_summary_delivery"].as_str(), Some("sequential_cutoff") @@ -947,6 +948,17 @@ async fn responses_websocket_reconnects_after_account_switch() { if index == 5 || (index == 3 && new_turn_on_switch) { drop(client_session); client_session = harness.client.new_session(); + if index == 3 { + client_session + .preconnect_websocket( + &harness.model_info, + /*service_tier*/ None, + &harness.session_telemetry, + &websocket_connection_metadata(&harness), + ) + .await + .expect("new account should establish its own socket during preconnect"); + } } input.push(message_item(&format!("request {index}"))); stream_until_complete_with_model_info( @@ -1098,125 +1110,15 @@ async fn responses_websocket_sends_responses_lite_metadata_per_request() { server.shutdown().await; } +#[test_case::test_case(MODEL, Some("ultrafast"), Some("ultrafast"); "matching hint")] +#[test_case::test_case(MODEL, None, Some("ultrafast"); "tier selected after handshake")] +#[test_case::test_case("gpt-5.2", Some("ultrafast"), None; "model and tier changed after handshake")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn responses_websocket_preconnect_is_reused_even_with_header_changes() { - skip_if_no_network!(); - - let server = start_websocket_server(vec![vec![vec![ - ev_response_created("resp-1"), - ev_completed("resp-1"), - ]]]) - .await; - - let harness = websocket_harness(&server).await; - let mut client_session = harness.client.new_session(); - let preconnect_metadata = websocket_connection_metadata(&harness); - client_session - .preconnect_websocket( - &harness.model_info, - &harness.session_telemetry, - &preconnect_metadata, - ) - .await - .expect("websocket preconnect failed"); - let prompt = prompt_with_input(vec![message_item("hello")]); - let responses_metadata = turn_metadata(&harness, /*turn_id*/ None); - let mut stream = client_session - .stream( - &prompt, - &harness.model_info, - &harness.session_telemetry, - harness.effort.clone(), - harness.summary, - /*service_tier*/ None, - &responses_metadata, - &codex_rollout_trace::InferenceTraceContext::disabled(), - ) - .await - .expect("websocket stream failed"); - - while let Some(event) = stream.next().await { - if matches!(event, Ok(ResponseEvent::Completed { .. })) { - break; - } - } - - assert_eq!(server.handshakes().len(), 1); - assert_eq!(server.single_connection().len(), 1); - - server.shutdown().await; -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn responses_websocket_request_prewarm_is_reused_even_with_header_changes() { - skip_if_no_network!(); - - let server = start_websocket_server(vec![vec![ - vec![ev_response_created("warm-1"), ev_completed("warm-1")], - vec![ev_response_created("resp-1"), ev_completed("resp-1")], - ]]) - .await; - - let harness = websocket_harness_with_options(&server, /*runtime_metrics_enabled*/ true).await; - let mut client_session = harness.client.new_session(); - let prompt = prompt_with_input(vec![message_item("hello")]); - let prewarm_responses_metadata = prewarm_metadata(&harness, /*turn_id*/ None); - client_session - .prewarm_websocket( - &prompt, - &harness.model_info, - &harness.session_telemetry, - harness.effort.clone(), - harness.summary, - /*service_tier*/ None, - &prewarm_responses_metadata, - ) - .await - .expect("websocket prewarm failed"); - let responses_metadata = turn_metadata(&harness, /*turn_id*/ None); - let mut stream = client_session - .stream( - &prompt, - &harness.model_info, - &harness.session_telemetry, - harness.effort.clone(), - harness.summary, - /*service_tier*/ None, - &responses_metadata, - &codex_rollout_trace::InferenceTraceContext::disabled(), - ) - .await - .expect("websocket stream failed"); - - while let Some(event) = stream.next().await { - if matches!(event, Ok(ResponseEvent::Completed { .. })) { - break; - } - } - - assert_eq!(server.handshakes().len(), 1); - let connection = server.single_connection(); - assert_eq!(connection.len(), 2); - let warmup = connection - .first() - .expect("missing warmup request") - .body_json(); - let follow_up = connection - .get(1) - .expect("missing follow-up request") - .body_json(); - assert_eq!(warmup["type"].as_str(), Some("response.create")); - assert_eq!(warmup["generate"].as_bool(), Some(false)); - assert_eq!(warmup["tools"], serde_json::json!([])); - assert_eq!(follow_up["type"].as_str(), Some("response.create")); - assert_eq!(follow_up["previous_response_id"].as_str(), Some("warm-1")); - assert_eq!(follow_up["input"], serde_json::json!([])); - - server.shutdown().await; -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn responses_websocket_prewarm_includes_model_and_tier_routing_hint() { +async fn responses_websocket_prewarm_reuses_advisory_model_and_tier_routing_hint( + warmup_model: &str, + preconnect_tier: Option<&str>, + warmup_tier: Option<&str>, +) { skip_if_no_network!(); let server = start_websocket_server(vec![vec![vec![ @@ -1227,15 +1129,36 @@ async fn responses_websocket_prewarm_includes_model_and_tier_routing_hint() { let harness = websocket_harness_for_codex_backend(&server).await; let mut model_info = harness.model_info.clone(); - let service_tier = ServiceTier::Fast.request_value(); model_info.service_tiers.push(ModelServiceTier { - id: service_tier.to_string(), - name: "Fast".to_string(), - description: "Priority processing".to_string(), + id: "ultrafast".to_string(), + name: "Ultrafast".to_string(), + description: "Ultrafast processing".to_string(), }); let mut client_session = harness.client.new_session(); let prompt = prompt_with_input(vec![message_item("hello")]); let responses_metadata = prewarm_metadata(&harness, /*turn_id*/ None); + client_session + .preconnect_websocket( + &model_info, + preconnect_tier.map(str::to_string), + &harness.session_telemetry, + &responses_metadata, + ) + .await + .expect("eager handshake should connect"); + let handshake = server.single_handshake(); + assert_eq!( + handshake.header(X_CODEX_ROUTING_HINT_HEADER), + Some(match preconnect_tier { + Some(tier) => format!("model={MODEL};tier={tier}"), + None => format!("model={MODEL}"), + }) + ); + assert_eq!( + handshake.header(OPENAI_BETA_HEADER), + Some(WS_V2_BETA_HEADER_VALUE.to_string()) + ); + model_info.slug = warmup_model.to_string(); client_session .prewarm_websocket( &prompt, @@ -1243,73 +1166,12 @@ async fn responses_websocket_prewarm_includes_model_and_tier_routing_hint() { &harness.session_telemetry, harness.effort.clone(), harness.summary, - Some(service_tier.to_string()), + warmup_tier.map(str::to_string), &responses_metadata, ) .await .expect("websocket prewarm failed"); - let handshake = server.single_handshake(); - assert_eq!( - handshake.header(X_CODEX_ROUTING_HINT_HEADER), - Some(format!("model={MODEL};tier={service_tier}")) - ); - assert_eq!( - handshake.header(OPENAI_BETA_HEADER), - Some(WS_V2_BETA_HEADER_VALUE.to_string()) - ); - let connection = server.single_connection(); - let prewarm = connection - .first() - .expect("missing prewarm request") - .body_json(); - assert_eq!(prewarm["service_tier"].as_str(), Some(service_tier)); - - server.shutdown().await; -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn responses_websocket_prewarm_uses_v2_when_provider_supports_websockets() { - skip_if_no_network!(); - - let server = start_websocket_server(vec![vec![vec![ - ev_response_created("resp-1"), - ev_completed("resp-1"), - ]]]) - .await; - - let harness = websocket_harness_with_options(&server, /*runtime_metrics_enabled*/ false).await; - let mut client_session = harness.client.new_session(); - let prompt = prompt_with_input(vec![message_item("hello")]); - let responses_metadata = prewarm_metadata(&harness, /*turn_id*/ None); - client_session - .prewarm_websocket( - &prompt, - &harness.model_info, - &harness.session_telemetry, - harness.effort.clone(), - harness.summary, - /*service_tier*/ None, - &responses_metadata, - ) - .await - .expect("websocket prewarm failed"); - - // V2 prewarm issues a request on the websocket connection. - assert_eq!(server.handshakes().len(), 1); - assert_eq!(server.single_connection().len(), 1); - - let handshake = server.single_handshake(); - let openai_beta_header = handshake - .header(OPENAI_BETA_HEADER) - .expect("missing OpenAI-Beta header"); - assert!( - openai_beta_header - .split(',') - .map(str::trim) - .any(|value| value == WS_V2_BETA_HEADER_VALUE) - ); - stream_until_complete(&mut client_session, &harness, &prompt).await; assert_eq!(server.handshakes().len(), 1); let connection = server.single_connection(); assert_eq!(connection.len(), 1); @@ -1317,69 +1179,15 @@ async fn responses_websocket_prewarm_uses_v2_when_provider_supports_websockets() .first() .expect("missing prewarm request") .body_json(); - assert_eq!(prewarm["type"].as_str(), Some("response.create")); - assert_eq!( - prewarm["input"], - serde_json::to_value(&prompt.input).unwrap() - ); + assert_eq!(prewarm["type"], "response.create"); + assert_eq!(prewarm["generate"], false); + assert_eq!(prewarm["model"], warmup_model); + assert_eq!(prewarm["service_tier"].as_str(), warmup_tier); + assert_eq!(prewarm.get("previous_response_id"), None); server.shutdown().await; } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn responses_websocket_preconnect_runs_when_only_v2_feature_enabled() { - skip_if_no_network!(); - - let server = start_websocket_server(vec![vec![vec![ - ev_response_created("resp-1"), - ev_completed("resp-1"), - ]]]) - .await; - - let harness = websocket_harness_with_options(&server, /*runtime_metrics_enabled*/ true).await; - let mut client_session = harness.client.new_session(); - let responses_metadata = websocket_connection_metadata(&harness); - client_session - .preconnect_websocket( - &harness.model_info, - &harness.session_telemetry, - &responses_metadata, - ) - .await - .expect("websocket preconnect failed"); - - assert_eq!(server.handshakes().len(), 1); - assert_eq!(server.single_connection().len(), 0); - assert_eq!( - server.single_handshake().header("x-codex-turn-metadata"), - None - ); - assert_eq!( - server - .single_handshake() - .header(X_CODEX_ROUTING_HINT_HEADER), - None - ); - - let prompt = prompt_with_input(vec![message_item("hello")]); - stream_until_complete(&mut client_session, &harness, &prompt).await; - - assert_eq!(server.handshakes().len(), 1); - assert_eq!(server.single_connection().len(), 1); - - let handshake = server.single_handshake(); - let openai_beta_header = handshake - .header(OPENAI_BETA_HEADER) - .expect("missing OpenAI-Beta header"); - assert!( - openai_beta_header - .split(',') - .map(str::trim) - .any(|value| value == WS_V2_BETA_HEADER_VALUE) - ); - server.shutdown().await; -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_websocket_v2_requests_use_v2_when_provider_supports_websockets() { skip_if_no_network!(); diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index 760d45bc02..9b47049ce6 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -262,6 +262,9 @@ async fn guardian_session_inherits_parent_http_fallback( guardian_request.header("x-codex-guardian").as_deref(), credits_enabled.then_some("reviewer") ); + if credits_enabled { + assert_eq!(guardian_request.header("x-codex-routing-hint"), None); + } let body = guardian_request.body_json(); assert_eq!( ( diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index bad5c1135b..11e6bdf92c 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -26,6 +26,7 @@ use codex_config::types::OAuthCredentialsStoreMode; use codex_core::EnvironmentConfig; use codex_core::EnvironmentMcpPolicy; use codex_core::TurnInputRequest; +use codex_core::X_CODEX_ROUTING_HINT_HEADER; use codex_core::config::Config; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_exec_server::CreateDirectoryOptions; @@ -45,6 +46,7 @@ use codex_history::RolloutItem; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::config_types::ServiceTier; use codex_protocol::config_types::Settings; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::mcp_policy::McpServerIdentity; @@ -60,9 +62,11 @@ use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::ConfirmationPolicies; use codex_protocol::openai_models::InputModality; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ModelServiceTier; use codex_protocol::openai_models::ModelVisibility; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ReasoningEffortPreset; +use codex_protocol::openai_models::ToolMode; use codex_protocol::openai_models::TruncationPolicyConfig; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; @@ -1562,6 +1566,18 @@ async fn shutdown_cancels_startup_prewarm_waiting_for_mcp_startup() -> anyhow::R }, TestMcpServerOptions::default(), ); + // Wait for discovery instead of skipping the pending optional server. + // Its startup timeout exceeds every bounded wait in this test. + config.mcp_optional_startup_grace = Duration::ZERO; + let mut servers = config.mcp_servers.get().clone(); + let pending = servers + .get_mut("shutdown_prewarm") + .expect("pending MCP server"); + pending.startup_timeout_sec = Some(Duration::from_secs(300)); + config + .mcp_servers + .set(servers) + .expect("valid MCP configuration"); }) .build_with_websocket_server(&server) .await?; @@ -1570,19 +1586,218 @@ async fn shutdown_cancels_startup_prewarm_waiting_for_mcp_startup() -> anyhow::R tokio::time::timeout(Duration::from_secs(5), pending_mcp_listener.accept()) .await .context("startup prewarm should start the MCP connection")??; + assert!( + server + .wait_for_connections(/*expected*/ 1, Duration::from_secs(2)) + .await, + "websocket handshake should overlap pending MCP tool startup" + ); + assert!( + server.single_connection().is_empty(), + "prewarm should wait for tools before sending generate=false" + ); tokio::time::timeout(Duration::from_secs(2), fixture.codex.shutdown_and_wait()) .await .context("shutdown should not wait for startup prewarm MCP startup")??; - tokio::time::sleep(Duration::from_millis(100)).await; + // Dropping the thread closes its cached socket. Wait for the server to drain preceding frames. + drop(fixture); assert!( - server.connections().is_empty(), - "startup prewarm should not send a websocket request after shutdown" + server + .wait_for_closed_connections(/*expected*/ 1, Duration::from_secs(3)) + .await, + "websocket server should drain pending frames after startup cancellation" + ); + assert_eq!( + server.single_connection().len(), + 0, + "blocked MCP discovery must prevent any startup prewarm request" ); server.shutdown().await; Ok(()) } +#[derive(Clone, Copy, PartialEq)] +enum PrewarmMcpScenario { + Responses, + ResponsesLite, + EagerSocketCloses, +} + +#[test_case(PrewarmMcpScenario::Responses; "responses")] +#[test_case(PrewarmMcpScenario::ResponsesLite; "responses lite")] +#[test_case(PrewarmMcpScenario::EagerSocketCloses; "first turn 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, +) -> anyhow::Result<()> { + skip_if_wine_exec!( + Ok(()), + "requires a Windows test_stdio_server in the Wine-exec environment" + ); + skip_if_no_network!(Ok(())); + let use_responses_lite = scenario == PrewarmMcpScenario::ResponsesLite; + let close_eager_socket = scenario == PrewarmMcpScenario::EagerSocketCloses; + + let warmup_events = vec![ + responses::ev_response_created("warm-1"), + responses::ev_completed("warm-1"), + ]; + let turn_events = vec![ + responses::ev_response_created("turn-1"), + responses::ev_assistant_message("message-1", "done"), + responses::ev_completed("turn-1"), + ]; + let connections = if close_eager_socket { + // A connection without scripted requests closes as soon as it is accepted. + vec![Vec::new(), vec![turn_events]] + } else { + vec![vec![warmup_events, turn_events]] + }; + let server = responses::start_websocket_server(connections).await; + let command = stdio_server_bin()?; + let service_tier = ServiceTier::Fast.request_value(); + let fixture = test_codex() + .with_model("gpt-5.4") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_info_override("gpt-5.4", move |model| { + model.supports_search_tool = false; + model.use_responses_lite = use_responses_lite; + if use_responses_lite { + model.tool_mode = Some(ToolMode::CodeMode); + } + model.service_tiers.push(ModelServiceTier { + id: service_tier.to_string(), + name: "Fast".to_string(), + description: "Priority processing".to_string(), + }); + }) + .with_config(move |config| { + config + .features + .enable(Feature::FastMode) + .expect("enable Fast Mode for routing hint"); + config.service_tier = Some(ServiceTier::Fast.request_value().to_string()); + config.tool_registry.turn_metadata_includes_tool_info = use_responses_lite; + config.mcp_optional_startup_grace = Duration::ZERO; + let barrier_file = config.cwd.join("allow-prewarm-mcp-initialize"); + insert_mcp_server( + config, + "delayed_prewarm", + stdio_transport( + command, + Some(HashMap::from([( + "MCP_TEST_INITIALIZE_BARRIER_FILE".to_string(), + barrier_file.to_string_lossy().to_string(), + )])), + Vec::new(), + ), + TestMcpServerOptions::default(), + ); + }) + .build_with_websocket_server(&server) + .await?; + + assert!( + server + .wait_for_connections(/*expected*/ 1, Duration::from_secs(5)) + .await, + "websocket handshake should finish while MCP discovery is blocked" + ); + assert_eq!( + server + .single_handshake() + .header(X_CODEX_ROUTING_HINT_HEADER), + Some(format!("model=gpt-5.4;tier={service_tier}")) + ); + assert!( + server.single_connection().is_empty(), + "prewarm should wait for delayed MCP discovery" + ); + if close_eager_socket { + assert!( + server + .wait_for_closed_connections(/*expected*/ 1, Duration::from_secs(5)) + .await, + "the eager socket must close before MCP discovery is released" + ); + } + std::fs::write( + fixture.workspace_path("allow-prewarm-mcp-initialize"), + b"ready", + )?; + 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), + ) + .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_eq!(prewarm["service_tier"].as_str(), Some(service_tier)); + let tool_body = if use_responses_lite { + &prewarm["input"][0] + } else { + &prewarm + }; + assert!( + responses::namespace_child_tool(tool_body, namespace, "echo").is_some(), + "prewarm should contain the discovered MCP echo tool: {prewarm:?}" + ); + let metadata: Value = serde_json::from_str( + prewarm["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .context("prewarm should include turn metadata")?, + )?; + if use_responses_lite { + assert!( + metadata["tool_namespaces_info"][namespace]["functions"]["echo"].is_object(), + "Responses Lite prewarm metadata should reflect discovered MCP tools: {metadata:?}" + ); + } + + 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), + ) + .await + .context("first turn should reuse the prewarmed socket")? + .body_json(); + assert_eq!(first_turn["previous_response_id"], "warm-1"); + assert_eq!(server.handshakes().len(), 1); + + fixture.codex.shutdown_and_wait().await?; + server.shutdown().await; + Ok(()) +} + #[derive(Clone, Copy)] enum InterruptedMcpStartupPhase { PreSamplingCompaction,