mirror of
https://github.com/openai/codex.git
synced 2026-05-03 10:56:37 +00:00
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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user