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,

View File

@@ -225,6 +225,7 @@ async fn model_change_from_image_to_text_strips_prior_image_content() -> Result<
visibility: ModelVisibility::List,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
priority: 1,
upgrade: None,
base_instructions: "base instructions".to_string(),

View File

@@ -351,5 +351,6 @@ fn test_remote_model(slug: &str, priority: i32) -> ModelInfo {
effective_context_window_percent: 95,
experimental_supported_tools: Vec::new(),
input_modalities: default_input_modalities(),
prefer_websockets: false,
}
}

View File

@@ -613,6 +613,7 @@ async fn ignores_remote_personality_if_remote_models_disabled() -> anyhow::Resul
effective_context_window_percent: 95,
experimental_supported_tools: Vec::new(),
input_modalities: default_input_modalities(),
prefer_websockets: false,
};
let _models_mock = mount_models_once(
@@ -729,6 +730,7 @@ async fn remote_model_friendly_personality_instructions_with_feature() -> anyhow
effective_context_window_percent: 95,
experimental_supported_tools: Vec::new(),
input_modalities: default_input_modalities(),
prefer_websockets: false,
};
let _models_mock = mount_models_once(
@@ -840,6 +842,7 @@ async fn user_turn_personality_remote_model_template_includes_update_message() -
effective_context_window_percent: 95,
experimental_supported_tools: Vec::new(),
input_modalities: default_input_modalities(),
prefer_websockets: false,
};
let _models_mock = mount_models_once(

View File

@@ -141,6 +141,7 @@ async fn remote_models_remote_model_uses_unified_exec() -> Result<()> {
visibility: ModelVisibility::List,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
priority: 1,
upgrade: None,
base_instructions: "base instructions".to_string(),
@@ -379,6 +380,7 @@ async fn remote_models_apply_remote_base_instructions() -> Result<()> {
visibility: ModelVisibility::List,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
priority: 1,
upgrade: None,
base_instructions: remote_base.to_string(),
@@ -862,6 +864,7 @@ fn test_remote_model_with_policy(
visibility,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
priority,
upgrade: None,
base_instructions: "base instructions".to_string(),

View File

@@ -409,6 +409,7 @@ async fn stdio_image_responses_are_sanitized_for_text_only_model() -> anyhow::Re
effective_context_window_percent: 95,
experimental_supported_tools: Vec::new(),
input_modalities: vec![InputModality::Text],
prefer_websockets: false,
}],
},
)

View File

@@ -560,6 +560,7 @@ async fn view_image_tool_returns_unsupported_message_for_text_only_model() -> an
visibility: ModelVisibility::List,
supported_in_api: true,
input_modalities: vec![InputModality::Text],
prefer_websockets: false,
priority: 1,
upgrade: None,
base_instructions: "base instructions".to_string(),