From 4ec8f2c6e8b6a3767814791183d1eb09e43d37ad Mon Sep 17 00:00:00 2001 From: luawl <252236154+luawl@users.noreply.github.com> Date: Tue, 14 Jul 2026 18:19:59 -0400 Subject: [PATCH 1/4] fix(proxy): fail over dead account proxy routes before upstream dispatch An account-bound upstream proxy can stop accepting connections while the account stays administratively active. Movable Responses requests selected onto that account then failed with a terminal sanitized 502 (or a bridge startup error) instead of failing over, and client retries could select the same dead route again (#1314). Rebuilt from PR #1322 onto current main, which already carries the sanitized pre-dispatch provenance (`retryable_same_contract` + `failure_phase == "connect"`) this change originally introduced as a dispatch-state enum: - `is_confirmed_pre_dispatch_transport_error` is the single predicate that authorizes cross-account replay; host-wide network loss keeps its account-neutral process recovery path and TLS verification failures stay non-replayable. - A confirmed pre-dispatch connect failure may try the next endpoint in the same resolved proxy pool even for a non-idempotent POST. - Raw HTTP/SSE streaming, native Responses WebSocket connects, and HTTP bridge session startup release the failed account's stream lease, record the bounded transient backoff floor (`record_error_backoff`, shared `ERROR_BACKOFF_THRESHOLD`), exclude the account, and retry another eligible account within the existing attempt and deadline budgets. - Hard previous-response/turn-state/file/single-account ownership fails closed on the original sanitized failure without crossing accounts, and selection exhaustion preserves that failure instead of generating `no_accounts`. - The API-key reservation stays request-scoped and singular across internal failover; ambiguous POST dispatch outcomes are surfaced without replay. Fixes #1314 Supersedes #1322 Co-Authored-By: Claude Fable 5 --- app/core/balancer/__init__.py | 2 + app/core/balancer/logic.py | 7 +- app/core/clients/codex.py | 15 +- app/core/clients/proxy.py | 18 ++ app/core/clients/proxy_websocket.py | 15 +- .../proxy/_service/http_bridge/helpers.py | 73 +++++ .../proxy/_service/http_bridge/mixin.py | 123 +++----- .../_service/http_bridge/proxy_failover.py | 47 +++ app/modules/proxy/_service/streaming/retry.py | 114 ++++++- app/modules/proxy/_service/websocket/mixin.py | 29 +- app/modules/proxy/load_balancer.py | 15 +- .../design.md | 71 +++++ .../proposal.md | 46 +++ .../specs/upstream-proxy-routing/spec.md | 48 +++ .../tasks.md | 9 + .../integration/test_http_responses_bridge.py | 84 +++++ tests/integration/test_proxy_responses.py | 52 ++++ .../test_proxy_websocket_responses.py | 113 +++++++ tests/unit/test_codex_client.py | 61 ++++ tests/unit/test_codex_upstream_paths.py | 60 ++++ tests/unit/test_load_balancer_concurrency.py | 19 ++ tests/unit/test_proxy_http_bridge.py | 181 +++++++++++ tests/unit/test_proxy_utils.py | 292 ++++++++++++++++++ tests/unit/test_proxy_websocket_client.py | 54 +++- 24 files changed, 1449 insertions(+), 99 deletions(-) create mode 100644 app/modules/proxy/_service/http_bridge/proxy_failover.py create mode 100644 openspec/changes/retry-account-proxy-connect-failures/design.md create mode 100644 openspec/changes/retry-account-proxy-connect-failures/proposal.md create mode 100644 openspec/changes/retry-account-proxy-connect-failures/specs/upstream-proxy-routing/spec.md create mode 100644 openspec/changes/retry-account-proxy-connect-failures/tasks.md diff --git a/app/core/balancer/__init__.py b/app/core/balancer/__init__.py index 2e29cad9db..c8d4dc525b 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, @@ -36,6 +37,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 a3caf85afa..96d9dc4a8f 100644 --- a/app/core/balancer/logic.py +++ b/app/core/balancer/logic.py @@ -91,6 +91,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" @@ -481,8 +482,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 @@ -519,7 +520,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 efe34dc805..9a60fb92e4 100644 --- a/app/core/clients/codex.py +++ b/app/core/clients/codex.py @@ -212,11 +212,20 @@ async def request_with_route_metadata( retryable_same_contract=False, ) from None return CodexRequestResult(response, candidate, index > 0) - except CodexTransportError: - if index == len(endpoints) - 1 or not allow_fallback: + except CodexTransportError as exc: + # A confirmed pre-dispatch connect failure proves the request + # never left for upstream, so trying the next endpoint in the + # same resolved pool is safe even for a non-idempotent POST. + # TLS verification failures are stable endpoint configuration + # errors rather than transient connect losses; they keep the + # idempotent-only rule. + if index == len(endpoints) - 1 or not ( + allow_fallback or (exc.retryable_same_contract and not exc.is_tls_verification_failure) + ): raise except Exception as exc: - if index == len(endpoints) - 1 or not allow_fallback: + pre_dispatch = is_pre_dispatch_connection_failure(exc) and not isinstance(exc, aiohttp.ClientSSLError) + if index == len(endpoints) - 1 or not (allow_fallback or pre_dispatch): raise _transport_error( "request", endpoint.id, diff --git a/app/core/clients/proxy.py b/app/core/clients/proxy.py index 71b0442369..d39f068154 100644 --- a/app/core/clients/proxy.py +++ b/app/core/clients/proxy.py @@ -492,6 +492,24 @@ def __init__( self.retry_after_seconds = retry_after_seconds +def is_confirmed_pre_dispatch_transport_error(exc: ProxyResponseError) -> bool: + """Return whether the transport proved the upstream request never dispatched. + + Only this provenance authorizes replaying a movable request on another + account: a typed connector failure while reaching the account's routed + proxy endpoint, before any request bytes could leave for upstream. + Host-wide network loss (``proxy_network_unavailable``) stays on its + account-neutral process recovery path instead of penalizing the selected + account, and ambiguous dispatch outcomes remain non-replayable. + """ + + if not (exc.retryable_same_contract and exc.failure_phase == "connect"): + return False + error = exc.payload.get("error") + error_code = error.get("code") if isinstance(error, dict) else None + return error_code != PROCESS_NETWORK_UNAVAILABLE_CODE + + def _process_network_failure_error( message: str, exc: Exception, diff --git a/app/core/clients/proxy_websocket.py b/app/core/clients/proxy_websocket.py index 0cec919728..962e0c5cf1 100644 --- a/app/core/clients/proxy_websocket.py +++ b/app/core/clients/proxy_websocket.py @@ -802,9 +802,22 @@ async def _connect_upstream_websocket( status_code if policy.preserve_handshake_status else 502, openai_error(error_code, message, error_type="server_error"), failure_phase="connect", + # Carry the client's dispatch provenance across the sanitizing + # boundary: a typed connector failure against the routed proxy + # proves no ``response.create`` frame could have reached + # upstream, so service-level failover may replay the request + # on another account. TLS verification failures are stable + # endpoint configuration errors and stay non-replayable. retryable_same_contract=( - policy.retry_routed_network_errors and error_code == PROCESS_NETWORK_UNAVAILABLE_CODE + (policy.retry_routed_network_errors and error_code == PROCESS_NETWORK_UNAVAILABLE_CODE) + or (exc.retryable_same_contract and not exc.is_tls_verification_failure) ), + failure_detail=( + "proxy_connect_pre_dispatch" + if exc.retryable_same_contract and not exc.is_tls_verification_failure + else "transport_error" + ), + failure_exception_type=type(exc).__name__, ) 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 146c10fa42..d0226c6983 100644 --- a/app/modules/proxy/_service/http_bridge/helpers.py +++ b/app/modules/proxy/_service/http_bridge/helpers.py @@ -757,6 +757,79 @@ def _http_bridge_session_has_visible_requests(session: "_HTTPBridgeSession") -> ) +async def _close_http_bridge_session( + service: Any, + session: "_HTTPBridgeSession", + *, + turn_state_lock_held: bool = False, + release_durable_session: bool = True, +) -> None: + session.closed = True + if turn_state_lock_held: + service._unregister_http_bridge_turn_states_locked(session) + service._unregister_http_bridge_previous_response_ids_locked(session) + else: + await service._unregister_http_bridge_turn_states(session) + await service._unregister_http_bridge_previous_response_ids(session) + account_lease = getattr(session, "account_lease", None) + try: + await service._load_balancer.release_account_lease(account_lease) + except Exception: + logger.warning("Failed to release HTTP bridge account lease during close", exc_info=True) + finally: + session.account_lease = None + if release_durable_session and _http_bridge_durable_release_allowed(service, session): + try: + await service._durable_bridge.release_live_session( + session_id=session.durable_session_id, + instance_id=_service_get_settings().http_responses_session_bridge_instance_id, + owner_epoch=session.durable_owner_epoch, + draining=shutdown_state.is_bridge_drain_active(), + ) + except Exception: + logger.warning("Failed to release durable HTTP bridge session", exc_info=True) + upstream_reader = session.upstream_reader + if upstream_reader is not None: + if upstream_reader is asyncio.current_task(): + session.upstream_reader = None + else: + await _await_cancelled_task( + upstream_reader, + label="http bridge upstream reader", + cleanup_tasks=service._background_cleanup_tasks, + ) + if session.upstream_reader is upstream_reader: + session.upstream_reader = None + try: + await session.upstream.close() + except Exception: + logger.debug("Failed to close HTTP bridge upstream websocket", exc_info=True) + pending_requests = getattr(session, "pending_requests", None) + pending_lock = getattr(session, "pending_lock", None) + response_create_gate = getattr(session, "response_create_gate", None) + if pending_requests is not None and pending_lock is not None: + async with pending_lock: + session.queued_request_count = 0 + await service._fail_pending_websocket_requests( + account=session.account, + account_id_value=session.account.id, + pending_requests=pending_requests, + pending_lock=pending_lock, + error_code="stream_incomplete", + error_message="HTTP bridge session closed before response.completed", + api_key=None, + response_create_gate=response_create_gate, + ) + _log_http_bridge_event( + "close", + session.key, + account_id=session.account.id, + model=session.request_model, + cache_key_family=session.key.affinity_kind, + model_class=_extract_model_class(session.request_model) if session.request_model else None, + ) + + async def _close_http_bridge_session_bounded( service: Any, session: "_HTTPBridgeSession", diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index 17e1877620..79d9000c6b 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -5,6 +5,7 @@ import logging from collections import deque from collections.abc import Collection +from dataclasses import replace from typing import Any, Literal, TypeVar, overload from uuid import uuid4 @@ -80,7 +81,6 @@ _http_bridge_can_single_instance_prompt_cache_takeover_without_anchor, _http_bridge_compatible, _http_bridge_continuity_lost_error_envelope, - _http_bridge_durable_release_allowed, _http_bridge_endpoint_matches_current_instance, _http_bridge_eviction_priority, _http_bridge_has_durable_recovery_anchor, @@ -119,8 +119,12 @@ _require_http_bridge_bound_account_not_excluded, _reserve_http_bridge_unanchored_handoff, ) +from app.modules.proxy._service.http_bridge.helpers import ( + _close_http_bridge_session as _helpers_close_http_bridge_session, +) 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, @@ -1637,77 +1641,7 @@ def _prune_http_bridge_sessions_locked(self) -> list["_HTTPBridgeSession"]: sessions_to_close.append(session) return sessions_to_close - async def _close_http_bridge_session( - self, - session: "_HTTPBridgeSession", - *, - turn_state_lock_held: bool = False, - release_durable_session: bool = True, - ) -> None: - session.closed = True - if turn_state_lock_held: - self._unregister_http_bridge_turn_states_locked(session) - self._unregister_http_bridge_previous_response_ids_locked(session) - else: - await self._unregister_http_bridge_turn_states(session) - await self._unregister_http_bridge_previous_response_ids(session) - account_lease = getattr(session, "account_lease", None) - try: - await self._load_balancer.release_account_lease(account_lease) - except Exception: - logger.warning("Failed to release HTTP bridge account lease during close", exc_info=True) - finally: - session.account_lease = None - if release_durable_session and _http_bridge_durable_release_allowed(self, session): - try: - await self._durable_bridge.release_live_session( - session_id=session.durable_session_id, - instance_id=_service_get_settings().http_responses_session_bridge_instance_id, - owner_epoch=session.durable_owner_epoch, - draining=shutdown_state.is_bridge_drain_active(), - ) - except Exception: - logger.warning("Failed to release durable HTTP bridge session", exc_info=True) - upstream_reader = session.upstream_reader - if upstream_reader is not None: - if upstream_reader is asyncio.current_task(): - session.upstream_reader = None - else: - await _await_cancelled_task( - upstream_reader, - label="http bridge upstream reader", - cleanup_tasks=self._background_cleanup_tasks, - ) - if session.upstream_reader is upstream_reader: - session.upstream_reader = None - try: - await session.upstream.close() - except Exception: - logger.debug("Failed to close HTTP bridge upstream websocket", exc_info=True) - pending_requests = getattr(session, "pending_requests", None) - pending_lock = getattr(session, "pending_lock", None) - response_create_gate = getattr(session, "response_create_gate", None) - if pending_requests is not None and pending_lock is not None: - async with pending_lock: - session.queued_request_count = 0 - await self._fail_pending_websocket_requests( - account=session.account, - account_id_value=session.account.id, - pending_requests=pending_requests, - pending_lock=pending_lock, - error_code="stream_incomplete", - error_message="HTTP bridge session closed before response.completed", - api_key=None, - response_create_gate=response_create_gate, - ) - _log_http_bridge_event( - "close", - session.key, - account_id=session.account.id, - model=session.request_model, - cache_key_family=session.key.affinity_kind, - model_class=_extract_model_class(session.request_model) if session.request_model else None, - ) + _close_http_bridge_session = _helpers_close_http_bridge_session async def _create_http_bridge_session( self, @@ -1750,7 +1684,11 @@ async def _create_http_bridge_session( if require_preferred_account: fallback_on_preferred_account_unavailable = False 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 = { @@ -1758,14 +1696,18 @@ async def _create_http_bridge_session( "kind": "http_bridge", "request_stage": request_stage, "api_key": api_key, - "affinity_policy": affinity, + "affinity_policy": ( + replace(affinity, reallocate_sticky=True) + if proxy_connect_failover.reallocate_sticky and not affinity.reallocate_sticky + else affinity + ), "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, "preferred_account_is_continuity_owner": preferred_account_is_continuity_owner, "lease_kind": "stream", "estimated_lease_tokens": _estimated_lease_tokens_from_request_usage_budget(request_usage_budget), @@ -1781,6 +1723,11 @@ 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: + # No eligible replacement exists after a confirmed + # pre-dispatch route failure: preserve the original + # sanitized failure instead of generating ``no_accounts``. + raise proxy_connect_failover.last_error is_local_account_cap = _is_local_account_cap_code(selection.error_code) if ( require_preferred_account @@ -1838,6 +1785,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 @@ -1863,6 +1819,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 @@ -1873,7 +1838,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 @@ -1885,7 +1850,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 @@ -1910,7 +1875,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 @@ -1949,7 +1914,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..7a7e453b30 --- /dev/null +++ b/app/modules/proxy/_service/http_bridge/proxy_failover.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from app.core.clients.proxy import ProxyResponseError, is_confirmed_pre_dispatch_transport_error +from app.db.models import Account +from app.modules.proxy._service.http_bridge.protocol import _HTTPBridgeServiceProtocol +from app.modules.proxy.load_balancer import AccountLease + + +@dataclass +class _HTTPBridgePreDispatchFailover: + """Bridge-session startup failover for confirmed dead account proxy routes. + + Only a transport failure that proves the upstream request never dispatched + may move a session to another account. The failed account's stream lease + is released before the bounded transient backoff floor is recorded, and a + hard-required account fails closed on the original sanitized failure. The + preserved ``last_error`` keeps that failure authoritative when selection + cannot produce a replacement, instead of a generated ``no_accounts``. + """ + + excluded_account_ids: set[str] + preferred_account_id: str | None + reallocate_sticky: bool + last_error: ProxyResponseError | None = None + + async def handle( + self, + service: _HTTPBridgeServiceProtocol, + 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 51c044c422..8f073ef8a6 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -14,7 +14,12 @@ from app.core.auth.refresh import RefreshError, is_transient_refresh_contention, refresh_contention_kind 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_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.resilience.network_recovery import ( @@ -351,6 +356,7 @@ async def _stream_with_retry( network_recovery = ProcessNetworkRecovery(transport="stream", request_id=request_id) settlement = _StreamSettlement() last_transient_exc: ProxyResponseError | None = None + last_pre_dispatch_transport_error: ProxyResponseError | None = None last_account_model_rejection: ProxyResponseError | None = None last_account_model_rejection_account_id: str | None = None account_model_replacement_account_id: str | None = None @@ -389,6 +395,29 @@ async def _release_tracked_stream_lease(lease: AccountLease | None) -> None: pass await proxy._load_balancer.release_account_lease(lease) + def _render_dispatch_transport_error(exc: ProxyResponseError) -> str: + # Terminal render of the preserved sanitized transport failure: + # the client sees the original upstream-unavailable error instead + # of a misleading generated ``no_accounts`` response. + 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 _settle_stream_usage_before_pending_penalty( current_settlement: _StreamSettlement, ) -> bool: @@ -1075,6 +1104,13 @@ async def _retry_account_model_rejection( selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES or not (propagate_http_errors and last_transient_exc is not None) ) + and ( + selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES + # A preserved confirmed pre-dispatch failure is + # terminal for this request: waiting for capacity + # recovery cannot resurrect the dead proxy route. + or last_pre_dispatch_transport_error is None + ) and ( selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES or (last_retryable_stream_error is None and last_security_work_retry_error is None) @@ -1110,6 +1146,15 @@ async def _retry_account_model_rejection( account_id=last_account_model_rejection_account_id, ) return + if last_pre_dispatch_transport_error is not None: + # No eligible replacement exists: preserve the original + # sanitized upstream-unavailable failure instead of + # generating a misleading ``no_accounts`` response. + await _drain_pending_post_refresh_penalty_on_terminal(settlement) + 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: await _drain_pending_post_refresh_penalty_on_terminal(settlement) no_accounts_msg = selection.error_message or "Local account capacity is exhausted" @@ -1889,6 +1934,43 @@ async def _retry_account_model_rejection( _facade()._raise_proxy_budget_exhausted() if _facade()._is_account_neutral_error_code(code): raise + if is_confirmed_pre_dispatch_transport_error(tex): + # The transport proved the request never + # dispatched: this account's proxy route is + # dead. Release the account's stream lease + # before recording health so its slot never + # outlives the failed route, then jump + # straight to the bounded transient backoff + # floor so independent requests stop + # rediscovering the dead route one generic + # error at a time. + await _release_tracked_stream_lease(current_account_lease) + current_account_lease = None + 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: + # Hard account ownership or exhausted + # attempts: fail closed on the original + # sanitized failure without crossing + # accounts. + raise + last_transient_exc = tex + last_pre_dispatch_transport_error = tex + transient_failed_account_id = account.id + excluded_account_ids.add(account.id) + affinity = replace(affinity, reallocate_sticky=True) + _facade().logger.info( + "Retrying stream after confirmed pre-dispatch proxy connect failure " + "request_id=%s account_id=%s attempt=%d", + request_id, + account.id, + attempt + 1, + ) + break classified = await proxy._handle_stream_error( account, _upstream_error_from_openai(error), @@ -2440,12 +2522,24 @@ async def _retry_account_model_rejection( action, ) if action == "failover_next": - await proxy._handle_stream_error( - account, - current_error_payload, - current_error_code, - http_status=retry_exc.status_code, - ) + if is_confirmed_pre_dispatch_transport_error(retry_exc): + # Confirmed dead proxy route on the + # post-refresh attempt: release the lease + # first, then apply the bounded transient + # backoff floor instead of a single generic + # error, and preserve the sanitized failure + # for terminal rendering. + await _release_tracked_stream_lease(current_account_lease) + current_account_lease = None + await proxy._load_balancer.record_error_backoff(account) + last_pre_dispatch_transport_error = retry_exc + else: + await proxy._handle_stream_error( + account, + current_error_payload, + current_error_code, + http_status=retry_exc.status_code, + ) last_transient_exc = retry_exc await _release_tracked_stream_lease(current_account_lease) current_account_lease = None @@ -2580,6 +2674,12 @@ async def _retry_account_model_rejection( return if propagate_http_errors and last_transient_exc is not None: raise last_transient_exc + if last_pre_dispatch_transport_error is not None: + # Attempt budget exhausted after confirmed pre-dispatch route + # failures: surface the original sanitized failure rather than + # a generated ``no_accounts`` response. + yield _render_dispatch_transport_error(last_pre_dispatch_transport_error) + return if last_retryable_stream_error is not None: retries_exhausted_msg = str(last_retryable_stream_error.error.get("message") or "Upstream error") event = response_failed_event( diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index c90b6e80b5..f5e831f8da 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -45,6 +45,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, @@ -3164,6 +3165,7 @@ async def _connect_proxy_websocket( last_failover_account = account continue except ProxyResponseError as exc: + confirmed_pre_dispatch = is_confirmed_pre_dispatch_transport_error(exc) if selected_account_model_replacement: # The account/model retry budget selected this replacement; # its connection failure must be surfaced rather than @@ -3177,9 +3179,16 @@ 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": + # Release the dead route's stream lease before recording + # the backoff so its concurrency slot never outlives the + # failed connection attempt. 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) @@ -3189,6 +3198,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, @@ -3902,13 +3913,25 @@ 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: + # A proven pre-dispatch proxy connect failure is account-local + # transient evidence. The caller applies the bounded transient + # backoff floor once the failed lease is released, so the generic + # single-error health write is skipped here. Hard account + # ownership fails closed on the original sanitized failure. + 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 27244b522a..57b73c6261 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, @@ -1510,7 +1511,17 @@ async def mark_permanent_failure(self, account: Account, error_code: str) -> boo 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 @@ -1518,7 +1529,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..ddf55b8cc5 --- /dev/null +++ b/openspec/changes/retry-account-proxy-connect-failures/design.md @@ -0,0 +1,71 @@ +# Design + +## Transport provenance + +The codebase already carries dispatch provenance on `CodexTransportError` and +`ProxyResponseError`: `retryable_same_contract` is true only when a typed +connector failure proves the upstream request was never dispatched, and +`failure_phase == "connect"` records where the attempt died. This change +reuses that single source of truth instead of introducing a parallel state. + +`is_confirmed_pre_dispatch_transport_error` is the one predicate that +authorizes cross-account replay. It requires proven pre-dispatch connect +provenance and explicitly excludes host-wide network loss +(`proxy_network_unavailable`), which keeps its account-neutral process +network recovery path (transport rotation and bounded same-account retry) +rather than penalizing the selected account. TLS verification failures are +stable endpoint configuration errors and never authorize replay. + +The native Responses WebSocket connect boundary previously collapsed routed +`CodexTransportError` provenance to the process-network case only; it now +carries the pre-dispatch provenance across the sanitizing conversion. No +proxy URL, credentials, or raw exception text is retained or exposed. + +## Retry order + +For an HTTP POST, a confirmed pre-dispatch connect 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 idempotent-only 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 (attempt budget, +request deadline, and a monotonically growing exclusion set). The original +sanitized 502 is retained and returned if no replacement exists, instead of +a generated `no_accounts` response. + +## 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 with the original sanitized failure. Soft prompt-cache and +process-session affinity may move; the failed account's sticky binding is +reallocated so the replacement selection does not immediately loop back. + +## 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 (`record_error_backoff`, shared `ERROR_BACKOFF_THRESHOLD`): +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, and it replaces — not stacks onto — the generic single-error health +write for the same failure. + +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; downstream-visible output always forbids replay. +- Do not treat idle disconnects or downstream idle timeouts as account + health evidence. +- 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..cef8812e45 --- /dev/null +++ b/openspec/changes/retry-account-proxy-connect-failures/proposal.md @@ -0,0 +1,46 @@ +# Retry confirmed account-proxy connect failures + +## Why + +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 because a confirmed +dead route only accrues generic single errors (#1314). + +## What Changes + +- Reuse the existing sanitized dispatch provenance + (`retryable_same_contract` + `failure_phase == "connect"`) as the single + predicate that authorizes cross-account replay, excluding host-wide + network loss and TLS verification failures. +- When the transport proves the connection to the selected proxy failed + before request dispatch, try another endpoint in the same proxy pool even + for a non-idempotent POST, then another eligible account for movable + Responses requests across raw HTTP/SSE, native Responses WebSocket, and + HTTP bridge session startup. +- Release the failed account's stream lease before recording bounded + transient account backoff (`record_error_backoff` jumps directly to the + existing 30s floor, capped at 300s) so independent requests stop + rediscovering the dead route. +- Preserve the original sanitized upstream-unavailable failure when no + replacement account exists. Ambiguous failures and hard continuity, file, + or required-account ownership remain non-replayable and fail closed. + +## Capabilities + +### New Capabilities + +(none) + +### Modified Capabilities + +- `upstream-proxy-routing` + +## Impact + +No API surface, schema, or dependency changes. Selection health-state gains +a minimum-error-count floor used only by the confirmed dead-route path; the +generic transient-error accounting, process network recovery, and +downstream-visible fail-closed semantics are unchanged. 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..a05dbba559 --- /dev/null +++ b/openspec/changes/retry-account-proxy-connect-failures/specs/upstream-proxy-routing/spec.md @@ -0,0 +1,48 @@ +# 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 468dfe6b12..24a14f1354 100644 --- a/tests/integration/test_http_responses_bridge.py +++ b/tests/integration/test_http_responses_bridge.py @@ -1603,6 +1603,90 @@ 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]] = [] + backed_off_accounts: list[str] = [] + handle_stream_error = AsyncMock() + + async def fake_select_account_with_budget(self, deadline, *, exclude_account_ids=None, **kwargs): + del self, deadline, kwargs + excluded = set(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", + retryable_same_contract=True, + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + ) + return upstream + + async def fake_record_error_backoff(self, account): + del self + backed_off_accounts.append(account.id) + + 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.LoadBalancer, "record_error_backoff", fake_record_error_backoff) + 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 backed_off_accounts == [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 1bac6a33a8..35fd86a45c 100644 --- a/tests/integration/test_proxy_responses.py +++ b/tests/integration/test_proxy_responses.py @@ -597,6 +597,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", + retryable_same_contract=True, + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + ) + 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_elides_required_tool_image_and_streams_item( async_client, diff --git a/tests/integration/test_proxy_websocket_responses.py b/tests/integration/test_proxy_websocket_responses.py index 8fbcd45288..3a2a9e0902 100644 --- a/tests/integration/test_proxy_websocket_responses.py +++ b/tests/integration/test_proxy_websocket_responses.py @@ -9,6 +9,7 @@ from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any, cast +from unittest.mock import AsyncMock import pytest from fastapi.testclient import TestClient @@ -569,6 +570,118 @@ def _websocket_response_create(text: str) -> dict[str, object]: } +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", + retryable_same_contract=True, + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + ) + + 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] + # The confirmed dead route skips the generic single-error health write. + 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 71974bf70e..6b6eb915ff 100644 --- a/tests/unit/test_codex_client.py +++ b/tests/unit/test_codex_client.py @@ -214,6 +214,67 @@ async def test_non_idempotent_request_failure_does_not_fallback(route: ResolvedU assert session.calls[0]["proxy"] == runtime_basic_auth_url("u", "p", "proxy.test:8080") +def _proxy_connect_error() -> aiohttp.ClientProxyConnectionError: + key = ConnectionKey("proxy.test", 8080, False, False, None, None, 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"}) + + +@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] == [ + runtime_basic_auth_url("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.retryable_same_contract is True + assert exc_info.value.failure_phase == "connect" + assert "ClientProxyConnectionError" in str(exc_info.value) + 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 b98bd2567f..50c52ef424 100644 --- a/tests/unit/test_codex_upstream_paths.py +++ b/tests/unit/test_codex_upstream_paths.py @@ -17,6 +17,7 @@ UpstreamProxyRouteTrace, codex_control_request, compact_responses, + is_confirmed_pre_dispatch_transport_error, stream_responses, thread_goal_request, transcribe_audio, @@ -699,6 +700,65 @@ async def test_stream_responses_marks_typed_routed_connector_failure_replay_safe assert exc_info.value.failed_session is None +@pytest.mark.asyncio +async def test_stream_responses_propagates_confirmed_pre_dispatch_failure_for_status_retry( + route: ResolvedUpstreamRoute, +) -> None: + key = ConnectionKey("proxy.test", 8080, False, False, None, None, None) + network_error = aiohttp.ClientProxyConnectionError(key, ConnectionRefusedError("connection refused")) + client = CodexClient(_NetworkFailureSession(network_error)) + 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=client, + raise_for_status=True, + ): + pass + + assert exc_info.value.status_code == 502 + assert exc_info.value.payload["error"]["code"] == "upstream_unavailable" + assert exc_info.value.retryable_same_contract is True + assert exc_info.value.failure_phase == "connect" + assert is_confirmed_pre_dispatch_transport_error(exc_info.value) is True + assert "connection refused" not in str(exc_info.value.payload["error"]["message"]) + + +@pytest.mark.asyncio +async def test_stream_responses_ambiguous_transport_failure_stays_terminal_without_retry_authorization( + route: ResolvedUpstreamRoute, +) -> None: + payload = ResponsesRequest(model="gpt-5.2", instructions="Reply.", input="hello", stream=True) + + events = [ + event + async for event 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, + ) + ] + + # Dispatch is unknown, so even ``raise_for_status`` callers receive the + # terminal downstream event rather than a replay-authorizing exception. + combined = "".join(events) + assert '"type":"response.failed"' in combined or '"type": "response.failed"' in combined + assert '"code":"upstream_unavailable"' in combined or '"code": "upstream_unavailable"' in combined + + @pytest.mark.asyncio async def test_codex_client_buffered_body_network_failure_is_neutral_but_unsafe( route: ResolvedUpstreamRoute, diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index 520f3c4e81..abc05e6e76 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -14,6 +14,7 @@ import app.modules.proxy.load_balancer as load_balancer_module from app.core.balancer import ( + ERROR_BACKOFF_THRESHOLD, HEALTH_TIER_DRAINING, HEALTH_TIER_HEALTHY, HEALTH_TIER_PROBING, @@ -402,6 +403,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 == ERROR_BACKOFF_THRESHOLD + assert runtime.last_error_at is not None + + await balancer.record_error_backoff(account) + + assert runtime.error_count == ERROR_BACKOFF_THRESHOLD + 1 + + @pytest.mark.asyncio async def test_successful_force_probes_promote_probing_account_to_healthy() -> None: account = _make_account("acc-force-probe-success") diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index fd1a43d1e5..3f8a5a9f5a 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -6242,6 +6242,187 @@ async def select_account(_deadline: float, **kwargs: object) -> proxy_service.Ac assert selection_kwargs[0]["prefer_earlier_reset_window"] == "primary" +def _pre_dispatch_proxy_error(message: str = "sanitized proxy connect failure") -> ProxyResponseError: + return ProxyResponseError( + 502, + openai_error("upstream_unavailable", message), + failure_phase="connect", + retryable_same_contract=True, + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + ) + + +def _bridge_selection_settings() -> SimpleNamespace: + return SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + + +@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())) + 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]] = [] + reallocate_flags: list[bool] = [] + released_leases: list[proxy_service.AccountLease] = [] + backed_off_accounts: list[object] = [] + upstream = cast(Any, 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) + affinity_policy = cast(proxy_service._AffinityPolicy, kwargs["affinity_policy"]) + reallocate_flags.append(affinity_policy.reallocate_sticky) + 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: + backed_off_accounts.append(account) + # The dead route's stream lease must settle before the health write. + 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=_bridge_selection_settings())), + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", 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._load_balancer, "record_error", AsyncMock()) + 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 reallocate_flags == [False, True] + assert backed_off_accounts == [account_a] + assert lease_a in released_leases + assert lease_b not 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())) + 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=_bridge_selection_settings())), + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", 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, + ) + + # The hard-required owner fails closed on the original sanitized failure. + assert exc_info.value is original_error + select_account.assert_awaited_once() + record_error_backoff.assert_awaited_once_with(account) + assert release_account_lease.await_args_list[0].args == (lease,) + + +@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())) + 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=_bridge_selection_settings())), + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", 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-no-replacement", None), + headers={}, + affinity=proxy_service._AffinityPolicy(key="sid-proxy-no-replacement"), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + ) + + # The original sanitized failure is preserved instead of ``no_accounts``. + 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 8745be1f00..afbf77ec77 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -14155,6 +14155,298 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, assert request_logs.calls[-1]["error_code"] == "stream_idle_timeout" +def _pre_dispatch_proxy_connect_error( + message: str = "sanitized proxy connect failure", +) -> proxy_service.ProxyResponseError: + return proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", message), + failure_phase="connect", + retryable_same_contract=True, + failure_detail="proxy_connect_pre_dispatch", + failure_exception_type="ClientProxyConnectionError", + ) + + +@pytest.mark.asyncio +async def test_stream_responses_confirmed_proxy_connect_failure_fails_over_after_lease_release(monkeypatch): + settings = _make_proxy_settings() + 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 = _make_api_key_data("key_proxy_connect_failover") + 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]] = [] + seen_reallocate_sticky: list[object] = [] + record_error = AsyncMock() + record_success = AsyncMock() + settle_usage = AsyncMock(return_value=True) + release_unsettled_usage = AsyncMock() + record_error_backoff_accounts: list[Account] = [] + + 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) + seen_reallocate_sticky.append(kwargs.get("reallocate_sticky")) + 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: + record_error_backoff_accounts.append(account) + # Leases must settle before the error-health write for the dead route. + 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 _pre_dispatch_proxy_connect_error() + 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 all('"type":"response.failed"' not in chunk for chunk in chunks) + assert stream_accounts == [account_a.chatgpt_account_id, account_b.chatgpt_account_id] + assert seen_excluded_account_ids == [set(), {account_a.id}] + assert seen_reallocate_sticky[-1] is True + assert record_error_backoff_accounts == [account_a] + 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() + 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() + record_error_backoff_accounts: list[Account] = [] + + 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: + record_error_backoff_accounts.append(failed_account) + assert lease in released_leases + + async def fake_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise _pre_dispatch_proxy_connect_error("original sanitized proxy failure") + 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}] + assert record_error_backoff_accounts == [account] + 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() + 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]] = [] + record_error_backoff = 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: + selections.append(dict(kwargs)) + selected = owner if len(selections) == 1 else alternate + return AccountSelection(account=selected, error_message=None, lease=lease if selected is owner else None) + + async def fake_stream(*_args: object, **_kwargs: object) -> AsyncIterator[str]: + raise _pre_dispatch_proxy_connect_error("file owner proxy unavailable") + yield "" + + 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", record_error_backoff) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(return_value=owner)) + 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_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" + # The file-pinned request must never cross to the alternate account. + assert len(selections) == 1 + record_error_backoff.assert_awaited_once_with(owner) + + +@pytest.mark.asyncio +async def test_stream_responses_ambiguous_dispatch_transport_failure_is_not_replayed(monkeypatch): + settings = _make_proxy_settings() + 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_backoff = AsyncMock() + + 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_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]: + # Dispatch is unknown, so the sanitized failure is yielded as a + # terminal downstream event instead of authorizing a replay. + yield ( + 'data: {"type":"response.failed","response":{"error":' + '{"code":"upstream_unavailable","message":"dispatch outcome is unknown"}}}\n\n' + ) + + 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[-1].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_backoff.assert_not_awaited() + + @pytest.mark.asyncio async def test_stream_responses_retries_hard_owner_after_transient_exclusion(monkeypatch): settings = _make_proxy_settings() diff --git a/tests/unit/test_proxy_websocket_client.py b/tests/unit/test_proxy_websocket_client.py index 6dcd3af3f5..38af67011e 100644 --- a/tests/unit/test_proxy_websocket_client.py +++ b/tests/unit/test_proxy_websocket_client.py @@ -17,7 +17,7 @@ import app.core.clients.proxy_websocket as proxy_websocket_module from app.core.clients.codex import CodexTransportError, CodexWebSocketResult -from app.core.clients.proxy import ProxyResponseError +from app.core.clients.proxy import ProxyResponseError, is_confirmed_pre_dispatch_transport_error from app.core.clients.proxy_websocket import ( CodexUpstreamWebSocket, RealtimeWebSocketProtocol, @@ -683,6 +683,58 @@ 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 "") + # An ambiguous routed transport failure must not authorize replay. + assert exc_info.value.retryable_same_contract is False + + +@pytest.mark.asyncio +async def test_connect_responses_websocket_routed_pre_dispatch_failure_carries_provenance(monkeypatch): + route = ResolvedUpstreamRoute( + mode="account_bound", + pool_id="pool_1", + endpoint=ResolvedProxyEndpoint("ep_1", "http", "proxy.test", 8080), + ) + monkeypatch.setattr( + proxy_websocket_module, + "get_settings", + lambda: SimpleNamespace( + upstream_base_url="https://chatgpt.com/backend-api", + upstream_connect_timeout_seconds=7.0, + max_sse_event_bytes=4321, + upstream_websocket_trust_env=False, + ), + ) + + class _PreDispatchFailingCodexClient(_FailingCodexClient): + async def open_ws_with_route_metadata( + self, + url: str, + *, + route: ResolvedUpstreamRoute, + **kwargs: object, + ) -> CodexWebSocketResult: + del url, route, kwargs + raise CodexTransportError( + "Codex upstream websocket failed via proxy endpoint ep_1: ClientProxyConnectionError", + failure_phase="connect", + retryable_same_contract=True, + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await connect_responses_websocket( + {"openai-beta": "responses_websockets=2026-02-06"}, + "access-token", + "account-123", + route=route, + codex_client=cast(Any, _PreDispatchFailingCodexClient()), + ) + + assert exc_info.value.status_code == 502 + assert _proxy_error_code(exc_info.value) == "upstream_unavailable" + assert exc_info.value.retryable_same_contract is True + assert exc_info.value.failure_phase == "connect" + assert exc_info.value.failure_detail == "proxy_connect_pre_dispatch" + assert is_confirmed_pre_dispatch_transport_error(exc_info.value) is True @pytest.mark.asyncio From 2f663af4ca69b7e4b4f3b106dcc7237d5daa8f79 Mon Sep 17 00:00:00 2001 From: Soju06 Date: Thu, 30 Jul 2026 08:44:24 +0000 Subject: [PATCH 2/4] test(proxy): guard replay-authorization boundaries for dead-route failover Maintainer follow-up to the #1322 takeover: pin the negative boundaries of the confirmed pre-dispatch predicate at the externally observable surfaces. - TLS verification connect failures do not authorize same-pool POST fallback or cross-account websocket replay. - Host-wide network loss (proxy_network_unavailable) and post-dispatch body-read failures never classify as confirmed pre-dispatch. - An idle bridge disconnect does not exclude the account or record the transient backoff floor, so healthy accounts stay healthy. Co-Authored-By: Claude Fable 5 --- tests/unit/test_codex_client.py | 27 +++++++++++++ tests/unit/test_proxy_http_bridge.py | 47 +++++++++++++++++++++++ tests/unit/test_proxy_utils.py | 26 +++++++++++++ tests/unit/test_proxy_websocket_client.py | 47 +++++++++++++++++++++++ 4 files changed, 147 insertions(+) diff --git a/tests/unit/test_codex_client.py b/tests/unit/test_codex_client.py index 6b6eb915ff..cf3389530c 100644 --- a/tests/unit/test_codex_client.py +++ b/tests/unit/test_codex_client.py @@ -275,6 +275,33 @@ async def test_exhausted_proxy_connect_failures_preserve_pre_dispatch_provenance assert "credentials must stay private" not in str(exc_info.value) +@pytest.mark.asyncio +async def test_non_idempotent_tls_verification_connect_failure_does_not_fallback( + route: ResolvedUpstreamRoute, +) -> None: + connection_key = ConnectionKey("proxy.test", 8080, True, True, None, None, None) + + class _TLSFailSession(_Session): + async def request(self, method: str, url: str, **kwargs: Any) -> _Response: + self.calls.append({"method": method, "url": url, **kwargs}) + raise aiohttp.ClientConnectorCertificateError(connection_key, ValueError("certificate verify failed")) + + session = _TLSFailSession() + client = CodexClient(session) + + 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 len(session.calls) == 1 + assert exc_info.value.is_tls_verification_failure is True + + @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_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index 3f8a5a9f5a..bbf6cc1907 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -6423,6 +6423,53 @@ async def select_account(_deadline: float, **kwargs: object) -> proxy_service.Ac assert selections == [set(), {account.id}] +@pytest.mark.asyncio +async def test_create_http_bridge_session_idle_close_error_is_not_treated_as_dead_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + account = cast(Any, SimpleNamespace(id="acc-proxy-idle", status=AccountStatus.ACTIVE, plan_type="plus")) + lease = proxy_service.AccountLease("lease-bridge-idle", account.id, "stream", time.monotonic()) + record_error_backoff = AsyncMock() + idle_error = ProxyResponseError( + 502, + openai_error("upstream_unavailable", "Upstream websocket closed while idle"), + failure_phase="upstream", + failure_detail="stream_idle_timeout", + ) + + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=_bridge_selection_settings())), + ) + monkeypatch.setattr( + service, + "_select_account_with_budget_compatible", + AsyncMock(return_value=proxy_service.AccountSelection(account=account, error_message=None, lease=lease)), + ) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(return_value=account)) + monkeypatch.setattr(service, "_open_upstream_websocket_with_budget", AsyncMock(side_effect=idle_error)) + monkeypatch.setattr(service._load_balancer, "release_account_lease", AsyncMock()) + 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("session_header", "sid-proxy-idle", None), + headers={}, + affinity=proxy_service._AffinityPolicy(key="sid-proxy-idle"), + api_key=None, + request_model="gpt-5.4", + idle_ttl_seconds=120.0, + ) + + # An idle disconnect is not provable pre-dispatch evidence: no account + # exclusion, no transient-backoff health write. + assert exc_info.value is idle_error + record_error_backoff.assert_not_awaited() + + @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 afbf77ec77..ffffd2e363 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -14168,6 +14168,32 @@ def _pre_dispatch_proxy_connect_error( ) +def test_is_confirmed_pre_dispatch_transport_error_requires_provable_connect_provenance(): + assert proxy_module.is_confirmed_pre_dispatch_transport_error(_pre_dispatch_proxy_connect_error()) is True + ambiguous = proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "dispatch outcome is unknown"), + failure_phase="upstream", + failure_detail="transport_error", + ) + assert proxy_module.is_confirmed_pre_dispatch_transport_error(ambiguous) is False + non_connect = proxy_service.ProxyResponseError( + 502, + openai_error("upstream_unavailable", "body read failed"), + failure_phase="body_read", + retryable_same_contract=True, + ) + assert proxy_module.is_confirmed_pre_dispatch_transport_error(non_connect) is False + process_network = proxy_service.ProxyResponseError( + 502, + openai_error("proxy_network_unavailable", "host lost DNS"), + failure_phase="connect", + retryable_same_contract=True, + ) + # Host-wide network loss keeps its account-neutral recovery path. + assert proxy_module.is_confirmed_pre_dispatch_transport_error(process_network) is False + + @pytest.mark.asyncio async def test_stream_responses_confirmed_proxy_connect_failure_fails_over_after_lease_release(monkeypatch): settings = _make_proxy_settings() diff --git a/tests/unit/test_proxy_websocket_client.py b/tests/unit/test_proxy_websocket_client.py index 38af67011e..23b2511b63 100644 --- a/tests/unit/test_proxy_websocket_client.py +++ b/tests/unit/test_proxy_websocket_client.py @@ -737,6 +737,53 @@ async def open_ws_with_route_metadata( assert is_confirmed_pre_dispatch_transport_error(exc_info.value) is True +@pytest.mark.asyncio +async def test_connect_responses_websocket_routed_tls_verification_failure_is_not_replayable(monkeypatch): + route = ResolvedUpstreamRoute( + mode="account_bound", + pool_id="pool_1", + endpoint=ResolvedProxyEndpoint("ep_1", "http", "proxy.test", 8080), + ) + monkeypatch.setattr( + proxy_websocket_module, + "get_settings", + lambda: SimpleNamespace( + upstream_base_url="https://chatgpt.com/backend-api", + upstream_connect_timeout_seconds=7.0, + max_sse_event_bytes=4321, + upstream_websocket_trust_env=False, + ), + ) + + class _TLSFailingCodexClient(_FailingCodexClient): + async def open_ws_with_route_metadata( + self, + url: str, + *, + route: ResolvedUpstreamRoute, + **kwargs: object, + ) -> CodexWebSocketResult: + del url, route, kwargs + raise CodexTransportError( + "Codex upstream websocket failed via proxy endpoint ep_1: ClientConnectorCertificateError", + failure_phase="connect", + retryable_same_contract=True, + is_tls_verification_failure=True, + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await connect_responses_websocket( + {"openai-beta": "responses_websockets=2026-02-06"}, + "access-token", + "account-123", + route=route, + codex_client=cast(Any, _TLSFailingCodexClient()), + ) + + assert exc_info.value.retryable_same_contract is False + assert is_confirmed_pre_dispatch_transport_error(exc_info.value) is False + + @pytest.mark.asyncio async def test_connect_responses_websocket_appends_required_beta_header(monkeypatch): fake_connection = _FakeConnection() From 6bfc2f6566e555bec06436315c2294bec30531a0 Mon Sep 17 00:00:00 2001 From: SSY <234955825+mereyabdenbekuly-ctrl@users.noreply.github.com> Date: Thu, 30 Jul 2026 17:12:09 +0500 Subject: [PATCH 3/4] fix(proxy): settle keyed failover health writes safely --- app/modules/proxy/_service/api_key_usage.py | 26 + .../proxy/_service/http_bridge/mixin.py | 17 + .../proxy/_service/http_bridge/protocol.py | 1 + .../_service/http_bridge/proxy_failover.py | 15 +- .../proxy/_service/http_bridge/streaming.py | 128 +++- .../proxy/_service/streaming/protocol.py | 1 + app/modules/proxy/_service/streaming/retry.py | 161 +++-- app/modules/proxy/_service/support.py | 20 + app/modules/proxy/_service/websocket/mixin.py | 31 +- .../proxy/_service/websocket/protocol.py | 1 + .../design.md | 30 +- .../specs/upstream-proxy-routing/spec.md | 3 +- .../tasks.md | 2 + .../test_proxy_websocket_responses.py | 31 + tests/unit/test_proxy_http_bridge.py | 236 ++++++- tests/unit/test_proxy_utils.py | 619 +++++++++++++++++- 16 files changed, 1234 insertions(+), 88 deletions(-) diff --git a/app/modules/proxy/_service/api_key_usage.py b/app/modules/proxy/_service/api_key_usage.py index f9fa125890..c9bbca703f 100644 --- a/app/modules/proxy/_service/api_key_usage.py +++ b/app/modules/proxy/_service/api_key_usage.py @@ -14,6 +14,7 @@ from app.core.exceptions import ProxyAuthError, ProxyRateLimitError from app.core.openai.models import CompactResponsePayload from app.core.utils.request_id import get_request_id +from app.db.models import Account from app.modules.api_keys.service import ( ApiKeyData, ApiKeyInvalidError, @@ -64,6 +65,7 @@ class _ApiKeyUsageServiceProtocol(Protocol): _repo_factory: ProxyRepoFactory _background_cleanup_tasks: set[asyncio.Task[None]] _stream_api_key_release_retry_semaphore: asyncio.Semaphore + _load_balancer: Any def _normalize_service_tier_value(value: Any) -> str | None: @@ -135,6 +137,30 @@ async def _release_websocket_request_state_reservation( ) -> None: self._cancel_request_state_api_key_reservation_heartbeat(request_state) await self._release_websocket_reservation(request_state.api_key_reservation) + request_state.api_key_reservation = None + lifecycle = request_state.deferred_account_backoff_lifecycle + if lifecycle is not None: + lifecycle.settlement_confirmed = True + pending_backoffs = ( + lifecycle.pending_backoffs if lifecycle is not None else request_state.deferred_account_error_backoffs + ) + if pending_backoffs: + await self._drain_deferred_account_error_backoffs(pending_backoffs) + + async def _drain_deferred_account_error_backoffs( + self, + pending_backoffs: dict[str, Account], + ) -> None: + if not pending_backoffs: + return + proxy = cast(_ApiKeyUsageServiceProtocol, self) + while pending_backoffs: + account_id, account = pending_backoffs.popitem() + try: + await proxy._load_balancer.record_error_backoff(account) + except BaseException: + pending_backoffs.setdefault(account_id, account) + raise async def _maybe_touch_api_key_reservation( self, diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index 79d9000c6b..13a92e54aa 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -156,6 +156,7 @@ _clear_websocket_precreated_replay_fallback, _complete_http_bridge_handoff, _copy_websocket_route_metadata_to_session, + _DeferredAccountBackoffLifecycle, _HTTPBridgeOwnerForward, _HTTPBridgeSession, _HTTPBridgeSessionKey, @@ -350,6 +351,8 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession": ... @overload async def _get_or_create_http_bridge_session( @@ -381,6 +384,8 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession | _HTTPBridgeOwnerForward": ... async def _get_or_create_http_bridge_session( self, @@ -411,6 +416,8 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession | _HTTPBridgeOwnerForward": settings = _service_get_settings() request_scope_id = ensure_request_scope_id() @@ -1489,6 +1496,8 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: "request_usage_budget": request_usage_budget, "request_deadline": request_deadline, "exclude_account_ids": exclude_account_ids, + "deferred_account_backoff_lifecycle": deferred_account_backoff_lifecycle, + "defer_account_health_writes": defer_account_health_writes, } try: create_signature = inspect.signature(create_session) @@ -1505,6 +1514,8 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: "request_deadline", "exclude_account_ids", "preferred_account_is_continuity_owner", + "deferred_account_backoff_lifecycle", + "defer_account_health_writes", ): if optional_kwarg not in create_signature.parameters: create_kwargs.pop(optional_kwarg, None) @@ -1661,6 +1672,8 @@ async def _create_http_bridge_session( request_usage_budget: ApiKeyRequestUsageBudget | None = None, request_deadline: float | None = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession": request_state = _WebSocketRequestState( request_id=f"http_bridge_connect_{uuid4().hex}", @@ -1791,6 +1804,8 @@ async def _create_http_bridge_session( selected_account_lease, exc, required_account=require_preferred_account and selected_is_preferred, + deferred_account_backoff_lifecycle=deferred_account_backoff_lifecycle, + defer_account_health_write=defer_account_health_writes, ): selected_account_lease = None continue @@ -1825,6 +1840,8 @@ async def _create_http_bridge_session( selected_account_lease, retry_exc, required_account=require_preferred_account and selected_is_preferred, + deferred_account_backoff_lifecycle=deferred_account_backoff_lifecycle, + defer_account_health_write=defer_account_health_writes, ): selected_account_lease = None continue diff --git a/app/modules/proxy/_service/http_bridge/protocol.py b/app/modules/proxy/_service/http_bridge/protocol.py index b88edce03e..841d02c014 100644 --- a/app/modules/proxy/_service/http_bridge/protocol.py +++ b/app/modules/proxy/_service/http_bridge/protocol.py @@ -102,4 +102,5 @@ async def _maybe_touch_request_state_api_key_reservation(self, *args: Any, **kwa async def _reserve_websocket_api_key_usage(self, *args: Any, **kwargs: Any) -> Any: ... async def _release_websocket_reservation(self, *args: Any, **kwargs: Any) -> None: ... async def _release_websocket_request_state_reservation(self, *args: Any, **kwargs: Any) -> None: ... + async def _drain_deferred_account_error_backoffs(self, *args: Any, **kwargs: Any) -> None: ... def _schedule_cancel_safe_cleanup(self, *args: Any, **kwargs: Any) -> None: ... diff --git a/app/modules/proxy/_service/http_bridge/proxy_failover.py b/app/modules/proxy/_service/http_bridge/proxy_failover.py index 7a7e453b30..c8a474302d 100644 --- a/app/modules/proxy/_service/http_bridge/proxy_failover.py +++ b/app/modules/proxy/_service/http_bridge/proxy_failover.py @@ -5,6 +5,7 @@ from app.core.clients.proxy import ProxyResponseError, is_confirmed_pre_dispatch_transport_error from app.db.models import Account from app.modules.proxy._service.http_bridge.protocol import _HTTPBridgeServiceProtocol +from app.modules.proxy._service.support import _DeferredAccountBackoffLifecycle from app.modules.proxy.load_balancer import AccountLease @@ -14,7 +15,8 @@ class _HTTPBridgePreDispatchFailover: Only a transport failure that proves the upstream request never dispatched may move a session to another account. The failed account's stream lease - is released before the bounded transient backoff floor is recorded, and a + is released before the bounded transient backoff floor is recorded; keyed + requests defer that write until their singular reservation settles. A hard-required account fails closed on the original sanitized failure. The preserved ``last_error`` keeps that failure authoritative when selection cannot produce a replacement, instead of a generated ``no_accounts``. @@ -33,11 +35,20 @@ async def handle( exc: ProxyResponseError, *, required_account: bool, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_write: bool = False, ) -> 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 ( + defer_account_health_write + and deferred_account_backoff_lifecycle is not None + and not deferred_account_backoff_lifecycle.settlement_confirmed + ): + deferred_account_backoff_lifecycle.pending_backoffs.setdefault(account.id, account) + else: + await service._load_balancer.record_error_backoff(account) if required_account: raise exc self.last_error = exc diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index ec2c3c9923..da677d8c32 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -151,6 +151,8 @@ _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS, # noqa: F401 _account_capacity_wait_payload, _account_selection_recovery_sleep_seconds_from_message, + _DeferredAccountBackoffLifecycle, + _DeferredAccountBackoffTracker, _event_type_from_payload, _HTTPBridgeOwnerForward, _HTTPBridgeSession, @@ -859,6 +861,7 @@ async def _stream_http_bridge_or_retry( return request_scope_id = ensure_request_scope_id() + deferred_account_backoff_tracker = _DeferredAccountBackoffTracker() try: async for line in self._stream_via_http_bridge( payload, @@ -886,14 +889,40 @@ async def _stream_http_bridge_or_retry( enforce_openai_sdk_contract=enforce_openai_sdk_contract, capacity_startup_wait_event=capacity_startup_wait_event, capacity_startup_ready_event=capacity_startup_ready_event, + deferred_account_backoff_tracker=deferred_account_backoff_tracker, ): yield line finally: with anyio.CancelScope(shield=True): - await _release_http_bridge_unanchored_handoffs_for_request( - self, - request_scope_id=request_scope_id, - ) + try: + lifecycle = deferred_account_backoff_tracker.current_lifecycle + if lifecycle is not None: + pending_backoffs = lifecycle.pending_backoffs + if lifecycle.settlement_confirmed: + await self._drain_deferred_account_error_backoffs(pending_backoffs) + elif not lifecycle.settlement_owned and ( + pending_backoffs or lifecycle.reservation != api_key_reservation + ): + # Session creation can fail before the request is + # submitted. Until submit returns, this wrapper owns + # the current lifecycle and may release exactly that + # reservation. Once ownership transfers, the request + # finalizer is the only safe settlement owner. + try: + await self._release_websocket_reservation(lifecycle.reservation) + except Exception: + logger.warning( + "Failed to release HTTP bridge API key reservation before deferred backoff", + exc_info=True, + ) + else: + lifecycle.settlement_confirmed = True + await self._drain_deferred_account_error_backoffs(pending_backoffs) + finally: + await _release_http_bridge_unanchored_handoffs_for_request( + self, + request_scope_id=request_scope_id, + ) async def _stream_via_http_bridge( self: Any, @@ -923,11 +952,14 @@ async def _stream_via_http_bridge( enforce_openai_sdk_contract: bool = True, capacity_startup_wait_event: asyncio.Event | None = None, capacity_startup_ready_event: asyncio.Event | None = None, + deferred_account_backoff_tracker: _DeferredAccountBackoffTracker | None = None, ) -> AsyncIterator[str]: del suppress_text_done_events request_id = ensure_request_id() dashboard_settings = await _service_get_settings_cache().get() runtime_config = _http_bridge_runtime_config(dashboard_settings, _service_get_settings()) + if deferred_account_backoff_tracker is None: + deferred_account_backoff_tracker = _DeferredAccountBackoffTracker() bridge_payload = payload.to_payload() bridge_client_metadata = _response_create_client_metadata( bridge_payload, @@ -940,6 +972,33 @@ async def _stream_via_http_bridge( if bridge_client_metadata is not None or "client_metadata" in bridge_payload: payload = payload.model_copy(update={"client_metadata": bridge_client_metadata}) + def begin_bridge_lifecycle( + reservation: ApiKeyUsageReservationData | None, + ) -> _DeferredAccountBackoffLifecycle: + previous_lifecycle = deferred_account_backoff_tracker.current_lifecycle + same_reservation = bool( + previous_lifecycle is not None + and ( + previous_lifecycle.reservation is reservation + or ( + previous_lifecycle.reservation is not None + and reservation is not None + and previous_lifecycle.reservation.reservation_id == reservation.reservation_id + ) + ) + ) + pending_backoffs = ( + previous_lifecycle.pending_backoffs + if previous_lifecycle is not None and not previous_lifecycle.settlement_owned and same_reservation + else {} + ) + lifecycle = _DeferredAccountBackoffLifecycle( + reservation=reservation, + pending_backoffs=pending_backoffs, + ) + deferred_account_backoff_tracker.current_lifecycle = lifecycle + return lifecycle + def prepare_bridge_request( request_payload: ResponsesRequest, *, @@ -966,8 +1025,25 @@ def prepare_bridge_request( ) request_state.capacity_startup_wait_event = capacity_startup_wait_event request_state.capacity_startup_ready_event = capacity_startup_ready_event + lifecycle = begin_bridge_lifecycle(request_state.api_key_reservation) + request_state.deferred_account_error_backoffs = lifecycle.pending_backoffs + request_state.deferred_account_backoff_tracker = deferred_account_backoff_tracker + request_state.deferred_account_backoff_lifecycle = lifecycle return request_state, text_data + async def release_unowned_bridge_lifecycle( + lifecycle: _DeferredAccountBackoffLifecycle | None, + request_state: _WebSocketRequestState | None, + ) -> None: + if lifecycle is None or lifecycle.settlement_owned: + return + if request_state is not None: + await self._release_websocket_request_state_reservation(request_state) + return + await self._release_websocket_reservation(lifecycle.reservation) + lifecycle.settlement_confirmed = True + await self._drain_deferred_account_error_backoffs(lifecycle.pending_backoffs) + incoming_turn_state_header = _sticky_key_from_turn_state_header(headers) if not forwarded_request else None incoming_session_header = _sticky_key_from_session_header(headers) if not forwarded_request else None explicit_prompt_cache_key = _prompt_cache_key_from_request_model(payload) @@ -1749,6 +1825,8 @@ def switch_to_account_neutral_replay() -> None: request_deadline=request_deadline, session_header_fallback_key=session_header_fallback_key, exclude_account_ids=fresh_replay_excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as exc: if not owner_unavailable_allows_account_neutral_replay(exc): @@ -2023,6 +2101,8 @@ def switch_to_account_neutral_replay() -> None: session_header_fallback_key=session_header_fallback_key, request_deadline=request_deadline, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: if owner_unavailable_allows_account_neutral_replay(capacity_exc): @@ -2130,6 +2210,7 @@ def switch_to_account_neutral_replay() -> None: request_scope_id=owner_recovery_scope_id, ) retry_request_state: _WebSocketRequestState | None = None + retry_unowned_lifecycle: _DeferredAccountBackoffLifecycle | None = None try: retry_api_key_reservation = api_key_reservation retry_reservation_reacquired = False @@ -2143,6 +2224,7 @@ def switch_to_account_neutral_replay() -> None: request_usage_budget=estimate_api_key_request_usage(recovery_payload), ) retry_reservation_reacquired = True + retry_unowned_lifecycle = begin_bridge_lifecycle(retry_api_key_reservation) retry_request_state, retry_text_data = prepare_bridge_request( recovery_payload, @@ -2185,7 +2267,18 @@ def switch_to_account_neutral_replay() -> None: yield event_block except BaseException: if retry_reservation_reacquired and retry_api_key_reservation is not None: - await self._release_websocket_reservation(retry_api_key_reservation) + retry_lifecycle = ( + retry_request_state.deferred_account_backoff_lifecycle + if retry_request_state is not None + else retry_unowned_lifecycle + ) + try: + await release_unowned_bridge_lifecycle(retry_lifecycle, retry_request_state) + except Exception: + logger.warning( + "Failed to release owner-recovery HTTP bridge reservation", + exc_info=True, + ) raise finally: if owner_recovery_scope_id is not None: @@ -2508,6 +2601,8 @@ async def rollback_pre_dispatch_recovery_claim() -> None: request_deadline=request_deadline, session_header_fallback_key=session_header_fallback_key, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: wait_plan = _http_bridge_capacity_wait_plan(capacity_exc, request_deadline=request_deadline) @@ -2612,6 +2707,8 @@ async def rollback_pre_dispatch_recovery_claim() -> None: request_usage_budget=request_state.request_usage_budget, request_deadline=request_deadline, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: wait_plan = _http_bridge_capacity_wait_plan(capacity_exc, request_deadline=request_deadline) @@ -2860,6 +2957,8 @@ async def rollback_pre_dispatch_recovery_claim() -> None: request_usage_budget=estimate_api_key_request_usage(retry_payload), request_deadline=request_deadline, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: wait_plan = _http_bridge_capacity_wait_plan(capacity_exc, request_deadline=request_deadline) @@ -2896,6 +2995,8 @@ async def rollback_pre_dispatch_recovery_claim() -> None: session, request_scope_id=local_recovery_scope_id, ) + retry_request_state: _WebSocketRequestState | None = None + retry_unowned_lifecycle: _DeferredAccountBackoffLifecycle | None = None try: retry_api_key_reservation = api_key_reservation retry_reservation_reacquired = False @@ -2909,6 +3010,7 @@ async def rollback_pre_dispatch_recovery_claim() -> None: request_usage_budget=estimate_api_key_request_usage(retry_payload), ) retry_reservation_reacquired = True + retry_unowned_lifecycle = begin_bridge_lifecycle(retry_api_key_reservation) retry_request_state, retry_text_data = prepare_bridge_request( retry_payload, @@ -2950,7 +3052,18 @@ async def rollback_pre_dispatch_recovery_claim() -> None: except BaseException: await rollback_pre_dispatch_recovery_claim() if retry_reservation_reacquired and retry_api_key_reservation is not None: - await self._release_websocket_reservation(retry_api_key_reservation) + retry_lifecycle = ( + retry_request_state.deferred_account_backoff_lifecycle + if retry_request_state is not None + else retry_unowned_lifecycle + ) + try: + await release_unowned_bridge_lifecycle(retry_lifecycle, retry_request_state) + except Exception: + logger.warning( + "Failed to release local-recovery HTTP bridge reservation", + exc_info=True, + ) raise finally: if local_recovery_scope_id is not None: @@ -3151,6 +3264,9 @@ async def startup_continuity_cooldown_terminal_event() -> str | None: text_data=text_data, queue_limit=queue_limit, ) + lifecycle = request_state.deferred_account_backoff_lifecycle + if lifecycle is not None: + lifecycle.settlement_owned = True except ProxyResponseError as exc: if request_state.bridge_soft_capacity_reroute_allowed: raise diff --git a/app/modules/proxy/_service/streaming/protocol.py b/app/modules/proxy/_service/streaming/protocol.py index 8e3d17a6b8..279a59390a 100644 --- a/app/modules/proxy/_service/streaming/protocol.py +++ b/app/modules/proxy/_service/streaming/protocol.py @@ -6,6 +6,7 @@ class _StreamingServiceProtocol(Protocol): _acquire_account_response_create_lease_or_overload: Any _cancel_api_key_reservation_heartbeat_task: Any + _drain_deferred_account_error_backoffs: Any _encryptor: Any _ensure_fresh_with_budget: Any _get_work_admission: Any diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index 8f073ef8a6..46a610b37b 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -373,6 +373,7 @@ async def _stream_with_retry( require_preferred_account = False last_retryable_stream_error: _RetryableStreamError | None = None pending_post_refresh_transient_penalties: list[tuple[Account, UpstreamError, str, int, int]] = [] + deferred_account_error_backoffs: dict[str, Account] = {} post_refresh_transient_replacement_selected = False require_security_work_authorized = False account_leases: list[AccountLease] = [] @@ -424,18 +425,21 @@ async def _settle_stream_usage_before_pending_penalty( apply_pending_penalty = post_refresh_transient_replacement_selected and bool( pending_post_refresh_transient_penalties ) + wait_for_health_write = apply_pending_penalty or bool(deferred_account_error_backoffs) + settle_kwargs = {"wait_for_settlement": True} if wait_for_health_write else {} + settled_result = await proxy._settle_stream_api_key_usage( + api_key, + api_key_reservation, + current_settlement, + request_id, + **settle_kwargs, + ) + if not settled_result: + return False + await proxy._drain_deferred_account_error_backoffs(deferred_account_error_backoffs) if apply_pending_penalty: - settled_result = await proxy._settle_stream_api_key_usage( - api_key, - api_key_reservation, - current_settlement, - request_id, - wait_for_settlement=True, - ) pending_penalties = list(pending_post_refresh_transient_penalties) pending_post_refresh_transient_penalties.clear() - if not settled_result: - return False for pending_penalty in pending_penalties: ( failed_account, @@ -452,23 +456,24 @@ async def _settle_stream_usage_before_pending_penalty( ) if transient_retry_count > 1: await proxy._load_balancer.record_errors(failed_account, transient_retry_count - 1) - return True - return await proxy._settle_stream_api_key_usage( - api_key, - api_key_reservation, - current_settlement, - request_id, - ) + return settled_result + + async def _record_or_defer_confirmed_route_backoff(account: Account) -> None: + if api_key is not None and api_key_reservation is not None: + deferred_account_error_backoffs.setdefault(account.id, account) + return + await proxy._load_balancer.record_error_backoff(account) async def _drain_pending_post_refresh_penalty_on_terminal( current_settlement: _StreamSettlement, ) -> bool: nonlocal post_refresh_transient_replacement_selected, settled - if pending_post_refresh_transient_penalties: + if pending_post_refresh_transient_penalties or deferred_account_error_backoffs: # A failed replacement selection still ends the request. Mark # it as terminal so the deferred failure is settled and # recorded before this path returns or re-raises. - post_refresh_transient_replacement_selected = True + if pending_post_refresh_transient_penalties: + post_refresh_transient_replacement_selected = True settled = await _settle_stream_usage_before_pending_penalty(current_settlement) return settled return True @@ -580,6 +585,12 @@ async def _iter_stream_once() -> AsyncIterator[str]: ): yield line except ProxyResponseError as exc: + if is_confirmed_pre_dispatch_transport_error(exc): + # Keep dispatch provenance intact for the outer account + # failover handler. Converting this into the generic + # transient wrapper would authorize same-account replay + # and lose the confirmed dead-route backoff semantics. + raise error = _parse_openai_error(exc.payload) error_code = _normalize_error_code( error.code if error else None, @@ -1349,6 +1360,13 @@ async def _retry_account_model_rejection( post_refresh_transient_replacement_selected = True account_id_value = account.id + if last_pre_dispatch_transport_error is not None: + # The preserved connect failure is only authoritative when + # replacement selection is empty. Once another account is + # actually attempted, its terminal outcome takes precedence. + if last_transient_exc is last_pre_dispatch_transport_error: + last_transient_exc = None + last_pre_dispatch_transport_error = None if last_account_model_rejection is not None and account.id != last_account_model_rejection_account_id: # The original 400 is only the fallback when account # selection cannot produce a replacement. Once this @@ -1946,7 +1964,7 @@ async def _retry_account_model_rejection( # error at a time. await _release_tracked_stream_lease(current_account_lease) current_account_lease = None - await proxy._load_balancer.record_error_backoff(account) + await _record_or_defer_confirmed_route_backoff(account) can_try_other_account = ( not require_preferred_account and account.id != file_preferred_account_id @@ -2493,6 +2511,56 @@ async def _retry_account_model_rejection( if _facade()._is_account_neutral_error_code(error_code): await _drain_pending_post_refresh_penalty_on_terminal(settlement) raise + if is_confirmed_pre_dispatch_transport_error(retry_exc): + # This retry still failed before dispatch. Handle + # the proven dead route before the generic + # failover policy: that policy does not know + # about hard account ownership and may otherwise + # cross a previous-response, turn-state, file, or + # single-account boundary. + await _release_tracked_stream_lease(current_account_lease) + current_account_lease = None + await _record_or_defer_confirmed_route_backoff(account) + last_transient_exc = retry_exc + last_pre_dispatch_transport_error = retry_exc + + verified_owner_replay_moved = False + if ( + attempt < max_attempts - 1 + and routing_strategy != "single_account" + and file_preferred_account_id is None + and turn_state_owner_account_id is None + ): + verified_owner_replay_moved = _move_verified_fresh_replay_from_owner( + account_id=account.id, + outcome="owner_post_refresh_proxy_connect_failure", + ) + + can_try_other_account = bool( + attempt < max_attempts - 1 + and routing_strategy != "single_account" + and file_preferred_account_id is None + and turn_state_owner_account_id is None + and not require_preferred_account + ) + if can_try_other_account: + excluded_account_ids.add(account.id) + if not verified_owner_replay_moved: + affinity = replace(affinity, reallocate_sticky=True) + _facade().logger.info( + "Retrying post-refresh stream after confirmed pre-dispatch proxy " + "connect failure request_id=%s account_id=%s attempt=%d", + request_id, + account.id, + attempt + 1, + ) + continue + + # Hard ownership or an exhausted attempt budget: + # stop here. The shared terminal path settles the + # reservation, drains the deferred backoff floor, + # and preserves this sanitized failure. + break current_error_payload = _upstream_error_from_openai(error) current_error_code = error_code or "upstream_error" classified = classify_upstream_failure( @@ -2522,24 +2590,12 @@ async def _retry_account_model_rejection( action, ) if action == "failover_next": - if is_confirmed_pre_dispatch_transport_error(retry_exc): - # Confirmed dead proxy route on the - # post-refresh attempt: release the lease - # first, then apply the bounded transient - # backoff floor instead of a single generic - # error, and preserve the sanitized failure - # for terminal rendering. - await _release_tracked_stream_lease(current_account_lease) - current_account_lease = None - await proxy._load_balancer.record_error_backoff(account) - last_pre_dispatch_transport_error = retry_exc - else: - await proxy._handle_stream_error( - account, - current_error_payload, - current_error_code, - http_status=retry_exc.status_code, - ) + await proxy._handle_stream_error( + account, + current_error_payload, + current_error_code, + http_status=retry_exc.status_code, + ) last_transient_exc = retry_exc await _release_tracked_stream_lease(current_account_lease) current_account_lease = None @@ -2574,7 +2630,12 @@ async def _retry_account_model_rejection( failed_account is account for failed_account, *_rest in pending_post_refresh_transient_penalties ) - health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement) + ordered_settlement_required = bool( + pending_post_refresh_transient_penalties or deferred_account_error_backoffs + ) + health_write_allowed = True + if ordered_settlement_required: + health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement) if ( health_write_allowed and settlement.account_health_error @@ -2587,7 +2648,11 @@ async def _retry_account_model_rejection( ) elif health_write_allowed and settlement.record_success: await proxy._load_balancer.record_success(account) - if not settled and not settlement.usage_settlement_transferred: + if ( + not settled + and not ordered_settlement_required + and not settlement.usage_settlement_transferred + ): settled = await _settle_stream_usage_before_pending_penalty(settlement) upstream_transport_metric_status = settlement.status _record_upstream_transport_metric_once(settlement.status) @@ -2774,11 +2839,17 @@ async def _retry_account_model_rejection( and api_key is not None and api_key_reservation is not None ): - release_coro = proxy._release_unsettled_stream_api_key_usage( - api_key=api_key, - api_key_reservation=api_key_reservation, - request_id=request_id, - ) + + async def _release_reservation_then_drain_backoffs() -> None: + released = await proxy._release_unsettled_stream_api_key_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + request_id=request_id, + ) + if released: + await proxy._drain_deferred_account_error_backoffs(deferred_account_error_backoffs) + + release_coro = _release_reservation_then_drain_backoffs() current_task = asyncio.current_task() if current_task is not None and current_task.cancelling(): proxy._schedule_cancel_safe_cleanup( @@ -2788,3 +2859,5 @@ async def _retry_account_model_rejection( ) else: await release_coro + elif settled and deferred_account_error_backoffs: + await proxy._drain_deferred_account_error_backoffs(deferred_account_error_backoffs) diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index dfa151ece4..017050f258 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -753,6 +753,19 @@ class _HTTPBridgeCompletedDeliveryScope: terminal_enqueued: bool = False +@dataclass(slots=True) +class _DeferredAccountBackoffLifecycle: + reservation: ApiKeyUsageReservationData | None + pending_backoffs: dict[str, Account] = field(default_factory=dict) + settlement_owned: bool = False + settlement_confirmed: bool = False + + +@dataclass(slots=True) +class _DeferredAccountBackoffTracker: + current_lifecycle: _DeferredAccountBackoffLifecycle | None = None + + @dataclass class _WebSocketRequestState: request_id: str @@ -920,6 +933,13 @@ class _WebSocketRequestState: client_ip: str | None = None downstream_visible: bool = False last_downstream_sequence_number: int | None = None + # Confirmed pre-dispatch account-route failures must not mutate account + # health while this request's API-key reservation is still live. The + # account objects are keyed by id so repeated connect attempts cannot + # stack the same backoff floor more than once before settlement. + deferred_account_error_backoffs: dict[str, Account] = field(default_factory=dict) + deferred_account_backoff_tracker: _DeferredAccountBackoffTracker | None = None + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None deferred_reasoning_downstream_texts: list[str] = field(default_factory=list) suppress_next_created_downstream: bool = False replay_downstream_response_id: str | None = None diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index f5e831f8da..e6680e7680 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -3005,6 +3005,13 @@ async def _connect_proxy_websocket( ) -> tuple[Account | None, UpstreamWebSocket | None]: proxy = cast(_WebSocketServiceProtocol, self) _ = proxy + + async def _record_or_defer_confirmed_route_backoff(account: Account) -> None: + if request_state.api_key_reservation is not None: + request_state.deferred_account_error_backoffs.setdefault(account.id, account) + return + await proxy._load_balancer.record_error_backoff(account) + if ( request_state.useragent is None and request_state.useragent_group is None @@ -3188,7 +3195,7 @@ async def _connect_proxy_websocket( 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 _record_or_defer_confirmed_route_backoff(account) last_failover_exc = exc last_failover_account = account excluded_account_ids.add(account.id) @@ -3199,7 +3206,7 @@ async def _connect_proxy_websocket( 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 _record_or_defer_confirmed_route_backoff(account) await proxy._emit_websocket_connect_failure( websocket, client_send_lock=client_send_lock, @@ -5394,8 +5401,7 @@ async def _finalize_websocket_request_state( if request_state.draining_until_terminal: await _release_websocket_response_create_gate(request_state, response_create_gate) - await proxy._release_websocket_reservation(request_state.api_key_reservation) - request_state.api_key_reservation = None + await proxy._release_websocket_request_state_reservation(request_state) # The reservation is settled; clear any terminal-bookkeeping # settlement claim so abort handling does not settle it again. request_state.terminal_settlement_phase = None @@ -5489,6 +5495,7 @@ async def _finalize_websocket_request_state( # persistence. The health write remains ordered below. upstream_control.reconnect_requested = True upstream_control.retire_after_drain = True + lifecycle = request_state.deferred_account_backoff_lifecycle settlement_committed = await proxy._settle_stream_api_key_usage( api_key, request_state.api_key_reservation, @@ -5496,13 +5503,27 @@ async def _finalize_websocket_request_state( response_id, # The reservation must be settled before the load-balancer # health write below (settlement-ordering invariant). - wait_for_settlement=settlement.account_health_error or settlement.record_success, + wait_for_settlement=( + lifecycle is not None + or settlement.account_health_error + or settlement.record_success + or bool(request_state.deferred_account_error_backoffs) + ), ) # Settlement responsibility has transferred (the settle path tracks # its own failure/cancellation fallback release). Clear any terminal- # bookkeeping settlement claim so a later abort of the surrounding # continuation does not race a duplicate release against it. request_state.terminal_settlement_phase = None + if settlement_committed: + request_state.api_key_reservation = None + if lifecycle is not None: + lifecycle.settlement_confirmed = True + pending_backoffs = ( + lifecycle.pending_backoffs if lifecycle is not None else request_state.deferred_account_error_backoffs + ) + if pending_backoffs: + await proxy._drain_deferred_account_error_backoffs(pending_backoffs) latency_ms = int((time.monotonic() - request_state.started_at) * 1000) cached_input_tokens = usage.input_tokens_details.cached_tokens if usage and usage.input_tokens_details else None reasoning_tokens = ( diff --git a/app/modules/proxy/_service/websocket/protocol.py b/app/modules/proxy/_service/websocket/protocol.py index 9eaed0669e..519cd7c868 100644 --- a/app/modules/proxy/_service/websocket/protocol.py +++ b/app/modules/proxy/_service/websocket/protocol.py @@ -12,6 +12,7 @@ class _WebSocketServiceProtocol(Protocol): _connect_proxy_websocket: Any _decide_websocket_failover_action: Any _downstream_websocket_is_idle: Any + _drain_deferred_account_error_backoffs: Any _emit_pending_websocket_keepalive: Any _emit_websocket_connect_failure: Any _emit_websocket_connect_timeout: Any diff --git a/openspec/changes/retry-account-proxy-connect-failures/design.md b/openspec/changes/retry-account-proxy-connect-failures/design.md index ddf55b8cc5..004f1b8e34 100644 --- a/openspec/changes/retry-account-proxy-connect-failures/design.md +++ b/openspec/changes/retry-account-proxy-connect-failures/design.md @@ -48,18 +48,32 @@ reallocated so the replacement selection does not immediately loop back. 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 (`record_error_backoff`, shared `ERROR_BACKOFF_THRESHOLD`): -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, and it replaces — not stacks onto — the generic single-error health -write for the same failure. +floor (`record_error_backoff`, shared `ERROR_BACKOFF_THRESHOLD`): 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, and it +replaces — not stacks onto — the generic single-error health write for the +same failure. 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. +reservation alive instead of releasing and racing to reacquire it. Because +account-health mutation must not race a live reservation, keyed requests queue +the dead-route floor and apply it only after the normal terminal finalizer has +settled or released that singular reservation. Startup failure and cancellation +paths release the reservation before draining the same queue. A failed terminal +settlement may fall back to release, but only a confirmed settlement or release +authorizes the queued health write; if both fail, the reservation and backoff +remain pending rather than racing an account-health mutation. + +The HTTP bridge transfers settlement ownership to a request state only after +its upstream submit call returns successfully. Outer startup cleanup releases +only the current unowned lifecycle. Each lifecycle generation owns its own +backoff queue, and terminal finalizers drain only that generation after its +settlement is confirmed. A newer retry therefore cannot drain an older failed +settlement's queue, and a late older finalizer cannot drain the newer queue. +Queue entries are claimed before the asynchronous health write so concurrent +terminal and outer cleanup cannot apply the same floor twice. ## Non-goals 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 index a05dbba559..0fc995b1b4 100644 --- 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 @@ -6,7 +6,7 @@ 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. +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, MUST settle or release at the terminal request outcome before the account-health write, and MUST NOT be reacquired solely for the internal failover. If neither settlement nor fallback release can be confirmed, the service MUST leave the health write unapplied. HTTP-bridge startup cleanup MUST release only an unowned current request lifecycle, and each reservation lifecycle MUST drain only its own health writes after confirmed settlement or release. 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. @@ -23,6 +23,7 @@ The service MUST NOT replay a request when dispatch is unknown or when the reque - **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 settles or releases any request-scoped API-key reservation before the account-health write - **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 diff --git a/openspec/changes/retry-account-proxy-connect-failures/tasks.md b/openspec/changes/retry-account-proxy-connect-failures/tasks.md index 25b68de4b2..6c41d2c2e6 100644 --- a/openspec/changes/retry-account-proxy-connect-failures/tasks.md +++ b/openspec/changes/retry-account-proxy-connect-failures/tasks.md @@ -5,5 +5,7 @@ - [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] Defer dead-route account-health writes until request-scoped API-key settlement. +- [x] Track HTTP-bridge settlement ownership per reservation generation and fail closed on unconfirmed release. - [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_proxy_websocket_responses.py b/tests/integration/test_proxy_websocket_responses.py index 3a2a9e0902..639d39bf71 100644 --- a/tests/integration/test_proxy_websocket_responses.py +++ b/tests/integration/test_proxy_websocket_responses.py @@ -606,7 +606,14 @@ def test_backend_responses_websocket_fails_over_confirmed_proxy_connect_before_d selection_exclusions: list[set[str]] = [] connect_accounts: list[str] = [] backed_off_accounts: list[str] = [] + settlement_order: list[str] = [] + settlement_wait_flags: list[bool] = [] handle_connect_error = AsyncMock() + reservation = proxy_module.ApiKeyUsageReservationData( + reservation_id="resv_ws_proxy_failover", + key_id="key_ws_proxy_failover", + model="gpt-5.4", + ) proxy_connect_error = proxy_module.ProxyResponseError( 502, proxy_module.openai_error("upstream_unavailable", "sanitized websocket proxy failure"), @@ -641,8 +648,20 @@ async def fake_try_open_websocket_connect_attempt(self, account, headers, **kwar async def fake_record_error_backoff(self, account): del self + assert settlement_order == ["settle"] + settlement_order.append("backoff") backed_off_accounts.append(account.id) + async def fake_reserve_websocket_api_key_usage(self, *_args, **_kwargs): + del self + return reservation + + async def fake_settle_stream_api_key_usage(self, *_args, **kwargs): + del self + settlement_wait_flags.append(bool(kwargs.get("wait_for_settlement"))) + settlement_order.append("settle") + return True + 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()) @@ -657,6 +676,16 @@ async def fake_record_error_backoff(self, account): fake_try_open_websocket_connect_attempt, ) monkeypatch.setattr(proxy_module.ProxyService, "_handle_websocket_connect_error", handle_connect_error) + monkeypatch.setattr( + proxy_module.ProxyService, + "_reserve_websocket_api_key_usage", + fake_reserve_websocket_api_key_usage, + ) + monkeypatch.setattr( + proxy_module.ProxyService, + "_settle_stream_api_key_usage", + fake_settle_stream_api_key_usage, + ) monkeypatch.setattr(proxy_module.LoadBalancer, "record_error_backoff", fake_record_error_backoff) request_payload = { @@ -678,6 +707,8 @@ async def fake_record_error_backoff(self, account): assert connect_accounts == [accounts[0].id, accounts[1].id] assert selection_exclusions == [set(), {accounts[0].id}] assert backed_off_accounts == [accounts[0].id] + assert settlement_order == ["settle", "backoff"] + assert settlement_wait_flags == [True] # The confirmed dead route skips the generic single-error health write. handle_connect_error.assert_not_awaited() diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index bbf6cc1907..b1c31c3f03 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -524,6 +524,60 @@ async def test_submit_http_bridge_request_early_failure_releases_published_hando assert session.unanchored_reservation_id is None +@pytest.mark.asyncio +@pytest.mark.parametrize("submit_succeeds", [True, False]) +async def test_http_bridge_submit_transfers_settlement_ownership_only_after_success( + monkeypatch: pytest.MonkeyPatch, + submit_succeeds: bool, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv-bridge-submit-owner", + key_id="key-bridge-submit-owner", + model="gpt-5.6-sol", + ) + lifecycle = proxy_support_module._DeferredAccountBackoffLifecycle(reservation=reservation) + request_state = proxy_service._WebSocketRequestState( + request_id="req-bridge-submit-owner", + model="gpt-5.6-sol", + service_tier=None, + reasoning_effort=None, + api_key_reservation=reservation, + started_at=time.monotonic(), + event_queue=asyncio.Queue(), + deferred_account_backoff_lifecycle=lifecycle, + ) + + async def submit(*_args: object, **_kwargs: object) -> None: + if not submit_succeeds: + raise RuntimeError("send failed before submit returned") + assert request_state.event_queue is not None + request_state.event_queue.put_nowait(None) + + monkeypatch.setattr(service, "_submit_http_bridge_request", submit) + monkeypatch.setattr(service, "_detach_http_bridge_request", AsyncMock(return_value=False)) + monkeypatch.setattr(service, "_maybe_release_idle_http_bridge_session_lease", AsyncMock(return_value=False)) + stream = service._stream_http_bridge_session_events( + session, + request_state=request_state, + text_data="{}", + queue_limit=4, + propagate_http_errors=True, + downstream_turn_state=None, + request_deadline=time.monotonic() + 10, + ) + + if submit_succeeds: + assert [event async for event in stream] == [] + else: + with pytest.raises(RuntimeError, match="send failed"): + async for _ in stream: + pass + + assert lifecycle.settlement_owned is submit_succeeds + + @pytest.mark.asyncio async def test_http_bridge_request_cleanup_releases_pre_submit_handoff( monkeypatch: pytest.MonkeyPatch, @@ -579,6 +633,163 @@ async def fail_before_submit(*args: object, **kwargs: object): assert session.unanchored_reservation_id is None +@pytest.mark.asyncio +async def test_http_bridge_owned_terminal_lifecycle_skips_outer_startup_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv-bridge-owned", + key_id="key-bridge-owned", + model="gpt-5.6-sol", + ) + account = cast(Any, SimpleNamespace(id="acc-bridge-owned")) + runtime_config = SimpleNamespace( + enabled=True, + idle_ttl_seconds=120.0, + codex_idle_ttl_seconds=1800.0, + max_sessions=8, + queue_limit=4, + prompt_cache_idle_ttl_seconds=120.0, + ) + + async def terminal_before_finalizer(*_args: object, **kwargs: object): + tracker = cast( + proxy_support_module._DeferredAccountBackoffTracker, + kwargs["deferred_account_backoff_tracker"], + ) + lifecycle = proxy_support_module._DeferredAccountBackoffLifecycle( + reservation=reservation, + pending_backoffs={account.id: account}, + settlement_owned=True, + ) + tracker.current_lifecycle = lifecycle + yield 'data: {"type":"response.completed"}\n\n' + + monkeypatch.setattr( + http_bridge_streaming_module, + "_service_get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=SimpleNamespace())), + ) + monkeypatch.setattr(http_bridge_streaming_module, "_service_get_settings", _make_app_settings) + monkeypatch.setattr(http_bridge_streaming_module, "_http_bridge_runtime_config", lambda *_args: runtime_config) + monkeypatch.setattr(service, "_resolve_file_account_for_responses", AsyncMock(return_value=None)) + monkeypatch.setattr(service, "_stream_via_http_bridge", terminal_before_finalizer) + release_reservation = AsyncMock() + drain_backoffs = AsyncMock() + monkeypatch.setattr(service, "_release_websocket_reservation", release_reservation) + monkeypatch.setattr(service, "_drain_deferred_account_error_backoffs", drain_backoffs) + payload = proxy_service.ResponsesRequest.model_validate( + {"model": "gpt-5.6-sol", "instructions": "test", "input": "hello"} + ) + + chunks = [ + chunk + async for chunk in service._stream_http_bridge_or_retry( + payload, + {}, + codex_session_affinity=True, + propagate_http_errors=True, + openai_cache_affinity=False, + api_key=None, + api_key_reservation=reservation, + suppress_text_done_events=False, + ) + ] + + assert chunks == ['data: {"type":"response.completed"}\n\n'] + release_reservation.assert_not_awaited() + drain_backoffs.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_http_bridge_startup_fallback_releases_current_lifecycle_before_backoff( + monkeypatch: pytest.MonkeyPatch, + release_fails: bool, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + original_reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv-bridge-original", + key_id="key-bridge-startup", + model="gpt-5.6-sol", + ) + current_reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv-bridge-current", + key_id="key-bridge-startup", + model="gpt-5.6-sol", + ) + account = cast(Any, SimpleNamespace(id="acc-bridge-startup")) + runtime_config = SimpleNamespace( + enabled=True, + idle_ttl_seconds=120.0, + codex_idle_ttl_seconds=1800.0, + max_sessions=8, + queue_limit=4, + prompt_cache_idle_ttl_seconds=120.0, + ) + tracker_seen: proxy_support_module._DeferredAccountBackoffTracker | None = None + order: list[str] = [] + + async def fail_before_submit(*_args: object, **kwargs: object): + nonlocal tracker_seen + tracker_seen = cast( + proxy_support_module._DeferredAccountBackoffTracker, + kwargs["deferred_account_backoff_tracker"], + ) + tracker_seen.current_lifecycle = proxy_support_module._DeferredAccountBackoffLifecycle( + reservation=current_reservation, + pending_backoffs={account.id: account}, + ) + raise RuntimeError("startup failed") + yield "" + + async def release_reservation(reservation: object) -> None: + assert reservation is current_reservation + order.append("release") + if release_fails: + raise RuntimeError("release failed") + + async def drain_backoffs(pending: dict[str, object]) -> None: + assert order == ["release"] + order.append("backoff") + pending.clear() + + monkeypatch.setattr( + http_bridge_streaming_module, + "_service_get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=SimpleNamespace())), + ) + monkeypatch.setattr(http_bridge_streaming_module, "_service_get_settings", _make_app_settings) + monkeypatch.setattr(http_bridge_streaming_module, "_http_bridge_runtime_config", lambda *_args: runtime_config) + monkeypatch.setattr(service, "_resolve_file_account_for_responses", AsyncMock(return_value=None)) + monkeypatch.setattr(service, "_stream_via_http_bridge", fail_before_submit) + monkeypatch.setattr(service, "_release_websocket_reservation", release_reservation) + monkeypatch.setattr(service, "_drain_deferred_account_error_backoffs", drain_backoffs) + payload = proxy_service.ResponsesRequest.model_validate( + {"model": "gpt-5.6-sol", "instructions": "test", "input": "hello"} + ) + + with pytest.raises(RuntimeError, match="startup failed"): + async for _ in service._stream_http_bridge_or_retry( + payload, + {}, + codex_session_affinity=True, + propagate_http_errors=True, + openai_cache_affinity=False, + api_key=None, + api_key_reservation=original_reservation, + suppress_text_done_events=False, + ): + pass + + assert tracker_seen is not None + assert order == (["release"] if release_fails else ["release", "backoff"]) + assert tracker_seen.current_lifecycle is not None + assert tracker_seen.current_lifecycle.settlement_confirmed is not release_fails + assert bool(tracker_seen.current_lifecycle.pending_backoffs) is release_fails + + @pytest.mark.asyncio async def test_durable_turn_state_fence_rejection_rolls_back_local_alias( monkeypatch: pytest.MonkeyPatch, @@ -6262,7 +6473,7 @@ def _bridge_selection_settings() -> SimpleNamespace: @pytest.mark.asyncio -async def test_create_http_bridge_session_fails_over_confirmed_proxy_connect_after_lease_release( +async def test_create_http_bridge_session_defers_confirmed_proxy_backoff_until_reservation_release( monkeypatch: pytest.MonkeyPatch, ) -> None: service = proxy_service.ProxyService(cast(Any, nullcontext())) @@ -6274,6 +6485,13 @@ async def test_create_http_bridge_session_fails_over_confirmed_proxy_connect_aft reallocate_flags: list[bool] = [] released_leases: list[proxy_service.AccountLease] = [] backed_off_accounts: list[object] = [] + settlement_order: list[str] = [] + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv-http-bridge-proxy-failover", + key_id="key-http-bridge-proxy-failover", + model="gpt-5.4", + ) + lifecycle = proxy_support_module._DeferredAccountBackoffLifecycle(reservation=reservation) upstream = cast(Any, SimpleNamespace(response_header=lambda _name: None, close=AsyncMock())) async def select_account(_deadline: float, **kwargs: object) -> proxy_service.AccountSelection: @@ -6293,6 +6511,12 @@ async def record_error_backoff(account: object) -> None: backed_off_accounts.append(account) # The dead route's stream lease must settle before the health write. assert lease_a in released_leases + assert settlement_order == ["settle"] + settlement_order.append("backoff") + + async def release_reservation(candidate: object) -> None: + assert candidate is reservation + settlement_order.append("settle") monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) monkeypatch.setattr( @@ -6310,6 +6534,7 @@ async def record_error_backoff(account: object) -> None: 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._load_balancer, "record_error", AsyncMock()) + monkeypatch.setattr(service, "_release_websocket_reservation", release_reservation) monkeypatch.setattr(service, "_relay_http_bridge_upstream_messages", AsyncMock()) session = await service._create_http_bridge_session( @@ -6319,12 +6544,21 @@ async def record_error_backoff(account: object) -> None: api_key=None, request_model="gpt-5.4", idle_ttl_seconds=120.0, + deferred_account_backoff_lifecycle=lifecycle, + defer_account_health_writes=True, ) + assert backed_off_accounts == [] + assert lifecycle.pending_backoffs == {account_a.id: account_a} + await service._release_websocket_reservation(reservation) + lifecycle.settlement_confirmed = True + await service._drain_deferred_account_error_backoffs(lifecycle.pending_backoffs) + assert session.account is account_b assert selections == [set(), {account_a.id}] assert reallocate_flags == [False, True] assert backed_off_accounts == [account_a] + assert settlement_order == ["settle", "backoff"] assert lease_a in released_leases assert lease_b not in released_leases diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index ffffd2e363..03f94b3e6d 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -12843,11 +12843,9 @@ async def fake_stream_once(*_args: object, **_kwargs: object): @pytest.mark.asyncio -@pytest.mark.parametrize("post_refresh_failure", ["model_capacity", "connect"]) @pytest.mark.parametrize("upstream_transport", ["http", "websocket"]) -async def test_stream_with_retry_post_refresh_transport_failure_retries_same_account( +async def test_stream_with_retry_post_refresh_model_capacity_retries_same_account( monkeypatch, - post_refresh_failure: str, upstream_transport: str, ): settings = _make_proxy_settings() @@ -12880,13 +12878,6 @@ async def fake_stream_once(*_args: object, **kwargs: object): ) if stream_once_calls == 2: assert kwargs["allow_transient_retry"] is True - if post_refresh_failure == "connect": - raise proxy_module.ProxyResponseError( - 502, - proxy_module.openai_error("upstream_unavailable", "Server disconnected"), - failure_phase="connect", - retryable_same_contract=True, - ) raise proxy_module.ProxyResponseError( 400, proxy_module.openai_error( @@ -12922,6 +12913,403 @@ async def fake_stream_once(*_args: object, **kwargs: object): assert sleeps +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_transport", "settlement_confirmed"), + [ + ("http", True), + ("websocket", True), + ("http", False), + ("websocket", False), + ], +) +async def test_stream_with_retry_post_refresh_confirmed_proxy_connect_failure_fails_over_after_settlement( + monkeypatch, + upstream_transport: str, + settlement_confirmed: bool, +): + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account_a = _make_account("acc_post_refresh_proxy_connect_a") + account_b = _make_account("acc_post_refresh_proxy_connect_b") + api_key = _make_api_key_data("key_post_refresh_proxy_connect") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv_post_refresh_proxy_connect", + key_id=api_key.id, + model="gpt-5.1", + ) + selection_exclusions: list[set[str]] = [] + stream_account_ids: list[str] = [] + settlement_order: list[str] = [] + settlement_wait_flags: list[bool] = [] + sleeps: list[float] = [] + handle_stream_error = AsyncMock() + record_success = AsyncMock() + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service, "_STREAM_MAX_ACCOUNT_ATTEMPTS", 2) + + async def fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr(streaming_retry_module.asyncio, "sleep", fake_sleep) + + async def select_account(_deadline: float, **kwargs: object) -> AccountSelection: + excluded = set(cast(set[str], kwargs["exclude_account_ids"])) + selection_exclusions.append(excluded) + return AccountSelection( + account=account_b if account_a.id in excluded else account_a, + error_message=None, + ) + + async def fake_stream_once(account: Account, *_args: object, **kwargs: object): + assert kwargs["upstream_stream_transport"] == upstream_transport + stream_account_ids.append(account.id) + if stream_account_ids == [account_a.id]: + raise proxy_module.ProxyResponseError( + 401, + proxy_module.openai_error("invalid_api_key", "expired", error_type="invalid_request_error"), + ) + if stream_account_ids == [account_a.id, account_a.id]: + raise _pre_dispatch_proxy_connect_error("post-refresh proxy route unavailable") + yield 'data: {"type":"response.completed","response":{"id":"resp_post_refresh_proxy_ok"}}\n\n' + + async def settle_usage(*_args: object, **kwargs: object) -> bool: + settlement = cast(proxy_service._StreamSettlement, _args[2]) + settlement.usage_settlement_transferred = True + settlement_wait_flags.append(bool(kwargs.get("wait_for_settlement"))) + settlement_order.append("settle") + return settlement_confirmed + + async def record_error_backoff(account: Account) -> None: + assert account is account_a + assert settlement_order == ["settle"] + settlement_order.append("backoff") + + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(side_effect=lambda account, **_k: account)) + monkeypatch.setattr(service, "_stream_once", fake_stream_once) + monkeypatch.setattr(service, "_settle_stream_api_key_usage", settle_usage) + monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service._load_balancer, "record_success", record_success) + + 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-post-refresh-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", + upstream_stream_transport_override=upstream_transport, + ) + ] + + terminal = json.loads(chunks[-1].split("data: ", 1)[1]) + assert terminal["response"]["id"] == "resp_post_refresh_proxy_ok" + assert stream_account_ids == [account_a.id, account_a.id, account_b.id] + assert selection_exclusions == [set(), {account_a.id}] + assert settlement_order == (["settle", "backoff"] if settlement_confirmed else ["settle"]) + assert settlement_wait_flags == [True] + assert sleeps == [] + assert all(call.args[2] != "upstream_unavailable" for call in handle_stream_error.await_args_list) + if settlement_confirmed: + record_success.assert_awaited_once_with(account_b) + else: + record_success.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stream_with_retry_post_refresh_confirmed_proxy_connect_failure_moves_verified_fresh_replay( + monkeypatch, +): + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + owner = _make_account("acc_post_refresh_verified_owner") + replacement = _make_account("acc_post_refresh_verified_replacement") + previous_response_id = "resp_post_refresh_verified_owner" + original_payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.1", + "instructions": "hi", + "input": [{"role": "user", "content": "full replay"}], + "previous_response_id": previous_response_id, + "stream": True, + } + ) + fresh_payload = original_payload.model_copy(update={"previous_response_id": None}) + streamed_previous_response_ids: list[str | None] = [] + selection_exclusions: list[set[str]] = [] + handle_stream_error = AsyncMock() + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service, "_STREAM_MAX_ACCOUNT_ATTEMPTS", 2) + monkeypatch.setattr(service, "_resolve_websocket_previous_response_owner", AsyncMock(return_value=owner.id)) + monkeypatch.setattr( + streaming_retry_module, + "_verified_cross_transport_fresh_replay", + lambda *_args, **_kwargs: fresh_payload, + ) + + async def select_account(_deadline: float, **kwargs: object) -> AccountSelection: + excluded = set(cast(set[str], kwargs["exclude_account_ids"])) + selection_exclusions.append(excluded) + return AccountSelection(account=replacement if owner.id in excluded else owner, error_message=None) + + async def fake_stream_once(account: Account, payload: ResponsesRequest, *_args: object, **_kwargs: object): + streamed_previous_response_ids.append(payload.previous_response_id) + if len(streamed_previous_response_ids) == 1: + raise proxy_module.ProxyResponseError( + 401, + proxy_module.openai_error("invalid_api_key", "expired", error_type="invalid_request_error"), + ) + if account is owner: + raise _pre_dispatch_proxy_connect_error("verified owner proxy route unavailable") + yield 'data: {"type":"response.completed","response":{"id":"resp_post_refresh_verified_ok"}}\n\n' + + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(side_effect=lambda account, **_k: account)) + monkeypatch.setattr(service, "_stream_once", fake_stream_once) + monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_success", AsyncMock()) + + chunks = [ + chunk + async for chunk in service._stream_with_retry( + original_payload, + {"session_id": "sid-post-refresh-verified-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", + upstream_stream_transport_override="http", + ) + ] + + assert json.loads(chunks[-1].split("data: ", 1)[1])["response"]["id"] == "resp_post_refresh_verified_ok" + assert streamed_previous_response_ids == [previous_response_id, previous_response_id, None] + assert selection_exclusions == [set(), {owner.id}] + assert all(call.args[2] != "upstream_unavailable" for call in handle_stream_error.await_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ownership", ["previous_response", "turn_state", "file", "single_account"]) +async def test_stream_with_retry_post_refresh_confirmed_proxy_connect_failure_keeps_hard_owner( + monkeypatch, + ownership: str, +): + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + owner = _make_account(f"acc_post_refresh_proxy_owner_{ownership}") + alternate = _make_account(f"acc_post_refresh_proxy_alternate_{ownership}") + selections: list[set[str]] = [] + stream_account_ids: list[str] = [] + record_error_backoff = AsyncMock() + handle_stream_error = AsyncMock() + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service, "_STREAM_MAX_ACCOUNT_ATTEMPTS", 2) + if ownership == "previous_response": + monkeypatch.setattr(service, "_resolve_websocket_previous_response_owner", AsyncMock(return_value=owner.id)) + elif ownership == "turn_state": + monkeypatch.setattr(service, "_resolve_compact_turn_state_owner", AsyncMock(return_value=owner.id)) + elif ownership == "single_account": + settings.routing_strategy = "single_account" + + async def select_account(_deadline: float, **kwargs: object) -> AccountSelection: + excluded = set(cast(set[str], kwargs["exclude_account_ids"])) + selections.append(excluded) + return AccountSelection(account=owner if len(selections) == 1 else alternate, error_message=None) + + async def fake_stream_once(account: Account, *_args: object, **_kwargs: object): + stream_account_ids.append(account.id) + if len(stream_account_ids) == 1: + raise proxy_module.ProxyResponseError( + 401, + proxy_module.openai_error("invalid_api_key", "expired", error_type="invalid_request_error"), + ) + raise _pre_dispatch_proxy_connect_error("hard owner proxy route unavailable") + yield "" + + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(side_effect=lambda account, **_k: account)) + monkeypatch.setattr(service, "_stream_once", fake_stream_once) + monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + + payload_data: dict[str, object] = { + "model": "gpt-5.1", + "instructions": "hi", + "input": [], + "stream": True, + } + if ownership == "previous_response": + payload_data["previous_response_id"] = "resp_post_refresh_proxy_owner" + payload = ResponsesRequest.model_validate(payload_data) + headers = {"session_id": f"sid-post-refresh-proxy-owner-{ownership}"} + if ownership == "turn_state": + headers["x-codex-turn-state"] = "turn-post-refresh-proxy-owner" + chunks = [ + chunk + async for chunk in service._stream_with_retry( + payload, + headers, + 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", + upstream_stream_transport_override="http", + rewritten_file_account_id=owner.id if ownership == "file" else None, + ) + ] + + failed = json.loads(chunks[-1].split("data: ", 1)[1]) + assert failed["response"]["error"]["message"] == "hard owner proxy route unavailable" + assert selections == [set()] + assert stream_account_ids == [owner.id, owner.id] + record_error_backoff.assert_awaited_once_with(owner) + assert all(call.args[2] != "upstream_unavailable" for call in handle_stream_error.await_args_list) + + +@pytest.mark.asyncio +async def test_stream_with_retry_post_refresh_confirmed_proxy_connect_failure_records_floor_without_attempts( + monkeypatch, +): + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_post_refresh_proxy_no_attempts") + stream_calls = 0 + record_error_backoff = AsyncMock() + handle_stream_error = AsyncMock() + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service, "_STREAM_MAX_ACCOUNT_ATTEMPTS", 1) + monkeypatch.setattr( + service, + "_select_account_with_budget_compatible", + AsyncMock(return_value=AccountSelection(account=account, error_message=None)), + ) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(side_effect=lambda target, **_k: target)) + + async def fake_stream_once(*_args: object, **_kwargs: object): + nonlocal stream_calls + stream_calls += 1 + if stream_calls == 1: + raise proxy_module.ProxyResponseError( + 401, + proxy_module.openai_error("invalid_api_key", "expired", error_type="invalid_request_error"), + ) + raise _pre_dispatch_proxy_connect_error("post-refresh proxy exhausted") + yield "" + + monkeypatch.setattr(service, "_stream_once", fake_stream_once) + monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + + 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-post-refresh-proxy-no-attempts"}, + 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", + upstream_stream_transport_override="http", + ) + ] + + failed = json.loads(chunks[-1].split("data: ", 1)[1]) + assert failed["response"]["error"]["message"] == "post-refresh proxy exhausted" + record_error_backoff.assert_awaited_once_with(account) + assert all(call.args[2] != "upstream_unavailable" for call in handle_stream_error.await_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("propagate_http_errors", [False, True]) +async def test_stream_with_retry_replacement_failure_supersedes_prior_proxy_connect_error( + monkeypatch, + propagate_http_errors: bool, +): + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account_a = _make_account("acc_proxy_connect_stale_a") + account_b = _make_account("acc_proxy_connect_stale_b") + stream_account_ids: list[str] = [] + + monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) + monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr(proxy_service, "_STREAM_MAX_ACCOUNT_ATTEMPTS", 2) + + async def select_account(_deadline: float, **kwargs: object) -> AccountSelection: + excluded = cast(set[str], kwargs["exclude_account_ids"]) + return AccountSelection(account=account_b if account_a.id in excluded else account_a, error_message=None) + + async def fake_stream_once(account: Account, *_args: object, **_kwargs: object): + stream_account_ids.append(account.id) + if stream_account_ids == [account_a.id]: + raise proxy_module.ProxyResponseError( + 401, + proxy_module.openai_error("invalid_api_key", "expired", error_type="invalid_request_error"), + ) + if account is account_a: + raise _pre_dispatch_proxy_connect_error("first account proxy route unavailable") + raise proxy_service._RetryableStreamError( + "replacement_failed", + cast(UpstreamError, {"message": "replacement account failed"}), + exclude_account=True, + ) + yield "" + + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", AsyncMock(side_effect=lambda account, **_k: account)) + monkeypatch.setattr(service, "_stream_once", fake_stream_once) + monkeypatch.setattr(service, "_handle_stream_error", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", AsyncMock()) + + 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-stale"}, + codex_session_affinity=False, + propagate_http_errors=propagate_http_errors, + openai_cache_affinity=False, + api_key=None, + api_key_reservation=None, + suppress_text_done_events=False, + request_transport="http", + upstream_stream_transport_override="http", + ) + ] + + failed = json.loads(chunks[-1].split("data: ", 1)[1]) + assert failed["response"]["error"]["message"] == "replacement account failed" + assert stream_account_ids == [account_a.id, account_a.id, account_b.id] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("replacement_outcome", "settlement_confirmed"), @@ -13153,7 +13541,7 @@ async def fake_stream_once(account: Account, *_args: object, **_kwargs: object): @pytest.mark.asyncio -async def test_stream_with_retry_post_refresh_connect_exhaustion_terminal_counts_every_retry(monkeypatch): +async def test_stream_with_retry_post_refresh_ambiguous_connect_exhaustion_counts_every_retry(monkeypatch): settings = _make_proxy_settings() service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) account = _make_account("acc_post_refresh_transient_terminal") @@ -13203,7 +13591,7 @@ async def fake_stream_once(*_args: object, **_kwargs: object): 502, proxy_module.openai_error("upstream_unavailable", "proxy cannot connect"), failure_phase="connect", - retryable_same_contract=True, + retryable_same_contract=False, ) yield "" # pragma: no cover @@ -14217,9 +14605,10 @@ async def test_stream_responses_confirmed_proxy_connect_failure_fails_over_after seen_reallocate_sticky: list[object] = [] record_error = AsyncMock() record_success = AsyncMock() - settle_usage = AsyncMock(return_value=True) release_unsettled_usage = AsyncMock() record_error_backoff_accounts: list[Account] = [] + settlement_order: list[str] = [] + settlement_wait_flags: list[bool] = [] monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) @@ -14241,6 +14630,13 @@ async def record_error_backoff(account: Account) -> None: # Leases must settle before the error-health write for the dead route. assert lease_a in released_leases assert response_lease_a in released_leases + assert settlement_order == ["settle"] + settlement_order.append("backoff") + + async def settle_usage(*_args: object, **kwargs: object) -> bool: + settlement_wait_flags.append(bool(kwargs.get("wait_for_settlement"))) + settlement_order.append("settle") + return True async def fake_stream( payload: ResponsesRequest, @@ -14300,10 +14696,12 @@ async def fake_stream( assert seen_excluded_account_ids == [set(), {account_a.id}] assert seen_reallocate_sticky[-1] is True assert record_error_backoff_accounts == [account_a] + assert settlement_order == ["settle", "backoff"] + assert settlement_wait_flags == [True] 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() + assert settlement_wait_flags == [True] release_unsettled_usage.assert_not_awaited() @@ -21635,10 +22033,19 @@ async def test_fail_pending_websocket_requests_penalizes_upstream_stream_drop(mo request_logs = _RequestLogsRecorder() service = proxy_service.ProxyService(_repo_factory(request_logs)) account = _make_account("acc_ws_drop") - handle_stream_error = AsyncMock() + order: list[str] = [] + + async def record_stream_error(*_args: object, **_kwargs: object) -> None: + assert order == ["release"] + order.append("health") + + async def release_request_state(*_args: object, **_kwargs: object) -> None: + order.append("release") + + handle_stream_error = AsyncMock(side_effect=record_stream_error) monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) - monkeypatch.setattr(service, "_release_websocket_reservation", AsyncMock()) + monkeypatch.setattr(service, "_release_websocket_request_state_reservation", release_request_state) request_state = proxy_service._WebSocketRequestState( request_id="ws_req_drop", @@ -21666,6 +22073,7 @@ async def test_fail_pending_websocket_requests_penalizes_upstream_stream_drop(mo {"message": "Upstream websocket closed before response.completed"}, "stream_incomplete", ) + assert order == ["release", "health"] assert list(pending_requests) == [] assert await service.drain_persistence_tasks(timeout_seconds=1) assert len(request_logs.calls) == 1 @@ -22673,6 +23081,169 @@ def cancel_settlement_before_start( handle_stream_error.assert_awaited_once() +@pytest.mark.asyncio +async def test_finalize_websocket_request_state_keeps_health_deferred_when_settlement_is_unconfirmed(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_ws_unconfirmed_settlement") + api_key = _make_api_key_data("key_ws_unconfirmed_settlement") + reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv_ws_unconfirmed_settlement", + key_id=api_key.id, + model="gpt-5.1", + ) + lifecycle = proxy_support._DeferredAccountBackoffLifecycle( + reservation=reservation, + pending_backoffs={account.id: account}, + settlement_owned=True, + ) + tracker = proxy_support._DeferredAccountBackoffTracker(current_lifecycle=lifecycle) + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_unconfirmed_settlement", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=reservation, + started_at=time.monotonic(), + deferred_account_error_backoffs=lifecycle.pending_backoffs, + deferred_account_backoff_tracker=tracker, + deferred_account_backoff_lifecycle=lifecycle, + ) + settle_usage = AsyncMock(return_value=False) + record_error_backoff = AsyncMock() + handle_stream_error = AsyncMock() + record_success = AsyncMock() + monkeypatch.setattr(service, "_settle_stream_api_key_usage", settle_usage) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) + monkeypatch.setattr(service._load_balancer, "record_success", record_success) + failed_payload: dict[str, JsonValue] = { + "type": "response.failed", + "response": { + "id": "resp_ws_unconfirmed_settlement", + "error": {"code": "server_error", "message": "upstream failed"}, + }, + } + failed_event = parse_sse_event(f"data: {json.dumps(failed_payload)}\n\n") + assert failed_event is not None + upstream_control = proxy_service._WebSocketUpstreamControl() + + await service._finalize_websocket_request_state( + request_state, + account=account, + account_id_value=account.id, + event=failed_event, + event_type="response.failed", + payload=failed_payload, + api_key=api_key, + upstream_control=upstream_control, + response_create_gate=asyncio.Semaphore(1), + ) + + assert request_state.api_key_reservation is reservation + assert lifecycle.settlement_confirmed is False + assert lifecycle.pending_backoffs == {account.id: account} + settle_args = settle_usage.await_args + assert settle_args is not None + assert settle_args.kwargs["wait_for_settlement"] is True + record_error_backoff.assert_not_awaited() + handle_stream_error.assert_not_awaited() + record_success.assert_not_awaited() + assert upstream_control.reconnect_requested is True + + +@pytest.mark.asyncio +async def test_finalize_websocket_request_state_does_not_drain_older_unconfirmed_bridge_lifecycle(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_bridge_old_finalizer") + api_key = _make_api_key_data("key_bridge_old_finalizer") + old_reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv_bridge_old_finalizer", + key_id=api_key.id, + model="gpt-5.1", + ) + new_reservation = proxy_service.ApiKeyUsageReservationData( + reservation_id="resv_bridge_new_generation", + key_id=api_key.id, + model="gpt-5.1", + ) + old_lifecycle = proxy_support._DeferredAccountBackoffLifecycle( + reservation=old_reservation, + pending_backoffs={account.id: account}, + settlement_owned=True, + ) + new_lifecycle = proxy_support._DeferredAccountBackoffLifecycle( + reservation=new_reservation, + settlement_owned=True, + ) + tracker = proxy_support._DeferredAccountBackoffTracker(current_lifecycle=new_lifecycle) + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_bridge_old_finalizer", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=new_reservation, + started_at=time.monotonic(), + deferred_account_error_backoffs=new_lifecycle.pending_backoffs, + deferred_account_backoff_tracker=tracker, + deferred_account_backoff_lifecycle=new_lifecycle, + ) + monkeypatch.setattr(service, "_settle_stream_api_key_usage", AsyncMock(return_value=True)) + record_error_backoff = AsyncMock() + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service._load_balancer, "record_success", AsyncMock()) + completed_payload: dict[str, JsonValue] = { + "type": "response.completed", + "response": {"id": "resp_bridge_old_finalizer", "usage": {"input_tokens": 1, "output_tokens": 1}}, + } + completed_event = parse_sse_event(f"data: {json.dumps(completed_payload)}\n\n") + assert completed_event is not None + + await service._finalize_websocket_request_state( + request_state, + account=account, + account_id_value=account.id, + event=completed_event, + event_type="response.completed", + payload=completed_payload, + api_key=api_key, + upstream_control=proxy_service._WebSocketUpstreamControl(), + response_create_gate=asyncio.Semaphore(1), + ) + + assert request_state.api_key_reservation is None + assert old_lifecycle.settlement_confirmed is False + assert new_lifecycle.settlement_confirmed is True + assert old_lifecycle.pending_backoffs == {account.id: account} + assert new_lifecycle.pending_backoffs == {} + record_error_backoff.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_deferred_account_backoff_concurrent_drains_claim_once(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_deferred_backoff_claim") + pending = {account.id: account} + started = asyncio.Event() + release = asyncio.Event() + calls: list[str] = [] + + async def record_error_backoff(candidate: Account) -> None: + calls.append(candidate.id) + started.set() + await release.wait() + + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + first = asyncio.create_task(service._drain_deferred_account_error_backoffs(pending)) + await started.wait() + second = asyncio.create_task(service._drain_deferred_account_error_backoffs(pending)) + await asyncio.sleep(0) + release.set() + await asyncio.gather(first, second) + + assert calls == [account.id] + assert pending == {} + + @pytest.mark.asyncio async def test_finalize_websocket_empty_prewarm_does_not_store_continuity_anchor(monkeypatch): request_logs = _RequestLogsRecorder() @@ -26899,7 +27470,12 @@ async def release_usage_reservation(self, reservation_id: str) -> None: @pytest.mark.asyncio -async def test_stream_api_key_settlement_wait_option_reports_confirmed_fallback(monkeypatch): +@pytest.mark.parametrize(("release_fails", "expected_result"), [(False, True), (True, False)]) +async def test_stream_api_key_settlement_wait_option_reports_fallback_release_result( + monkeypatch, + release_fails: bool, + expected_result: bool, +): released: list[str] = [] repo = SimpleNamespace(api_keys=object()) @@ -26913,10 +27489,12 @@ def __init__(self, api_keys_repository: object) -> None: async def finalize_usage_reservation(self, reservation_id: str, **kwargs: object) -> None: del reservation_id, kwargs - raise RuntimeError("temporary finalize failure") + raise RuntimeError("settlement unavailable") async def release_usage_reservation(self, reservation_id: str) -> None: released.append(reservation_id) + if release_fails: + raise RuntimeError("release unavailable") monkeypatch.setattr(proxy_service, "ApiKeysService", FakeApiKeysService) @@ -26946,15 +27524,14 @@ async def release_usage_reservation(self, reservation_id: str) -> None: output_tokens=2, ) - committed = await service._settle_stream_api_key_usage( + result = await service._settle_stream_api_key_usage( api_key, reservation, settlement, request_id="req_stream_wait_failure", wait_for_settlement=True, ) - - assert committed is True + assert result is expected_result assert await service.drain_persistence_tasks(timeout_seconds=1) assert released == ["resv_stream_wait_failure"] From 5bf458dbdb49c611cd89b5132e28f3801093c54c Mon Sep 17 00:00:00 2001 From: Soju06 Date: Thu, 6 Aug 2026 05:01:10 +0000 Subject: [PATCH 4/4] fix(proxy): move verified fresh replay off dead pinned owner before pre-dispatch hard-ownership raise A confirmed pre-dispatch connect failure guarantees zero upstream bytes, so the confirmed dead-route branch can safely run the verified-owner replay step (same file/turn-state/single-account guards as the post-refresh path) ahead of the require_preferred_account fail-closed raise, moving a locally verified full-resend to a fresh account instead of erroring. Co-Authored-By: Claude Fable 5 --- app/modules/proxy/_service/streaming/retry.py | 23 ++++- tests/unit/test_proxy_utils.py | 88 +++++++++++++++++++ 2 files changed, 109 insertions(+), 2 deletions(-) diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index 46a610b37b..a2a0f6ff46 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -1965,7 +1965,25 @@ async def _retry_account_model_rejection( await _release_tracked_stream_lease(current_account_lease) current_account_lease = None await _record_or_defer_confirmed_route_backoff(account) - can_try_other_account = ( + # A confirmed pre-dispatch connect failure + # guarantees zero upstream bytes, so a + # locally verified full-resend replay is + # safe to move off the dead owner before + # the hard-ownership raise (same file/ + # turn-state/single-account guards as the + # post-visible failover path below). + verified_owner_replay_moved = False + if ( + attempt < max_attempts - 1 + and routing_strategy != "single_account" + and file_preferred_account_id is None + and turn_state_owner_account_id is None + ): + verified_owner_replay_moved = _move_verified_fresh_replay_from_owner( + account_id=account.id, + outcome="owner_pre_dispatch_proxy_connect_failure", + ) + can_try_other_account = verified_owner_replay_moved or ( not require_preferred_account and account.id != file_preferred_account_id and attempt < max_attempts - 1 @@ -1980,7 +1998,8 @@ async def _retry_account_model_rejection( last_pre_dispatch_transport_error = tex transient_failed_account_id = account.id excluded_account_ids.add(account.id) - affinity = replace(affinity, reallocate_sticky=True) + if not verified_owner_replay_moved: + affinity = replace(affinity, reallocate_sticky=True) _facade().logger.info( "Retrying stream after confirmed pre-dispatch proxy connect failure " "request_id=%s account_id=%s attempt=%d", diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 03f94b3e6d..08afbe68b0 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -33731,6 +33731,94 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, assert handle_stream_error.await_args.args[0] == owner_account +@pytest.mark.asyncio +async def test_stream_verified_fresh_replay_moves_off_owner_after_pre_dispatch_connect_failure(monkeypatch): + # Regression: a confirmed pre-dispatch dead-route failure on a pinned owner + # guarantees zero upstream bytes, so a locally verified full-resend replay + # must move to a fresh account instead of failing closed under hard + # ownership. + settings = _make_proxy_settings() + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + owner_account = _make_account("acc_stream_replay_dead_route_owner") + replacement_account = _make_account("acc_stream_replay_dead_route_replacement") + session_id = "sid_stream_verified_dead_route_replay" + previous_response_id = "resp_stream_verified_dead_route_owner" + request_logs.response_owner_by_id[(previous_response_id, None, session_id)] = owner_account.id + initial_input: list[JsonValue] = [{"role": "user", "content": "first turn"}] + full_input: list[JsonValue] = [ + *initial_input, + {"role": "user", "content": "fresh full resend after dead route"}, + ] + service._websocket_continuity_index[(session_id, None)] = proxy_service._WebSocketContinuityState( + last_completed_response_id=previous_response_id, + last_completed_input_count=len(initial_input), + last_completed_input_prefix_fingerprint=proxy_service._fingerprint_input_items(initial_input), + ) + selection_calls: list[dict[str, object]] = [] + streamed_payloads: list[ResponsesRequest] = [] + streamed_account_ids: list[str | None] = [] + record_error_backoff = AsyncMock() + record_error = AsyncMock() + + async def fake_select_account(**kwargs): + selection_calls.append(dict(kwargs)) + if kwargs.get("required_account_id") == owner_account.id: + return AccountSelection(account=owner_account, error_message=None) + assert kwargs.get("required_account_id") is None + assert kwargs.get("exclude_account_ids") == {owner_account.id} + assert kwargs.get("reallocate_sticky") is True + return AccountSelection(account=replacement_account, error_message=None) + + async def fake_stream(payload, headers, access_token, account_id, base_url=None, raise_for_status=False, **kwargs): + del headers, access_token, base_url, raise_for_status, kwargs + streamed_payloads.append(payload) + streamed_account_ids.append(account_id) + if account_id == owner_account.chatgpt_account_id: + raise _pre_dispatch_proxy_connect_error("pinned owner proxy route unavailable") + assert account_id == replacement_account.chatgpt_account_id + yield ( + 'data: {"type":"response.completed","response":{"id":"resp_dead_route_replay_ok",' + '"status":"completed","usage":{"input_tokens":1,"output_tokens":1,' + '"total_tokens":2}}}\n\n' + ) + + 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", AsyncMock(side_effect=fake_select_account)) + monkeypatch.setattr(service._load_balancer, "record_success", AsyncMock()) + monkeypatch.setattr(service._load_balancer, "record_error", record_error) + monkeypatch.setattr(service._load_balancer, "record_error_backoff", record_error_backoff) + monkeypatch.setattr(service, "_handle_stream_error", AsyncMock(return_value={"failure_class": "transient"})) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(side_effect=lambda account, **kwargs: account)) + monkeypatch.setattr(service, "_settle_stream_api_key_usage", AsyncMock(return_value=True)) + monkeypatch.setattr(proxy_service, "core_stream_responses", fake_stream) + + payload = ResponsesRequest.model_validate( + { + "model": "gpt-5.6-sol", + "instructions": "test verified full replay after dead route", + "input": full_input, + "previous_response_id": previous_response_id, + "stream": True, + } + ) + + chunks = [chunk async for chunk in service.stream_responses(payload, {"session_id": session_id})] + + assert json.loads(chunks[-1].split("data: ", 1)[1])["type"] == "response.completed" + assert all('"type":"response.failed"' not in chunk for chunk in chunks) + assert len(selection_calls) >= 2 + assert streamed_account_ids == [ + owner_account.chatgpt_account_id, + replacement_account.chatgpt_account_id, + ] + assert [streamed.previous_response_id for streamed in streamed_payloads] == [previous_response_id, None] + assert streamed_payloads[1].input == full_input + record_error_backoff.assert_awaited_once_with(owner_account) + record_error.assert_not_awaited() + + def test_cross_transport_fresh_replay_requires_matching_ws_continuity_prefix(): service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) first_input: list[JsonValue] = [