diff --git a/.all-contributorsrc b/.all-contributorsrc index 021a069c36..ecfb8f6a12 100644 --- a/.all-contributorsrc +++ b/.all-contributorsrc @@ -1193,6 +1193,16 @@ "code", "test" ] + }, + { + "login": "glopyglerky", + "name": "glopyglerky", + "avatar_url": "https://avatars.githubusercontent.com/u/189872235?v=4", + "profile": "https://github.com/glopyglerky", + "contributions": [ + "code", + "test" + ] } ], "contributorsPerLine": 7, diff --git a/app/core/balancer/__init__.py b/app/core/balancer/__init__.py index c8d4dc525b..5431263f32 100644 --- a/app/core/balancer/__init__.py +++ b/app/core/balancer/__init__.py @@ -12,6 +12,7 @@ ROUTING_POLICY_PRESERVE, TRAFFIC_CLASS_FOREGROUND, TRAFFIC_CLASS_OPPORTUNISTIC, + USAGE_LIMIT_REACHED, AccountState, FailoverAction, ResetPreferenceWindow, @@ -29,6 +30,7 @@ handle_quota_exceeded, handle_rate_limit, plausible_rate_limit_reset_at, + pool_usage_exhaustion, select_account, ) @@ -54,6 +56,7 @@ "RoutingStrategy", "TrafficClass", "SelectionResult", + "USAGE_LIMIT_REACHED", "UsageWeightedOrder", "account_status_for_permanent_failure", "configure_replica_salt", @@ -63,5 +66,6 @@ "handle_quota_exceeded", "handle_rate_limit", "plausible_rate_limit_reset_at", + "pool_usage_exhaustion", "select_account", ] diff --git a/app/core/balancer/logic.py b/app/core/balancer/logic.py index 96d9dc4a8f..a4d7b14725 100644 --- a/app/core/balancer/logic.py +++ b/app/core/balancer/logic.py @@ -146,6 +146,90 @@ class AccountState: class SelectionResult: account: AccountState | None error_message: str | None + error_code: str | None = None + resets_at: int | None = None + + +USAGE_LIMIT_REACHED = "usage_limit_reached" + + +def pool_usage_exhaustion( + states: Iterable[AccountState], + *, + current: float, + ignore_standard_quota: bool = False, + ignore_standard_quota_account_ids: Collection[str] | None = None, +) -> SelectionResult | None: + """Describe pool-wide subscription exhaustion without parsing retry text.""" + if ignore_standard_quota: + return None + ignored_account_ids = set(ignore_standard_quota_account_ids or ()) + + def _primary_usage_evidence(state: AccountState) -> float | None: + return state.priority_used_percent if state.priority_used_percent is not None else state.used_percent + + def _secondary_usage_evidence(state: AccountState) -> float | None: + if state.limit_scoped_usage and state.priority_secondary_used_percent is None: + return _primary_usage_evidence(state) + return ( + state.priority_secondary_used_percent + if state.priority_secondary_used_percent is not None + else state.secondary_used_percent + ) + + def _usage_exhausted(state: AccountState) -> bool: + if state.status not in (AccountStatus.QUOTA_EXCEEDED, AccountStatus.RATE_LIMITED): + return False + usage_values = (_primary_usage_evidence(state), _secondary_usage_evidence(state)) + return any(value is not None and float(value) >= 100.0 for value in usage_values) + + def _usage_exhausted_reset_at(state: AccountState) -> float | None: + if state.status not in (AccountStatus.QUOTA_EXCEEDED, AccountStatus.RATE_LIMITED): + return None + candidates: list[float] = [] + primary_evidence = _primary_usage_evidence(state) + secondary_evidence = _secondary_usage_evidence(state) + if primary_evidence is not None and float(primary_evidence) >= 100.0 and state.primary_reset_at is not None: + candidates.append(float(state.primary_reset_at)) + if ( + secondary_evidence is not None + and float(secondary_evidence) >= 100.0 + and state.secondary_reset_at is not None + ): + candidates.append(float(state.secondary_reset_at)) + if not candidates: + return None + return max(candidates) + + eligible = [ + state + for state in states + if not state.ignore_standard_quota + and state.account_id not in ignored_account_ids + and state.status + not in ( + AccountStatus.PAUSED, + AccountStatus.REAUTH_REQUIRED, + AccountStatus.DEACTIVATED, + ) + ] + if not eligible or any(not _usage_exhausted(state) for state in eligible): + return None + + # Only usage-proven exhausted accounts reach this branch; surface the + # earliest reset for an actually exhausted window so the structured 429 + # does not retry a secondary-window exhaustion at the primary reset time. + reset_candidates = [reset_at for state in eligible if (reset_at := _usage_exhausted_reset_at(state)) is not None] + resets_at = int(min(reset_candidates)) if reset_candidates else None + message = "Usage limit reached" + if resets_at is not None: + message = _format_retry_hint(max(0.0, resets_at - current)) + return SelectionResult( + account=None, + error_message=message, + error_code=USAGE_LIMIT_REACHED, + resets_at=resets_at, + ) @dataclass(frozen=True, slots=True) @@ -380,6 +464,8 @@ def select_account( primary_first_usage_weighted: bool = False, routing_costs: RoutingCostsByAccount | None = None, replica_salt: str | None = None, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> SelectionResult: """Select an eligible account by applying availability checks and routing strategy. @@ -447,6 +533,7 @@ def select_account( available: list[AccountState] = [] in_error_backoff: list[AccountState] = [] all_states = list(states) + usage_exhaustion_state_list = list(usage_exhaustion_states) if usage_exhaustion_states is not None else all_states bypass_account_ids = None if bypass_quota_exceeded_account_ids is None else set(bypass_quota_exceeded_account_ids) for state in all_states: @@ -530,6 +617,15 @@ def _backoff_expires_at(s: AccountState) -> float: return SelectionResult(None, f"opportunistic burn window closed: {reason}") available = opportunistic_available else: + if allow_usage_exhaustion_error: + usage_exhaustion = pool_usage_exhaustion( + usage_exhaustion_state_list, + current=current, + ignore_standard_quota=ignore_standard_quota or bypass_quota_exceeded, + ignore_standard_quota_account_ids=bypass_account_ids, + ) + if usage_exhaustion is not None: + return usage_exhaustion reauth_required = [s for s in all_states if s.status == AccountStatus.REAUTH_REQUIRED] deactivated = [s for s in all_states if s.status == AccountStatus.DEACTIVATED] paused = [s for s in all_states if s.status == AccountStatus.PAUSED] diff --git a/app/core/errors.py b/app/core/errors.py index 269d142d4c..395470dbef 100644 --- a/app/core/errors.py +++ b/app/core/errors.py @@ -47,8 +47,17 @@ class ResponseFailedEvent(TypedDict): PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE = "Previous response was not found; retry without previous_response_id." -def openai_error(code: str, message: str, error_type: str = "server_error") -> OpenAIErrorEnvelope: - return {"error": {"message": message, "type": error_type, "code": code}} +def openai_error( + code: str, + message: str, + error_type: str = "server_error", + *, + resets_at: int | float | None = None, +) -> OpenAIErrorEnvelope: + detail: OpenAIErrorDetail = {"message": message, "type": error_type, "code": code} + if resets_at is not None: + detail["resets_at"] = int(resets_at) + return {"error": detail} def dashboard_error(code: str, message: str) -> DashboardErrorEnvelope: @@ -105,9 +114,10 @@ def response_failed_event( response_id: str | None = None, created_at: int | None = None, error_param: str | None = None, + resets_at: int | float | None = None, incomplete_details: dict[str, str] | None = None, ) -> ResponseFailedEvent: - error = openai_error(code, message, error_type)["error"] + error = openai_error(code, message, error_type, resets_at=resets_at)["error"] if error_param: error["param"] = error_param if created_at is None: diff --git a/app/modules/proxy/_load_balancer/sticky_selection.py b/app/modules/proxy/_load_balancer/sticky_selection.py index 810399f251..199c015a59 100644 --- a/app/modules/proxy/_load_balancer/sticky_selection.py +++ b/app/modules/proxy/_load_balancer/sticky_selection.py @@ -170,6 +170,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback: bool, traffic_class: TrafficClass, ignore_standard_quota: bool, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> _StickySelectionOutcome: ... async def release_account_lease(self, lease: AccountLease | None) -> None: ... @@ -204,6 +206,7 @@ class StickySelectionRequest(Generic[SelectionInputsT]): selection_inputs: SelectionInputsT reload_inputs: Callable[[], Awaitable[SelectionInputsT]] record_account_cap_rejection: AccountCapRejectionCallback + allow_usage_exhaustion_error: bool = True @dataclass(frozen=True, slots=True) @@ -226,6 +229,7 @@ class StickySelectionOutcome(Generic[SelectionInputsT]): selected_lease: AccountLease | None error_message: str | None error_code: str | None + resets_at: int | None = None disposition: StickySelectionDisposition = "shared_result" @@ -261,11 +265,13 @@ async def run_sticky_selection_path( redact_sensitive_details = request.redact_sensitive_details load_selection_inputs = request.reload_inputs _record_account_cap_rejection = request.record_account_cap_rejection + allow_usage_exhaustion_error = request.allow_usage_exhaustion_error selected_snapshot: Account | None = None selected_lease: AccountLease | None = None error_message: str | None = None selection_error_code: str | None = None + selection_resets_at: int | None = None def _direct_error( *, @@ -425,9 +431,11 @@ def _direct_error( sticky_outcome = _StickySelectionOutcome(selection=SelectionResult(None, None)) if hard_sticky and not selection_states: selection_error_code = "hard_affinity_saturated" + selection_resets_at = None result = SelectionResult(None, "Hard affinity owner account is unavailable") elif not selection_states and states: selection_error_code = _account_cap_error_code(lease_kind) + selection_resets_at = None result = SelectionResult(None, _account_cap_error_message(lease_kind, caps)) logger.warning( "Account cap exhausted during sticky selection lease_kind=%s reason=%s candidates=%s", @@ -452,17 +460,22 @@ def _direct_error( traffic_class=traffic_class, ignore_standard_quota=False, routing_costs_by_account_id=effective_routing_costs, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=states, ) if result.account is None: selection_error_code = "hard_affinity_saturated" + selection_resets_at = None result = SelectionResult( None, result.error_message or "Hard affinity owner account is unavailable", ) else: selection_error_code = None + selection_resets_at = None else: selection_error_code = None + selection_resets_at = None try: async with owner._repo_factory() as repos: sticky_outcome = await owner._select_with_stickiness( @@ -485,8 +498,25 @@ def _direct_error( traffic_class=traffic_class, ignore_standard_quota=False, routing_costs_by_account_id=effective_routing_costs, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=states, ) result = sticky_outcome.selection + if ( + result.account is None + and result.error_code is None + and lease_kind is not None + and len(selection_states) < len(states) + and any( + state.status == AccountStatus.ACTIVE for state in states if state not in selection_states + ) + ): + selection_error_code = _account_cap_error_code(lease_kind) + result = SelectionResult( + None, + _account_cap_error_message(lease_kind, caps), + error_code=selection_error_code, + ) except BaseException: async with owner._runtime_lock: owner._release_due_probe_reservation_locked(probe_reservation) @@ -545,6 +575,8 @@ def _direct_error( probe_reservation_invalidated = True if result.account is None: error_message = result.error_message + selection_error_code = result.error_code or selection_error_code + selection_resets_at = result.resets_at or selection_resets_at elif probe_reservation_invalidated: selected = None else: @@ -876,6 +908,7 @@ def _direct_error( selected_lease=selected_lease, error_message=error_message, error_code=selection_error_code, + resets_at=selection_resets_at, ) @@ -900,6 +933,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback: bool = False, traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> _StickySelectionOutcome: if not sticky_key or not sticky_repo: return _StickySelectionOutcome( @@ -914,6 +949,8 @@ async def _select_with_stickiness( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs_by_account_id=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) ) if sticky_kind is None: @@ -1039,6 +1076,8 @@ def finish_selection( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs_by_account_id=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) pool_also_exhausted = pool_best.account is not None and ( pool_best.account.account_id == pinned.account_id @@ -1125,6 +1164,8 @@ def finish_selection( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs_by_account_id=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) if persist_fallback and chosen.account is not None and chosen.account.account_id in account_map: return finish_selection(chosen, persist_account_id=chosen.account.account_id) @@ -1328,6 +1369,8 @@ def _select_account_preferring_budget_safe( traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, routing_costs_by_account_id: RoutingCostsByAccount | None = None, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> SelectionResult: state_list = list(states) if routing_strategy not in ("sequential_drain", "reset_drain", "single_account"): @@ -1346,6 +1389,7 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=False, ) if recovery_probe.account is not None: return recovery_probe @@ -1378,6 +1422,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) best_health_states = _best_health_tier_states(state_list) @@ -1395,6 +1441,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) if burn_first.account is not None: return burn_first @@ -1418,6 +1466,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) if preferred.account is not None: return preferred @@ -1435,6 +1485,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) return select_account( state_list, @@ -1448,6 +1500,8 @@ def _select_account_preferring_budget_safe( traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, routing_costs=routing_costs_by_account_id, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) diff --git a/app/modules/proxy/_load_balancer/unbound_selection.py b/app/modules/proxy/_load_balancer/unbound_selection.py index 5731d018ee..11c5a5e7b3 100644 --- a/app/modules/proxy/_load_balancer/unbound_selection.py +++ b/app/modules/proxy/_load_balancer/unbound_selection.py @@ -14,7 +14,7 @@ SelectionResult, TrafficClass, ) -from app.db.models import Account +from app.db.models import Account, AccountStatus from app.modules.proxy._load_balancer.sticky_selection import ( SelectionInputsProtocol, StickySelectionOwner, @@ -70,6 +70,7 @@ class UnboundSelectionRequest(Generic[SelectionInputsT]): selection_inputs: SelectionInputsT reload_inputs: Callable[[], Awaitable[SelectionInputsT]] record_account_cap_rejection: AccountCapRejectionCallback + allow_usage_exhaustion_error: bool = True @dataclass(frozen=True, slots=True) @@ -79,6 +80,7 @@ class UnboundSelectionOutcome(Generic[SelectionInputsT]): selected_lease: AccountLease | None error_message: str | None error_code: str | None + resets_at: int | None = None disposition: str = "shared_result" @@ -105,11 +107,13 @@ async def run_unbound_selection_path( redact_sensitive_details = request.redact_sensitive_details load_selection_inputs = request.reload_inputs _record_account_cap_rejection = request.record_account_cap_rejection + allow_usage_exhaustion_error = request.allow_usage_exhaustion_error selected_snapshot: Account | None = None selected_lease: AccountLease | None = None error_message: str | None = None selection_error_code: str | None = None + selection_resets_at: int | None = None def _direct_error( *, @@ -161,6 +165,7 @@ def _direct_error( ) if not selection_states and states: selection_error_code = _account_cap_error_code(lease_kind) + selection_resets_at = None error_message = _account_cap_error_message(lease_kind, caps) result = SelectionResult(None, error_message) logger.warning( @@ -172,6 +177,7 @@ def _direct_error( _record_account_cap_rejection(lease_kind) else: selection_error_code = None + selection_resets_at = None result = _select_account_preferring_budget_safe( selection_states, prefer_earlier_reset=prefer_earlier_reset_accounts, @@ -184,7 +190,22 @@ def _direct_error( traffic_class=traffic_class, ignore_standard_quota=False, routing_costs_by_account_id=effective_routing_costs, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=states, ) + if ( + result.account is None + and result.error_code is None + and lease_kind is not None + and len(selection_states) < len(states) + and any(state.status == AccountStatus.ACTIVE for state in states if state not in selection_states) + ): + selection_error_code = _account_cap_error_code(lease_kind) + result = SelectionResult( + None, + _account_cap_error_message(lease_kind, caps), + error_code=selection_error_code, + ) probing_result_requires_reservation = _probing_result_requires_recovery_reservation( selection_states, result.account, @@ -266,6 +287,8 @@ def _direct_error( selected_snapshot.reset_at = selected_reset_at elif result.account is None: error_message = result.error_message + selection_error_code = result.error_code or selection_error_code + selection_resets_at = result.resets_at or selection_resets_at if probe_reservation_invalidated: selected_snapshot = None @@ -442,4 +465,5 @@ def _direct_error( selected_lease=selected_lease, error_message=error_message, error_code=selection_error_code, + resets_at=selection_resets_at, ) diff --git a/app/modules/proxy/_service/codex_control.py b/app/modules/proxy/_service/codex_control.py index 1e768bf761..6879c7c3c4 100644 --- a/app/modules/proxy/_service/codex_control.py +++ b/app/modules/proxy/_service/codex_control.py @@ -36,6 +36,7 @@ from app.modules.proxy.affinity import _AffinityPolicy, _sticky_key_for_codex_control_request from app.modules.proxy.helpers import _header_account_id, _normalize_error_code, _parse_openai_error from app.modules.proxy.load_balancer import AccountSelection, effective_account_concurrency_caps +from app.modules.proxy.selection_errors import selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -338,10 +339,8 @@ async def _finalize_success( if account is None: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call_control(target: Account) -> CodexControlResponse: diff --git a/app/modules/proxy/_service/compact.py b/app/modules/proxy/_service/compact.py index f8eea368f6..f6a87d32a7 100644 --- a/app/modules/proxy/_service/compact.py +++ b/app/modules/proxy/_service/compact.py @@ -58,6 +58,7 @@ AccountSelection, effective_account_concurrency_caps, ) +from app.modules.proxy.selection_errors import selection_failure_response from app.modules.proxy.work_admission import AdmissionLease, WorkAdmissionController logger = logging.getLogger("app.modules.proxy.service") @@ -897,14 +898,10 @@ async def _call_compact( else: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - status_code = 429 if log_error_code == "account_response_create_cap" else 503 + status_code, error_payload = selection_failure_response(selection) raise ProxyResponseError( status_code, - openai_error( - log_error_code, - log_error_message, - error_type="rate_limit_error" if status_code == 429 else "server_error", - ), + error_payload, ) assert account is not None account_id_value = account.id diff --git a/app/modules/proxy/_service/file_ops.py b/app/modules/proxy/_service/file_ops.py index b07696ef35..acc0e39a13 100644 --- a/app/modules/proxy/_service/file_ops.py +++ b/app/modules/proxy/_service/file_ops.py @@ -36,6 +36,7 @@ ) from app.modules.proxy.helpers import _header_account_id, _normalize_error_code, _parse_openai_error from app.modules.proxy.load_balancer import AccountSelection +from app.modules.proxy.selection_errors import selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -455,10 +456,8 @@ async def _proxy_files_call( if not account: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call(target: Account) -> dict[str, JsonValue]: diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index 13a92e54aa..b5d7b3389c 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -130,7 +130,6 @@ _await_cancelled_task, _call_with_supported_optional_kwargs, _estimated_lease_tokens_from_request_usage_budget, - _is_local_account_cap_code, _prefer_earlier_reset_window, _proxy_admission_wait_timeout_seconds, _raise_proxy_unavailable, @@ -212,6 +211,7 @@ DurableBridgeLookup, ) from app.modules.proxy.load_balancer import CONTINUITY_OWNER_UNAVAILABLE, AccountLease +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED, selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -1741,7 +1741,6 @@ async def _create_http_bridge_session( # 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 and preferred_account_id is not None @@ -1749,16 +1748,8 @@ async def _create_http_bridge_session( and selection.error_code == CONTINUITY_OWNER_UNAVAILABLE ): raise _http_bridge_previous_response_owner_unavailable_error() - status_code = 429 if is_local_account_cap else 503 - error_type = "rate_limit_error" if status_code == 429 else "server_error" - raise ProxyResponseError( - status_code, - openai_error( - selection.error_code or "no_accounts", - selection.error_message or "No active accounts available", - error_type=error_type, - ), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) if require_preferred_account and preferred_account_id is not None and account.id != preferred_account_id: await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None @@ -2182,6 +2173,16 @@ def require_bound_account() -> None: ): preferred_candidate_id = None continue + if selection.error_code == USAGE_LIMIT_REACHED and ( + required_preferred_account_id is not None or hard_close_account_bound + ): + complete_failed_handoff() + raise _http_bridge_previous_response_owner_unavailable_error() + if selection.error_code == USAGE_LIMIT_REACHED: + record_selected_account_takeover(None) + status_code, error_payload = selection_failure_response(selection) + complete_failed_handoff() + raise ProxyResponseError(status_code, error_payload) try: should_retry_selection = await _sleep_for_account_selection_recovery( selection, @@ -2220,16 +2221,9 @@ def require_bound_account() -> None: preferred_candidate_id = None continue record_selected_account_takeover(None) - status_code = 429 if _is_local_account_cap_code(selection.error_code) else 503 + status_code, error_payload = selection_failure_response(selection) complete_failed_handoff() - raise ProxyResponseError( - status_code, - openai_error( - selection.error_code or "no_accounts", - selection.error_message or "No active accounts available", - error_type="rate_limit_error" if status_code == 429 else "server_error", - ), - ) + raise ProxyResponseError(status_code, error_payload) if required_preferred_account_id is not None and account.id != required_preferred_account_id: if selection.lease is not None: selected_account_lease = selection.lease diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index a2a0f6ff46..f09346a162 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -76,6 +76,7 @@ is_upstream_model_capacity_error, ) from app.modules.proxy.load_balancer import AccountLease, AccountSelection +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED, selection_failure_response _REQUEST_TRANSPORT_HTTP = "http" _REQUEST_TRANSPORT_WEBSOCKET = "websocket" @@ -1111,6 +1112,7 @@ async def _retry_account_model_rejection( continue if ( not account + and selection.error_code != USAGE_LIMIT_REACHED and ( selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES or not (propagate_http_errors and last_transient_exc is not None) @@ -1166,6 +1168,40 @@ async def _retry_account_model_rejection( raise last_pre_dispatch_transport_error yield _render_dispatch_transport_error(last_pre_dispatch_transport_error) return + if selection.error_code == USAGE_LIMIT_REACHED: + await _drain_pending_post_refresh_penalty_on_terminal(settlement) + no_accounts_msg = selection.error_message or "Usage limit reached" + status_code, error_payload = selection_failure_response(selection) + await proxy._write_request_log( + account_id=None, + api_key=api_key, + request_id=request_id, + model=payload.model, + latency_ms=int((time.monotonic() - start) * 1000), + status="error", + error_code=USAGE_LIMIT_REACHED, + error_message=no_accounts_msg, + reasoning_effort=payload.reasoning.effort if payload.reasoning else None, + transport=request_transport, + upstream_transport=upstream_stream_transport, + service_tier=payload.service_tier, + requested_service_tier=payload.service_tier, + useragent=useragent, + useragent_group=useragent_group, + client_ip=client_ip, + ) + if propagate_http_errors: + raise ProxyResponseError(status_code, error_payload) + yield format_sse_event( + response_failed_event( + USAGE_LIMIT_REACHED, + no_accounts_msg, + error_type=USAGE_LIMIT_REACHED, + response_id=request_id, + resets_at=selection.resets_at, + ) + ) + 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" diff --git a/app/modules/proxy/_service/transcribe.py b/app/modules/proxy/_service/transcribe.py index e2abbe4c1b..f27d08f4e6 100644 --- a/app/modules/proxy/_service/transcribe.py +++ b/app/modules/proxy/_service/transcribe.py @@ -30,6 +30,7 @@ from app.modules.proxy._service.support import _request_log_client_fields, _RequestLogFailureMetadata from app.modules.proxy.helpers import _header_account_id, _normalize_error_code, _parse_openai_error from app.modules.proxy.load_balancer import AccountSelection +from app.modules.proxy.selection_errors import selection_failure_response logger = logging.getLogger("app.modules.proxy.service") T = TypeVar("T") @@ -202,10 +203,8 @@ async def transcribe( if not account: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call_transcribe(target: Account) -> dict[str, JsonValue]: diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index e6680e7680..50308cbc78 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -79,7 +79,6 @@ ProcessNetworkRecovery, process_network_error_code, ) -from app.core.resilience.overload import is_local_overload_error_code from app.core.types import JsonValue from app.core.upstream_proxy import UpstreamProxyRouteError from app.core.utils.request_id import get_request_id, reset_request_id, set_request_id @@ -473,6 +472,7 @@ openai_validation_error, validate_model_access, ) +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED, selection_failure_response from app.modules.proxy.tool_call_dedupe import ( mark_duplicate_tool_call_downstream_event, rewrite_parallel_tool_call_text, @@ -3317,6 +3317,8 @@ async def _select_websocket_connect_account( account = selection.account if account is not None: break + if selection.error_code == USAGE_LIMIT_REACHED: + break async def _heartbeat(remaining_seconds: float) -> None: event = _account_capacity_wait_payload( @@ -3423,18 +3425,15 @@ async def _heartbeat(remaining_seconds: float) -> None: return None if require_preferred_account and preferred_account_id is not None: if _facade()._is_local_account_cap_code(error_code): + status_code, error_payload = selection_failure_response(selection) await proxy._emit_websocket_connect_failure( websocket, client_send_lock=client_send_lock, account_id=preferred_account_id, api_key=api_key, request_state=request_state, - status_code=429, - payload=openai_error( - error_code, - error_message, - error_type="rate_limit_error", - ), + status_code=status_code, + payload=error_payload, error_code=error_code, error_message=error_message, ) @@ -3475,7 +3474,7 @@ async def _heartbeat(remaining_seconds: float) -> None: len(exclude_account_ids), api_key is not None, ) - status_code = 429 if is_local_overload_error_code(error_code) else 503 + status_code, error_payload = selection_failure_response(selection) await proxy._emit_websocket_connect_failure( websocket, client_send_lock=client_send_lock, @@ -3483,11 +3482,7 @@ async def _heartbeat(remaining_seconds: float) -> None: api_key=api_key, request_state=request_state, status_code=status_code, - payload=openai_error( - error_code, - error_message, - error_type="rate_limit_error" if status_code == 429 else "server_error", - ), + payload=error_payload, error_code=error_code, error_message=error_message, ) diff --git a/app/modules/proxy/api.py b/app/modules/proxy/api.py index e1f37338af..c50e0e810d 100644 --- a/app/modules/proxy/api.py +++ b/app/modules/proxy/api.py @@ -264,6 +264,7 @@ WarmupSkippedAccount, WarmupSubmittedAccount, ) +from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED from app.modules.proxy.types import ( CreditStatusDetailsData, RateLimitResetCreditsData, @@ -6444,6 +6445,17 @@ async def _opportunistic_admission_denial( ) if selection.account is not None: return None + if selection.error_code == USAGE_LIMIT_REACHED: + return _logged_error_json_response( + request, + 429, + openai_error( + USAGE_LIMIT_REACHED, + selection.error_message or "Usage limit reached", + error_type=USAGE_LIMIT_REACHED, + resets_at=selection.resets_at, + ), + ) message = selection.error_message or "opportunistic burn window closed" if not message.startswith("opportunistic burn window closed"): message = f"opportunistic burn window closed: {message}" diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index 57b73c6261..824588b05d 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -23,6 +23,7 @@ ROUTING_POLICY_PRESERVE, TRAFFIC_CLASS_FOREGROUND, TRAFFIC_CLASS_OPPORTUNISTIC, + USAGE_LIMIT_REACHED, AccountState, ResetPreferenceWindow, RoutingCostsByAccount, @@ -184,6 +185,7 @@ class AccountSelection: account: Account | None error_message: str | None error_code: str | None = None + resets_at: int | None = None lease: AccountLease | None = None catalog_omission_quota_admission: CatalogOmissionQuotaAdmission | None = None @@ -441,6 +443,7 @@ async def select_account( traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, concurrency_caps: AccountConcurrencyCaps | None = None, redact_sensitive_details: bool = False, + allow_usage_exhaustion_error: bool = True, ) -> AccountSelection: if (required_account_is_ownership_constraint or required_continuity_owner) and required_account_id is None: raise ValueError("required account ownership flags require required_account_id") @@ -575,6 +578,7 @@ async def load_selection_inputs() -> _SelectionInputs: error_message: str | None = None selected_lease: AccountLease | None = None selection_error_code: str | None = None + selection_resets_at: int | None = None legacy_existing_account_id: str | None = None if sticky_source == "session_header" and legacy_sticky_key is not None: async with self._repo_factory() as repos: @@ -639,6 +643,7 @@ async def load_selection_inputs() -> _SelectionInputs: selection_inputs=selection_inputs, reload_inputs=load_selection_inputs, record_account_cap_rejection=_record_account_cap_rejection, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, ), ) selection_inputs = unbound_outcome.selection_inputs @@ -646,6 +651,7 @@ async def load_selection_inputs() -> _SelectionInputs: selected_lease = unbound_outcome.selected_lease error_message = unbound_outcome.error_message selection_error_code = unbound_outcome.error_code + selection_resets_at = unbound_outcome.resets_at if unbound_outcome.disposition == "direct_error": return AccountSelection( account=None, @@ -683,6 +689,7 @@ async def load_selection_inputs() -> _SelectionInputs: selection_inputs=selection_inputs, reload_inputs=load_selection_inputs, record_account_cap_rejection=_record_account_cap_rejection, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, ), ) selection_inputs = sticky_outcome.selection_inputs @@ -690,6 +697,7 @@ async def load_selection_inputs() -> _SelectionInputs: selected_lease = sticky_outcome.selected_lease error_message = sticky_outcome.error_message selection_error_code = sticky_outcome.error_code + selection_resets_at = sticky_outcome.resets_at if sticky_outcome.disposition == "direct_error": return AccountSelection( account=None, @@ -736,7 +744,12 @@ async def load_selection_inputs() -> _SelectionInputs: and (selection_inputs.accounts or selection_inputs.error_code is not None) ): set_normal() - return AccountSelection(account=None, error_message=error_message, error_code=selection_error_code) + return AccountSelection( + account=None, + error_message=error_message, + error_code=selection_error_code, + resets_at=selection_resets_at, + ) if not circuit_breaker_open: set_normal() logger.info( @@ -1189,8 +1202,16 @@ async def check_opportunistic_admission( deterministic_probe=True, traffic_class=TRAFFIC_CLASS_OPPORTUNISTIC, ignore_standard_quota=False, + usage_exhaustion_states=states, ) if result.account is None: + if result.error_code == USAGE_LIMIT_REACHED: + return AccountSelection( + account=None, + error_message=result.error_message, + error_code=result.error_code, + resets_at=result.resets_at, + ) return AccountSelection( account=None, error_message=result.error_message, @@ -1411,6 +1432,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback: bool = False, traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND, ignore_standard_quota: bool = False, + allow_usage_exhaustion_error: bool = True, + usage_exhaustion_states: Iterable[AccountState] | None = None, ) -> _StickySelectionOutcome: return await _run_select_with_stickiness( states=states, @@ -1432,6 +1455,8 @@ async def _select_with_stickiness( preserve_existing_mapping_on_fallback=preserve_existing_mapping_on_fallback, traffic_class=traffic_class, ignore_standard_quota=ignore_standard_quota, + allow_usage_exhaustion_error=allow_usage_exhaustion_error, + usage_exhaustion_states=usage_exhaustion_states, ) _persist_sticky_mutation = staticmethod(_persist_sticky_mutation) @@ -2278,6 +2303,7 @@ def _state_from_account( pressure_pct = inflight_pressure_pct + leased_token_pressure_pct effective_used_percent = None if used_percent is None else min(100.0, used_percent + pressure_pct) effective_secondary_used_percent = None if secondary_used is None else min(100.0, secondary_used + pressure_pct) + usage_exhaustion_evidence_status = status in (AccountStatus.QUOTA_EXCEEDED, AccountStatus.RATE_LIMITED) return AccountState( account_id=account.id, @@ -2297,6 +2323,8 @@ def _state_from_account( plan_type=account.plan_type, capacity_credits=capacity_credits, health_tier=new_tier, + priority_used_percent=used_percent if usage_exhaustion_evidence_status else None, + priority_secondary_used_percent=secondary_used if usage_exhaustion_evidence_status else None, inflight_response_creates=runtime.inflight_response_creates, inflight_streams=runtime.inflight_streams, leased_tokens=runtime.leased_tokens, diff --git a/app/modules/proxy/selection_errors.py b/app/modules/proxy/selection_errors.py new file mode 100644 index 0000000000..cdcfcd9e39 --- /dev/null +++ b/app/modules/proxy/selection_errors.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from typing import Protocol + +from app.core.errors import OpenAIErrorEnvelope, openai_error +from app.core.resilience.overload import is_local_overload_error_code + +USAGE_LIMIT_REACHED = "usage_limit_reached" + + +class SelectionFailure(Protocol): + error_message: str | None + error_code: str | None + resets_at: int | None + + +def selection_failure_response(selection: SelectionFailure) -> tuple[int, OpenAIErrorEnvelope]: + """Map an account-selection failure to its externally visible HTTP response. + + The ``usage_limit_reached`` mapping is strictly for upstream usage/quota + exhaustion of the whole eligible pool. Local capacity codes (account caps, + admission gates, fair-share throttles) resolve against the canonical + ``LOCAL_OVERLOAD_CODES`` registry so they keep their stable 429 + ``rate_limit_error`` contract and are never reclassified as upstream usage + exhaustion or collapsed into a generic 503. + """ + code = selection.error_code or "no_accounts" + message = selection.error_message or "No active accounts available" + if code == USAGE_LIMIT_REACHED: + return ( + 429, + openai_error( + code, + message, + error_type=USAGE_LIMIT_REACHED, + resets_at=selection.resets_at, + ), + ) + if is_local_overload_error_code(code): + return 429, openai_error(code, message, error_type="rate_limit_error") + return 503, openai_error(code, message) diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index 2ad683dfb2..848f41c6c3 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -751,6 +751,7 @@ from app.modules.proxy.ring_membership import ( RingMembershipService, ) +from app.modules.proxy.selection_errors import selection_failure_response from app.modules.proxy.work_admission import WorkAdmissionController logger = logging.getLogger(__name__) @@ -1033,10 +1034,8 @@ async def thread_goal_request( if account is None: log_error_code = selection.error_code or "no_accounts" log_error_message = selection.error_message or "No active accounts available" - raise ProxyResponseError( - 503, - openai_error(log_error_code, log_error_message), - ) + status_code, error_payload = selection_failure_response(selection) + raise ProxyResponseError(status_code, error_payload) account_id_value = account.id async def _call_goal(target: Account) -> dict[str, JsonValue]: @@ -1898,6 +1897,7 @@ def log_account_id(account_id: str | None) -> str | None: traffic_class=effective_traffic_class, concurrency_caps=concurrency_caps, redact_sensitive_details=redact_sensitive_details, + allow_usage_exhaustion_error=not required_preferred_account, ) if preferred_selection.account is not None: logger.info( diff --git a/openspec/changes/report-pool-usage-exhaustion/.openspec.yaml b/openspec/changes/report-pool-usage-exhaustion/.openspec.yaml new file mode 100644 index 0000000000..0bd76e6186 --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-18 diff --git a/openspec/changes/report-pool-usage-exhaustion/proposal.md b/openspec/changes/report-pool-usage-exhaustion/proposal.md new file mode 100644 index 0000000000..41367d46a4 --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/proposal.md @@ -0,0 +1,35 @@ +## Why + +When every account eligible for a Responses request is exhausted by known pool +usage windows, codex-lb can currently collapse the selection failure into a +generic no-account/server-unavailable response. That hides the user-actionable +upstream condition from Codex/OpenAI-compatible clients and makes agents treat a +quota window as infrastructure failure. + +## What Changes + +- Preserve the stable `usage_limit_reached` code from account selection when the + whole eligible pool is exhausted by usage windows. +- Return HTTP `429` with an OpenAI-style error envelope whose + `error.code` and `error.type` are both `usage_limit_reached`. +- Preserve the selected reset hint as `error.resets_at` when account selection + has one, and use the same contract across HTTP, streaming, bridge, and + WebSocket selection-failure paths. + +## Capabilities + +### New Capabilities + +None. + +### Modified Capabilities + +- `responses-api-compat`: define the externally visible Responses error + contract for pool-wide usage exhaustion. + +## Impact + +- Affected code: account selection failure mapping and Responses proxy surfaces. +- Affected APIs: failure status/body for pool-wide usage exhaustion changes from + generic unavailable/no-account semantics to HTTP 429 `usage_limit_reached`. +- Configuration and schema: no changes. diff --git a/openspec/changes/report-pool-usage-exhaustion/specs/responses-api-compat/spec.md b/openspec/changes/report-pool-usage-exhaustion/specs/responses-api-compat/spec.md new file mode 100644 index 0000000000..65482726e5 --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/specs/responses-api-compat/spec.md @@ -0,0 +1,63 @@ +## ADDED Requirements + +### Requirement: Pool usage exhaustion is reported as a usage-limit error + +The proxy MUST report pool-wide Responses usage exhaustion as a usage-limit +error. When every account eligible for a Responses request is exhausted by known +usage windows, the proxy MUST reject the request with HTTP `429` and an +OpenAI-style error envelope whose `error.code` and `error.type` are both +`usage_limit_reached`. If account selection has an authoritative upstream reset +timestamp for the exhausted pool, the response envelope MUST include that +timestamp as `error.resets_at`; the proxy MUST NOT expose the capped +human-facing retry hint or a synthesized fallback as `error.resets_at`. The +proxy MUST NOT collapse this condition into generic `no_accounts`, +`server_error`, or HTTP `503` semantics. Exhaustion classification MUST be +based on structured account state after the same eligibility filtering as +ordinary selection, and MUST NOT reclassify local capacity or overload codes +(account caps, admission gates, fair-share throttles) as usage exhaustion. + +#### Scenario: Public Responses request exhausts the eligible usage pool + +- **WHEN** account selection for a public `/v1/responses` or + `/backend-api/codex/responses` request finds only usage-exhausted eligible + accounts +- **THEN** the response status is HTTP `429` +- **AND** the response body has `error.code = "usage_limit_reached"` +- **AND** the response body has `error.type = "usage_limit_reached"` +- **AND** any selected pool reset timestamp is surfaced as `error.resets_at` + +#### Scenario: Streaming selection failure preserves usage-limit semantics + +- **WHEN** a streaming Responses request cannot select an account because every + eligible account is usage-exhausted before downstream-visible output +- **THEN** the terminal error event uses `usage_limit_reached` +- **AND** clients do not receive a generic no-account/server-unavailable error + +#### Scenario: Usage-limit selection failures are terminal, not waitable + +- **WHEN** account selection fails with `usage_limit_reached` on a streaming, + HTTP-bridge, or WebSocket Responses path +- **THEN** the proxy reports the structured usage-limit failure immediately +- **AND** it does not enter an account-capacity recovery wait for the + remaining request budget before reporting it + +#### Scenario: Local capacity codes keep their rate-limit contract + +- **WHEN** account selection fails with a local capacity or overload code such + as `account_stream_cap` or `account_response_create_cap` +- **THEN** the response keeps HTTP `429` with `error.type = "rate_limit_error"` + and the stable local error code +- **AND** the response is not reported as `usage_limit_reached` + +#### Scenario: Unusable non-exhausted pools keep existing semantics + +- **WHEN** every account is paused, deactivated, or requires re-authentication + and no eligible account is exhausted by a known usage window +- **THEN** the pre-existing `no_accounts` failure semantics are preserved + +#### Scenario: Owner-scoped exhaustion preserves continuity semantics + +- **WHEN** a request is pinned to a previous-response or file owner account and + only that owner is usage-exhausted while the wider eligible pool is usable +- **THEN** the proxy keeps the existing continuity-owner failure semantics +- **AND** it does not report pool-wide `usage_limit_reached` diff --git a/openspec/changes/report-pool-usage-exhaustion/tasks.md b/openspec/changes/report-pool-usage-exhaustion/tasks.md new file mode 100644 index 0000000000..9dd87603ff --- /dev/null +++ b/openspec/changes/report-pool-usage-exhaustion/tasks.md @@ -0,0 +1,26 @@ +## 1. Error contract + +- [x] Preserve `usage_limit_reached` from pool-wide account selection failures. +- [x] Map pool-wide usage exhaustion to HTTP 429 with OpenAI-style + `error.code = "usage_limit_reached"` and + `error.type = "usage_limit_reached"`. +- [x] Preserve `error.resets_at` when account selection provides a reset hint. + +## 2. Proxy surfaces + +- [x] Apply the same selection-failure response helper across HTTP, streaming, + bridge, compact, file, transcription, WebSocket, and Codex-control paths. +- [x] Keep local capacity cap errors as 429 `rate_limit_error` responses rather + than weakening their existing contract. + +## 3. Regression coverage + +- [x] Add unit coverage for pool usage exhaustion selection and response mapping. +- [x] Add externally routed HTTP/streaming regressions for the 429 envelope. + +## 4. Validation + +- [x] Run focused pytest for selection, load balancer, and Responses proxy + regressions. +- [x] Run lint/type checks for touched Python files. +- [x] Validate the OpenSpec change strictly. diff --git a/tests/integration/test_http_responses_bridge.py b/tests/integration/test_http_responses_bridge.py index 24a14f1354..41ee741d2a 100644 --- a/tests/integration/test_http_responses_bridge.py +++ b/tests/integration/test_http_responses_bridge.py @@ -6660,6 +6660,39 @@ async def test_backend_responses_http_bridge_startup_error_omits_turn_state_head assert "x-codex-turn-state" not in response.headers +@pytest.mark.asyncio +async def test_backend_responses_http_bridge_pool_usage_exhaustion_returns_429(async_client, monkeypatch): + _install_bridge_settings(monkeypatch, enabled=True) + + async def fake_select_account_with_budget(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="Usage limit reached", + error_code="usage_limit_reached", + ) + + monkeypatch.setattr( + proxy_module.ProxyService, + "_select_account_with_budget", + fake_select_account_with_budget, + ) + + response = await async_client.post( + "/backend-api/codex/responses", + json={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "hello", + "stream": True, + }, + ) + + assert response.status_code == 429 + assert response.json()["error"]["type"] == "usage_limit_reached" + assert response.json()["error"]["code"] == "usage_limit_reached" + assert "x-codex-turn-state" not in response.headers + + @pytest.mark.asyncio async def test_v1_responses_http_bridge_startup_error_omits_turn_state_header(async_client, monkeypatch): _install_bridge_settings(monkeypatch, enabled=True) diff --git a/tests/integration/test_proxy_api_extended.py b/tests/integration/test_proxy_api_extended.py index 2ed0938fd1..45cf44e34b 100644 --- a/tests/integration/test_proxy_api_extended.py +++ b/tests/integration/test_proxy_api_extended.py @@ -514,6 +514,32 @@ async def fake_select(*_args, **_kwargs): assert response.json()["error"]["code"] == "no_accounts" +@pytest.mark.asyncio +async def test_thread_goal_get_maps_pool_usage_exhaustion_for_codex(async_client, monkeypatch): + async def fake_select(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="Usage limit reached", + error_code="usage_limit_reached", + ) + + monkeypatch.setattr(proxy_module.ProxyService, "_select_account_with_budget", fake_select) + + response = await async_client.post( + "/backend-api/codex/thread/goal/get", + json={"threadId": "019debd9-2372-7f23-92b9-9f34002a6355"}, + ) + + assert response.status_code == 429 + assert response.json() == { + "error": { + "message": "Usage limit reached", + "type": "usage_limit_reached", + "code": "usage_limit_reached", + } + } + + @pytest.mark.asyncio async def test_thread_goal_set_propagates_upstream_errors(async_client, monkeypatch): await _import_account(async_client, "acc_goal_set_error", "goal-set-error@example.com") diff --git a/tests/integration/test_proxy_responses.py b/tests/integration/test_proxy_responses.py index 35fd86a45c..ecbf93a8af 100644 --- a/tests/integration/test_proxy_responses.py +++ b/tests/integration/test_proxy_responses.py @@ -175,6 +175,104 @@ async def test_proxy_responses_no_accounts(async_client): assert event["response"]["error"]["code"] == "no_accounts" +def _install_usage_limited_selection(monkeypatch, *, resets_at: int | None = 1_700_003_600) -> None: + async def fake_select_account(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 300s", + error_code="usage_limit_reached", + resets_at=resets_at, + ) + + monkeypatch.setattr( + "app.modules.proxy.load_balancer.LoadBalancer.select_account", + fake_select_account, + ) + + +@pytest.mark.asyncio +async def test_v1_responses_pool_usage_exhaustion_returns_429(async_client, monkeypatch): + _install_usage_limited_selection(monkeypatch) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + + response = await async_client.post("/v1/responses", json=payload) + + assert response.status_code == 429 + error = response.json()["error"] + assert error["type"] == "usage_limit_reached" + assert error["code"] == "usage_limit_reached" + assert error["resets_at"] == 1_700_003_600 + + +@pytest.mark.asyncio +async def test_v1_responses_pool_usage_exhaustion_omits_unknown_reset(async_client, monkeypatch): + _install_usage_limited_selection(monkeypatch, resets_at=None) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + + response = await async_client.post("/v1/responses", json=payload) + + assert response.status_code == 429 + error = response.json()["error"] + assert error["type"] == "usage_limit_reached" + assert error["code"] == "usage_limit_reached" + assert "resets_at" not in error + + +@pytest.mark.asyncio +async def test_backend_responses_pool_usage_exhaustion_returns_429(async_client, monkeypatch): + _install_usage_limited_selection(monkeypatch) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + request_id = "req_stream_usage_limited" + + response = await async_client.post( + "/backend-api/codex/responses", + json=payload, + headers={"x-request-id": request_id}, + ) + + # Codex only classifies a terminal response as usage-limited when it sees + # both HTTP 429 and error.type == "usage_limit_reached" (#1246). + assert response.status_code == 429 + error = response.json()["error"] + assert error["type"] == "usage_limit_reached" + assert error["code"] == "usage_limit_reached" + assert error["resets_at"] == 1_700_003_600 + + +@pytest.mark.asyncio +async def test_v1_responses_mixed_unusable_pool_keeps_no_accounts_semantics(async_client, monkeypatch): + # Paused/deactivated/reauth-only pools must keep the pre-existing + # no_accounts semantics; only usage/quota exhaustion of the whole + # eligible pool may surface the new usage_limit_reached contract. + async def fake_select_account(*_args, **_kwargs): + return proxy_module.AccountSelection( + account=None, + error_message="All accounts are paused, deactivated, or require re-authentication", + error_code=None, + ) + + monkeypatch.setattr( + "app.modules.proxy.load_balancer.LoadBalancer.select_account", + fake_select_account, + ) + payload = {"model": "gpt-5.4", "instructions": "hi", "input": [], "stream": True} + + async with async_client.stream("POST", "/v1/responses", json=payload) as resp: + assert resp.status_code == 200 + lines = [line async for line in resp.aiter_lines() if line] + + # The synthetic failure keeps the #1479 SDK stream contract: a sequenced + # synthetic response.created precedes the sequenced response.failed. + created = _extract_first_raw_event(lines) + assert created["type"] == "response.created" + assert created["sequence_number"] == 0 + failed = _extract_first_event(lines) + assert failed["type"] == "response.failed" + assert failed["sequence_number"] == 1 + assert failed["response"]["error"]["code"] == "no_accounts" + assert failed["response"]["error"]["type"] == "server_error" + + @pytest.mark.asyncio async def test_backend_responses_prohibits_fast_model_alias_priority_tier(async_client, monkeypatch): raw_account_id = "acc_prohibit_fast_mode" diff --git a/tests/unit/test_load_balancer.py b/tests/unit/test_load_balancer.py index f2d67e1d9e..f6b8c12662 100644 --- a/tests/unit/test_load_balancer.py +++ b/tests/unit/test_load_balancer.py @@ -676,6 +676,343 @@ def test_select_account_skips_rate_limited_until_reset(): assert result.account.account_id == "b" +def test_select_account_reports_pool_wide_usage_exhaustion_structurally(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + reset_at=int(now + 600), + primary_reset_at=int(now + 60), + ), + AccountState( + "b", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 7200), + primary_reset_at=int(now + 3600), + ), + AccountState("paused", AccountStatus.PAUSED), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Rate limit exceeded. Try again in 60s" + assert result.resets_at == int(now + 60) + + +def test_select_account_fails_over_when_one_account_remains_usable(): + now = 1_700_000_000.0 + states = [ + AccountState("exhausted", AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=int(now + 3600)), + AccountState("usable", AccountStatus.ACTIVE, used_percent=40.0), + ] + + result = select_account(states, now=now) + + assert result.account is not None + assert result.account.account_id == "usable" + assert result.error_code is None + assert result.error_message is None + + +def test_select_account_reports_secondary_usage_exhaustion_reset(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=10.0, + secondary_used_percent=100.0, + reset_at=int(now + 60), + secondary_reset_at=int(now + 3600), + ) + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Rate limit exceeded. Try again in 300s" + assert result.resets_at == int(now + 3600) + + +def test_select_account_omits_synthesized_primary_usage_reset(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + reset_at=int(now + 60), + primary_reset_at=None, + ) + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Usage limit reached" + assert result.resets_at is None + + +def test_select_account_waits_for_latest_exhausted_window_per_account(): + now = 1_700_000_000.0 + states = [ + AccountState( + "a", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + secondary_used_percent=100.0, + reset_at=int(now + 600), + primary_reset_at=int(now + 60), + secondary_reset_at=int(now + 3600), + ), + AccountState( + "b", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 7200), + primary_reset_at=int(now + 7200), + ), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.error_message == "Rate limit exceeded. Try again in 300s" + assert result.resets_at == int(now + 3600) + + +def test_select_account_requires_usage_window_evidence_for_quota_exhaustion(): + now = 1_700_000_000.0 + states = [ + AccountState("a", AccountStatus.QUOTA_EXCEEDED, reset_at=int(now + 3600)), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_select_account_can_disable_pool_usage_exhaustion_for_owner_scope(): + now = 1_700_000_000.0 + states = [ + AccountState( + "owner", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 3600), + ) + ] + + result = select_account(states, now=now, allow_usage_exhaustion_error=False) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_budget_safe_selection_uses_full_scope_for_usage_exhaustion() -> None: + now = time.time() + cap_filtered_states = [ + AccountState( + "exhausted", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 3600), + ) + ] + full_scope_states = [ + AccountState( + "capped-but-usable", + AccountStatus.ACTIVE, + used_percent=50.0, + reset_at=int(now + 3600), + ), + *cap_filtered_states, + ] + + result = _select_account_preferring_budget_safe( + cap_filtered_states, + prefer_earlier_reset=False, + routing_strategy="usage_weighted", + budget_threshold_pct=95.0, + usage_exhaustion_states=full_scope_states, + ) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_budget_safe_capacity_selection_forwards_usage_exhaustion_controls() -> None: + now = time.time() + owner_scope = [ + AccountState( + "owner", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 600), + primary_reset_at=int(now + 3600), + ) + ] + full_scope = [ + *owner_scope, + AccountState( + "pool-usable", + AccountStatus.ACTIVE, + used_percent=10.0, + ), + ] + + result = _select_account_preferring_budget_safe( + owner_scope, + prefer_earlier_reset=False, + routing_strategy="capacity_weighted", + budget_threshold_pct=95.0, + allow_usage_exhaustion_error=False, + usage_exhaustion_states=full_scope, + ) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_opportunistic_budget_safe_selection_uses_full_scope_for_usage_exhaustion() -> None: + now = time.time() + cap_filtered_states = [ + AccountState( + "exhausted", + AccountStatus.QUOTA_EXCEEDED, + used_percent=100.0, + reset_at=int(now + 3600), + ) + ] + full_scope_states = [ + AccountState( + "capped-but-usable", + AccountStatus.ACTIVE, + used_percent=50.0, + reset_at=int(now + 3600), + ), + *cap_filtered_states, + ] + + result = _select_account_preferring_budget_safe( + cap_filtered_states, + prefer_earlier_reset=False, + routing_strategy="usage_weighted", + budget_threshold_pct=95.0, + traffic_class="opportunistic", + usage_exhaustion_states=full_scope_states, + ) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "Rate limit exceeded. Try again in 300s" + + +def test_select_account_uses_raw_priority_usage_for_exhaustion_evidence() -> None: + now = 1_700_000_000.0 + states = [ + AccountState( + "pressure-adjusted", + AccountStatus.RATE_LIMITED, + used_percent=100.0, + priority_used_percent=98.0, + reset_at=int(now + 60), + ) + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "No available accounts" + + +def test_select_account_does_not_treat_generic_rate_limit_as_usage_exhaustion(): + now = 1_700_000_000.0 + states = [ + AccountState("a", AccountStatus.RATE_LIMITED, used_percent=5.0, reset_at=int(now + 60)), + AccountState("b", AccountStatus.RATE_LIMITED, secondary_used_percent=10.0, reset_at=int(now + 120)), + ] + + result = select_account(states, now=now) + + assert result.account is None + assert result.error_code is None + assert result.error_message == "No available accounts" + + +def test_select_account_does_not_misclassify_transient_backoff_as_usage_exhaustion(): + now = 1_700_000_000.0 + states = [ + AccountState("quota", AccountStatus.QUOTA_EXCEEDED, reset_at=int(now + 3600)), + AccountState( + "transient", + AccountStatus.ACTIVE, + error_count=3, + last_error_at=now, + ), + ] + + result = select_account(states, now=now, allow_backoff_fallback=False) + + assert result.account is None + assert result.error_code is None + + +def test_select_account_does_not_report_ignored_standard_quota_as_pool_exhaustion(): + now = 1_700_000_000.0 + state = AccountState("quota", AccountStatus.QUOTA_EXCEEDED, reset_at=int(now + 3600)) + + result = select_account([state], now=now, ignore_standard_quota=True) + + assert result.account is not None + assert result.account.account_id == "quota" + + +def test_select_account_excludes_per_account_standard_quota_bypass_from_pool_exhaustion(): + now = 1_700_000_000.0 + state = AccountState( + "quota", + AccountStatus.QUOTA_EXCEEDED, + reset_at=int(now + 3600), + cooldown_until=now + 30, + ignore_standard_quota=True, + ) + + result = select_account([state], now=now) + + assert result.account is None + assert result.error_code is None + + +def test_select_account_excludes_scoped_standard_quota_bypass_from_pool_exhaustion(): + now = 1_700_000_000.0 + state = AccountState( + "quota", + AccountStatus.QUOTA_EXCEEDED, + reset_at=int(now + 3600), + cooldown_until=now + 30, + ) + + result = select_account([state], now=now, bypass_quota_exceeded_account_ids={"quota"}) + + assert result.account is None + assert result.error_code is None + + def test_select_account_reports_paused_and_deactivated_without_reauth_reason(): states = [ AccountState("paused", AccountStatus.PAUSED, used_percent=5.0), @@ -1291,12 +1628,14 @@ def test_select_account_caps_quota_exceeded_retry_hint(): AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=far_future_reset, + primary_reset_at=far_future_reset, ), AccountState( "b", AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=int(now + 271_819), + primary_reset_at=int(now + 271_819), ), ] result = select_account(states, now=now) @@ -1315,6 +1654,7 @@ def test_select_account_preserves_short_quota_exceeded_retry_hint(): AccountStatus.QUOTA_EXCEEDED, used_percent=100.0, reset_at=int(now + 60), + primary_reset_at=int(now + 60), ), ] result = select_account(states, now=now) @@ -1878,6 +2218,52 @@ def test_state_from_account_keeps_active_account_selectable_when_primary_usage_s assert selection.account.account_id == state.account_id +def test_state_from_account_keeps_raw_usage_evidence_separate_from_pressure(monkeypatch): + now = 1_700_000_000.0 + future_reset = int(now + 300) + monkeypatch.setattr("app.modules.proxy.load_balancer.time.time", lambda: now) + monkeypatch.setattr("app.core.usage.quota.time.time", lambda: now) + + state = _state_from_account( + account=_make_test_account(status=AccountStatus.RATE_LIMITED, reset_at=future_reset, blocked_at=int(now)), + primary_entry=_make_test_usage( + window="primary", + used_percent=98.0, + reset_at=future_reset, + recorded_at=_epoch_to_naive_utc(now - 30), + ), + secondary_entry=None, + runtime=RuntimeState(inflight_streams=1), + ) + + assert state.used_percent == 100.0 + assert state.priority_used_percent == 98.0 + + +def test_state_from_account_preserves_pressure_for_active_routing(monkeypatch): + now = 1_700_000_000.0 + future_reset = int(now + 300) + monkeypatch.setattr("app.modules.proxy.load_balancer.time.time", lambda: now) + monkeypatch.setattr("app.core.usage.quota.time.time", lambda: now) + + state = _state_from_account( + account=_make_test_account(status=AccountStatus.ACTIVE), + primary_entry=_make_test_usage( + window="primary", + used_percent=94.0, + reset_at=future_reset, + recorded_at=_epoch_to_naive_utc(now - 30), + ), + secondary_entry=None, + runtime=RuntimeState(inflight_streams=1), + ) + + assert state.status == AccountStatus.ACTIVE + assert state.used_percent == 96.5 + assert state.priority_used_percent is None + assert _state_above_sticky_budget_threshold(state, 95.0) is True + + def test_state_from_account_clears_stale_advisory_account_reset_for_active_account(monkeypatch): now = 1_700_000_000.0 future_reset = int(now + 300) diff --git a/tests/unit/test_load_balancer_concurrency.py b/tests/unit/test_load_balancer_concurrency.py index abc05e6e76..d8bd90ecea 100644 --- a/tests/unit/test_load_balancer_concurrency.py +++ b/tests/unit/test_load_balancer_concurrency.py @@ -122,6 +122,28 @@ async def test_account_lease_uses_explicit_dashboard_cap_snapshot_not_startup_en assert third is None +@pytest.mark.asyncio +async def test_opportunistic_selection_preserves_usage_limit_exhaustion_error() -> None: + account = _make_account("acc-opportunistic-usage-exhausted") + account.status = AccountStatus.QUOTA_EXCEEDED + reset_at = int(time.time() + 300) + account.reset_at = reset_at + usage_repo = _StubUsageRepository( + {account.id: _usage_row_with_percent(1, account.id, used_percent=100.0, reset_at=reset_at)}, + {}, + ) + balancer = LoadBalancer(lambda: _repo_factory(_StubAccountsRepository([account]), usage_repo)) + + result = await balancer.select_account( + routing_strategy="usage_weighted", + traffic_class=load_balancer_module.TRAFFIC_CLASS_OPPORTUNISTIC, + ) + + assert result.account is None + assert result.error_code == "usage_limit_reached" + assert result.resets_at == reset_at + + class _StubAccountsRepository: def __init__(self, accounts: list[Account]) -> None: self._accounts = accounts @@ -834,6 +856,57 @@ async def test_account_stream_cap_returns_stable_local_reason_until_released() - assert recovered.lease is not None +@pytest.mark.asyncio +async def test_stream_cap_takes_precedence_over_remaining_quota_exhausted_account() -> None: + now_epoch = int(datetime.now(tz=timezone.utc).timestamp()) + capped = _make_account("acc-stream-cap-mixed-capped") + exhausted = _make_account("acc-stream-cap-mixed-exhausted") + exhausted.status = AccountStatus.QUOTA_EXCEEDED + exhausted.reset_at = now_epoch + 3600 + accounts_repo = _StubAccountsRepository([capped, exhausted]) + usage_repo = _StubUsageRepository( + primary={ + capped.id: _usage_row_with_percent( + 203, + capped.id, + used_percent=50.0, + reset_at=now_epoch + 300, + ), + exhausted.id: _usage_row_with_percent( + 204, + exhausted.id, + used_percent=100.0, + reset_at=now_epoch + 3600, + ), + }, + secondary={}, + ) + balancer = LoadBalancer(lambda: _repo_factory(accounts_repo, usage_repo)) + leases = [ + ( + await balancer.select_account( + routing_strategy="usage_weighted", + lease_kind="stream", + ) + ).lease + for _ in range(8) + ] + + selected = await balancer.select_account( + routing_strategy="usage_weighted", + lease_kind="stream", + ) + + assert selected.account is None + assert selected.error_code == "account_stream_cap" + assert selected.error_message is not None + assert "Account stream capacity is exhausted" in selected.error_message + assert selected.resets_at is None + + for lease in leases: + await balancer.release_account_lease(lease) + + @pytest.mark.asyncio async def test_account_stream_recovery_reserve_keeps_last_slot_for_reattach() -> None: now_epoch = int(datetime.now(tz=timezone.utc).timestamp()) diff --git a/tests/unit/test_openai_errors.py b/tests/unit/test_openai_errors.py index 2534ead7c7..70ec2d136f 100644 --- a/tests/unit/test_openai_errors.py +++ b/tests/unit/test_openai_errors.py @@ -29,6 +29,18 @@ def test_response_failed_event_accepts_incomplete_details(): assert response.get("incomplete_details") == {"reason": "max_output_tokens"} +def test_response_failed_event_preserves_reset_hint(): + event = response_failed_event( + "usage_limit_reached", + "Rate limit exceeded. Try again in 1h", + error_type="usage_limit_reached", + response_id="resp_1", + resets_at=1_700_003_600, + ) + + assert event["response"]["error"]["resets_at"] == 1_700_003_600 + + def test_previous_response_not_found_classifier_covers_openai_shapes(): assert is_previous_response_not_found_error( code="previous_response_not_found", diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index b1c31c3f03..b0a61903b6 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -6855,6 +6855,112 @@ async def sleep_for_recovery(*_args: object, **kwargs: object) -> bool: assert sleep_calls[0]["max_sleep_seconds"] == pytest.approx(119.5) +@pytest.mark.asyncio +async def test_reconnect_http_bridge_session_skips_capacity_wait_for_usage_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + + async def select_account(_deadline: float, **_kwargs: object) -> proxy_service.AccountSelection: + return proxy_service.AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + + request_state = proxy_service._WebSocketRequestState( + request_id="req-reconnect-usage-limit-now", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=100.0, + ) + monkeypatch.setattr(proxy_service.time, "monotonic", lambda: 100.5) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=settings)), + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr( + http_bridge_mixin_module, + "_sleep_for_account_selection_recovery", + lambda *_args, **_kwargs: pytest.fail("usage_limit_reached must not enter recovery wait"), + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await service._reconnect_http_bridge_session(session, request_state=request_state) + + assert exc_info.value.status_code == 429 + assert exc_info.value.payload["error"]["code"] == "usage_limit_reached" + assert exc_info.value.payload["error"]["type"] == "usage_limit_reached" + assert exc_info.value.payload["error"]["resets_at"] == 1_700_003_600 + + +@pytest.mark.asyncio +async def test_reconnect_http_bridge_session_preserves_owner_error_for_owner_usage_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + ) + + async def select_account(_deadline: float, **_kwargs: object) -> proxy_service.AccountSelection: + return proxy_service.AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + + request_state = proxy_service._WebSocketRequestState( + request_id="req-reconnect-owner-usage-limit", + model="gpt-5.4", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=100.0, + preferred_account_id=session.account.id, + ) + monkeypatch.setattr(proxy_service.time, "monotonic", lambda: 100.5) + monkeypatch.setattr(proxy_service, "get_settings", lambda: _make_app_settings()) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=settings)), + ) + monkeypatch.setattr(service, "_select_account_with_budget_compatible", select_account) + monkeypatch.setattr( + http_bridge_mixin_module, + "_sleep_for_account_selection_recovery", + lambda *_args, **_kwargs: pytest.fail("owner-only usage_limit_reached must not enter recovery wait"), + ) + + with pytest.raises(ProxyResponseError) as exc_info: + await service._reconnect_http_bridge_session( + session, + request_state=request_state, + require_preferred_account=True, + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.payload["error"]["code"] == "previous_response_owner_unavailable" + assert exc_info.value.payload["error"]["type"] == "server_error" + + @pytest.mark.asyncio async def test_reconnect_http_bridge_session_preserves_exclusions_after_capacity_wait( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index 08afbe68b0..0822c0f75a 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -2971,6 +2971,49 @@ async def test_opportunistic_admission_uses_api_key_enforced_model(): ) +@pytest.mark.asyncio +async def test_opportunistic_admission_preserves_usage_limit_denial(): + api_key = ApiKeyData( + id="key_opportunistic_usage_limit", + name="opportunistic usage limit", + key_prefix="sk-opportunistic", + allowed_models=None, + enforced_model=None, + enforced_reasoning_effort=None, + enforced_service_tier=None, + traffic_class=proxy_api.TRAFFIC_CLASS_OPPORTUNISTIC, + expires_at=None, + is_active=True, + created_at=utcnow(), + last_used_at=None, + ) + selection = AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + service = SimpleNamespace(check_opportunistic_admission=AsyncMock(return_value=selection)) + context = SimpleNamespace(service=service) + request = Request({"type": "http", "method": "GET", "path": "/v1/opportunistic/admission", "headers": []}) + + response = await proxy_api._opportunistic_admission_denial( + request, + cast(proxy_api.ProxyContext, context), + api_key, + model="gpt-5.1", + ) + + assert response is not None + assert response.status_code == 429 + body = json.loads(bytes(response.body)) + assert body["error"]["code"] == "usage_limit_reached" + assert body["error"]["type"] == "usage_limit_reached" + assert body["error"]["message"] == "Rate limit exceeded. Try again in 1h" + assert body["error"]["resets_at"] == 1_700_003_600 + assert "Retry-After" not in response.headers + + @pytest.mark.asyncio async def test_opportunistic_admission_scopes_single_account_to_selected_account(monkeypatch): settings = _make_proxy_settings() @@ -10574,6 +10617,11 @@ async def test_service_compact_passes_chatgpt_account_id_to_core(monkeypatch): monkeypatch.setattr(proxy_service, "get_settings_cache", lambda: _SettingsCache(settings)) monkeypatch.setattr(proxy_service, "get_settings", lambda: settings) + monkeypatch.setattr( + streaming_retry_module, + "_account_selection_recovery_sleep_seconds", + lambda _selection: pytest.fail("usage_limit_reached must not enter recovery wait"), + ) monkeypatch.setattr( service._load_balancer, "select_account", @@ -11865,6 +11913,46 @@ async def test_stream_responses_propagates_selection_error_code(monkeypatch): assert request_logs.calls[0]["error_code"] == "additional_quota_data_unavailable" +@pytest.mark.asyncio +async def test_stream_responses_preserves_usage_limit_reset_hint(monkeypatch): + settings = _make_proxy_settings() + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + + 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( + return_value=AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ), + ) + + 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-usage-limit"})] + + event = json.loads(chunks[0].split("data: ", 1)[1]) + assert event["response"]["error"]["code"] == "usage_limit_reached" + assert event["response"]["error"]["type"] == "usage_limit_reached" + assert event["response"]["error"]["resets_at"] == 1_700_003_600 + assert await service.drain_persistence_tasks(timeout_seconds=1) + assert request_logs.calls[0]["error_code"] == "usage_limit_reached" + + @pytest.mark.asyncio async def test_stream_with_retry_keeps_sse_alive_while_account_capacity_recovers(monkeypatch): settings = _make_proxy_settings() @@ -18895,6 +18983,62 @@ async def test_select_websocket_connect_account_requires_preferred_account_for_p assert select_account.await_args.kwargs["request_stage"] == "reattach" +@pytest.mark.asyncio +async def test_select_websocket_connect_account_preserves_continuity_for_owner_usage_limit(monkeypatch): + request_logs = _RequestLogsRecorder() + service = proxy_service.ProxyService(_repo_factory(request_logs)) + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_prev_owner_usage_limit", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + request_stage="reattach", + ) + emit_connect_failure = AsyncMock() + select_account = AsyncMock( + return_value=AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ) + + monkeypatch.setattr(service, "_select_account_with_budget", select_account) + monkeypatch.setattr(service, "_emit_websocket_connect_failure", emit_connect_failure) + + result = await service._select_websocket_connect_account( + time.monotonic() + 10_000.0, + sticky_key=None, + sticky_kind=None, + prefer_earlier_reset=False, + prefer_earlier_reset_window="secondary", + routing_strategy="usage_weighted", + model="gpt-5.1", + request_state=request_state, + api_key=None, + client_send_lock=anyio.Lock(), + websocket=cast(WebSocket, SimpleNamespace()), + reallocate_sticky=False, + sticky_max_age_seconds=None, + exclude_account_ids=set(), + preferred_account_id="acc_owner", + require_preferred_account=True, + ) + + assert result is None + emit_connect_failure.assert_awaited_once() + call = emit_connect_failure.await_args + assert call is not None + assert call.kwargs["status_code"] == 502 + assert call.kwargs["error_code"] == "previous_response_owner_unavailable" + assert call.kwargs["account_id"] == "acc_owner" + assert call.kwargs["payload"]["error"]["code"] == "previous_response_owner_unavailable" + assert call.kwargs["payload"]["error"]["type"] == "server_error" + + @pytest.mark.asyncio async def test_select_websocket_connect_account_records_fail_closed_for_preferred_account_mismatch( monkeypatch, @@ -19173,6 +19317,65 @@ async def fake_sleep_for_account_selection_recovery(*_args: object, **kwargs: ob assert sent_payload["request_id"] == "ws_req_capacity_wait" +@pytest.mark.asyncio +async def test_select_websocket_connect_account_skips_capacity_wait_for_usage_limit(monkeypatch): + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + request_state = proxy_service._WebSocketRequestState( + request_id="ws_req_usage_limit_now", + model="gpt-5.1", + service_tier=None, + reasoning_effort=None, + api_key_reservation=None, + started_at=0.0, + ) + websocket_send = AsyncMock() + + monkeypatch.setattr( + service, + "_select_account_with_budget", + AsyncMock( + return_value=AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 1h", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ), + ) + monkeypatch.setattr( + websocket_mixin_module, + "_sleep_for_account_selection_recovery", + lambda *_args, **_kwargs: pytest.fail("usage_limit_reached must not enter recovery wait"), + ) + + result = await service._select_websocket_connect_account( + time.monotonic() + 10_000.0, + sticky_key=None, + sticky_kind=None, + prefer_earlier_reset=False, + routing_strategy="usage_weighted", + model="gpt-5.1", + request_state=request_state, + api_key=None, + client_send_lock=anyio.Lock(), + websocket=cast(WebSocket, SimpleNamespace(send_text=websocket_send)), + reallocate_sticky=False, + sticky_max_age_seconds=None, + exclude_account_ids=set(), + preferred_account_id=None, + require_preferred_account=False, + ) + + assert result is None + await_args = websocket_send.await_args + assert await_args is not None + sent_payload = json.loads(await_args.args[0]) + assert sent_payload["status"] == 429 + assert sent_payload["error"]["code"] == "usage_limit_reached" + assert sent_payload["error"]["type"] == "usage_limit_reached" + assert sent_payload["error"]["resets_at"] == 1_700_003_600 + + @pytest.mark.asyncio @pytest.mark.parametrize( ("preferred_account_id", "require_preferred_account", "file_required", "defer_no_account_error"), diff --git a/tests/unit/test_selection_errors.py b/tests/unit/test_selection_errors.py new file mode 100644 index 0000000000..8a670dad33 --- /dev/null +++ b/tests/unit/test_selection_errors.py @@ -0,0 +1,72 @@ +import pytest + +from app.core.resilience.overload import LOCAL_OVERLOAD_CODES +from app.modules.proxy.load_balancer import AccountSelection +from app.modules.proxy.selection_errors import selection_failure_response + + +def test_pool_usage_exhaustion_is_codex_compatible_429(): + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="Usage limit reached", + error_code="usage_limit_reached", + ) + ) + + assert status == 429 + assert payload == { + "error": { + "message": "Usage limit reached", + "type": "usage_limit_reached", + "code": "usage_limit_reached", + } + } + + +def test_unusable_pool_remains_no_accounts_503(): + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="All accounts require re-authentication", + error_code=None, + ) + ) + + assert status == 503 + assert payload["error"]["type"] == "server_error" + assert payload["error"]["code"] == "no_accounts" + + +def test_pool_usage_exhaustion_preserves_authoritative_reset(): + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="Rate limit exceeded. Try again in 300s", + error_code="usage_limit_reached", + resets_at=1_700_003_600, + ) + ) + + assert status == 429 + assert payload["error"]["resets_at"] == 1_700_003_600 + + +@pytest.mark.parametrize("local_code", sorted(LOCAL_OVERLOAD_CODES)) +def test_local_overload_codes_keep_rate_limit_contract(local_code: str): + # Covers every canonical local capacity code, including codes added later + # (e.g. api_key_stream_fair_share): local overload must stay a 429 + # rate_limit_error and never be reclassified as upstream usage exhaustion + # or a 503. + status, payload = selection_failure_response( + AccountSelection( + account=None, + error_message="Local capacity is exhausted", + error_code=local_code, + ) + ) + + assert status == 429 + assert payload["error"]["type"] == "rate_limit_error" + assert payload["error"]["code"] == local_code + assert "resets_at" not in payload["error"]