diff --git a/app/core/balancer/__init__.py b/app/core/balancer/__init__.py index 75f3648374..f06baff11b 100644 --- a/app/core/balancer/__init__.py +++ b/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", diff --git a/app/core/balancer/logic.py b/app/core/balancer/logic.py index 2501eac230..53be3f5179 100644 --- a/app/core/balancer/logic.py +++ b/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)) diff --git a/app/core/clients/codex.py b/app/core/clients/codex.py index 0dbdd5efc1..18eb5a9b86 100644 --- a/app/core/clients/codex.py +++ b/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): diff --git a/app/core/clients/proxy.py b/app/core/clients/proxy.py index 76af69ec1b..1de0134024 100644 --- a/app/core/clients/proxy.py +++ b/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()), ) diff --git a/app/core/clients/proxy_websocket.py b/app/core/clients/proxy_websocket.py index 0d3b0b9a12..ca894c3c2d 100644 --- a/app/core/clients/proxy_websocket.py +++ b/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: diff --git a/app/modules/proxy/_service/http_bridge/helpers.py b/app/modules/proxy/_service/http_bridge/helpers.py index ac421da4b6..575a4087f9 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/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, diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index e99dd36115..7f7cf839e8 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/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,7 +1848,7 @@ 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), @@ -1874,7 +1856,7 @@ async def _create_http_bridge_session( "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 diff --git a/app/modules/proxy/_service/http_bridge/proxy_failover.py b/app/modules/proxy/_service/http_bridge/proxy_failover.py new file mode 100644 index 0000000000..e09686fbf8 --- /dev/null +++ b/app/modules/proxy/_service/http_bridge/proxy_failover.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from app.core.clients.proxy import ProxyResponseError, is_confirmed_pre_dispatch_transport_error +from app.db.models import Account +from app.modules.proxy.load_balancer import AccountLease + + +@dataclass +class _HTTPBridgePreDispatchFailover: + excluded_account_ids: set[str] + preferred_account_id: str | None + reallocate_sticky: bool + last_error: ProxyResponseError | None = None + + async def handle( + self, + service: Any, + account: Account, + lease: AccountLease | None, + exc: ProxyResponseError, + *, + required_account: bool, + ) -> bool: + if not is_confirmed_pre_dispatch_transport_error(exc): + return False + await service._load_balancer.release_account_lease(lease) + await service._load_balancer.record_error_backoff(account) + if required_account: + raise exc + self.last_error = exc + self.excluded_account_ids.add(account.id) + self.preferred_account_id = None + self.reallocate_sticky = True + return True diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index 84da68415d..68f9373ded 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -14,7 +14,13 @@ from app.core.auth.refresh import RefreshError from app.core.balancer import failover_decision from app.core.balancer.types import UpstreamError -from app.core.clients.proxy import ProxyResponseError, _resolve_stream_transport, pop_stream_timeout_overrides +from app.core.clients.proxy import ( + ProxyResponseError, + _resolve_stream_transport, + is_ambiguous_dispatch_transport_error, + is_confirmed_pre_dispatch_transport_error, + pop_stream_timeout_overrides, +) from app.core.errors import openai_error, response_failed_event from app.core.openai.requests import ResponsesRequest, extract_input_file_ids from app.core.upstream_proxy import UpstreamProxyRouteError @@ -326,6 +332,8 @@ async def _stream_with_retry( upstream_transport_metric_recorded = False settlement = _StreamSettlement() last_transient_exc: ProxyResponseError | None = None + last_pre_dispatch_transport_error: ProxyResponseError | None = None + handled_dispatch_transport_error: ProxyResponseError | None = None last_security_work_retry_error: _RetryableStreamError | None = None excluded_account_ids: set[str] = set() deferred_capacity_account: Account | None = None @@ -336,6 +344,7 @@ async def _stream_with_retry( last_retryable_stream_error: _RetryableStreamError | None = None require_security_work_authorized = False account_leases: list[AccountLease] = [] + current_account_lease: AccountLease | None = None estimated_lease_tokens = _facade()._estimated_lease_tokens_from_request_usage_budget( estimate_api_key_request_usage(payload) ) @@ -465,6 +474,71 @@ def _record_upstream_transport_metric_once(status: str) -> None: status=status, ) + def _render_dispatch_transport_error(exc: ProxyResponseError) -> str: + error = _parse_openai_error(exc.payload) + error_code = ( + _normalize_error_code( + error.code if error else None, + error.type if error else None, + ) + or "upstream_unavailable" + ) + error_message = error.message if error and error.message else "Upstream transport failed" + event = response_failed_event( + error_code, + error_message, + error_type=(error.type if error else None) or "server_error", + response_id=request_id, + error_param=error.param if error else None, + ) + _apply_error_metadata(event["response"]["error"], error) + return format_sse_event(event) + + async def _handle_dispatch_transport_error( + exc: ProxyResponseError, + account: Account, + *, + outcome: str, + ) -> bool | None: + nonlocal affinity, current_account_lease + nonlocal handled_dispatch_transport_error, last_pre_dispatch_transport_error + confirmed_pre_dispatch = is_confirmed_pre_dispatch_transport_error(exc) + ambiguous_dispatch = is_ambiguous_dispatch_transport_error(exc) + if not confirmed_pre_dispatch and not ambiguous_dispatch: + return None + handled_dispatch_transport_error = exc + await _release_tracked_stream_lease(current_account_lease) + current_account_lease = None + if ambiguous_dispatch: + error = _parse_openai_error(exc.payload) + error_code = _normalize_error_code( + error.code if error else None, + error.type if error else None, + ) + await proxy._handle_stream_error( + account, + _upstream_error_from_openai(error), + error_code, + http_status=exc.status_code, + ) + return False + last_pre_dispatch_transport_error = exc + await proxy._load_balancer.record_error_backoff(account) + can_try_other_account = ( + not require_preferred_account and account.id != file_preferred_account_id and attempt < max_attempts - 1 + ) + if not can_try_other_account: + return False + excluded_account_ids.add(account.id) + affinity = replace(affinity, reallocate_sticky=True) + logger.info( + "Retrying stream after confirmed pre-dispatch proxy failure request_id=%s account_id=%s phase=%s", + request_id, + account.id, + outcome, + ) + return True + try: if payload.previous_response_id is not None: previous_response_lookup_session_id = _owner_lookup_session_id_from_headers(headers) @@ -700,6 +774,7 @@ def _record_upstream_transport_metric_once(status: str) -> None: deferred_capacity_lease = None if ( not account + and last_pre_dispatch_transport_error is None and ( selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES or not (propagate_http_errors and last_transient_exc is not None) @@ -730,6 +805,11 @@ def _record_upstream_transport_metric_once(status: str) -> None: continue break if not account: + if last_pre_dispatch_transport_error is not None: + if propagate_http_errors: + raise last_pre_dispatch_transport_error + yield _render_dispatch_transport_error(last_pre_dispatch_transport_error) + return if selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES: no_accounts_msg = selection.error_message or "Local account capacity is exhausted" error_code = selection.error_code @@ -1224,6 +1304,19 @@ def _record_upstream_transport_metric_once(status: str) -> None: request_id, ) return + if isinstance(tex, ProxyResponseError): + dispatch_transport_retry = await _handle_dispatch_transport_error( + tex, + account, + outcome="previsible", + ) + if dispatch_transport_retry is not None: + if dispatch_transport_retry: + break + if propagate_http_errors: + raise tex + yield _render_dispatch_transport_error(tex) + return if isinstance(tex, ProxyResponseError) and tex.status_code != 500: error = _parse_openai_error(tex.payload) code = _normalize_error_code( @@ -1458,6 +1551,22 @@ def _record_upstream_transport_metric_once(status: str) -> None: await proxy._handle_stream_error(account, exc.error, exc.code) return except ProxyResponseError as exc: + if is_confirmed_pre_dispatch_transport_error(exc) or is_ambiguous_dispatch_transport_error(exc): + dispatch_transport_retry = ( + False + if handled_dispatch_transport_error is exc and current_account_lease is None + else await _handle_dispatch_transport_error( + exc, + account, + outcome="outer", + ) + ) + if dispatch_transport_retry: + continue + if propagate_http_errors: + raise + yield _render_dispatch_transport_error(exc) + return if exc.status_code == 401: remaining_budget = _facade()._remaining_budget_seconds(deadline) if remaining_budget <= 0: @@ -1631,6 +1740,18 @@ def _record_upstream_transport_metric_once(status: str) -> None: request_id, ) return + dispatch_transport_retry = await _handle_dispatch_transport_error( + retry_exc, + account, + outcome="post_refresh", + ) + if dispatch_transport_retry is not None: + if dispatch_transport_retry: + continue + if propagate_http_errors: + raise retry_exc + yield _render_dispatch_transport_error(retry_exc) + return error = _parse_openai_error(retry_exc.payload) error_code = _normalize_error_code( error.code if error else None, @@ -1787,6 +1908,11 @@ def _record_upstream_transport_metric_once(status: str) -> None: return # When HTTP error propagation is enabled and the last failure was # a transient 500, re-raise to preserve the upstream status/payload. + if last_pre_dispatch_transport_error is not None: + if propagate_http_errors: + raise last_pre_dispatch_transport_error + yield _render_dispatch_transport_error(last_pre_dispatch_transport_error) + return if propagate_http_errors and last_transient_exc is not None: raise last_transient_exc if last_retryable_stream_error is not None: diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index dc988fafe3..6975955259 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -40,6 +40,7 @@ apply_codex_installation_headers, apply_codex_installation_metadata, filter_inbound_headers, + is_confirmed_pre_dispatch_transport_error, pop_compact_timeout_overrides, pop_stream_timeout_overrides, pop_transcribe_timeout_overrides, @@ -1869,6 +1870,7 @@ async def _connect_proxy_websocket( force_refresh=forced_refresh_account_id == account.id, ) except ProxyResponseError as exc: + confirmed_pre_dispatch = is_confirmed_pre_dispatch_transport_error(exc) action = await proxy._decide_websocket_failover_action( account=account, exc=exc, @@ -1876,9 +1878,13 @@ async def _connect_proxy_websocket( attempt=attempt + 1, max_attempts=max_attempts, deterministic_failover_enabled=getattr(base_settings, "deterministic_failover_enabled", True), + require_preferred_account=require_preferred_account, ) if action == "failover_next": await proxy._load_balancer.release_account_lease(selected_stream_lease) + selected_stream_lease = None + if confirmed_pre_dispatch: + await proxy._load_balancer.record_error_backoff(account) last_failover_exc = exc last_failover_account = account excluded_account_ids.add(account.id) @@ -1888,6 +1894,8 @@ async def _connect_proxy_websocket( error_message = error.message if error else None await proxy._load_balancer.release_account_lease(selected_stream_lease) selected_stream_lease = None + if confirmed_pre_dispatch: + await proxy._load_balancer.record_error_backoff(account) await proxy._emit_websocket_connect_failure( websocket, client_send_lock=client_send_lock, @@ -2396,13 +2404,20 @@ async def _decide_websocket_failover_action( attempt: int, max_attempts: int, deterministic_failover_enabled: bool, + require_preferred_account: bool = False, ) -> str: proxy = cast(_WebSocketServiceProtocol, self) _ = proxy - classified = await proxy._handle_websocket_connect_error(account, exc) - failure_class = classified["failure_class"] if isinstance(classified, dict) else "non_retryable" + confirmed_pre_dispatch = is_confirmed_pre_dispatch_transport_error(exc) + if confirmed_pre_dispatch: + failure_class = "retryable_transient" + else: + classified = await proxy._handle_websocket_connect_error(account, exc) + failure_class = classified["failure_class"] if isinstance(classified, dict) else "non_retryable" candidates_remaining = max_attempts - attempt - if exc.status_code == 401 and candidates_remaining > 0: + if confirmed_pre_dispatch: + action = "surface" if require_preferred_account or candidates_remaining <= 0 else "failover_next" + elif exc.status_code == 401 and candidates_remaining > 0: action = "failover_next" elif deterministic_failover_enabled: action = failover_decision( diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index c87f0de6f6..6395d341b4 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -13,6 +13,7 @@ from app.core import usage as usage_core from app.core.balancer import ( + ERROR_BACKOFF_THRESHOLD, HEALTH_TIER_DRAINING, HEALTH_TIER_HEALTHY, HEALTH_TIER_PROBING, @@ -1535,7 +1536,17 @@ async def mark_permanent_failure(self, account: Account, error_code: str) -> Non async def record_error(self, account: Account) -> None: await self.record_errors(account, 1) - async def record_errors(self, account: Account, count: int) -> None: + async def record_error_backoff(self, account: Account) -> None: + """Record one error and immediately enter the bounded transient backoff.""" + await self.record_errors(account, 1, minimum_error_count=ERROR_BACKOFF_THRESHOLD) + + async def record_errors( + self, + account: Account, + count: int, + *, + minimum_error_count: int = 0, + ) -> None: """Record *count* transient errors in a single lock acquisition.""" if count < 1: return @@ -1543,7 +1554,7 @@ async def record_errors(self, account: Account, count: int) -> None: async with lock: account_snapshot = _clone_account(account) state = self._state_for(account) - state.error_count += count + state.error_count = max(state.error_count + count, minimum_error_count) state.last_error_at = time.time() self._sync_runtime_state(account, state) runtime = self._runtime.get(account.id) diff --git a/openspec/changes/retry-account-proxy-connect-failures/design.md b/openspec/changes/retry-account-proxy-connect-failures/design.md new file mode 100644 index 0000000000..5546a6dee8 --- /dev/null +++ b/openspec/changes/retry-account-proxy-connect-failures/design.md @@ -0,0 +1,55 @@ +# Design + +## Transport provenance + +`CodexTransportError` will carry a typed dispatch state and the sanitized +transport exception class. The only replay-safe state is `not_dispatched`, +which is assigned when the client library explicitly reports failure to +connect to the configured HTTP or SOCKS proxy. Generic connector errors, +timeouts, response-body failures, websocket receive failures, and unknown +exceptions remain `unknown`. + +The state is copied onto `ProxyResponseError` when websocket/bridge startup or +HTTP stream startup crosses the core-client boundary. No proxy URL, +credentials, or raw exception text is retained or exposed. + +## Retry order + +For an HTTP POST, a confirmed `not_dispatched` failure may try the next endpoint +in the already-resolved proxy pool. This is safe despite the method being +non-idempotent because no request reached upstream. Ambiguous POST failures +retain the existing no-fallback rule. + +If every same-pool endpoint fails before dispatch, raw HTTP/SSE, native +Responses WebSocket, and the HTTP responses bridge exclude the selected +account and use the existing bounded account-selection loop. The original +sanitized 502 is retained and returned if no replacement exists. + +## Ownership boundary + +Only movable requests can cross accounts. A client or proxy continuation that +requires its owner, an account-scoped uploaded file, a forced single-account +route, or another hard preferred-account contract fails closed on the original +account. Soft prompt-cache and process-session affinity may move. + +## Account backoff and resource ordering + +A confirmed dead account route is stronger evidence than a generic transient +stream error. It raises the account to the existing transient error-backoff +floor immediately: 30 seconds at the floor, exponentially bounded by the +existing 300-second cap. It does not pause, deactivate, rate-limit, or +quota-penalize the account. + +Per-account response-create and stream leases are released before recording +the backoff. The downstream API-key reservation is request-scoped rather than +account-scoped, so an internal pre-dispatch failover keeps that single +reservation alive instead of releasing and racing to reacquire it. The normal +terminal finalizer settles or releases it exactly once after the replacement +attempt or final failure. + +## Non-goals + +- Do not replay ambiguous failures after proxy acceptance, header wait, or + response-body processing. +- Do not add endpoint-health persistence or change proxy-pool membership. +- Do not broaden generic `upstream_unavailable` retry classification. diff --git a/openspec/changes/retry-account-proxy-connect-failures/proposal.md b/openspec/changes/retry-account-proxy-connect-failures/proposal.md new file mode 100644 index 0000000000..b43c8931e2 --- /dev/null +++ b/openspec/changes/retry-account-proxy-connect-failures/proposal.md @@ -0,0 +1,15 @@ +# Retry confirmed account-proxy connect failures + +An account-bound upstream proxy can remain administratively active after its +listener stops accepting connections. Responses requests selected onto that +account then fail before upstream sees the request, but the proxy currently +turns the transport failure into a terminal stream event or a bridge startup +error. Client retries may select the same account again. + +Add explicit, sanitized dispatch provenance to account-routed transport +failures. When the transport library proves that the connection to the +selected proxy failed before request dispatch, try another endpoint in the +same proxy pool and then another eligible account for movable Responses +requests. Apply bounded transient account backoff so independent requests do +not immediately rediscover the dead route. Ambiguous failures and hard +continuity or file ownership remain non-replayable and fail closed. diff --git a/openspec/changes/retry-account-proxy-connect-failures/specs/upstream-proxy-routing/spec.md b/openspec/changes/retry-account-proxy-connect-failures/specs/upstream-proxy-routing/spec.md new file mode 100644 index 0000000000..65712fd4c1 --- /dev/null +++ b/openspec/changes/retry-account-proxy-connect-failures/specs/upstream-proxy-routing/spec.md @@ -0,0 +1,67 @@ +# upstream-proxy-routing Delta Specification + +## ADDED Requirements + +### Requirement: Confirmed account-proxy connection failures fail over safely + +When an account-routed transport reports that it could not connect to the +selected proxy endpoint and proves that the upstream request was not +dispatched, the service MUST classify the failure with sanitized structured +pre-dispatch provenance. For a route with another usable endpoint in the same +proxy pool, the client MUST try that endpoint before moving accounts, including +for a non-idempotent request. If the pool cannot connect, movable Responses +requests MUST exclude the failed account and retry another eligible account +within the existing request budget and attempt limits. + +This behavior MUST cover raw HTTP/SSE, native Responses WebSocket, and the HTTP +responses bridge. Before recording transient account backoff, the service MUST +release response-create and stream leases held for the failed account. A +request-scoped API-key reservation MUST remain singular across an internal +pre-dispatch failover and MUST settle or release exactly once at the terminal +request outcome. The confirmed failure MUST place the account at the existing +bounded transient error-backoff floor, but MUST NOT pause, deactivate, +rate-limit, or quota-penalize it. + +The service MUST NOT replay a request when dispatch is unknown or when the +request depends on hard previous-response, turn-state, uploaded-file, +single-account, or other required account ownership. If no eligible +replacement account exists, the service MUST preserve the original sanitized +upstream-unavailable failure instead of replacing it with a generated +`no_accounts` error. + +#### Scenario: POST uses a healthy endpoint from the same proxy pool + +- **GIVEN** a non-idempotent Responses POST is routed through a proxy pool with two endpoints +- **AND** connecting to the first endpoint fails before request dispatch +- **WHEN** the second endpoint is reachable +- **THEN** the service sends the request through the second endpoint +- **AND** it does not move the request to another account + +#### Scenario: movable request retries another account + +- **GIVEN** two eligible accounts and the first account's complete proxy route refuses connections before dispatch +- **WHEN** a fresh Responses request has no hard account ownership +- **THEN** the service releases the first account's response-create and stream leases +- **AND** it records bounded transient backoff for the first account +- **AND** it excludes the first account and completes through the second account +- **AND** no failure event from the first attempt is forwarded downstream + +#### Scenario: hard account ownership fails closed + +- **GIVEN** a Responses request depends on a previous-response owner or an account-scoped uploaded file +- **AND** the required account's proxy refuses the connection before dispatch +- **WHEN** another account is otherwise eligible +- **THEN** the service does not send the request to the other account +- **AND** it returns the sanitized upstream-unavailable failure for the required account + +#### Scenario: ambiguous transport failure is not replayed + +- **WHEN** a POST transport failure cannot prove that request dispatch was impossible +- **THEN** the service does not use that failure as authorization to retry another proxy endpoint or account + +#### Scenario: empty replacement pool preserves the original failure + +- **GIVEN** a movable request has a confirmed pre-dispatch proxy connection failure +- **AND** no other eligible account can be selected +- **THEN** the client receives the original sanitized upstream-unavailable failure +- **AND** the failure is not replaced with `no_accounts` diff --git a/openspec/changes/retry-account-proxy-connect-failures/tasks.md b/openspec/changes/retry-account-proxy-connect-failures/tasks.md new file mode 100644 index 0000000000..25b68de4b2 --- /dev/null +++ b/openspec/changes/retry-account-proxy-connect-failures/tasks.md @@ -0,0 +1,9 @@ +# Tasks + +- [x] Add sanitized pre-dispatch provenance to account-routed transport errors. +- [x] Permit same-pool fallback for non-idempotent requests only when dispatch is proven impossible. +- [x] Add movable-account failover for raw HTTP/SSE and HTTP bridge startup. +- [x] Preserve native WebSocket failover while applying the same confirmed-route backoff. +- [x] Preserve hard continuity/file pins and the original failure when no replacement exists. +- [x] Add core-client, load-balancer, HTTP/SSE, native WebSocket, and HTTP-bridge regressions. +- [x] Run focused tests, static checks, strict OpenSpec validation, and the relevant broader suites. diff --git a/tests/integration/test_http_responses_bridge.py b/tests/integration/test_http_responses_bridge.py index a31042d85d..33332fa411 100644 --- a/tests/integration/test_http_responses_bridge.py +++ b/tests/integration/test_http_responses_bridge.py @@ -19,6 +19,7 @@ from sqlalchemy import select import app.modules.proxy.service as proxy_module +from app.core.clients.codex import CodexTransportDispatchState from app.core.config.settings import Settings from app.core.utils.request_id import ( reset_request_id, @@ -1476,6 +1477,83 @@ def _make_api_key_data( ) +@pytest.mark.asyncio +async def test_v1_responses_http_bridge_fails_over_confirmed_proxy_connect_before_dispatch( + async_client, + monkeypatch, +): + _install_bridge_settings(monkeypatch, enabled=True) + first_account_id = await _import_account( + async_client, + "acc_http_bridge_proxy_connect_a", + "http-bridge-proxy-connect-a@example.com", + ) + second_account_id = await _import_account( + async_client, + "acc_http_bridge_proxy_connect_b", + "http-bridge-proxy-connect-b@example.com", + ) + first_account = await _get_account(first_account_id) + second_account = await _get_account(second_account_id) + upstream = _FakeBridgeUpstreamWebSocket() + connect_calls: list[str | None] = [] + selection_exclusions: list[set[str]] = [] + handle_stream_error = AsyncMock() + + async def fake_select_account_with_budget(self, deadline, **kwargs): + del self, deadline + excluded = set(cast(set[str], kwargs.get("exclude_account_ids") or set())) + selection_exclusions.append(excluded) + account = second_account if first_account.id in excluded else first_account + return AccountSelection(account=account, error_message=None, error_code=None) + + async def fake_ensure_fresh_with_budget(self, target, *, force=False, timeout_seconds): + del self, force, timeout_seconds + return target + + async def fake_connect_responses_websocket( + headers, + access_token, + account_id_header, + **kwargs, + ): + del headers, access_token, kwargs + connect_calls.append(account_id_header) + if len(connect_calls) == 1: + raise proxy_module.ProxyResponseError( + 502, + proxy_module.openai_error("upstream_unavailable", "sanitized bridge proxy failure"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + return upstream + + monkeypatch.setattr(proxy_module.ProxyService, "_select_account_with_budget", fake_select_account_with_budget) + monkeypatch.setattr(proxy_module.ProxyService, "_ensure_fresh_with_budget", fake_ensure_fresh_with_budget) + monkeypatch.setattr(proxy_module.ProxyService, "_handle_stream_error", handle_stream_error) + monkeypatch.setattr(proxy_module, "connect_responses_websocket", fake_connect_responses_websocket) + + events = await _collect_sse_events( + async_client, + "/v1/responses", + json_body={ + "model": "gpt-5.4", + "instructions": "Return exactly OK.", + "input": "hello", + "prompt_cache_key": "http-bridge-proxy-connect-failover-key", + "stream": True, + }, + ) + + _assert_created_text_delta_completed(events) + assert len(connect_calls) == 2 + assert selection_exclusions == [set(), {first_account.id}] + assert len(upstream.sent_text) == 1 + handle_stream_error.assert_not_awaited() + + @pytest.mark.asyncio async def test_v1_responses_http_bridge_codex_session_uses_extended_idle_ttl(async_client, app_instance, monkeypatch): _install_bridge_settings_with_limits(monkeypatch, enabled=True, codex_idle_ttl_seconds=600.0) diff --git a/tests/integration/test_proxy_responses.py b/tests/integration/test_proxy_responses.py index c64089ef18..0f905936bd 100644 --- a/tests/integration/test_proxy_responses.py +++ b/tests/integration/test_proxy_responses.py @@ -14,6 +14,7 @@ import app.modules.proxy.api as proxy_api_module import app.modules.proxy.service as proxy_module from app.core.auth import generate_unique_account_id +from app.core.clients.codex import CodexTransportDispatchState from app.core.config.settings import Settings from app.core.openai.models import CompactResponsePayload from app.core.types import JsonValue @@ -502,6 +503,58 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, assert captured_account_ids[1] != invalidated_account_id +@pytest.mark.asyncio +async def test_proxy_responses_confirmed_proxy_connect_failure_fails_over_before_dispatch(async_client, monkeypatch): + for raw_account_id, email in ( + ("acc_stream_proxy_connect_a", "stream-proxy-connect-a@example.com"), + ("acc_stream_proxy_connect_b", "stream-proxy-connect-b@example.com"), + ): + auth_json = _make_auth_json(raw_account_id, email) + files = {"auth_json": ("auth.json", json.dumps(auth_json), "application/json")} + response = await async_client.post("/api/accounts/import", files=files) + assert response.status_code == 200 + + captured_account_ids: list[str | None] = [] + failed_account_id: str | None = None + + async def fake_stream(payload, headers, access_token, account_id, base_url=None, raise_for_status=False, **kwargs): + del payload, headers, access_token, base_url, kwargs + nonlocal failed_account_id + assert raise_for_status is True + if failed_account_id is None: + failed_account_id = account_id + captured_account_ids.append(account_id) + if account_id == failed_account_id: + raise proxy_module.ProxyResponseError( + 502, + proxy_module.openai_error("upstream_unavailable", "sanitized account proxy failure"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + yield ( + 'data: {"type":"response.completed","response":{"id":"resp_stream_proxy_failover",' + '"object":"response","status":"completed","usage":{"input_tokens":2,"output_tokens":1}}}\n\n' + ) + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + + async with async_client.stream( + "POST", + "/backend-api/codex/responses", + json={"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True}, + ) as resp: + assert resp.status_code == 200 + lines = [line async for line in resp.aiter_lines() if line] + + event = _extract_first_event(lines) + assert event["type"] == "response.completed" + assert event["response"]["id"] == "resp_stream_proxy_failover" + assert captured_account_ids[0] == failed_account_id + assert captured_account_ids[1] != failed_account_id + + @pytest.mark.asyncio async def test_proxy_responses_compaction_trigger_streams_single_compaction_item(async_client, monkeypatch): email = "compact-trigger@example.com" diff --git a/tests/integration/test_proxy_websocket_responses.py b/tests/integration/test_proxy_websocket_responses.py index e8034285e6..ea1c3e88fd 100644 --- a/tests/integration/test_proxy_websocket_responses.py +++ b/tests/integration/test_proxy_websocket_responses.py @@ -5,6 +5,7 @@ from collections import deque from types import SimpleNamespace from typing import Any, cast +from unittest.mock import AsyncMock import pytest from fastapi.testclient import TestClient @@ -12,6 +13,7 @@ import app.modules.proxy.api as proxy_api_module import app.modules.proxy.service as proxy_module +from app.core.clients.codex import CodexTransportDispatchState from app.core.utils.request_id import get_request_id pytestmark = pytest.mark.integration @@ -164,6 +166,117 @@ def _websocket_settings(**overrides): return SimpleNamespace(**values) +def test_backend_responses_websocket_fails_over_confirmed_proxy_connect_before_dispatch( + app_instance, + monkeypatch, +): + upstream = _FakeUpstreamWebSocket( + [ + _FakeUpstreamMessage( + "text", + text=json.dumps( + { + "type": "response.created", + "response": {"id": "resp_ws_proxy_failover", "status": "in_progress"}, + }, + separators=(",", ":"), + ), + ), + _FakeUpstreamMessage( + "text", + text=json.dumps( + { + "type": "response.completed", + "response": { + "id": "resp_ws_proxy_failover", + "status": "completed", + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + }, + separators=(",", ":"), + ), + ), + ] + ) + accounts = [SimpleNamespace(id="acct_ws_proxy_a"), SimpleNamespace(id="acct_ws_proxy_b")] + selection_exclusions: list[set[str]] = [] + connect_accounts: list[str] = [] + backed_off_accounts: list[str] = [] + handle_connect_error = AsyncMock() + proxy_connect_error = proxy_module.ProxyResponseError( + 502, + proxy_module.openai_error("upstream_unavailable", "sanitized websocket proxy failure"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + + class _FakeSettingsCache: + async def get(self): + return _websocket_settings() + + async def allow_firewall(_websocket): + return None + + async def allow_proxy_api_key(_authorization: str | None, *, request: object | None = None): + return None + + async def fake_select_websocket_connect_account(self, deadline, **kwargs): + del self, deadline + excluded = set(cast(set[str], kwargs["exclude_account_ids"])) + selection_exclusions.append(excluded) + return accounts[1] if accounts[0].id in excluded else accounts[0] + + async def fake_try_open_websocket_connect_attempt(self, account, headers, **kwargs): + del self, headers, kwargs + connect_accounts.append(account.id) + if account.id == accounts[0].id: + raise proxy_connect_error + return account, upstream + + async def fake_record_error_backoff(self, account): + del self + backed_off_accounts.append(account.id) + + monkeypatch.setattr(proxy_api_module, "_websocket_firewall_denial_response", allow_firewall) + monkeypatch.setattr(proxy_api_module, "validate_proxy_api_key_authorization", allow_proxy_api_key) + monkeypatch.setattr(proxy_module, "get_settings_cache", lambda: _FakeSettingsCache()) + monkeypatch.setattr( + proxy_module.ProxyService, + "_select_websocket_connect_account", + fake_select_websocket_connect_account, + ) + monkeypatch.setattr( + proxy_module.ProxyService, + "_try_open_websocket_connect_attempt", + fake_try_open_websocket_connect_attempt, + ) + monkeypatch.setattr(proxy_module.ProxyService, "_handle_websocket_connect_error", handle_connect_error) + monkeypatch.setattr(proxy_module.LoadBalancer, "record_error_backoff", fake_record_error_backoff) + + request_payload = { + "type": "response.create", + "model": "gpt-5.4", + "instructions": "", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "retry safely"}]}], + "stream": True, + } + + with TestClient(app_instance) as client: + with client.websocket_connect("/backend-api/codex/responses") as websocket: + websocket.send_text(json.dumps(request_payload)) + first_event = json.loads(websocket.receive_text()) + second_event = json.loads(websocket.receive_text()) + + assert first_event["type"] == "response.created" + assert second_event["type"] == "response.completed" + assert connect_accounts == [accounts[0].id, accounts[1].id] + assert selection_exclusions == [set(), {accounts[0].id}] + assert backed_off_accounts == [accounts[0].id] + handle_connect_error.assert_not_awaited() + + def test_backend_responses_websocket_session_ended_auth_failure_fails_over_before_visible_output( app_instance, monkeypatch, diff --git a/tests/unit/test_codex_client.py b/tests/unit/test_codex_client.py index 4a05ac3e21..b1f394ad76 100644 --- a/tests/unit/test_codex_client.py +++ b/tests/unit/test_codex_client.py @@ -5,9 +5,15 @@ import aiohttp import pytest +from aiohttp.client_reqrep import ConnectionKey from python_socks import ProxyType -from app.core.clients.codex import CodexClient, require_route_or_direct_egress_opt_in +from app.core.clients.codex import ( + CodexClient, + CodexTransportDispatchState, + CodexTransportError, + require_route_or_direct_egress_opt_in, +) from app.core.upstream_proxy import ResolvedProxyEndpoint, ResolvedUpstreamRoute pytestmark = pytest.mark.unit @@ -42,6 +48,31 @@ async def ws_connect(self, url: str, **kwargs: Any) -> object: return object() +def _proxy_connect_error() -> aiohttp.ClientProxyConnectionError: + key = ConnectionKey( + host="proxy.test", + port=8080, + is_ssl=False, + ssl=False, + proxy=None, + proxy_auth=None, + proxy_headers_hash=None, + ) + return aiohttp.ClientProxyConnectionError(key, OSError("proxy credentials must stay private")) + + +class _ProxyConnectFailureSession(_Session): + def __init__(self, *, fail_all: bool = False) -> None: + super().__init__() + self.fail_all_proxy_connects = fail_all + + async def request(self, method: str, url: str, **kwargs: Any) -> _Response: + self.calls.append({"method": method, "url": url, **kwargs}) + if self.fail_all_proxy_connects or len(self.calls) == 1: + raise _proxy_connect_error() + return _Response(headers={"content-type": "application/json"}) + + class _HandshakeFailure(Exception): status = 426 @@ -211,6 +242,49 @@ async def test_non_idempotent_request_failure_does_not_fallback(route: ResolvedU assert session.calls[0]["proxy"] == "http://u:p@proxy.test:8080" +@pytest.mark.asyncio +async def test_non_idempotent_pre_dispatch_proxy_failure_uses_same_pool_fallback( + route: ResolvedUpstreamRoute, +) -> None: + session = _ProxyConnectFailureSession() + client = CodexClient(session) + + result = await client.request_with_route_metadata( + "POST", + "https://upstream.test", + route=route, + buffer_response=False, + json={"x": 1}, + ) + + assert result.fallback_used is True + assert result.route.endpoint_id == "ep_2" + assert [call["proxy"] for call in session.calls] == [ + "http://u:p@proxy.test:8080", + "http://proxy-two.test:8081", + ] + + +@pytest.mark.asyncio +async def test_exhausted_proxy_connect_failures_preserve_pre_dispatch_provenance( + route: ResolvedUpstreamRoute, +) -> None: + client = CodexClient(_ProxyConnectFailureSession(fail_all=True)) + + with pytest.raises(CodexTransportError) as exc_info: + await client.request_with_route_metadata( + "POST", + "https://upstream.test", + route=route, + buffer_response=False, + json={"x": 1}, + ) + + assert exc_info.value.dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED + assert exc_info.value.exception_type == "ClientProxyConnectionError" + assert "credentials must stay private" not in str(exc_info.value) + + @pytest.mark.asyncio async def test_transport_errors_do_not_expose_proxy_credentials(route: ResolvedUpstreamRoute) -> None: client = CodexClient(_Session(fail_all=True)) diff --git a/tests/unit/test_codex_upstream_paths.py b/tests/unit/test_codex_upstream_paths.py index 725b4f1a0e..d34690d6ad 100644 --- a/tests/unit/test_codex_upstream_paths.py +++ b/tests/unit/test_codex_upstream_paths.py @@ -6,7 +6,7 @@ import pytest import app.core.clients.proxy as proxy_module -from app.core.clients.codex import CodexTransportError, CodexWebSocketResult +from app.core.clients.codex import CodexTransportDispatchState, CodexTransportError, CodexWebSocketResult from app.core.clients.files import create_file, finalize_file from app.core.clients.proxy import ( ProxyResponseError, @@ -113,6 +113,23 @@ async def request(self, method: str, url: str, *, route: ResolvedUpstreamRoute, raise CodexTransportError("Codex upstream request failed via proxy endpoint ep_1: OSError") +class _PreDispatchTransportErrorCodexClient: + async def request_with_route_metadata( + self, + method: str, + url: str, + *, + route: ResolvedUpstreamRoute, + **kwargs: Any, + ) -> object: + del method, url, route, kwargs + raise CodexTransportError( + "Codex upstream request failed via proxy endpoint ep_1: ClientProxyConnectionError", + dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + exception_type="ClientProxyConnectionError", + ) + + class _FakeCodexWebSocket: def __init__(self, *, fail_receive: bool = False, fail_send: bool = False) -> None: self.sent: list[str | bytes] = [] @@ -591,6 +608,58 @@ async def test_stream_responses_routed_transport_errors_are_unavailable(route: R assert "ep_1" in combined +@pytest.mark.asyncio +async def test_stream_responses_propagates_confirmed_pre_dispatch_failure_for_status_retry( + route: ResolvedUpstreamRoute, +) -> None: + payload = ResponsesRequest(model="gpt-5.2", instructions="Reply.", input="hello", stream=True) + + with pytest.raises(ProxyResponseError) as exc_info: + async for _ in stream_responses( + payload, + {"user-agent": "codex"}, + "access", + "chatgpt_account", + session=cast(Any, object()), + upstream_stream_transport_override="http", + route=route, + codex_client=cast(Any, _PreDispatchTransportErrorCodexClient()), + raise_for_status=True, + ): + pass + + assert exc_info.value.status_code == 502 + assert exc_info.value.upstream_dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED + assert exc_info.value.failure_detail == "proxy_connect_pre_dispatch" + assert exc_info.value.failure_exception_type == "ClientProxyConnectionError" + + +@pytest.mark.asyncio +async def test_stream_responses_propagates_ambiguous_dispatch_failure_without_retry_authorization( + route: ResolvedUpstreamRoute, +) -> None: + payload = ResponsesRequest(model="gpt-5.2", instructions="Reply.", input="hello", stream=True) + + with pytest.raises(ProxyResponseError) as exc_info: + async for _ in stream_responses( + payload, + {"user-agent": "codex"}, + "access", + "chatgpt_account", + session=cast(Any, object()), + upstream_stream_transport_override="http", + route=route, + codex_client=cast(Any, _TransportErrorCodexClient()), + raise_for_status=True, + ): + pass + + assert exc_info.value.status_code == 502 + assert exc_info.value.upstream_dispatch_state is CodexTransportDispatchState.UNKNOWN + assert exc_info.value.retryable_same_contract is False + assert exc_info.value.failure_detail == "transport_error" + + @pytest.mark.asyncio async def test_responses_websocket_uses_codex_client_when_route_is_resolved(route: ResolvedUpstreamRoute) -> None: client = _WsCodexClient() diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index ddc52eb4a4..c6416db855 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -277,6 +277,24 @@ async def test_record_error_updates_are_atomic_with_per_account_lock() -> None: assert runtime.last_error_at is not None +@pytest.mark.asyncio +async def test_record_error_backoff_enters_floor_without_adding_full_threshold() -> None: + account = _make_account("acc-error-backoff") + accounts_repo = _StubAccountsRepository([account]) + usage_repo = _StubUsageRepository(primary={}, secondary={}) + balancer = LoadBalancer(lambda: _repo_factory(accounts_repo, usage_repo)) + + await balancer.record_error_backoff(account) + + runtime = balancer._runtime[account.id] + assert runtime.error_count == 3 + assert runtime.last_error_at is not None + + await balancer.record_error_backoff(account) + + assert runtime.error_count == 4 + + @pytest.mark.asyncio async def test_stale_reclaim_keeps_active_stream_lease_within_stream_budget( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index 76439e7365..854afc47ea 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -18,6 +18,7 @@ from fastapi import WebSocket from app.core.auth.refresh import RefreshError +from app.core.clients.codex import CodexTransportDispatchState from app.core.clients.proxy import CODEX_RESPONSES_LITE_WEBSOCKET_METADATA_KEY, ProxyResponseError from app.core.clients.proxy_websocket import ( CodexResponsesWebSocket, @@ -61,6 +62,17 @@ def _make_app_settings(*, bridge_enabled: bool = True, **overrides: Any) -> Sett return Settings(http_responses_session_bridge_enabled=bridge_enabled, **overrides) +def _pre_dispatch_proxy_error(message: str = "sanitized proxy connect failure") -> ProxyResponseError: + return ProxyResponseError( + 502, + openai_error("upstream_unavailable", message), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + + def _make_bridge_session( *, key: proxy_service._HTTPBridgeSessionKey | None = None, @@ -3518,6 +3530,173 @@ async def select_account(_deadline: float, **kwargs: object) -> proxy_service.Ac assert selection_kwargs[0]["prefer_earlier_reset_window"] == "primary" +@pytest.mark.asyncio +async def test_create_http_bridge_session_fails_over_confirmed_proxy_connect_after_lease_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + account_a = cast(Any, SimpleNamespace(id="acc-proxy-a", status=AccountStatus.ACTIVE, plan_type="plus")) + account_b = cast(Any, SimpleNamespace(id="acc-proxy-b", status=AccountStatus.ACTIVE, plan_type="plus")) + lease_a = proxy_service.AccountLease("lease-bridge-a", account_a.id, "stream", time.monotonic()) + lease_b = proxy_service.AccountLease("lease-bridge-b", account_b.id, "stream", time.monotonic()) + selections: list[set[str]] = [] + released_leases: list[proxy_service.AccountLease] = [] + upstream = cast( + UpstreamResponsesWebSocket, + SimpleNamespace(response_header=lambda _name: None, close=AsyncMock()), + ) + + async def select_account(_deadline: float, **kwargs: object) -> proxy_service.AccountSelection: + excluded = set(cast(set[str], kwargs["exclude_account_ids"])) + selections.append(excluded) + if not excluded: + return proxy_service.AccountSelection(account=account_a, error_message=None, lease=lease_a) + return proxy_service.AccountSelection(account=account_b, error_message=None, lease=lease_b) + + async def release_account_lease(lease: proxy_service.AccountLease | None) -> None: + if lease is not None: + released_leases.append(lease) + + async def record_error_backoff(account: object) -> None: + assert account is account_a + assert lease_a in released_leases + + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, "get_settings_cache", lambda: SimpleNamespace(get=AsyncMock(return_value=settings)) + ) + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(side_effect=[account_a, account_b])) + monkeypatch.setattr( + service, + "_open_upstream_websocket_with_budget", + AsyncMock(side_effect=[_pre_dispatch_proxy_error(), upstream]), + ) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_account_lease) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service, "_relay_http_bridge_upstream_messages", AsyncMock()) + + session = await service._create_http_bridge_session( + proxy_service._HTTPBridgeSessionKey("session_header", "sid-proxy-failover", None), + headers={}, + affinity=proxy_service._AffinityPolicy(key="sid-proxy-failover"), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + ) + + assert session.account is account_b + assert selections == [set(), {account_a.id}] + assert lease_a in released_leases + + +@pytest.mark.asyncio +async def test_create_http_bridge_session_confirmed_proxy_failure_keeps_hard_owner_pinned( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + account = cast(Any, SimpleNamespace(id="acc-proxy-owner", status=AccountStatus.ACTIVE, plan_type="plus")) + lease = proxy_service.AccountLease("lease-bridge-owner", account.id, "stream", time.monotonic()) + select_account = AsyncMock( + return_value=proxy_service.AccountSelection(account=account, error_message=None, lease=lease) + ) + release_account_lease = AsyncMock() + record_error_backoff = AsyncMock() + original_error = _pre_dispatch_proxy_error("owner proxy unavailable") + + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, "get_settings_cache", lambda: SimpleNamespace(get=AsyncMock(return_value=settings)) + ) + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(return_value=account)) + monkeypatch.setattr( + service, + "_open_upstream_websocket_with_budget", + AsyncMock(side_effect=original_error), + ) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_account_lease) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + + with pytest.raises(ProxyResponseError) as exc_info: + await service._create_http_bridge_session( + proxy_service._HTTPBridgeSessionKey("turn_state_header", "turn-owner", None, strength="hard"), + headers={"x-codex-turn-state": "turn-owner"}, + affinity=proxy_service._AffinityPolicy(key="turn-owner"), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + preferred_account_id=account.id, + require_preferred_account=True, + fallback_on_preferred_account_unavailable=False, + ) + + assert exc_info.value is original_error + select_account.assert_awaited_once() + release_account_lease.assert_awaited_once_with(lease) + record_error_backoff.assert_awaited_once_with(account) + + +@pytest.mark.asyncio +async def test_create_http_bridge_session_preserves_proxy_failure_when_no_replacement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + account = cast(Any, SimpleNamespace(id="acc-proxy-only", status=AccountStatus.ACTIVE, plan_type="plus")) + lease = proxy_service.AccountLease("lease-bridge-only", account.id, "stream", time.monotonic()) + selections: list[set[str]] = [] + original_error = _pre_dispatch_proxy_error("original bridge proxy failure") + + async def select_account(_deadline: float, **kwargs: object) -> proxy_service.AccountSelection: + excluded = set(cast(set[str], kwargs["exclude_account_ids"])) + selections.append(excluded) + if not excluded: + return proxy_service.AccountSelection(account=account, error_message=None, lease=lease) + return proxy_service.AccountSelection( + account=None, + error_message="No active accounts available", + error_code="no_accounts", + ) + + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, "get_settings_cache", lambda: SimpleNamespace(get=AsyncMock(return_value=settings)) + ) + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(return_value=account)) + monkeypatch.setattr(service, "_open_upstream_websocket_with_budget", AsyncMock(side_effect=original_error)) + monkeypatch.setattr(service._load_balancer, "release_account_lease", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", AsyncMock()) + + with pytest.raises(ProxyResponseError) as exc_info: + await service._create_http_bridge_session( + proxy_service._HTTPBridgeSessionKey("session_header", "sid-proxy-only", None), + headers={}, + affinity=proxy_service._AffinityPolicy(key="sid-proxy-only"), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + ) + + assert exc_info.value is original_error + assert selections == [set(), {account.id}] + + @pytest.mark.asyncio async def test_reconnect_http_bridge_session_passes_dashboard_reset_window_to_selection( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index d3d7791876..0abc63042c 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -24,6 +24,7 @@ import app.core.clients.proxy as proxy_module import app.modules.proxy.load_balancer as load_balancer_module +from app.core.clients.codex import CodexTransportDispatchState from app.core.clients.proxy import _build_upstream_headers, filter_inbound_headers from app.core.config.settings import Settings from app.core.crypto import TokenEncryptor @@ -10876,6 +10877,315 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, record_success.assert_awaited_once_with(account_b) +@pytest.mark.asyncio +async def test_stream_responses_confirmed_proxy_connect_failure_fails_over_once_after_lease_release(monkeypatch): + settings = _make_proxy_settings(log_proxy_service_tier_trace=False) + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + account_a = _make_account("acc_proxy_connect_stream_a") + account_b = _make_account("acc_proxy_connect_stream_b") + api_key = ApiKeyData( + id="key_proxy_connect_failover", + name="proxy-connect-failover", + key_prefix="sk-proxy", + allowed_models=None, + enforced_model=None, + enforced_reasoning_effort=None, + enforced_service_tier=None, + expires_at=None, + is_active=True, + created_at=utcnow(), + last_used_at=None, + ) + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv_proxy_connect_failover", + key_id=api_key.id, + model="gpt-5.1", + ) + lease_a = AccountLease("lease-proxy-a", account_a.id, "stream", time.monotonic()) + lease_b = AccountLease("lease-proxy-b", account_b.id, "stream", time.monotonic()) + response_lease_a = AccountLease("lease-proxy-response-a", account_a.id, "response_create", time.monotonic()) + response_lease_b = AccountLease("lease-proxy-response-b", account_b.id, "response_create", time.monotonic()) + released_leases: list[AccountLease] = [] + stream_accounts: list[str | None] = [] + seen_excluded_account_ids: list[set[str]] = [] + record_error = AsyncMock() + record_success = AsyncMock() + settle_usage = AsyncMock(return_value=True) + release_unsettled_usage = AsyncMock() + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + + async def select_account(**kwargs: object) -> AccountSelection: + excluded = set(cast(set[str] | None, kwargs.get("exclude_account_ids")) or set()) + seen_excluded_account_ids.append(excluded) + if not excluded: + return AccountSelection(account=account_a, error_message=None, lease=lease_a) + return AccountSelection(account=account_b, error_message=None, lease=lease_b) + + async def release_account_lease(lease: AccountLease | None) -> None: + if lease is not None: + released_leases.append(lease) + + async def record_error_backoff(account: Account) -> None: + assert account == account_a + assert lease_a in released_leases + assert response_lease_a in released_leases + + async def fake_stream( + payload: ResponsesRequest, + headers: Mapping[str, str], + access_token: str, + account_id: str | None, + base_url: str | None = None, + raise_for_status: bool = False, + **_kwargs: object, + ) -> AsyncIterator[str]: + del payload, headers, access_token, base_url + assert raise_for_status is True + stream_accounts.append(account_id) + if account_id == account_a.chatgpt_account_id: + raise proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "sanitized proxy connect failure"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + yield ( + 'data: {"type":"response.completed","response":{"id":"resp_proxy_connect_ok",' + '"usage":{"input_tokens":1,"output_tokens":2}}}\n\n' + ) + + monkeypatch.setattr(service._load_balancer, "select_account", select_account) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_account_lease) + monkeypatch.setattr(service._load_balancer, "record_error", record_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service._load_balancer, "record_success", record_success) + monkeypatch.setattr(service, "_settle_stream_api_key_usage", settle_usage) + monkeypatch.setattr(service, "_release_unsettled_stream_api_key_usage", release_unsettled_usage) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(side_effect=[account_a, account_b])) + monkeypatch.setattr( + service, + "_acquire_account_response_create_lease_or_overload", + AsyncMock(side_effect=[response_lease_a, response_lease_b]), + ) + monkeypatch.setattr(proxy_service, "core_stream_responses", fake_stream) + + payload = ResponsesRequest.model_validate({"model": "gpt-5.1", "instructions": "hi", "input": [], "stream": True}) + chunks = [ + chunk + async for chunk in service._stream_with_retry( + payload, + {"session_id": "sid-proxy-connect"}, + codex_session_affinity=False, + propagate_http_errors=False, + openai_cache_affinity=False, + api_key=api_key, + api_key_reservation=reservation, + suppress_text_done_events=False, + request_transport="http", + ) + ] + + event = json.loads(chunks[0].split("data: ", 1)[1]) + assert event["type"] == "response.completed" + assert event["response"]["id"] == "resp_proxy_connect_ok" + assert stream_accounts == [account_a.chatgpt_account_id, account_b.chatgpt_account_id] + assert seen_excluded_account_ids == [set(), {account_a.id}] + assert request_logs.calls[0]["failure_detail"] == "proxy_connect_pre_dispatch" + record_error.assert_not_awaited() + record_success.assert_awaited_once_with(account_b) + settle_usage.assert_awaited_once() + release_unsettled_usage.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stream_responses_confirmed_proxy_connect_failure_preserves_original_when_no_replacement(monkeypatch): + settings = _make_proxy_settings(log_proxy_service_tier_trace=False) + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + account = _make_account("acc_proxy_connect_stream_only") + lease = AccountLease("lease-proxy-only", account.id, "stream", time.monotonic()) + released_leases: list[AccountLease] = [] + selections: list[set[str]] = [] + record_error = AsyncMock() + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + + async def select_account(**kwargs: object) -> AccountSelection: + excluded = set(cast(set[str] | None, kwargs.get("exclude_account_ids")) or set()) + selections.append(excluded) + if not excluded: + return AccountSelection(account=account, error_message=None, lease=lease) + return AccountSelection(account=None, error_message="No active accounts available", error_code="no_accounts") + + async def release_account_lease(candidate: AccountLease | None) -> None: + if candidate is not None: + released_leases.append(candidate) + + async def record_error_backoff(failed_account: Account) -> None: + assert failed_account == account + assert lease in released_leases + + async def fake_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "original sanitized proxy failure"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + yield "" + + monkeypatch.setattr(service._load_balancer, "select_account", select_account) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_account_lease) + monkeypatch.setattr(service._load_balancer, "record_error", record_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(return_value=account)) + monkeypatch.setattr( + service, + "_acquire_account_response_create_lease_or_overload", + AsyncMock(return_value=None), + ) + monkeypatch.setattr(proxy_service, "core_stream_responses", fake_stream) + + payload = ResponsesRequest.model_validate({"model": "gpt-5.1", "instructions": "hi", "input": [], "stream": True}) + chunks = [chunk async for chunk in service.stream_responses(payload, {"session_id": "sid-proxy-only"})] + + event = json.loads(chunks[0].split("data: ", 1)[1]) + assert event["type"] == "response.failed" + assert event["response"]["error"]["code"] == "upstream_unavailable" + assert event["response"]["error"]["message"] == "original sanitized proxy failure" + assert selections == [set(), {account.id}] + record_error.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stream_responses_confirmed_proxy_connect_failure_keeps_file_owner_pinned(monkeypatch): + settings = _make_proxy_settings(log_proxy_service_tier_trace=False) + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + owner = _make_account("acc_proxy_connect_file_owner") + alternate = _make_account("acc_proxy_connect_file_alternate") + lease = AccountLease("lease-proxy-file-owner", owner.id, "stream", time.monotonic()) + selections: list[dict[str, object]] = [] + original_error = proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "file owner proxy unavailable"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + + async def select_account(**kwargs: object) -> AccountSelection: + selections.append(dict(kwargs)) + account = owner if len(selections) == 1 else alternate + return AccountSelection(account=account, error_message=None, lease=lease if account is owner else None) + + monkeypatch.setattr(service._load_balancer, "select_account", select_account) + monkeypatch.setattr(service._load_balancer, "release_account_lease", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_error", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", AsyncMock()) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(return_value=owner)) + monkeypatch.setattr( + service, + "_acquire_account_response_create_lease_or_overload", + AsyncMock(return_value=None), + ) + + async def fake_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise original_error + yield "" + + monkeypatch.setattr(proxy_service, "core_stream_responses", fake_stream) + payload = ResponsesRequest.model_validate({"model": "gpt-5.1", "instructions": "hi", "input": [], "stream": True}) + + chunks = [ + chunk + async for chunk in service._stream_with_retry( + payload, + {"session_id": "sid-file-owner"}, + codex_session_affinity=False, + propagate_http_errors=False, + openai_cache_affinity=False, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + request_transport="http", + rewritten_file_account_id=owner.id, + ) + ] + + event = json.loads(chunks[0].split("data: ", 1)[1]) + assert event["type"] == "response.failed" + assert event["response"]["error"]["message"] == "file owner proxy unavailable" + assert len(selections) == 1 + + +@pytest.mark.asyncio +async def test_stream_responses_ambiguous_dispatch_transport_failure_is_not_replayed(monkeypatch): + settings = _make_proxy_settings(log_proxy_service_tier_trace=False) + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + first_account = _make_account("acc_ambiguous_dispatch_a") + alternate = _make_account("acc_ambiguous_dispatch_b") + lease = AccountLease("lease-ambiguous-dispatch", first_account.id, "stream", time.monotonic()) + select_account = AsyncMock( + side_effect=[ + AccountSelection(account=first_account, error_message=None, lease=lease), + AccountSelection(account=alternate, error_message=None), + ] + ) + record_error = AsyncMock() + record_error_backoff = AsyncMock() + ambiguous_error = proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "dispatch outcome is unknown"), + failure_phase="upstream", + failure_detail="transport_error", + failure_exception_type="ClientConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.UNKNOWN, + ) + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(service._load_balancer, "select_account", select_account) + monkeypatch.setattr(service._load_balancer, "release_account_lease", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_error", record_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(return_value=first_account)) + monkeypatch.setattr( + service, + "_acquire_account_response_create_lease_or_overload", + AsyncMock(return_value=None), + ) + + async def fake_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise ambiguous_error + yield "" + + monkeypatch.setattr(proxy_service, "core_stream_responses", fake_stream) + payload = ResponsesRequest.model_validate({"model": "gpt-5.1", "instructions": "hi", "input": [], "stream": True}) + + chunks = [chunk async for chunk in service.stream_responses(payload, {"session_id": "sid-ambiguous"})] + + event = json.loads(chunks[0].split("data: ", 1)[1]) + assert event["type"] == "response.failed" + assert event["response"]["error"]["message"] == "dispatch outcome is unknown" + select_account.assert_awaited_once() + record_error.assert_awaited_once_with(first_account) + record_error_backoff.assert_not_awaited() + + @pytest.mark.asyncio async def test_stream_responses_first_event_upstream_unavailable_fails_over(monkeypatch): settings = _make_proxy_settings(log_proxy_service_tier_trace=False) @@ -13408,6 +13718,156 @@ async def test_connect_proxy_websocket_fails_over_after_upstream_connect_timeout release_reservation.assert_not_awaited() +@pytest.mark.asyncio +async def test_connect_proxy_websocket_confirmed_proxy_failure_backs_off_after_lease_release(monkeypatch): + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + base_settings = _make_proxy_settings(log_proxy_service_tier_trace=False) + base_settings.deterministic_failover_enabled = False + first_account = _make_account("acc_ws_proxy_connect") + second_account = _make_account("acc_ws_proxy_ok") + first_lease = AccountLease("lease-ws-proxy-a", first_account.id, "stream", time.monotonic()) + second_lease = AccountLease("lease-ws-proxy-b", second_account.id, "stream", time.monotonic()) + released_leases: list[AccountLease] = [] + upstream = SimpleNamespace() + record_error = AsyncMock() + + async def release_account_lease(lease: AccountLease | None) -> None: + if lease is not None: + released_leases.append(lease) + + async def record_error_backoff(account: Account) -> None: + assert account == first_account + assert first_lease in released_leases + + proxy_connect_error = proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "sanitized proxy connect failure"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + monkeypatch.setattr(proxy_service, "get_settings", lambda: base_settings) + monkeypatch.setattr( + service._load_balancer, + "select_account", + AsyncMock( + side_effect=[ + AccountSelection(account=first_account, error_message=None, lease=first_lease), + AccountSelection(account=second_account, error_message=None, lease=second_lease), + ] + ), + ) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_account_lease) + monkeypatch.setattr(service._load_balancer, "record_error", record_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(side_effect=[first_account, second_account])) + monkeypatch.setattr( + service, + "_open_upstream_websocket_with_budget", + AsyncMock(side_effect=[proxy_connect_error, upstream]), + ) + + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_proxy_connect", + model="gpt-5.1", + service_tier="fast", + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + ) + websocket_send = AsyncMock() + + selected_account, selected_upstream = await service._connect_proxy_websocket( + {}, + sticky_key=None, + sticky_kind=None, + prefer_earlier_reset=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + model="gpt-5.1", + request_state=request_state, + api_key=None, + client_send_lock=anyio.Lock(), + websocket=cast(WebSocket, SimpleNamespace(send_text=websocket_send)), + ) + + assert selected_account == second_account + assert selected_upstream == upstream + assert first_lease in released_leases + record_error.assert_not_awaited() + websocket_send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_connect_proxy_websocket_confirmed_proxy_failure_keeps_previous_response_owner_pinned(monkeypatch): + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + owner = _make_account("acc_ws_proxy_owner") + lease = AccountLease("lease-ws-proxy-owner", owner.id, "stream", time.monotonic()) + released_leases: list[AccountLease] = [] + select_account = AsyncMock(return_value=AccountSelection(account=owner, error_message=None, lease=lease)) + + async def release_account_lease(candidate: AccountLease | None) -> None: + if candidate is not None: + released_leases.append(candidate) + + async def record_error_backoff(account: Account) -> None: + assert account == owner + assert lease in released_leases + + proxy_connect_error = proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "owner proxy unavailable"), + failure_phase="connect", + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + upstream_dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + ) + monkeypatch.setattr(service._load_balancer, "select_account", select_account) + monkeypatch.setattr(service._load_balancer, "release_account_lease", release_account_lease) + monkeypatch.setattr(service._load_balancer, "record_error", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(return_value=owner)) + monkeypatch.setattr(service, "_open_upstream_websocket_with_budget", AsyncMock(side_effect=proxy_connect_error)) + + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_proxy_owner", + model="gpt-5.1", + service_tier="fast", + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + previous_response_id="resp_proxy_owner", + preferred_account_id=owner.id, + ) + websocket_send = AsyncMock() + + selected_account, selected_upstream = await service._connect_proxy_websocket( + {}, + sticky_key=None, + sticky_kind=None, + prefer_earlier_reset=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + model="gpt-5.1", + request_state=request_state, + api_key=None, + client_send_lock=anyio.Lock(), + websocket=cast(WebSocket, SimpleNamespace(send_text=websocket_send)), + ) + + assert selected_account is None + assert selected_upstream is None + select_account.assert_awaited_once() + send_args = websocket_send.await_args + assert send_args is not None + sent_payload = json.loads(send_args.args[0]) + assert sent_payload["status"] == 502 + assert sent_payload["error"]["message"] == "owner proxy unavailable" + + @pytest.mark.asyncio async def test_connect_proxy_websocket_surfaces_connect_timeout_when_no_failover_account(monkeypatch): request_logs = _RequestLogsRecorder() diff --git a/tests/unit/test_proxy_websocket_client.py b/tests/unit/test_proxy_websocket_client.py index edecf2cd64..ffa5e7547e 100644 --- a/tests/unit/test_proxy_websocket_client.py +++ b/tests/unit/test_proxy_websocket_client.py @@ -13,7 +13,7 @@ from websockets.http11 import Response import app.core.clients.proxy_websocket as proxy_websocket_module -from app.core.clients.codex import CodexTransportError, CodexWebSocketResult +from app.core.clients.codex import CodexTransportDispatchState, CodexTransportError, CodexWebSocketResult from app.core.clients.proxy import ProxyResponseError from app.core.clients.proxy_websocket import connect_responses_websocket from app.core.upstream_proxy import ResolvedProxyEndpoint, ResolvedUpstreamRoute @@ -174,7 +174,11 @@ async def open_ws_with_route_metadata( **kwargs: object, ) -> CodexWebSocketResult: del url, route, kwargs - raise CodexTransportError("Codex upstream websocket failed via proxy endpoint ep_1: OSError") + raise CodexTransportError( + "Codex upstream websocket failed via proxy endpoint ep_1: ClientProxyConnectionError", + dispatch_state=CodexTransportDispatchState.NOT_DISPATCHED, + exception_type="ClientProxyConnectionError", + ) async def close(self) -> None: self.closed = True @@ -310,6 +314,9 @@ async def test_connect_responses_websocket_routed_transport_error_maps_proxy_err assert exc_info.value.status_code == 502 assert _proxy_error_code(exc_info.value) == "upstream_unavailable" assert "ep_1" in (_proxy_error_message(exc_info.value) or "") + assert exc_info.value.upstream_dispatch_state is CodexTransportDispatchState.NOT_DISPATCHED + assert exc_info.value.failure_detail == "proxy_connect_pre_dispatch" + assert exc_info.value.failure_exception_type == "ClientProxyConnectionError" @pytest.mark.asyncio