fix(proxy): fail over dead account proxy routes

app/core/balancer/__init__.py

@@ -1,4 +1,5 @@
 from app.core.balancer.logic import (
+    ERROR_BACKOFF_THRESHOLD,
     HEALTH_TIER_DRAINING,
     HEALTH_TIER_HEALTHY,
     HEALTH_TIER_PROBING,
@@ -33,6 +34,7 @@
     "HEALTH_TIER_DRAINING",
     "HEALTH_TIER_HEALTHY",
     "HEALTH_TIER_PROBING",
+    "ERROR_BACKOFF_THRESHOLD",
     "REAUTH_REQUIRED_FAILURE_CODES",
     "AccountState",
     "RoutingCost",

app/core/balancer/logic.py

@@ -88,6 +88,7 @@
 DRAIN_SECONDARY_THRESHOLD_PCT = 90.0
 DRAIN_ERROR_WINDOW_SECONDS = 60.0
 DRAIN_ERROR_COUNT_THRESHOLD = 2
+ERROR_BACKOFF_THRESHOLD = 3
 PROBE_QUIET_SECONDS = 60.0
 PROBE_SUCCESS_STREAK_REQUIRED = 3
 ROUTING_POLICY_NORMAL = "normal"
@@ -467,8 +468,8 @@ def select_account(
             state.error_count = 0
         if state.cooldown_until and current < state.cooldown_until:
             continue
-        if state.error_count >= 3:
-            backoff = min(300, 30 * (2 ** (state.error_count - 3)))
+        if state.error_count >= ERROR_BACKOFF_THRESHOLD:
+            backoff = min(300, 30 * (2 ** (state.error_count - ERROR_BACKOFF_THRESHOLD)))
             if state.last_error_at and current - state.last_error_at < backoff:
                 in_error_backoff.append(state)
                 continue
@@ -505,7 +506,7 @@ def select_account(
         if allow_backoff_fallback and (len(in_error_backoff) > 1 or (in_error_backoff and hard_blocked_exists)):
 
             def _backoff_expires_at(s: AccountState) -> float:
-                backoff = min(300, 30 * (2 ** (s.error_count - 3)))
+                backoff = min(300, 30 * (2 ** (s.error_count - ERROR_BACKOFF_THRESHOLD)))
                 return (s.last_error_at or 0.0) + backoff
 
             available.append(min(in_error_backoff, key=_backoff_expires_at))

app/core/clients/codex.py

@@ -2,9 +2,11 @@
 
 import asyncio
 from dataclasses import dataclass
+from enum import StrEnum
 from typing import Any, Mapping
 
 import aiohttp
+from aiohttp_socks import ProxyConnectionError as SocksProxyConnectionError
 from aiohttp_socks import ProxyConnector
 from python_socks import ProxyType
 
@@ -13,6 +15,11 @@
 _RESERVED = frozenset({"akamai", "extra_fp", "impersonate", "ja3", "proxies", "proxy"})
 
 
+class CodexTransportDispatchState(StrEnum):
+    UNKNOWN = "unknown"
+    NOT_DISPATCHED = "not_dispatched"
+
+
 class CodexTransportError(RuntimeError):
     """Sanitized upstream transport failure.
 
@@ -21,9 +28,18 @@ class CodexTransportError(RuntimeError):
     instead of the original transport message.
     """
 
-    def __init__(self, message: str, *, status_code: int | None = None) -> None:
+    def __init__(
+        self,
+        message: str,
+        *,
+        status_code: int | None = None,
+        dispatch_state: CodexTransportDispatchState = CodexTransportDispatchState.UNKNOWN,
+        exception_type: str | None = None,
+    ) -> None:
         super().__init__(message)
         self.status_code = status_code
+        self.dispatch_state = dispatch_state
+        self.exception_type = exception_type
 
 
 def require_route_or_direct_egress_opt_in(
@@ -174,7 +190,8 @@ async def request_with_route_metadata(
                     response = await _buffer_response(response)
                 return CodexRequestResult(response, candidate, index > 0)
             except Exception as exc:
-                if index == len(endpoints) - 1 or not allow_fallback:
+                pre_dispatch = _transport_dispatch_state(exc) is CodexTransportDispatchState.NOT_DISPATCHED
+                if index == len(endpoints) - 1 or not (allow_fallback or pre_dispatch):
                     raise _transport_error("request", endpoint.id, exc) from None
         raise RuntimeError("unreachable Codex client fallback state")
 
@@ -384,9 +401,17 @@ def _transport_error(operation: str, endpoint_id: str, exc: Exception) -> CodexT
     return CodexTransportError(
         codex_transport_error_message(operation, endpoint_id, exc),
         status_code=_transport_error_status_code(exc),
+        dispatch_state=_transport_dispatch_state(exc),
+        exception_type=type(exc).__name__,
     )
 
 
+def _transport_dispatch_state(exc: Exception) -> CodexTransportDispatchState:
+    if isinstance(exc, (aiohttp.ClientProxyConnectionError, SocksProxyConnectionError)):
+        return CodexTransportDispatchState.NOT_DISPATCHED
+    return CodexTransportDispatchState.UNKNOWN
+
+
 def _transport_error_status_code(exc: Exception) -> int | None:
     response = getattr(exc, "response", None)
     for source in (response, exc):

app/core/clients/proxy.py

@@ -39,6 +39,7 @@
 
 from app.core.clients.codex import (
     CodexClient,
+    CodexTransportDispatchState,
     CodexTransportError,
     codex_transport_error_message,
     create_codex_session,
@@ -445,6 +446,7 @@ def __init__(
         failure_exception_type: str | None = None,
         upstream_status_code: int | None = None,
         upstream_error_code: str | None = None,
+        upstream_dispatch_state: CodexTransportDispatchState | None = None,
     ) -> None:
         super().__init__(f"Proxy response error ({status_code})")
         self.status_code = status_code
@@ -455,6 +457,15 @@ def __init__(
         self.failure_exception_type = failure_exception_type
         self.upstream_status_code = upstream_status_code
         self.upstream_error_code = upstream_error_code
+        self.upstream_dispatch_state = upstream_dispatch_state
+
+
+def is_confirmed_pre_dispatch_transport_error(exc: ProxyResponseError) -> bool:
+    return exc.upstream_dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED
+
+
+def is_ambiguous_dispatch_transport_error(exc: ProxyResponseError) -> bool:
+    return exc.upstream_dispatch_state is CodexTransportDispatchState.UNKNOWN
 
 
 @dataclass(frozen=True)
@@ -3033,7 +3044,18 @@ async def _stream_via_http_after_websocket_rejection(
         failure_phase = "upstream"
         failure_detail = "transport_error"
         failure_exception_type = type(exc).__name__
-        retryable_same_contract = True
+        retryable_same_contract = exc.dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED
+        if raise_for_status:
+            confirmed_pre_dispatch = exc.dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED
+            raise ProxyResponseError(
+                502,
+                openai_error("upstream_unavailable", response_error_message, error_type="server_error"),
+                failure_phase="connect" if confirmed_pre_dispatch else "upstream",
+                retryable_same_contract=confirmed_pre_dispatch,
+                failure_detail="proxy_connect_pre_dispatch" if confirmed_pre_dispatch else "transport_error",
+                failure_exception_type=exc.exception_type or type(exc).__name__,
+                upstream_dispatch_state=exc.dispatch_state,
+            ) from exc
         yield format_sse_event(
             response_failed_event("upstream_unavailable", response_error_message, response_id=get_request_id()),
         )

app/core/clients/proxy_websocket.py

@@ -22,6 +22,7 @@
 
 from app.core.clients.codex import (
     CodexClient,
+    CodexTransportDispatchState,
     CodexTransportError,
     codex_transport_error_message,
     create_codex_session,
@@ -454,6 +455,15 @@ async def connect_responses_websocket(
             raise ProxyResponseError(
                 502,
                 openai_error("upstream_unavailable", str(exc), error_type="server_error"),
+                failure_phase="connect",
+                retryable_same_contract=exc.dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED,
+                failure_detail=(
+                    "proxy_connect_pre_dispatch"
+                    if exc.dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED
+                    else "transport_error"
+                ),
+                failure_exception_type=exc.exception_type or type(exc).__name__,
+                upstream_dispatch_state=exc.dispatch_state,
             ) from exc
         except Exception:
             if owns_codex_client:

app/modules/proxy/_service/http_bridge/helpers.py

@@ -85,6 +85,7 @@
     _sticky_key_from_compact_payload as _sticky_key_from_compact_payload,
 )
 from app.modules.proxy._service.http_bridge.protocol import _HTTPBridgeServiceProtocol
+from app.modules.proxy._service.http_bridge.service_stubs import _headers_with_turn_state
 from app.modules.proxy._service.observability import (
     _hash_identifier as _hash_identifier,
 )
@@ -1064,6 +1065,23 @@ def _make_http_bridge_session_header_fallback_key(
     )
 
 
+def _promote_http_bridge_session_to_codex_affinity(
+    session: "_HTTPBridgeSession",
+    *,
+    turn_state: str,
+    settings: Settings,
+) -> None:
+    session.affinity = _AffinityPolicy(key=turn_state, kind=StickySessionKind.CODEX_SESSION)
+    session.codex_session = True
+    session.downstream_turn_state = turn_state
+    session.downstream_turn_state_aliases.add(turn_state)
+    session.idle_ttl_seconds = max(
+        session.idle_ttl_seconds,
+        float(settings.http_responses_session_bridge_codex_idle_ttl_seconds),
+    )
+    session.headers = _headers_with_turn_state(session.headers, turn_state)
+
+
 async def _http_bridge_should_wait_for_registration(
     self,
     key: _HTTPBridgeSessionKey,

app/modules/proxy/_service/http_bridge/mixin.py

@@ -34,7 +34,6 @@
 from app.core.clients.proxy import codex_control_request as core_codex_control_request  # noqa: F401
 from app.core.clients.proxy import compact_responses as core_compact_responses  # noqa: F401
 from app.core.clients.proxy import transcribe_audio as core_transcribe_audio  # noqa: F401
-from app.core.config.settings import Settings
 from app.core.errors import openai_error
 from app.core.metrics.prometheus import (
     PROMETHEUS_AVAILABLE,
@@ -111,6 +110,7 @@
     _persist_http_bridge_previous_response_alias,
     _persist_http_bridge_turn_state_alias,
     _preferred_http_bridge_reconnect_turn_state,
+    _promote_http_bridge_session_to_codex_affinity,
     _raise_http_bridge_incompatible_admission_handoff,
     _reconcile_durable_http_bridge_ownership,
     _record_bridge_drain_recovery_allowed,
@@ -124,6 +124,7 @@
 )
 from app.modules.proxy._service.http_bridge.owner_forwarding import _HTTPBridgeOwnerForwardingMixin
 from app.modules.proxy._service.http_bridge.protocol import _HTTPBridgeServiceProtocol
+from app.modules.proxy._service.http_bridge.proxy_failover import _HTTPBridgePreDispatchFailover
 from app.modules.proxy._service.http_bridge.request_submit import _HTTPBridgeRequestSubmitMixin
 from app.modules.proxy._service.http_bridge.service_stubs import (
     _await_cancelled_task,
@@ -1714,22 +1715,7 @@ def _unregister_http_bridge_previous_response_ids_locked(self, session: "_HTTPBr
         session.previous_response_ids.clear()
         getattr(session, "previous_response_alias_registration_generations", {}).clear()
 
-    def _promote_http_bridge_session_to_codex_affinity(
-        self,
-        session: "_HTTPBridgeSession",
-        *,
-        turn_state: str,
-        settings: Settings,
-    ) -> None:
-        session.affinity = _AffinityPolicy(key=turn_state, kind=StickySessionKind.CODEX_SESSION)
-        session.codex_session = True
-        session.downstream_turn_state = turn_state
-        session.downstream_turn_state_aliases.add(turn_state)
-        session.idle_ttl_seconds = max(
-            session.idle_ttl_seconds,
-            float(settings.http_responses_session_bridge_codex_idle_ttl_seconds),
-        )
-        session.headers = _headers_with_turn_state(session.headers, turn_state)
+    _promote_http_bridge_session_to_codex_affinity = staticmethod(_promote_http_bridge_session_to_codex_affinity)
 
     async def _claim_durable_http_bridge_session(
         self,
@@ -1809,15 +1795,8 @@ async def _claim_durable_http_bridge_session(
                 return
             raise
 
-    async def _refresh_durable_http_bridge_session(self, session: "_HTTPBridgeSession") -> None:
-        """Renew the durable lease; callers must hold ``self._http_bridge_lock``."""
-
-        await _renew_durable_http_bridge_lease(self, session)
-
-    async def reconcile_durable_http_bridge_ownership(self) -> int:
-        """Close local sessions whose durable row is owned by another instance/epoch."""
-
-        return await _reconcile_durable_http_bridge_ownership(self)
+    _refresh_durable_http_bridge_session = _renew_durable_http_bridge_lease
+    reconcile_durable_http_bridge_ownership = _reconcile_durable_http_bridge_ownership
 
     async def _create_http_bridge_session(
         self,
@@ -1856,8 +1835,11 @@ async def _create_http_bridge_session(
         settings = await _service_get_settings_cache().get()
         excluded_account_ids: set[str] = set()
         retry_same_account_once = preferred_account_id is not None
-        preferred_candidate_id = preferred_account_id
+        proxy_connect_failover = _HTTPBridgePreDispatchFailover(
+            excluded_account_ids, preferred_account_id, affinity.reallocate_sticky
+        )
         selected_account_lease: AccountLease | None = None
+
         while True:
             select_kwargs = {
                 "request_id": request_state.request_log_id or request_state.request_id,
@@ -1866,15 +1848,15 @@ async def _create_http_bridge_session(
                 "api_key": api_key,
                 "sticky_key": affinity.key,
                 "sticky_kind": affinity.kind,
-                "reallocate_sticky": affinity.reallocate_sticky,
+                "reallocate_sticky": proxy_connect_failover.reallocate_sticky,
                 "sticky_max_age_seconds": affinity.max_age_seconds,
                 "prefer_earlier_reset_accounts": settings.prefer_earlier_reset_accounts,
                 "prefer_earlier_reset_window": _prefer_earlier_reset_window(settings),
                 "routing_strategy": _routing_strategy(settings),
                 "model": request_model,
                 "service_tier": request_service_tier,
                 "exclude_account_ids": excluded_account_ids,
-                "preferred_account_id": preferred_candidate_id,
+                "preferred_account_id": proxy_connect_failover.preferred_account_id,
                 "lease_kind": "stream",
                 "estimated_lease_tokens": _estimated_lease_tokens_from_request_usage_budget(request_usage_budget),
                 "fallback_on_preferred_account_unavailable": fallback_on_preferred_account_unavailable,
@@ -1889,6 +1871,8 @@ async def _create_http_bridge_session(
                     preferred_account_id=preferred_account_id,
                     selected_account_id=None,
                 )
+                if proxy_connect_failover.last_error is not None:
+                    raise proxy_connect_failover.last_error
                 status_code = 429 if _is_local_account_cap_code(selection.error_code) else 503
                 error_type = "rate_limit_error" if status_code == 429 else "server_error"
                 raise ProxyResponseError(
@@ -1937,6 +1921,15 @@ async def _create_http_bridge_session(
                 )
                 break
             except ProxyResponseError as exc:
+                if await proxy_connect_failover.handle(
+                    self,
+                    account,
+                    selected_account_lease,
+                    exc,
+                    required_account=require_preferred_account and selected_is_preferred,
+                ):
+                    selected_account_lease = None
+                    continue
                 if exc.status_code != 401 or _remaining_budget_seconds(deadline) <= 0:
                     await self._load_balancer.release_account_lease(selected_account_lease)
                     selected_account_lease = None
@@ -1962,6 +1955,15 @@ async def _create_http_bridge_session(
                     )
                     break
                 except ProxyResponseError as retry_exc:
+                    if await proxy_connect_failover.handle(
+                        self,
+                        account,
+                        selected_account_lease,
+                        retry_exc,
+                        required_account=require_preferred_account and selected_is_preferred,
+                    ):
+                        selected_account_lease = None
+                        continue
                     if retry_exc.status_code != 401:
                         await self._load_balancer.release_account_lease(selected_account_lease)
                         selected_account_lease = None
@@ -1972,7 +1974,7 @@ async def _create_http_bridge_session(
                         selected_account_lease = None
                         raise
                     excluded_account_ids.add(account.id)
-                    preferred_candidate_id = None
+                    proxy_connect_failover.preferred_account_id = None
                     await self._load_balancer.release_account_lease(selected_account_lease)
                     selected_account_lease = None
                     continue
@@ -1984,7 +1986,7 @@ async def _create_http_bridge_session(
                         selected_account_lease = None
                         raise
                     excluded_account_ids.add(account.id)
-                    preferred_candidate_id = None
+                    proxy_connect_failover.preferred_account_id = None
                     await self._load_balancer.release_account_lease(selected_account_lease)
                     selected_account_lease = None
                     continue
@@ -2009,7 +2011,7 @@ async def _create_http_bridge_session(
                             ),
                         ) from exc
                     excluded_account_ids.add(account.id)
-                    preferred_candidate_id = None
+                    proxy_connect_failover.preferred_account_id = None
                     await self._load_balancer.release_account_lease(selected_account_lease)
                     selected_account_lease = None
                     continue
@@ -2048,7 +2050,7 @@ async def _create_http_bridge_session(
                             ),
                         ) from exc
                     excluded_account_ids.add(account.id)
-                    preferred_candidate_id = None
+                    proxy_connect_failover.preferred_account_id = None
                     await self._load_balancer.release_account_lease(selected_account_lease)
                     selected_account_lease = None
                     continue