Prefer websocket transport when model opts in (#11386)

Summary
- add a `prefer_websockets` field to `ModelInfo`, defaulting to `false`
in all fixtures and constructors
- wire the new flag into websocket selection so models that opt in
always use websocket transport even when the feature gate is off

Testing
- Not run (not requested)
This commit is contained in:
pakrym-oai
2026-02-10 18:50:48 -08:00
committed by GitHub
parent bfd4e2112c
commit c68999ee6d
14 changed files with 89 additions and 19 deletions

View File

@@ -105,7 +105,7 @@ async fn responses_websocket_preconnect_reuses_connection() {
let harness = websocket_harness(&server).await;
let mut client_session = harness.client.new_session();
client_session
.prewarm_websocket(&harness.otel_manager, None)
.prewarm_websocket(&harness.otel_manager, &harness.model_info, None)
.await
.expect("websocket prewarm failed");
let prompt = prompt_with_input(vec![message_item("hello")]);
@@ -130,7 +130,7 @@ async fn responses_websocket_preconnect_is_reused_even_with_header_changes() {
let harness = websocket_harness(&server).await;
let mut client_session = harness.client.new_session();
client_session
.prewarm_websocket(&harness.otel_manager, None)
.prewarm_websocket(&harness.otel_manager, &harness.model_info, None)
.await
.expect("websocket prewarm failed");
let prompt = prompt_with_input(vec![message_item("hello")]);
@@ -158,6 +158,36 @@ async fn responses_websocket_preconnect_is_reused_even_with_header_changes() {
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_prewarm_uses_model_preference_when_feature_disabled() {
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, false, false, false, true).await;
let mut client_session = harness.client.new_session();
client_session
.prewarm_websocket(&harness.otel_manager, &harness.model_info, None)
.await
.expect("websocket prewarm failed");
// Prewarm should only perform the handshake, not send response.create.
assert_eq!(server.handshakes().len(), 1);
assert_eq!(server.single_connection().len(), 0);
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);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[traced_test]
async fn responses_websocket_emits_websocket_telemetry_events() {
@@ -887,26 +917,32 @@ async fn websocket_harness_with_runtime_metrics(
server: &WebSocketTestServer,
runtime_metrics_enabled: bool,
) -> WebsocketTestHarness {
websocket_harness_with_options(server, runtime_metrics_enabled, false).await
websocket_harness_with_options(server, runtime_metrics_enabled, true, false, false).await
}
async fn websocket_harness_with_v2(
server: &WebSocketTestServer,
websocket_v2_enabled: bool,
) -> WebsocketTestHarness {
websocket_harness_with_options(server, false, websocket_v2_enabled).await
websocket_harness_with_options(server, false, true, websocket_v2_enabled, false).await
}
async fn websocket_harness_with_options(
server: &WebSocketTestServer,
runtime_metrics_enabled: bool,
websocket_enabled: bool,
websocket_v2_enabled: bool,
prefer_websockets: bool,
) -> WebsocketTestHarness {
let provider = websocket_provider(server);
let codex_home = TempDir::new().unwrap();
let mut config = load_default_config_for_test(&codex_home).await;
config.model = Some(MODEL.to_string());
config.features.enable(Feature::ResponsesWebsockets);
if websocket_enabled {
config.features.enable(Feature::ResponsesWebsockets);
} else {
config.features.disable(Feature::ResponsesWebsockets);
}
if runtime_metrics_enabled {
config.features.enable(Feature::RuntimeMetrics);
}
@@ -914,7 +950,8 @@ async fn websocket_harness_with_options(
config.features.enable(Feature::ResponsesWebsocketsV2);
}
let config = Arc::new(config);
let model_info = ModelsManager::construct_model_info_offline(MODEL, &config);
let mut model_info = ModelsManager::construct_model_info_offline(MODEL, &config);
model_info.prefer_websockets = prefer_websockets;
let conversation_id = ThreadId::new();
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let exporter = InMemoryMetricExporter::default();
@@ -944,7 +981,7 @@ async fn websocket_harness_with_options(
provider.clone(),
SessionSource::Exec,
config.model_verbosity,
true,
websocket_enabled,
websocket_v2_enabled,
false,
runtime_metrics_enabled,