Fallback to HTTP on websocket connect failures

weidapao · weidapao · commit 2284b7ee5f52 · 2026-04-27T21:57:37.000+08:00
diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs
@@ -1014,7 +1014,7 @@ impl ModelClientSession {
     pub async fn preconnect_websocket(
         &mut self,
         session_telemetry: &SessionTelemetry,
-        _model_info: &ModelInfo,
+        model_info: &ModelInfo,
     ) -> std::result::Result<(), ApiError> {
         if !self.client.responses_websocket_enabled() {
             return Ok(());
@@ -1033,7 +1033,7 @@ impl ModelClientSession {
             client_setup.api_auth.as_ref(),
             PendingUnauthorizedRetry::default(),
         );
-        let connection = self
+        let connection = match self
             .client
             .connect_websocket(
                 session_telemetry,
@@ -1044,7 +1044,15 @@ impl ModelClientSession {
                 auth_context,
                 RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT),
             )
-            .await?;
+            .await
+        {
+            Ok(connection) => connection,
+            Err(err) if should_fallback_to_http_after_websocket_connect_error(&err) => {
+                self.try_switch_fallback_transport(session_telemetry, model_info);
+                return Ok(());
+            }
+            Err(err) => return Err(err),
+        };
         self.websocket_session.connection = Some(connection);
         self.websocket_session
             .set_connection_reused(/*connection_reused*/ false);
@@ -1335,6 +1343,10 @@ impl ModelClientSession {
                 {
                     return Ok(WebsocketStreamOutcome::FallbackToHttp);
                 }
+                Err(err) if should_fallback_to_http_after_websocket_connect_error(&err) => {
+                    self.reset_websocket_session();
+                    return Ok(WebsocketStreamOutcome::FallbackToHttp);
+                }
                 Err(ApiError::Transport(
                     unauthorized_transport @ TransportError::Http { status, .. },
                 )) if status == StatusCode::UNAUTHORIZED => {
@@ -1905,6 +1917,13 @@ fn api_error_http_status(error: &ApiError) -> Option<u16> {
     }
 }
 
+fn should_fallback_to_http_after_websocket_connect_error(error: &ApiError) -> bool {
+    matches!(
+        error,
+        ApiError::Transport(TransportError::Timeout | TransportError::Network(_))
+    )
+}
+
 struct ApiTelemetry {
     session_telemetry: SessionTelemetry,
     auth_context: AuthRequestTelemetryContext,
diff --git a/codex-rs/core/tests/suite/websocket_fallback.rs b/codex-rs/core/tests/suite/websocket_fallback.rs
@@ -75,6 +75,57 @@ async fn websocket_fallback_switches_to_http_on_upgrade_required_connect() -> Re
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn websocket_fallback_switches_to_http_on_connect_timeout() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = responses::start_mock_server().await;
+    Mock::given(method("GET"))
+        .and(path_regex(".*/responses$"))
+        .respond_with(ResponseTemplate::new(426).set_delay(Duration::from_millis(500)))
+        .mount(&server)
+        .await;
+
+    let response_mock = mount_sse_once(
+        &server,
+        sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
+    )
+    .await;
+
+    let mut builder = test_codex().with_config({
+        let base_url = format!("{}/v1", server.uri());
+        move |config| {
+            config.model_provider.base_url = Some(base_url);
+            config.model_provider.wire_api = WireApi::Responses;
+            config.model_provider.supports_websockets = true;
+            config.model_provider.stream_max_retries = Some(5);
+            config.model_provider.request_max_retries = Some(0);
+            config.model_provider.websocket_connect_timeout_ms = Some(50);
+        }
+    });
+    let test = builder.build(&server).await?;
+
+    test.submit_turn("hello").await?;
+
+    let requests = server.received_requests().await.unwrap_or_default();
+    let websocket_attempts = requests
+        .iter()
+        .filter(|req| req.method == Method::GET && req.url.path().ends_with("/responses"))
+        .count();
+    let http_attempts = requests
+        .iter()
+        .filter(|req| req.method == Method::POST && req.url.path().ends_with("/responses"))
+        .count();
+
+    // Timeout during the WebSocket handshake is a strong signal that the local transport path
+    // cannot carry WebSockets, so fallback should activate before consuming stream retries.
+    assert_eq!(websocket_attempts, 1);
+    assert_eq!(http_attempts, 1);
+    assert_eq!(response_mock.requests().len(), 1);
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn websocket_fallback_switches_to_http_after_retries_exhausted() -> Result<()> {
     skip_if_no_network!(Ok(()));
