Prefer websockets when providers support them (#13592)

Remove all flags and model settings.

---------

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
pakrym-oai
2026-03-17 19:46:44 -07:00
committed by GitHub
parent d950543e65
commit 770616414a
34 changed files with 348 additions and 303 deletions

View File

@@ -717,6 +717,7 @@ async fn chatgpt_auth_sends_correct_request() {
let mut model_provider = built_in_model_providers(/* openai_base_url */ None)["openai"].clone();
model_provider.base_url = Some(format!("{}/api/codex", server.uri()));
model_provider.supports_websockets = false;
let mut builder = test_codex()
.with_auth(create_dummy_codex_auth())
.with_config(move |config| {
@@ -791,6 +792,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() {
let model_provider = ModelProviderInfo {
base_url: Some(format!("{}/v1", server.uri())),
supports_websockets: false,
..built_in_model_providers(/* openai_base_url */ None)["openai"].clone()
};
@@ -1832,7 +1834,6 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() {
config.model_verbosity,
false,
false,
false,
None,
);
let mut client_session = client.new_session();
@@ -1968,6 +1969,7 @@ async fn token_count_includes_rate_limits_snapshot() {
let mut provider = built_in_model_providers(/* openai_base_url */ None)["openai"].clone();
provider.base_url = Some(format!("{}/v1", server.uri()));
provider.supports_websockets = false;
let mut builder = test_codex()
.with_auth(CodexAuth::from_api_key("test"))

View File

@@ -8,7 +8,6 @@ use codex_core::ResponseEvent;
use codex_core::WireApi;
use codex_core::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER;
use codex_core::features::Feature;
use codex_core::ws_version_from_features;
use codex_otel::SessionTelemetry;
use codex_otel::TelemetryAuthMode;
use codex_otel::metrics::MetricsClient;
@@ -98,6 +97,28 @@ async fn responses_websocket_streams_request() {
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_streams_without_feature_flag_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, false).await;
let mut client_session = harness.client.new_session();
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)]
async fn responses_websocket_preconnect_reuses_connection() {
skip_if_no_network!();
@@ -133,7 +154,7 @@ async fn responses_websocket_request_prewarm_reuses_connection() {
]])
.await;
let harness = websocket_harness_with_options(&server, false, false, true, true).await;
let harness = websocket_harness_with_options(&server, true).await;
let mut client_session = harness.client.new_session();
let prompt = prompt_with_input(vec![message_item("hello")]);
client_session
@@ -252,7 +273,7 @@ async fn responses_websocket_request_prewarm_is_reused_even_with_header_changes(
]])
.await;
let harness = websocket_harness_with_options(&server, false, false, true, true).await;
let harness = websocket_harness_with_options(&server, true).await;
let mut client_session = harness.client.new_session();
let prompt = prompt_with_input(vec![message_item("hello")]);
client_session
@@ -308,7 +329,7 @@ async fn responses_websocket_request_prewarm_is_reused_even_with_header_changes(
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_prewarm_uses_v2_when_model_prefers_websockets_and_feature_disabled() {
async fn responses_websocket_prewarm_uses_v2_when_provider_supports_websockets() {
skip_if_no_network!();
let server = start_websocket_server(vec![vec![vec![
@@ -317,7 +338,7 @@ async fn responses_websocket_prewarm_uses_v2_when_model_prefers_websockets_and_f
]]])
.await;
let harness = websocket_harness_with_options(&server, false, false, false, true).await;
let harness = websocket_harness_with_options(&server, false).await;
let mut client_session = harness.client.new_session();
let prompt = prompt_with_input(vec![message_item("hello")]);
client_session
@@ -374,7 +395,7 @@ async fn responses_websocket_preconnect_runs_when_only_v2_feature_enabled() {
]]])
.await;
let harness = websocket_harness_with_options(&server, false, false, true, false).await;
let harness = websocket_harness_with_options(&server, true).await;
let mut client_session = harness.client.new_session();
client_session
.preconnect_websocket(&harness.session_telemetry, &harness.model_info)
@@ -404,7 +425,7 @@ async fn responses_websocket_preconnect_runs_when_only_v2_feature_enabled() {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_v2_requests_use_v2_when_model_prefers_websockets() {
async fn responses_websocket_v2_requests_use_v2_when_provider_supports_websockets() {
skip_if_no_network!();
let server = start_websocket_server(vec![vec![
@@ -417,7 +438,7 @@ async fn responses_websocket_v2_requests_use_v2_when_model_prefers_websockets()
]])
.await;
let harness = websocket_harness_with_options(&server, false, false, true, true).await;
let harness = websocket_harness_with_options(&server, true).await;
let mut client_session = harness.client.new_session();
let prompt_one = prompt_with_input(vec![message_item("hello")]);
let prompt_two = prompt_with_input(vec![
@@ -466,7 +487,7 @@ async fn responses_websocket_v2_incremental_requests_are_reused_across_turns() {
]])
.await;
let harness = websocket_harness_with_options(&server, false, false, true, true).await;
let harness = websocket_harness_with_options(&server, false).await;
let prompt_one = prompt_with_input(vec![message_item("hello")]);
let prompt_two = prompt_with_input(vec![
message_item("hello"),
@@ -510,7 +531,7 @@ async fn responses_websocket_v2_wins_when_both_features_enabled() {
]])
.await;
let harness = websocket_harness_with_options(&server, false, true, true, false).await;
let harness = websocket_harness_with_options(&server, false).await;
let mut client_session = harness.client.new_session();
let prompt_one = prompt_with_input(vec![message_item("hello")]);
let prompt_two = prompt_with_input(vec![
@@ -1534,69 +1555,39 @@ async fn websocket_harness_with_runtime_metrics(
server: &WebSocketTestServer,
runtime_metrics_enabled: bool,
) -> WebsocketTestHarness {
websocket_harness_with_options(server, runtime_metrics_enabled, true, false, false).await
websocket_harness_with_options(server, runtime_metrics_enabled).await
}
async fn websocket_harness_with_v2(
server: &WebSocketTestServer,
websocket_v2_enabled: bool,
runtime_metrics_enabled: bool,
) -> WebsocketTestHarness {
websocket_harness_with_options(server, false, true, websocket_v2_enabled, false).await
websocket_harness_with_options(server, runtime_metrics_enabled).await
}
async fn websocket_harness_with_options(
server: &WebSocketTestServer,
runtime_metrics_enabled: bool,
websocket_enabled: bool,
websocket_v2_enabled: bool,
prefer_websockets: bool,
) -> WebsocketTestHarness {
websocket_harness_with_provider_options(
websocket_provider(server),
runtime_metrics_enabled,
websocket_enabled,
websocket_v2_enabled,
prefer_websockets,
)
.await
websocket_harness_with_provider_options(websocket_provider(server), runtime_metrics_enabled)
.await
}
async fn websocket_harness_with_provider_options(
provider: ModelProviderInfo,
runtime_metrics_enabled: bool,
websocket_enabled: bool,
websocket_v2_enabled: bool,
prefer_websockets: bool,
) -> WebsocketTestHarness {
let codex_home = TempDir::new().unwrap();
let mut config = load_default_config_for_test(&codex_home).await;
config.model = Some(MODEL.to_string());
if websocket_enabled {
config
.features
.enable(Feature::ResponsesWebsockets)
.expect("test config should allow feature update");
} else {
config
.features
.disable(Feature::ResponsesWebsockets)
.expect("test config should allow feature update");
}
if runtime_metrics_enabled {
config
.features
.enable(Feature::RuntimeMetrics)
.expect("test config should allow feature update");
}
if websocket_v2_enabled {
config
.features
.enable(Feature::ResponsesWebsocketsV2)
.expect("test config should allow feature update");
}
let config = Arc::new(config);
let mut model_info = codex_core::test_support::construct_model_info_offline(MODEL, &config);
model_info.prefer_websockets = prefer_websockets;
let model_info = codex_core::test_support::construct_model_info_offline(MODEL, &config);
let conversation_id = ThreadId::new();
let auth_manager =
codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("Test API Key"));
@@ -1627,7 +1618,6 @@ async fn websocket_harness_with_provider_options(
provider.clone(),
SessionSource::Exec,
config.model_verbosity,
ws_version_from_features(&config),
false,
runtime_metrics_enabled,
None,

View File

@@ -96,6 +96,7 @@ fn non_openai_model_provider(server: &MockServer) -> ModelProviderInfo {
let mut provider = built_in_model_providers(/* openai_base_url */ None)["openai"].clone();
provider.name = "OpenAI (test)".into();
provider.base_url = Some(format!("{}/v1", server.uri()));
provider.supports_websockets = false;
provider
}

View File

@@ -53,7 +53,6 @@ fn test_model_info(
visibility: ModelVisibility::List,
supported_in_api: true,
input_modalities,
prefer_websockets: false,
used_fallback_model_metadata: false,
supports_search_tool: false,
priority: 1,
@@ -849,7 +848,6 @@ async fn model_switch_to_smaller_model_updates_token_context_window() -> Result<
visibility: ModelVisibility::List,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
used_fallback_model_metadata: false,
supports_search_tool: false,
priority: 1,

View File

@@ -351,7 +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,
used_fallback_model_metadata: false,
supports_search_tool: false,
}

View File

@@ -659,7 +659,6 @@ 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,
used_fallback_model_metadata: false,
supports_search_tool: false,
};
@@ -775,7 +774,6 @@ 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,
used_fallback_model_metadata: false,
supports_search_tool: false,
};

View File

@@ -289,7 +289,6 @@ 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,
used_fallback_model_metadata: false,
supports_search_tool: false,
priority: 1,
@@ -533,7 +532,6 @@ async fn remote_models_apply_remote_base_instructions() -> Result<()> {
visibility: ModelVisibility::List,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
used_fallback_model_metadata: false,
supports_search_tool: false,
priority: 1,
@@ -1001,7 +999,6 @@ fn test_remote_model_with_policy(
visibility,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
used_fallback_model_metadata: false,
supports_search_tool: false,
priority,

View File

@@ -419,7 +419,6 @@ 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,
used_fallback_model_metadata: false,
supports_search_tool: false,
}],

View File

@@ -64,7 +64,6 @@ fn test_model_info(
visibility,
supported_in_api: true,
input_modalities: default_input_modalities(),
prefer_websockets: false,
used_fallback_model_metadata: false,
supports_search_tool: false,
priority: 1,

View File

@@ -1270,7 +1270,6 @@ 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,
used_fallback_model_metadata: false,
supports_search_tool: false,
priority: 1,

View File

@@ -1,5 +1,4 @@
use anyhow::Result;
use codex_core::features::Feature;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
@@ -45,10 +44,7 @@ async fn websocket_fallback_switches_to_http_on_upgrade_required_connect() -> Re
move |config| {
config.model_provider.base_url = Some(base_url);
config.model_provider.wire_api = codex_core::WireApi::Responses;
config
.features
.enable(Feature::ResponsesWebsockets)
.expect("test config should allow feature update");
config.model_provider.supports_websockets = true;
// If we don't treat 426 specially, the sampling loop would retry the WebSocket
// handshake before switching to the HTTP transport.
config.model_provider.stream_max_retries = Some(2);
@@ -94,10 +90,7 @@ async fn websocket_fallback_switches_to_http_after_retries_exhausted() -> Result
move |config| {
config.model_provider.base_url = Some(base_url);
config.model_provider.wire_api = codex_core::WireApi::Responses;
config
.features
.enable(Feature::ResponsesWebsockets)
.expect("test config should allow feature update");
config.model_provider.supports_websockets = true;
config.model_provider.stream_max_retries = Some(2);
config.model_provider.request_max_retries = Some(0);
}
@@ -142,10 +135,7 @@ async fn websocket_fallback_hides_first_websocket_retry_stream_error() -> Result
move |config| {
config.model_provider.base_url = Some(base_url);
config.model_provider.wire_api = codex_core::WireApi::Responses;
config
.features
.enable(Feature::ResponsesWebsockets)
.expect("test config should allow feature update");
config.model_provider.supports_websockets = true;
config.model_provider.stream_max_retries = Some(2);
config.model_provider.request_max_retries = Some(0);
}
@@ -220,10 +210,7 @@ async fn websocket_fallback_is_sticky_across_turns() -> Result<()> {
move |config| {
config.model_provider.base_url = Some(base_url);
config.model_provider.wire_api = codex_core::WireApi::Responses;
config
.features
.enable(Feature::ResponsesWebsockets)
.expect("test config should allow feature update");
config.model_provider.supports_websockets = true;
config.model_provider.stream_max_retries = Some(2);
config.model_provider.request_max_retries = Some(0);
}