openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-14 16:38:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Overlap startup WebSocket preconnect with tool discovery (#47635)

Browse Source 
## Why

Startup prewarm waited for tool discovery before opening its WebSocket connection. Slow MCP startup therefore delayed connection setup even though the handshake could proceed independently.

## What changed

- Run the WebSocket handshake concurrently with tool capture, then send `generate=false` with the discovered tools and finalized request metadata.
- Include model and effective service-tier routing hints in preconnect, including the inherited tier for spawned subagents. Omit routing hints for guardian reviewer requests.
- Recover authentication failures during preconnect and switch to HTTP fallback on `426 Upgrade Required`.

## Testing

Add and extend regression coverage for handshake overlap with delayed MCP discovery, cancellation before warmup, finalized Responses Lite tool metadata, socket reuse and reconnection, service-tier changes and inheritance, and authentication recovery during preconnect.

GitOrigin-RevId: ecb53540765df8227e55e79910a13a05260a4c8d
This commit is contained in:
Ben Romano
2026-09-23 19:16:20 +00:00
committed by copyberry
parent dce62a50a2
commit ce7df36237
11 changed files with 727 additions and 440 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
+1 -5
View File
@@ -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";
codex-rs/core/src/client.rs
+95 -67
View File
@@ -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<HeaderValue> {
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<String>,
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:
codex-rs/core/src/session/mod.rs
+1 -1
View File
@@ -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<String>,
fast_mode_enabled: bool,
model_info: &ModelInfo,
codex-rs/core/src/session/session.rs
+1 -1
View File
@@ -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,
codex-rs/core/src/session_startup_prewarm.rs
+58 -12
View File
@@ -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<CodexResult<ModelClientSession>>,
@@ -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
codex-rs/core/tests/common/responses.rs
+26
View File
@@ -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<Mutex<Vec<Vec<WebSocketRequest>>>>,
handshakes: Arc<Mutex<Vec<WebSocketHandshake>>>,
request_log_updated: Arc<Notify>,
closed_connections: watch::Receiver<usize>,
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<WebSocketHandshake> {
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,
}
codex-rs/core/tests/suite/agent_websocket.rs
+184 -67
View File
@@ -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
codex-rs/core/tests/suite/client.rs
+48
View File
@@ -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,
codex-rs/core/tests/suite/client_websockets.rs
+92 -284
View File
@@ -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!();
codex-rs/core/tests/suite/guardian_review.rs
+3
View File
@@ -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!(
(
codex-rs/core/tests/suite/rmcp_client.rs
+218 -3
View File
@@ -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,