diff --git a/app/modules/api_keys/repository.py b/app/modules/api_keys/repository.py index 8fb9b77a83..4467ddbdad 100644 --- a/app/modules/api_keys/repository.py +++ b/app/modules/api_keys/repository.py @@ -106,6 +106,11 @@ class ApiKeysRepository: def __init__(self, session: AsyncSession) -> None: self._session = session + async def update_last_used(self, key_id: str, *, commit: bool = True) -> None: + await self._session.execute(update(ApiKey).where(ApiKey.id == key_id).values(last_used_at=utcnow())) + if commit: + await self._session.commit() + @staticmethod def _build_account_costs(rows: Sequence[object]) -> list[ApiKeyAccountCost]: account_costs: list[ApiKeyAccountCost] = [] diff --git a/app/modules/proxy/_service/http_bridge/protocol.py b/app/modules/proxy/_service/http_bridge/protocol.py index 841d02c014..29103b18e7 100644 --- a/app/modules/proxy/_service/http_bridge/protocol.py +++ b/app/modules/proxy/_service/http_bridge/protocol.py @@ -88,6 +88,7 @@ def _schedule_http_bridge_session_closes( async def _open_upstream_websocket_with_budget(self, *args: Any, **kwargs: Any) -> Any: ... async def _resolve_websocket_previous_response_owner(self, *args: Any, **kwargs: Any) -> Any: ... async def _acquire_request_state_response_create_admission(self, *args: Any, **kwargs: Any) -> None: ... + async def _reconnect_http_bridge_session(self, *args: Any, **kwargs: Any) -> Any: ... async def _handle_proxy_error(self, account: Account, exc: ProxyResponseError) -> None: ... async def _handle_stream_error( self, account: Account, error: Any, code: str, http_status: int | None = None diff --git a/app/modules/proxy/_service/http_bridge/request_submit.py b/app/modules/proxy/_service/http_bridge/request_submit.py index c5901c28ec..ce4c279002 100644 --- a/app/modules/proxy/_service/http_bridge/request_submit.py +++ b/app/modules/proxy/_service/http_bridge/request_submit.py @@ -55,6 +55,7 @@ reset_request_id, set_request_id, ) +from app.core.utils.retry import backoff_seconds from app.core.utils.sse import format_sse_event, parse_sse_data_json from app.modules.api_keys.service import ( ApiKeyData, @@ -247,6 +248,24 @@ async def _send_http_bridge_request_text_with_archive_id( reset_request_id(token) +def _prepare_http_bridge_terminal_capacity_replay(request_state: _WebSocketRequestState) -> str | None: + request_text = request_state.request_text + if not isinstance(request_text, str) or not request_text: + return None + if request_state.response_id is None or request_state.replay_count >= 1: + return None + if request_state.upstream_model_output_seen: + return None + request_state.replay_count += 1 + request_state.awaiting_response_created = True + request_state.response_id = None + request_state.response_event_count = 0 + request_state.latency_response_created_ms = None + request_state.upstream_model_output_seen = False + _clear_websocket_request_error_overrides(request_state) + return request_text + + def _text_with_account_installation_id(text_data: str, codex_installation_id: str | None) -> str: payload = json.loads(text_data) if not isinstance(payload, dict): @@ -1688,18 +1707,18 @@ async def _detach_http_bridge_request( request_state: _WebSocketRequestState, ) -> bool: detached = False - async with session.pending_lock: - if request_state in session.pending_requests and not request_state.draining_until_terminal: - request_state.draining_until_terminal = True - request_state.downstream_visible = False - session.queued_request_count = max(0, session.queued_request_count - 1) - session.upstream_control.reconnect_requested = True - session.upstream_control.retire_after_drain = True - detached = True - # Queue revocation and pending ownership use the same lock. A - # completed handler that wins first keeps its local queue reference; - # a detach that wins first leaves no queue for that handler to claim. - request_state.event_queue = None + # Revoke downstream delivery immediately; lifecycle ownership below + # may legitimately be held by a reconnect owner for an arbitrary wait. + request_state.event_queue = None + async with session.lifecycle_lock: + async with session.pending_lock: + if request_state in session.pending_requests and not request_state.draining_until_terminal: + request_state.draining_until_terminal = True + request_state.downstream_visible = False + session.queued_request_count = max(0, session.queued_request_count - 1) + session.upstream_control.reconnect_requested = True + session.upstream_control.retire_after_drain = True + detached = True await _release_websocket_response_create_gate(request_state, session.response_create_gate) if not detached: if request_state.terminal_settlement_phase == "abandoned": @@ -1921,6 +1940,119 @@ async def _retry_http_bridge_request_on_fresh_upstream( logger.warning("HTTP bridge retry on fresh upstream failed", exc_info=True) return False + async def _retry_http_bridge_terminal_capacity_request( + self: Any, + session: "_HTTPBridgeSession", + request_state: _WebSocketRequestState, + *, + error_code: str, + preserve_for_reader_failure: bool = False, + ) -> bool: + original_account_id = session.account.id + original_response_id = request_state.response_id + original_response_event_count = request_state.response_event_count + original_replay_count = request_state.replay_count + original_output_seen = request_state.upstream_model_output_seen + original_preferred_account_id = request_state.preferred_account_id + original_error_overrides = ( + request_state.error_code_override, + request_state.error_message_override, + request_state.error_type_override, + request_state.error_param_override, + request_state.error_http_status_override, + ) + async with session.pending_lock: + if ( + (session.pending_requests and session.pending_requests != deque([request_state])) + or request_state.replay_count >= 1 + or session.admission_waiter_count + ): + return False + if not session.pending_requests: + session.pending_requests.appendleft(request_state) + session.queued_request_count += 1 + + async def owns_request() -> bool: + if session.closed or session.upstream_control.retire_after_drain: + return False + async with session.pending_lock: + return len(session.pending_requests) == 1 and session.pending_requests[0] is request_state + + retry_sent = False + try: + await self._acquire_request_state_response_create_admission( + request_state, + response_create_gate=session.response_create_gate, + account_id=original_account_id, + surface="http_bridge_capacity_retry", + bridge_session=session, + ) + if not await owns_request(): + return False + delay = backoff_seconds(original_replay_count + 1) + if request_state.bridge_request_deadline is not None: + remaining = max(0.0, request_state.bridge_request_deadline - _service_time().monotonic()) + if remaining <= 0: + return False + delay = min(delay, remaining) + await asyncio.sleep(delay) + if ( + request_state.bridge_request_deadline is not None + and _service_time().monotonic() >= request_state.bridge_request_deadline + ): + return False + async with session.lifecycle_lock: + if not await owns_request(): + return False + request_state.preferred_account_id = original_account_id + await self._reconnect_http_bridge_session( + session, + request_state=request_state, + require_preferred_account=True, + ) + if session.account.id != original_account_id or not await owns_request(): + return False + request_text = _prepare_http_bridge_terminal_capacity_replay(request_state) + if request_text is None: + return False + request_text = self._http_bridge_text_with_account_installation_id(session, request_state, request_text) + await _send_http_bridge_request_text_with_archive_id(session, request_state, request_text) + session.last_used_at = _service_time().monotonic() + retry_sent = True + return True + except UpstreamWebSocketTransportError: + raise + except Exception: + logger.warning("HTTP bridge terminal capacity retry failed", exc_info=True) + return False + finally: + if not retry_sent and not preserve_for_reader_failure: + async with session.pending_lock: + if request_state in session.pending_requests: + session.pending_requests.remove(request_state) + session.queued_request_count = sum( + 1 + for pending in session.pending_requests + if _http_bridge_request_counts_against_queue(pending) + ) + if ( + request_state.response_create_gate_acquired + or request_state.account_response_create_lease is not None + ): + await _release_websocket_response_create_gate(request_state, session.response_create_gate) + request_state.response_id = original_response_id + request_state.response_event_count = original_response_event_count + request_state.replay_count = original_replay_count + request_state.upstream_model_output_seen = original_output_seen + request_state.preferred_account_id = original_preferred_account_id + ( + request_state.error_code_override, + request_state.error_message_override, + request_state.error_type_override, + request_state.error_param_override, + request_state.error_http_status_override, + ) = original_error_overrides + async def _retry_http_bridge_precreated_request( self: Any, session: "_HTTPBridgeSession", diff --git a/app/modules/proxy/_service/http_bridge/upstream_events.py b/app/modules/proxy/_service/http_bridge/upstream_events.py index e9fe0c6607..519296be6b 100644 --- a/app/modules/proxy/_service/http_bridge/upstream_events.py +++ b/app/modules/proxy/_service/http_bridge/upstream_events.py @@ -411,6 +411,137 @@ def _durable_pending_tool_call_manifest( "Upstream flagged this request as possible cybersecurity work. " "codex-lb is retrying on an account marked as authorized for security work." ) +_SECURITY_WORK_NO_AUTHORIZED_ACCOUNTS_MESSAGE = ( + "Upstream flagged this request as possible cybersecurity work, but no account is marked as authorized for " + "security work. codex-lb is continuing with normal account selection; the upstream request may still fail until " + "an account with Trusted Access for Cyber is marked as security-work-authorized." +) +_HTTP_BRIDGE_BACKGROUND_CLOSE_TIMEOUT_SECONDS = 5.0 +_HTTP_BRIDGE_BACKGROUND_CLEANUP_WARN_THRESHOLD = 100 +_HTTP_BRIDGE_TERMINAL_CAPACITY_RETRY_CODES = frozenset({"overloaded_error", "server_is_overloaded"}) +_HTTP_BRIDGE_TERMINAL_CAPACITY_RETRY_MESSAGES = ( + "selected model is at capacity", + "servers are currently overloaded", +) + + +def _http_bridge_terminal_payload_contains_output(payload: dict[str, JsonValue] | None) -> bool: + if not isinstance(payload, dict): + return False + candidates: list[JsonValue | None] = [payload.get("output")] + response = payload.get("response") + if isinstance(response, dict): + candidates.append(response.get("output")) + usage_candidates: list[JsonValue | None] = [payload.get("usage")] + if isinstance(response, dict): + usage_candidates.append(response.get("usage")) + for output in candidates: + if output is None: + continue + if isinstance(output, list): + if output: + return True + continue + return True + for usage in usage_candidates: + if not isinstance(usage, dict): + continue + output_tokens = usage.get("output_tokens") + if isinstance(output_tokens, (int, float)) and not isinstance(output_tokens, bool) and output_tokens > 0: + return True + output_token_details = usage.get("output_tokens_details") + if not isinstance(output_token_details, dict): + continue + reasoning_tokens = output_token_details.get("reasoning_tokens") + if ( + isinstance(reasoning_tokens, (int, float)) + and not isinstance(reasoning_tokens, bool) + and reasoning_tokens > 0 + ): + return True + return False + + +def _http_bridge_terminal_capacity_retry_message(message: str | None) -> bool: + if not isinstance(message, str): + return False + normalized = " ".join(message.casefold().split()) + return any(marker in normalized for marker in _HTTP_BRIDGE_TERMINAL_CAPACITY_RETRY_MESSAGES) + + +def _http_bridge_terminal_capacity_retry_error_code( + request_state: _WebSocketRequestState | None, + *, + event_type: str | None, + payload: dict[str, JsonValue] | None, + has_other_pending_requests: bool, +) -> str | None: + """Classify one output-free accepted overload that native Codex can replay.""" + if request_state is None or request_state.enforce_openai_sdk_contract: + return None + if has_other_pending_requests: + return None + if request_state.last_downstream_sequence_number is not None: + return None + if request_state.downstream_visible or request_state.upstream_model_output_seen: + return None + if request_state.pending_function_call_ids or request_state.pending_tool_call_types: + return None + if request_state.response_id is None or request_state.awaiting_response_created: + return None + if request_state.response_event_count < 1: + return None + if request_state.event_queue is None: + return None + if not request_state.request_text or request_state.replay_count >= 1: + return None + if event_type not in {"error", "response.failed"}: + return None + if _http_bridge_terminal_payload_contains_output(payload): + return None + error_code = _normalize_error_code( + _websocket_event_error_code(event_type, payload), + _websocket_event_error_type(event_type, payload), + ) + if error_code in _HTTP_BRIDGE_TERMINAL_CAPACITY_RETRY_CODES: + return error_code + if not _http_bridge_terminal_capacity_retry_message(_websocket_event_error_message(event_type, payload)): + return None + return error_code or "model_at_capacity" + + +def _http_bridge_transport_close_capacity_retry_error_code( + request_state: _WebSocketRequestState | None, + *, + has_other_pending_requests: bool, + error_code: str | None, + error_message: str | None, +) -> str | None: + """Classify output-free accepted disconnects that native Codex can replay.""" + if request_state is None or request_state.enforce_openai_sdk_contract: + return None + if has_other_pending_requests: + return None + if request_state.downstream_visible or request_state.upstream_model_output_seen: + return None + if request_state.pending_function_call_ids or request_state.pending_tool_call_types: + return None + if request_state.response_id is None or request_state.awaiting_response_created: + return None + if request_state.response_event_count < 1: + return None + if request_state.event_queue is None: + return None + if not request_state.request_text or request_state.replay_count >= 1: + return None + normalized_error_code = _normalize_error_code(error_code, None) + if normalized_error_code == "proxy_network_unavailable": + return None + if _http_bridge_terminal_capacity_retry_message(error_message): + return normalized_error_code or "model_at_capacity" + if normalized_error_code in {"stream_incomplete", "upstream_error", "upstream_unavailable"}: + return "stream_incomplete" + return None async def _wait_before_http_bridge_model_capacity_retry( @@ -1008,6 +1139,7 @@ async def _relay_http_bridge_upstream_messages( async with session.pending_lock: archive_request_state = session.pending_requests[0] if len(session.pending_requests) == 1 else None + has_other_pending_requests = len(session.pending_requests) != 1 response_events_seen = max( (request_state.response_event_count for request_state in session.pending_requests), default=0, @@ -1023,7 +1155,22 @@ async def _relay_http_bridge_upstream_messages( # or tool side effects. Clean websocket closes remain eligible # for the bounded pre-created retry path below. if message.error_code != "proxy_network_unavailable": - retried = await self._retry_http_bridge_precreated_request(session) + capacity_retry_code = _http_bridge_transport_close_capacity_retry_error_code( + archive_request_state, + has_other_pending_requests=has_other_pending_requests, + error_code=message.error_code, + error_message=message.error, + ) + capacity_retry_attempted = capacity_retry_code is not None and archive_request_state is not None + if capacity_retry_attempted and archive_request_state is not None: + retried = await self._retry_http_bridge_terminal_capacity_request( + session, + archive_request_state, + error_code=capacity_retry_code, + preserve_for_reader_failure=True, + ) + if not retried and not capacity_retry_attempted: + retried = await self._retry_http_bridge_precreated_request(session) if retried: continue close_classification = ( @@ -1031,6 +1178,22 @@ async def _relay_http_bridge_upstream_messages( if message.close_code is not None else None ) + # Keep this classification independent from retry eligibility: + # a blocked untyped close is still account-neutral. + close_is_untyped = message.error_code is None + close_is_clean = ( + message.kind == "close" + and _classify_upstream_close( + message.close_code, + response_events_seen=response_events_seen, + ) + == "clean" + ) + account_neutral_close = ( + message.error_code in {"proxy_network_unavailable", "upstream_keepalive_timeout"} + or close_is_untyped + or close_is_clean + ) async with session.lifecycle_lock: await self._fail_http_bridge_reader_and_maybe_retire( session, @@ -1043,18 +1206,7 @@ async def _relay_http_bridge_upstream_messages( if close_classification is not None else "websocket_transport_error" ), - penalize_account=( - message.error_code != "proxy_network_unavailable" - and message.error_code != "upstream_keepalive_timeout" - and not ( - message.kind == "close" - and _classify_upstream_close( - message.close_code, - response_events_seen=response_events_seen, - ) - == "clean" - ) - ), + penalize_account=not account_neutral_close, ) break except asyncio.CancelledError: @@ -2195,6 +2347,21 @@ async def _process_parsed_http_bridge_upstream_event( if retried: return + capacity_retry_code = _http_bridge_terminal_capacity_retry_error_code( + terminal_request_state, + event_type=settlement_event_type, + payload=settlement_payload, + has_other_pending_requests=has_other_pending_requests, + ) + if capacity_retry_code is not None: + retried = await self._retry_http_bridge_terminal_capacity_request( + session, + terminal_request_state, + error_code=capacity_retry_code, + ) + if retried: + return + matched_event_queue = ( completed_event_queue if completed_event_queue_claimed and matched_request_state is terminal_request_state diff --git a/app/modules/proxy/_service/streaming/mixin.py b/app/modules/proxy/_service/streaming/mixin.py index 470e7a9376..14e999c7e7 100644 --- a/app/modules/proxy/_service/streaming/mixin.py +++ b/app/modules/proxy/_service/streaming/mixin.py @@ -298,6 +298,7 @@ _RetryableStreamError, _StreamSettlement, _TerminalStreamError, + _TransientStreamError, _ttft_event_latency_ms, _WebSocketUpstreamControl, ) @@ -577,6 +578,8 @@ async def _stream_once( settlement.record_success = False settlement.account_health_error = True settlement.error = {"message": error_message} + if allow_transient_retry and payload.previous_response_id is None: + raise _TransientStreamError(error_code, settlement.error) yield format_sse_event( response_failed_event( error_code, @@ -727,14 +730,12 @@ async def _stream_once( settlement.error = {"message": error_message or "Upstream error"} settlement.record_success = False settlement.account_health_error = False - if event and event.type in ("response.completed", "response.incomplete"): usage = event.response.usage if event.response else None if event.response and event.response.id: response_id = event.response.id if event.type == "response.incomplete": status = "error" - if event_type in _facade()._TEXT_DELTA_EVENT_TYPES: saw_text_delta = True if not _facade()._should_suppress_text_done_event( diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index f09346a162..0d29147ca8 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -90,6 +90,21 @@ def _facade() -> Any: return sys.modules["app.modules.proxy.service"] +def _proxy_response_error_is_transient_stream_retry(exc: ProxyResponseError) -> bool: + error = _parse_openai_error(exc.payload) + code = _normalize_error_code(error.code if error else None, error.type if error else None) + message = error.message if error else None + return _facade()._should_retry_transient_stream_error(code, message) + + +def _transient_stream_retry_delay(deadline: float, retry_count: int) -> float | None: + """Return a backoff only when it leaves budget for the next attempt.""" + + remaining_budget_seconds = _facade()._remaining_budget_seconds(deadline) + delay = backoff_seconds(retry_count) + return delay if 0 <= delay < remaining_budget_seconds else None + + def _http_downstream_request_is_sticky(payload: ResponsesRequest, headers: Mapping[str, str]) -> bool: return ( payload.previous_response_id is not None @@ -357,6 +372,7 @@ async def _stream_with_retry( network_recovery = ProcessNetworkRecovery(transport="stream", request_id=request_id) settlement = _StreamSettlement() last_transient_exc: ProxyResponseError | None = None + last_exhausted_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 @@ -650,6 +666,12 @@ async def _iter_stream_once() -> AsyncIterator[str]: and not settlement.downstream_visible ): delay = backoff_seconds(transient_retries) + remaining_budget = _facade()._remaining_budget_seconds(deadline) + if remaining_budget <= delay: + raise ProxyResponseError( + 502, + openai_error("upstream_request_timeout", "Proxy request budget exhausted"), + ) from exc _facade().logger.info( "Transient post-refresh stream error, retrying same account " "request_id=%s account_id=%s retry=%s/%s delay=%.2fs code=%s", @@ -661,6 +683,11 @@ async def _iter_stream_once() -> AsyncIterator[str]: exc.code, ) await asyncio.sleep(delay) + if _facade()._remaining_budget_seconds(deadline) <= 0: + raise ProxyResponseError( + 502, + openai_error("upstream_request_timeout", "Proxy request budget exhausted"), + ) from exc continue error_message = str(exc.error.get("message") or "Upstream error") settlement.record_success = False @@ -1329,6 +1356,34 @@ async def _retry_account_model_rejection( # instead of returning a generic no_accounts event. if propagate_http_errors and last_transient_exc is not None: raise last_transient_exc + if last_retryable_stream_error is not None: + error_message = str(last_retryable_stream_error.error.get("message") or "Upstream error") + event = response_failed_event( + last_retryable_stream_error.code, + error_message, + response_id=request_id, + ) + yield format_sse_event(event) + 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=last_retryable_stream_error.code, + error_message=error_message, + 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, + conversation_id=conversation_id, + client_ip=client_ip, + ) + return if last_security_work_retry_error is not None: message = ( last_security_work_retry_error.error.get("message") @@ -1360,6 +1415,20 @@ async def _retry_account_model_rejection( client_ip=client_ip, ) return + if last_exhausted_transient_exc is not None: + error = _parse_openai_error(last_exhausted_transient_exc.payload) + error_code = _normalize_error_code(error.code if error else None, error.type if error else None) + error_message = error.message if error else None + event = response_failed_event( + error_code or "upstream_error", + error_message or "Upstream error", + 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) + yield format_sse_event(event) + return no_accounts_msg = selection.error_message or "No active accounts available" error_code = selection.error_code or "no_accounts" event = response_failed_event( @@ -1881,7 +1950,11 @@ async def _retry_account_model_rejection( settlement.error_code or "upstream_error", ) return - if isinstance(tex, ProxyResponseError) and tex.status_code != 500: + if ( + isinstance(tex, ProxyResponseError) + and tex.status_code != 500 + and not _proxy_response_error_is_transient_stream_retry(tex) + ): error = _parse_openai_error(tex.payload) code = _normalize_error_code( error.code if error else None, @@ -2079,12 +2152,19 @@ async def _retry_account_model_rejection( ) break raise - error_code = tex.code if isinstance(tex, _TransientStreamError) else "server_error" - error_payload: UpstreamError = ( - tex.error - if isinstance(tex, _TransientStreamError) - else _upstream_error_from_openai(_parse_openai_error(tex.payload)) - ) + if isinstance(tex, _TransientStreamError): + error_code = tex.code + error_payload: UpstreamError = tex.error + else: + parsed_error = _parse_openai_error(tex.payload) + error_code = ( + _normalize_error_code( + parsed_error.code if parsed_error else None, + parsed_error.type if parsed_error else None, + ) + or "server_error" + ) + error_payload = _upstream_error_from_openai(parsed_error) error_message = str(error_payload.get("message") or "") recovery_decision = await _wait_for_process_network_recovery( account, @@ -2106,19 +2186,20 @@ async def _retry_account_model_rejection( and _facade()._remaining_budget_seconds(deadline) > 0 and not settlement.downstream_visible ): - delay = backoff_seconds(transient_retries) - _facade().logger.info( - "Transient stream error, retrying same account " - "request_id=%s account_id=%s retry=%s/%s delay=%.2fs code=%s", - request_id, - account.id, - transient_retries, - _facade()._MAX_TRANSIENT_SAME_ACCOUNT_RETRIES, - delay, - error_code, - ) - await asyncio.sleep(delay) - continue # inner loop: retry same account + delay = _transient_stream_retry_delay(deadline, transient_retries) + if delay is not None: + _facade().logger.info( + "Transient stream error, retrying same account " + "request_id=%s account_id=%s retry=%s/%s delay=%.2fs code=%s", + request_id, + account.id, + transient_retries, + _facade()._MAX_TRANSIENT_SAME_ACCOUNT_RETRIES, + delay, + error_code, + ) + await asyncio.sleep(delay) + continue # inner loop: retry same account # Exhausted same-account retries — penalize and failover _facade().logger.warning( "Transient retries exhausted for account " @@ -2135,6 +2216,7 @@ async def _retry_account_model_rejection( # Preserve last ProxyResponseError for propagate_http_errors path. if isinstance(tex, ProxyResponseError): last_transient_exc = tex + last_exhausted_transient_exc = tex if isinstance(tex, _TransientStreamError) and ( tex.preserve_on_selection_exhausted or error_code == "stream_incomplete" ): @@ -2844,6 +2926,20 @@ async def _retry_account_model_rejection( _apply_error_metadata(event["response"]["error"], error) yield format_sse_event(event) return + if last_exhausted_transient_exc is not None: + error = _parse_openai_error(last_exhausted_transient_exc.payload) + error_code = _normalize_error_code(error.code if error else None, error.type if error else None) + error_message = error.message if error else None + event = response_failed_event( + error_code or "upstream_error", + error_message or "Upstream error", + 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) + yield format_sse_event(event) + return retries_exhausted_msg = "No available accounts after retries" _facade().logger.warning( diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index 7c26e6f81f..6ed24f7185 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -1255,6 +1255,8 @@ def _record_response_event(request_state: _WebSocketRequestState | None, event_t request_state.last_upstream_activity_at = time.monotonic() if event_type in {"response.failed", "response.incomplete"}: return + if event_type.startswith("response.output_"): + request_state.upstream_model_output_seen = True request_state.response_event_count += 1 diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index 50308cbc78..97fa751e0e 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -5480,7 +5480,8 @@ async def _finalize_websocket_request_state( if ( error_code == "stream_incomplete" and request_state.previous_response_id is not None - and error_message == "Upstream websocket closed before response.completed" + and isinstance(error_message, str) + and error_message.startswith("Upstream websocket closed before response.completed") ): settlement.account_health_error = False proxy._cancel_request_state_api_key_reservation_heartbeat(request_state) diff --git a/openspec/changes/retry-server-is-overloaded/design.md b/openspec/changes/retry-server-is-overloaded/design.md index c7630a10a9..dceb442a08 100644 --- a/openspec/changes/retry-server-is-overloaded/design.md +++ b/openspec/changes/retry-server-is-overloaded/design.md @@ -21,15 +21,31 @@ request. downstream-visible output and remains bounded by the existing request budget. 4. Cover the classifier plus both raw and HTTP-bridge routed streaming paths so the fix is not limited to a helper-only assertion. +5. For native Codex bridge requests only, distinguish lifecycle progress + (`response.created` and `response.in_progress`) from actual model output. If + either overload code terminates an accepted response before model output, + hold the response-create admission slot, use the existing transient + backoff, reconnect the same account, and replay the unchanged request once. +6. Preserve the retry response's real upstream ID. Codex ignores IDs on + `response.created` and derives its next `previous_response_id` from + `response.completed`; rewriting the successful ID to the failed attempt's + ID would corrupt the next turn's continuity. ## Risks and Mitigations - **Duplicate generation:** Existing stream settlement visibility checks prevent - replay after downstream-visible output. + replay after downstream-visible output. The bridge additionally records any + reasoning, item, or tool output before allowing accepted-response replay. - **Unbounded retry:** The change reuses existing retry counters and request deadlines; it adds no new loop. - **Unknown client errors:** Only the exact upstream overload code is added, so authentication and invalid-request failures remain non-retryable. +- **Public SDK lifecycle:** Accepted-response replay is disabled when the OpenAI + SDK contract is enforced, so public streams never receive a second + `response.created` event. +- **Continuity drift:** The retry stays on the account that accepted the parent + response, preserves the original `previous_response_id`, and exposes the + successful retry's actual completed ID for the following turn. ## Verification @@ -37,4 +53,10 @@ request. - Integration test a first-event `server_is_overloaded` envelope followed by a successful attempt through `/backend-api/codex/responses`, with and without the HTTP responses session bridge. +- Integration test the production sequence `response.created`, + `response.in_progress`, `server_is_overloaded`: assert backoff precedes a + same-account replay, the parent anchor is unchanged, and the next request + anchors on the retry's completed response ID. +- Unit-test fail-closed guards for public SDK streams, prior model output, and + exhausted replay budget. - Run Ruff, focused pytest, and strict OpenSpec validation. diff --git a/openspec/changes/retry-server-is-overloaded/proposal.md b/openspec/changes/retry-server-is-overloaded/proposal.md index 20d44ad401..fc0f907720 100644 --- a/openspec/changes/retry-server-is-overloaded/proposal.md +++ b/openspec/changes/retry-server-is-overloaded/proposal.md @@ -13,6 +13,9 @@ using its bounded pre-visible retry and failover behavior. accompanies the streamed error envelope. - Allow the streaming retry path to handle `server_is_overloaded` with the same bounded retry behavior as other transient server errors. +- Let the native Codex HTTP bridge wait briefly and retry once on the same + account when overload arrives after `response.created` but before any model + output, including continuation requests. - Add unit and externally routed integration coverage for the no-5xx SSE path. - Extend the Responses compatibility requirement to cover both known upstream overload codes. @@ -23,4 +26,6 @@ using its bounded pre-visible retry and failover behavior. code instead of stopping an agent mid-task. - Retry remains bounded by the existing stream retry budget and remains disabled once downstream-visible output makes replay unsafe. +- The accepted-response retry is limited to the native Codex stream contract; + public OpenAI SDK streams retain their single-response lifecycle. - No API, schema, or configuration changes are introduced. diff --git a/openspec/changes/retry-server-is-overloaded/specs/responses-api-compat/spec.md b/openspec/changes/retry-server-is-overloaded/specs/responses-api-compat/spec.md index 4872f5ddb7..ad6b5d735f 100644 --- a/openspec/changes/retry-server-is-overloaded/specs/responses-api-compat/spec.md +++ b/openspec/changes/retry-server-is-overloaded/specs/responses-api-compat/spec.md @@ -22,9 +22,51 @@ When `classify_upstream_failure` observes an upstream error envelope whose `code - **THEN** the returned `failure_class` is `retryable_transient` - **AND** the streaming retry layer is eligible to retry the request before surfacing the terminal overload event +#### Scenario: Non-streaming Responses retries status-level overload JSON + +- **GIVEN** a non-streaming `/v1/responses` request has not emitted any downstream output +- **WHEN** upstream returns an HTTP status error envelope with `code="server_is_overloaded"` +- **THEN** the retry layer MUST treat the envelope as a transient overload even when the HTTP status is not 500 +- **AND** it MUST retry the same request within the bounded transient retry budget before returning the overload error to the client + +#### Scenario: Streaming Responses retries initial output-free EOF + +- **GIVEN** a streaming `/v1/responses` request is not anchored to a previous response +- **WHEN** upstream closes before emitting the first SSE event +- **THEN** the retry layer MUST treat the close as a transient output-free failure +- **AND** it MUST retry the same request within the bounded transient retry budget before returning `stream_incomplete` +- **AND** it MUST NOT perform this retry for anchored continuations, after any downstream-visible output, or after the retry budget is exhausted + #### Scenario: HTTP bridge retries a pre-created overload event - **GIVEN** the HTTP responses session bridge is enabled - **WHEN** the first upstream `response.failed` or `error` event has `code="overloaded_error"` or `code="server_is_overloaded"` - **THEN** the bridge MUST retry the pre-created request before forwarding that terminal event - **AND** the bridge MUST preserve its existing no-replay behavior after downstream-visible output or for other fail-fast error codes + +#### Scenario: Native Codex bridge retries an accepted output-free overload + +- **GIVEN** the native Codex HTTP responses session bridge has accepted a continuation request on an account +- **AND** upstream has emitted `response.created` and optionally `response.in_progress`, but no reasoning, text, item, or tool output +- **WHEN** upstream terminates that response with `code="overloaded_error"` or `code="server_is_overloaded"` +- **THEN** the bridge MUST wait for a bounded transient backoff and replay the unchanged request exactly once on the same account +- **AND** the replay MUST preserve the original parent `previous_response_id` +- **AND** the bridge MUST expose the successful replay's actual `response.completed` ID so the next continuation anchors on the successful response +- **AND** the bridge MUST NOT perform this accepted-response replay for public OpenAI SDK streams, after any model output, while another request is pending, or after the replay budget is exhausted + +#### Scenario: Native Codex bridge retries accepted output-free abrupt closes + +- **GIVEN** the native Codex HTTP responses session bridge has accepted a continuation request on an account +- **AND** upstream has emitted only lifecycle events such as `response.created` or `response.in_progress` +- **WHEN** the upstream websocket closes before `response.completed` without a terminal error event +- **THEN** the bridge MUST wait for a bounded transient backoff and replay the unchanged request exactly once on the same account +- **AND** the replay MUST preserve the original parent `previous_response_id` +- **AND** the bridge MUST NOT perform this accepted-response replay for public OpenAI SDK streams, after any reasoning, text, item, or tool output, while another request is pending, or after the replay budget is exhausted +- **AND** when an untyped transport close is not replayed because those safety gates block replay or the replay budget is exhausted, the bridge MUST surface `stream_incomplete` without recording an upstream account-health penalty + +#### Scenario: Previous-response stream close suffix remains account-neutral + +- **GIVEN** a websocket Responses continuation is anchored by `previous_response_id` +- **WHEN** finalization receives `response.failed` with `code="stream_incomplete"` and an `Upstream websocket closed before response.completed` message that includes transport detail suffix text +- **THEN** the request log MUST record `stream_incomplete` +- **AND** the proxy MUST NOT record an upstream account-health penalty for that account diff --git a/openspec/changes/retry-server-is-overloaded/tasks.md b/openspec/changes/retry-server-is-overloaded/tasks.md index 52a306ee26..b810a02e58 100644 --- a/openspec/changes/retry-server-is-overloaded/tasks.md +++ b/openspec/changes/retry-server-is-overloaded/tasks.md @@ -3,4 +3,13 @@ - [x] Add unit coverage for classification without an HTTP 5xx status. - [x] Add integration coverage proving the public Responses route retries the first terminal overload event. - [x] Cover the production-default HTTP responses session bridge retry path. +- [x] Retry one native Codex continuation overload that arrives after + `response.created` but before model output, with same-account backoff and + preserved completed-response continuity. +- [x] Retry one native Codex continuation whose upstream websocket closes after + lifecycle-only events but before model output, with the same safety gates. +- [x] Add lifecycle negative controls for public SDK streams, prior model + output, and exhausted replay budget. +- [x] Keep unreplayable untyped bridge closes and previous-response close + suffixes account-neutral. - [x] Update the Responses compatibility requirement and validate OpenSpec. diff --git a/openspec/specs/responses-api-compat/spec.md b/openspec/specs/responses-api-compat/spec.md index 577c772514..a8464dd0a6 100644 --- a/openspec/specs/responses-api-compat/spec.md +++ b/openspec/specs/responses-api-compat/spec.md @@ -201,6 +201,21 @@ When `classify_upstream_failure` observes an upstream error envelope whose `code - **THEN** the returned `failure_class` is `retryable_transient` - **AND** the streaming retry layer is eligible to retry the request before surfacing the terminal overload event +#### Scenario: Non-streaming Responses retries status-level overload JSON + +- **GIVEN** a non-streaming `/v1/responses` request has not emitted any downstream output +- **WHEN** upstream returns an HTTP status error envelope with `code="server_is_overloaded"` +- **THEN** the retry layer MUST treat the envelope as a transient overload even when the HTTP status is not 500 +- **AND** it MUST retry the same request within the bounded transient retry budget before returning the overload error to the client + +#### Scenario: Streaming Responses retries initial output-free EOF + +- **GIVEN** a streaming `/v1/responses` request is not anchored to a previous response +- **WHEN** upstream closes before emitting the first SSE event +- **THEN** the retry layer MUST treat the close as a transient output-free failure +- **AND** it MUST retry the same request within the bounded transient retry budget before returning `stream_incomplete` +- **AND** it MUST NOT perform this retry for anchored continuations, after any downstream-visible output, or after the retry budget is exhausted + #### Scenario: HTTP bridge retries a pre-created overload event - **GIVEN** the HTTP responses session bridge is enabled @@ -208,6 +223,33 @@ When `classify_upstream_failure` observes an upstream error envelope whose `code - **THEN** the bridge MUST retry the pre-created request before forwarding that terminal event - **AND** the bridge MUST preserve its existing no-replay behavior after downstream-visible output or for other fail-fast error codes +#### Scenario: Native Codex bridge retries an accepted output-free overload + +- **GIVEN** the native Codex HTTP responses session bridge has accepted a continuation request on an account +- **AND** upstream has emitted `response.created` and optionally `response.in_progress`, but no reasoning, text, item, or tool output +- **WHEN** upstream terminates that response with `code="overloaded_error"` or `code="server_is_overloaded"` +- **THEN** the bridge MUST wait for a bounded transient backoff and replay the unchanged request exactly once on the same account +- **AND** the replay MUST preserve the original parent `previous_response_id` +- **AND** the bridge MUST expose the successful replay's actual `response.completed` ID so the next continuation anchors on the successful response +- **AND** the bridge MUST NOT perform this accepted-response replay for public OpenAI SDK streams, after any model output, while another request is pending, or after the replay budget is exhausted + +#### Scenario: Native Codex bridge retries accepted output-free abrupt closes + +- **GIVEN** the native Codex HTTP responses session bridge has accepted a continuation request on an account +- **AND** upstream has emitted only lifecycle events such as `response.created` or `response.in_progress` +- **WHEN** the upstream websocket closes before `response.completed` without a terminal error event +- **THEN** the bridge MUST wait for a bounded transient backoff and replay the unchanged request exactly once on the same account +- **AND** the replay MUST preserve the original parent `previous_response_id` +- **AND** the bridge MUST NOT perform this accepted-response replay for public OpenAI SDK streams, after any reasoning, text, item, or tool output, while another request is pending, or after the replay budget is exhausted +- **AND** when an untyped transport close is not replayed because those safety gates block replay or the replay budget is exhausted, the bridge MUST surface `stream_incomplete` without recording an upstream account-health penalty + +#### Scenario: Previous-response stream close suffix remains account-neutral + +- **GIVEN** a websocket Responses continuation is anchored by `previous_response_id` +- **WHEN** finalization receives `response.failed` with `code="stream_incomplete"` and an `Upstream websocket closed before response.completed` message that includes transport detail suffix text +- **THEN** the request log MUST record `stream_incomplete` +- **AND** the proxy MUST NOT record an upstream account-health penalty for that account + ### Requirement: Strict function tool parameter schemas are pre-validated The service MUST pre-validate the JSON schema attached to a function tool when that tool sets `strict: true`, before opening any upstream connection. The validation rules mirror OpenAI's Structured Outputs strict-mode policy (https://platform.openai.com/docs/guides/structured-outputs) and the existing `enforce_strict_text_format` policy for `text.format.json_schema`: diff --git a/tests/integration/test_http_responses_bridge.py b/tests/integration/test_http_responses_bridge.py index 41ee741d2a..791190d942 100644 --- a/tests/integration/test_http_responses_bridge.py +++ b/tests/integration/test_http_responses_bridge.py @@ -34,6 +34,7 @@ from app.db.session import SessionLocal from app.dependencies import get_proxy_service_for_app from app.modules.proxy._service import support as proxy_support +from app.modules.proxy._service.http_bridge import request_submit as http_bridge_request_submit_module from app.modules.proxy._service.http_bridge import streaming as http_bridge_streaming_module from app.modules.proxy._service.http_bridge.helpers import ( _release_http_bridge_unanchored_handoff, @@ -510,6 +511,109 @@ async def send_text(self, text: str) -> None: ) +class _AcceptedContinuationOverloadUpstreamWebSocket(_FakeBridgeUpstreamWebSocket): + def __init__( + self, + order: list[tuple[str, object]], + *, + error_code: str = "server_is_overloaded", + error_message: str = "Our servers are currently overloaded. Please try again later.", + ) -> None: + super().__init__("resp_capacity_parent") + self.order = order + self.error_code = error_code + self.error_message = error_message + + async def send_text(self, text: str) -> None: + if not self.sent_text: + self.order.append(("parent_send", None)) + await super().send_text(text) + return + self.sent_text.append(text) + self.order.append(("overload_send", None)) + response_id = "resp_capacity_failed" + await self._messages.put( + _FakeUpstreamMessage( + "text", + text=json.dumps( + { + "type": "response.created", + "response": {"id": response_id, "object": "response", "status": "in_progress"}, + }, + separators=(",", ":"), + ), + ) + ) + await self._messages.put( + _FakeUpstreamMessage( + "text", + text=json.dumps( + { + "type": "response.in_progress", + "response": {"id": response_id, "object": "response", "status": "in_progress"}, + }, + separators=(",", ":"), + ), + ) + ) + await self._messages.put( + _FakeUpstreamMessage( + "text", + text=json.dumps( + { + "type": "error", + "error": { + "type": "service_unavailable_error", + "code": self.error_code, + "message": self.error_message, + }, + }, + separators=(",", ":"), + ), + ) + ) + + +class _AcceptedContinuationAbruptCloseUpstreamWebSocket(_FakeBridgeUpstreamWebSocket): + def __init__(self, order: list[tuple[str, object]]) -> None: + super().__init__("resp_capacity_parent") + self.order = order + + async def send_text(self, text: str) -> None: + if not self.sent_text: + self.order.append(("parent_send", None)) + await super().send_text(text) + return + self.sent_text.append(text) + self.order.append(("close_send", None)) + response_id = "resp_capacity_closed" + await self._messages.put( + _FakeUpstreamMessage( + "text", + text=json.dumps( + { + "type": "response.created", + "response": {"id": response_id, "object": "response", "status": "in_progress"}, + }, + separators=(",", ":"), + ), + ) + ) + await self._messages.put( + _FakeUpstreamMessage( + "text", + text=json.dumps( + { + "type": "response.in_progress", + "response": {"id": response_id, "object": "response", "status": "in_progress"}, + }, + separators=(",", ":"), + ), + ) + ) + await self._messages.put(_FakeUpstreamMessage("close", error="no close frame received or sent")) + + class _CreatedOnlyUpstreamWebSocket(_FakeBridgeUpstreamWebSocket): async def send_text(self, text: str) -> None: self.sent_text.append(text) @@ -8804,6 +8908,337 @@ async def fail_legacy_stream(*args, **kwargs): assert connect_count == 2 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error_code", "error_message"), + [ + ("server_is_overloaded", "Our servers are currently overloaded. Please try again later."), + ("model_at_capacity", "Selected model is at capacity. Please try a different model."), + ], +) +async def test_backend_responses_http_bridge_waits_and_retries_accepted_output_free_server_overload( + async_client, + monkeypatch, + error_code, + error_message, +): + _install_bridge_settings(monkeypatch, enabled=True) + account_id = await _import_account( + async_client, + "acc_http_bridge_accepted_server_overload", + "http-bridge-accepted-server-overload@example.com", + ) + account = await _get_account(account_id) + order: list[tuple[str, object]] = [] + first_upstream = _AcceptedContinuationOverloadUpstreamWebSocket( + order, + error_code=error_code, + error_message=error_message, + ) + retry_upstream = _FakeBridgeUpstreamWebSocket("resp_capacity_retry") + upstreams = [first_upstream, retry_upstream] + connect_account_ids: list[str | None] = [] + connect_count = 0 + + class _RecordingAsyncio: + def __getattr__(self, name: str) -> object: + return getattr(asyncio, name) + + async def sleep(self, delay: float) -> None: + order.append(("sleep", delay)) + + async def fake_select_account_with_budget( + self, + deadline, + *, + request_id, + kind, + request_stage="first_turn", + sticky_key, + sticky_kind, + reallocate_sticky, + sticky_max_age_seconds, + prefer_earlier_reset_accounts, + routing_strategy, + model, + exclude_account_ids=None, + additional_limit_name=None, + api_key=None, + preferred_account_id=None, + ): + del preferred_account_id + del ( + self, + deadline, + request_id, + kind, + request_stage, + sticky_key, + sticky_kind, + reallocate_sticky, + sticky_max_age_seconds, + prefer_earlier_reset_accounts, + routing_strategy, + model, + exclude_account_ids, + additional_limit_name, + api_key, + ) + 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, + *, + base_url=None, + session=None, + ): + del headers, access_token, base_url, session + nonlocal connect_count + order.append(("connect", connect_count)) + connect_account_ids.append(account_id_header) + upstream = upstreams[connect_count] + connect_count += 1 + return upstream + + async def fail_legacy_stream(*args, **kwargs): + raise AssertionError("legacy core_stream_responses path must not be used when HTTP bridge is enabled") + + 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, "connect_responses_websocket", fake_connect_responses_websocket) + monkeypatch.setattr(proxy_module, "core_stream_responses", fail_legacy_stream) + monkeypatch.setattr(http_bridge_request_submit_module, "backoff_seconds", lambda attempt: 0.25) + monkeypatch.setattr(http_bridge_request_submit_module, "asyncio", _RecordingAsyncio()) + + prompt_cache_key = "accepted-server-overload-retry-key" + parent_events = await _collect_sse_events( + async_client, + "/backend-api/codex/responses", + json_body={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "parent", + "prompt_cache_key": prompt_cache_key, + "stream": True, + }, + ) + parent_response_id = parent_events[-1]["response"]["id"] + assert parent_response_id == "resp_capacity_parent_1" + + retry_events = await _collect_sse_events( + async_client, + "/backend-api/codex/responses", + json_body={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "retry accepted overload", + "previous_response_id": parent_response_id, + "prompt_cache_key": prompt_cache_key, + "stream": True, + }, + ) + + retry_created_ids = [event["response"]["id"] for event in retry_events if event["type"] == "response.created"] + assert retry_created_ids == ["resp_capacity_failed", "resp_capacity_retry_1"] + assert not any(event["type"] in {"error", "response.failed"} for event in retry_events) + retry_response_id = retry_events[-1]["response"]["id"] + assert retry_response_id == "resp_capacity_retry_1" + assert connect_count == 2 + assert connect_account_ids == [ + "acc_http_bridge_accepted_server_overload", + "acc_http_bridge_accepted_server_overload", + ] + assert order.index(("overload_send", None)) < order.index(("sleep", 0.25)) < order.index(("connect", 1)) + + first_attempt_payload = json.loads(first_upstream.sent_text[1]) + retry_payload = json.loads(retry_upstream.sent_text[0]) + assert first_attempt_payload["previous_response_id"] == parent_response_id + assert retry_payload["previous_response_id"] == parent_response_id + + follow_up_events = await _collect_sse_events( + async_client, + "/backend-api/codex/responses", + json_body={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "continue after retry", + "previous_response_id": retry_response_id, + "prompt_cache_key": prompt_cache_key, + "stream": True, + }, + ) + + follow_up_payload = json.loads(retry_upstream.sent_text[1]) + assert follow_up_payload["previous_response_id"] == retry_response_id + assert follow_up_events[-1]["response"]["id"] == "resp_capacity_retry_2" + assert connect_count == 2 + + +@pytest.mark.asyncio +async def test_backend_responses_http_bridge_waits_and_retries_accepted_output_free_abrupt_close( + async_client, + monkeypatch, +): + _install_bridge_settings(monkeypatch, enabled=True) + account_id = await _import_account( + async_client, + "acc_http_bridge_accepted_abrupt_close", + "http-bridge-accepted-abrupt-close@example.com", + ) + account = await _get_account(account_id) + order: list[tuple[str, object]] = [] + first_upstream = _AcceptedContinuationAbruptCloseUpstreamWebSocket(order) + retry_upstream = _FakeBridgeUpstreamWebSocket("resp_capacity_retry") + upstreams = [first_upstream, retry_upstream] + connect_account_ids: list[str | None] = [] + connect_count = 0 + + class _RecordingAsyncio: + def __getattr__(self, name: str) -> object: + return getattr(asyncio, name) + + async def sleep(self, delay: float) -> None: + order.append(("sleep", delay)) + + async def fake_select_account_with_budget( + self, + deadline, + *, + request_id, + kind, + request_stage="first_turn", + sticky_key, + sticky_kind, + reallocate_sticky, + sticky_max_age_seconds, + prefer_earlier_reset_accounts, + routing_strategy, + model, + exclude_account_ids=None, + additional_limit_name=None, + api_key=None, + preferred_account_id=None, + ): + del preferred_account_id + del ( + self, + deadline, + request_id, + kind, + request_stage, + sticky_key, + sticky_kind, + reallocate_sticky, + sticky_max_age_seconds, + prefer_earlier_reset_accounts, + routing_strategy, + model, + exclude_account_ids, + additional_limit_name, + api_key, + ) + 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, + *, + base_url=None, + session=None, + ): + del headers, access_token, base_url, session + nonlocal connect_count + order.append(("connect", connect_count)) + connect_account_ids.append(account_id_header) + upstream = upstreams[connect_count] + connect_count += 1 + return upstream + + async def fail_legacy_stream(*args, **kwargs): + raise AssertionError("legacy core_stream_responses path must not be used when HTTP bridge is enabled") + + 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, "connect_responses_websocket", fake_connect_responses_websocket) + monkeypatch.setattr(proxy_module, "core_stream_responses", fail_legacy_stream) + monkeypatch.setattr(http_bridge_request_submit_module, "backoff_seconds", lambda attempt: 0.25) + monkeypatch.setattr(http_bridge_request_submit_module, "asyncio", _RecordingAsyncio()) + + prompt_cache_key = "accepted-abrupt-close-retry-key" + parent_events = await _collect_sse_events( + async_client, + "/backend-api/codex/responses", + json_body={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "parent", + "prompt_cache_key": prompt_cache_key, + "stream": True, + }, + ) + parent_response_id = parent_events[-1]["response"]["id"] + assert parent_response_id == "resp_capacity_parent_1" + + retry_events = await _collect_sse_events( + async_client, + "/backend-api/codex/responses", + json_body={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "retry accepted abrupt close", + "previous_response_id": parent_response_id, + "prompt_cache_key": prompt_cache_key, + "stream": True, + }, + ) + + retry_created_ids = [event["response"]["id"] for event in retry_events if event["type"] == "response.created"] + assert retry_created_ids == ["resp_capacity_closed", "resp_capacity_retry_1"] + assert not any(event["type"] in {"error", "response.failed"} for event in retry_events) + retry_response_id = retry_events[-1]["response"]["id"] + assert retry_response_id == "resp_capacity_retry_1" + assert connect_count == 2 + assert connect_account_ids == [ + "acc_http_bridge_accepted_abrupt_close", + "acc_http_bridge_accepted_abrupt_close", + ] + assert order.index(("close_send", None)) < order.index(("sleep", 0.25)) < order.index(("connect", 1)) + + first_attempt_payload = json.loads(first_upstream.sent_text[1]) + retry_payload = json.loads(retry_upstream.sent_text[0]) + assert first_attempt_payload["previous_response_id"] == parent_response_id + assert retry_payload["previous_response_id"] == parent_response_id + + follow_up_events = await _collect_sse_events( + async_client, + "/backend-api/codex/responses", + json_body={ + "model": "gpt-5.1", + "instructions": "Return exactly OK.", + "input": "continue after retry", + "previous_response_id": retry_response_id, + "prompt_cache_key": prompt_cache_key, + "stream": True, + }, + ) + + follow_up_payload = json.loads(retry_upstream.sent_text[1]) + assert follow_up_payload["previous_response_id"] == retry_response_id + assert follow_up_events[-1]["response"]["id"] == "resp_capacity_retry_2" + assert connect_count == 2 + + @pytest.mark.asyncio async def test_v1_responses_http_bridge_rejects_oversized_response_create_before_upstream( async_client, diff --git a/tests/integration/test_proxy_api_extended.py b/tests/integration/test_proxy_api_extended.py index 44b00dee0d..83ee323702 100644 --- a/tests/integration/test_proxy_api_extended.py +++ b/tests/integration/test_proxy_api_extended.py @@ -2560,6 +2560,125 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, assert log.failure_detail == "upstream_eof_before_terminal_event" +@pytest.mark.asyncio +async def test_proxy_stream_retries_initial_upstream_eof_before_visible_output(async_client, monkeypatch): + expected_account_id = await _import_account( + async_client, + "acc_stream_initial_eof_retry", + "stream-initial-eof-retry@example.com", + ) + calls = 0 + + async def fake_stream(payload, headers, access_token, account_id, base_url=None, raise_for_status=False): + nonlocal calls + del payload, headers, access_token, account_id, base_url, raise_for_status + calls += 1 + if calls == 1: + if False: + yield "" + return + yield _sse_event({"type": "response.completed", "response": {"id": "resp_initial_eof_retry_ok"}}) + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + monkeypatch.setattr("app.modules.proxy._service.streaming.retry.backoff_seconds", lambda _attempt: 0.0) + + payload = {"model": "gpt-5.1", "instructions": "hi", "input": [], "stream": True} + async with async_client.stream( + "POST", + "/backend-api/codex/responses", + json=payload, + headers={"x-request-id": "req_stream_initial_eof_retry"}, + ) as resp: + assert resp.status_code == 200 + lines = [line async for line in resp.aiter_lines() if line] + + events = [ + json.loads(line[6:]) for line in lines if line.startswith("data: ") and not line.startswith("data: [DONE]") + ] + assert events[-1]["type"] == "response.completed" + assert events[-1]["response"]["id"] == "resp_initial_eof_retry_ok" + assert not [event for event in events if event.get("type") == "response.failed"] + assert calls == 2 + + async with SessionLocal() as session: + result = await session.execute( + select(RequestLog) + .where(RequestLog.account_id == expected_account_id) + .order_by(RequestLog.requested_at.desc()) + ) + logs = list(result.scalars().all()) + assert logs[0].status == "success" + initial_failure = next((item for item in logs if item.error_code == "stream_incomplete"), None) + assert initial_failure is not None + assert initial_failure.error_message == "Upstream websocket closed before response.completed" + assert initial_failure.latency_first_token_ms is None + + +@pytest.mark.asyncio +async def test_proxy_stream_does_not_retry_anchored_initial_upstream_eof(async_client, monkeypatch): + expected_account_id = await _import_account( + async_client, + "acc_stream_anchored_initial_eof", + "stream-anchored-initial-eof@example.com", + ) + calls = 0 + + async def fake_stream(payload, headers, access_token, account_id, base_url=None, raise_for_status=False): + nonlocal calls + del payload, headers, access_token, account_id, base_url, raise_for_status + calls += 1 + if False: + yield "" + return + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + monkeypatch.setattr("app.modules.proxy._service.streaming.retry.backoff_seconds", lambda _attempt: 0.0) + + payload = { + "model": "gpt-5.1", + "instructions": "hi", + "input": [], + "previous_response_id": "resp_parent", + "stream": True, + } + async with async_client.stream( + "POST", + "/backend-api/codex/responses", + json=payload, + headers={"x-request-id": "req_stream_anchored_initial_eof"}, + ) as resp: + assert resp.status_code == 200 + lines = [line async for line in resp.aiter_lines() if line] + + events = [ + event + for line in lines + if line.startswith("data: ") and not line.startswith("data: [DONE]") + if (event := json.loads(line[6:])).get("type") != "codex.keepalive" + ] + assert len(events) == 1 + assert events[0]["type"] == "response.failed" + assert events[0]["response"]["id"] == "req_stream_anchored_initial_eof" + assert events[0]["response"]["error"] == { + "code": "stream_incomplete", + "message": "Upstream websocket closed before response.completed", + "type": "server_error", + } + assert calls == 1 + + async with SessionLocal() as session: + result = await session.execute( + select(RequestLog) + .where(RequestLog.account_id == expected_account_id) + .order_by(RequestLog.requested_at.desc()) + ) + log = result.scalars().first() + assert log is not None + assert log.status == "error" + assert log.error_code == "stream_incomplete" + assert log.error_message == "Upstream websocket closed before response.completed" + + @pytest.mark.asyncio async def test_proxy_stream_classifies_core_generated_eof_failure(async_client, monkeypatch): expected_account_id = await _import_account(async_client, "acc_stream_core_eof", "stream-core-eof@example.com") @@ -2650,6 +2769,7 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, ][-1] assert event["type"] == "response.failed" assert event["response"]["error"]["code"] == "stream_incomplete" + assert event["response"]["error"]["message"] == "Upstream closed stream without completion" async with SessionLocal() as session: result = await session.execute( diff --git a/tests/integration/test_proxy_transient_retry.py b/tests/integration/test_proxy_transient_retry.py index 85171374a9..0227c793bb 100644 --- a/tests/integration/test_proxy_transient_retry.py +++ b/tests/integration/test_proxy_transient_retry.py @@ -357,8 +357,8 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, @pytest.mark.asyncio -async def test_stream_empty_upstream_body_surfaces_without_replay(async_client, monkeypatch): - """An untyped empty upstream stream may be post-dispatch, so it is not replayed.""" +async def test_stream_anchored_empty_upstream_body_surfaces_without_replay(async_client, monkeypatch): + """An anchored empty upstream stream may be post-dispatch, so it is not replayed.""" await _import_account(async_client, "acc_empty_body_no_replay", "empty-body-no-replay@example.com") call_count = 0 @@ -375,7 +375,13 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) - payload = {"model": "gpt-5.1", "instructions": "hi", "input": [], "stream": True} + payload = { + "model": "gpt-5.1", + "instructions": "hi", + "input": [], + "previous_response_id": "resp-empty-body-parent", + "stream": True, + } async with async_client.stream("POST", "/backend-api/codex/responses", json=payload) as resp: assert resp.status_code == 200 lines = [line async for line in resp.aiter_lines() if line] @@ -988,6 +994,40 @@ async def fake_stream(payload, headers, access_token, account_id, base_url=None, assert call_count == 3 +@pytest.mark.asyncio +async def test_v1_responses_non_streaming_server_overload_status_retries_same_account(async_client, monkeypatch): + """OpenAI SDK non-streaming responses can receive overloaded status JSON. + + The status may be non-500, but the normalized upstream code is still a + transient overload before any downstream response is visible. + """ + await _import_account(async_client, "acc_prop_overload", "prop-overload@example.com") + + call_count = 0 + seen_account_ids: list[str | None] = [] + + async def fake_stream(payload, headers, access_token, account_id, base_url=None, raise_for_status=False): + nonlocal call_count + call_count += 1 + seen_account_ids.append(account_id) + if call_count == 1: + raise ProxyResponseError( + 429, + openai_error("server_is_overloaded", "Our servers are currently overloaded. Please try again later."), + failure_phase="status", + ) + yield _success_sse_event("resp_nonstream_overload_retry_ok") + + monkeypatch.setattr(proxy_module, "core_stream_responses", fake_stream) + + payload = {"model": "gpt-5.1", "input": "hi"} + response = await async_client.post("/v1/responses", json=payload) + + assert response.status_code == 200 + assert response.json()["id"] == "resp_nonstream_overload_retry_ok" + assert seen_account_ids == ["acc_prop_overload", "acc_prop_overload"] + + # =========================================================================== # Compact — HTTP 500 retry # =========================================================================== diff --git a/tests/unit/test_proxy_http_bridge.py b/tests/unit/test_proxy_http_bridge.py index b0a61903b6..b14be709ac 100644 --- a/tests/unit/test_proxy_http_bridge.py +++ b/tests/unit/test_proxy_http_bridge.py @@ -371,6 +371,536 @@ def test_verified_replay_model_fork_preserves_recovery_kind() -> None: assert fork_key != replay_key +def _accepted_capacity_retry_state(**overrides: object) -> proxy_service._WebSocketRequestState: + values: dict[str, object] = { + "request_id": "req-accepted-capacity", + "model": "gpt-5.6-sol", + "service_tier": None, + "reasoning_effort": None, + "api_key_reservation": None, + "started_at": 1.0, + "transport": "http", + "enforce_openai_sdk_contract": False, + "request_text": json.dumps( + { + "type": "response.create", + "previous_response_id": "resp-parent", + "input": "continue", + } + ), + "previous_response_id": "resp-parent", + "response_id": "resp-capacity-failed", + "response_event_count": 2, + "event_queue": asyncio.Queue(), + } + values.update(overrides) + return proxy_service._WebSocketRequestState(**cast(Any, values)) + + +def _server_overloaded_event_payload() -> dict[str, object]: + return { + "type": "error", + "error": { + "type": "service_unavailable_error", + "code": "server_is_overloaded", + "message": "Our servers are currently overloaded. Please try again later.", + }, + } + + +def test_terminal_capacity_retry_accepts_native_output_free_continuation() -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_terminal_capacity_retry_error_code( + _accepted_capacity_retry_state(), + event_type="error", + payload=cast(Any, _server_overloaded_event_payload()), + has_other_pending_requests=False, + ) + == "server_is_overloaded" + ) + + +@pytest.mark.parametrize( + "payload", + [ + {**_server_overloaded_event_payload(), "usage": {"output_tokens": 1}}, + { + "type": "response.failed", + "response": { + "error": _server_overloaded_event_payload()["error"], + "usage": {"output_tokens_details": {"reasoning_tokens": 1}}, + }, + }, + ], +) +def test_terminal_capacity_retry_rejects_terminal_usage_output(payload: dict[str, object]) -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_terminal_capacity_retry_error_code( + _accepted_capacity_retry_state(), + event_type=cast(str, payload["type"]), + payload=cast(Any, payload), + has_other_pending_requests=False, + ) + is None + ) + + +def test_terminal_capacity_retry_rejects_generic_model_change_message() -> None: + payload = { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "invalid_request_error", + "message": "Please try a different model.", + }, + } + + assert ( + http_bridge_upstream_events_module._http_bridge_terminal_capacity_retry_error_code( + _accepted_capacity_retry_state(), + event_type="error", + payload=cast(Any, payload), + has_other_pending_requests=False, + ) + is None + ) + + +def test_transport_close_capacity_retry_accepts_native_output_free_continuation() -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_transport_close_capacity_retry_error_code( + _accepted_capacity_retry_state(), + has_other_pending_requests=False, + error_code=None, + error_message="no close frame received or sent", + ) + == "stream_incomplete" + ) + + +def test_transport_close_capacity_retry_accepts_lifecycle_only_downstream_sequence() -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_transport_close_capacity_retry_error_code( + _accepted_capacity_retry_state(last_downstream_sequence_number=1), + has_other_pending_requests=False, + error_code=None, + error_message="no close frame received or sent", + ) + == "stream_incomplete" + ) + + +def test_terminal_capacity_retry_rejects_another_pending_request() -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_terminal_capacity_retry_error_code( + _accepted_capacity_retry_state(), + event_type="error", + payload=cast(Any, _server_overloaded_event_payload()), + has_other_pending_requests=True, + ) + is None + ) + + +def test_transport_close_capacity_retry_rejects_another_pending_request() -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_transport_close_capacity_retry_error_code( + _accepted_capacity_retry_state(), + has_other_pending_requests=True, + error_code=None, + error_message="no close frame received or sent", + ) + is None + ) + + +@pytest.mark.parametrize( + ("unsafe_field", "unsafe_value"), + [ + ("enforce_openai_sdk_contract", True), + ("upstream_model_output_seen", True), + ("downstream_visible", True), + ("replay_count", 1), + ("response_id", None), + ("awaiting_response_created", True), + ], +) +def test_terminal_capacity_retry_rejects_unsafe_lifecycle( + unsafe_field: str, + unsafe_value: object, +) -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_terminal_capacity_retry_error_code( + _accepted_capacity_retry_state(**{unsafe_field: unsafe_value}), + event_type="error", + payload=cast(Any, _server_overloaded_event_payload()), + has_other_pending_requests=False, + ) + is None + ) + + +@pytest.mark.parametrize( + ("unsafe_field", "unsafe_value"), + [ + ("enforce_openai_sdk_contract", True), + ("upstream_model_output_seen", True), + ("downstream_visible", True), + ("replay_count", 1), + ("response_id", None), + ("awaiting_response_created", True), + ], +) +def test_transport_close_capacity_retry_rejects_unsafe_lifecycle( + unsafe_field: str, + unsafe_value: object, +) -> None: + assert ( + http_bridge_upstream_events_module._http_bridge_transport_close_capacity_retry_error_code( + _accepted_capacity_retry_state(**{unsafe_field: unsafe_value}), + has_other_pending_requests=False, + error_code=None, + error_message="no close frame received or sent", + ) + is None + ) + + +def test_response_event_tracking_distinguishes_lifecycle_from_model_output() -> None: + request_state = _accepted_capacity_retry_state(response_event_count=0) + + proxy_service._record_response_event(request_state, "response.created") + proxy_service._record_response_event(request_state, "response.in_progress") + + assert request_state.upstream_model_output_seen is False + + proxy_service._record_response_event(request_state, "response.output_item.added") + + assert request_state.upstream_model_output_seen is True + + +@pytest.mark.asyncio +async def test_detach_http_bridge_request_waits_for_lifecycle_reconnect_owner() -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = _accepted_capacity_retry_state() + session = _make_bridge_session( + pending_requests=deque([request_state]), + queued_request_count=1, + ) + + async with session.lifecycle_lock: + detach_task = asyncio.create_task(service._detach_http_bridge_request(session, request_state=request_state)) + await asyncio.sleep(0) + + assert request_state.event_queue is None + assert request_state.draining_until_terminal is False + assert detach_task.done() is False + + assert await detach_task is True + assert request_state.draining_until_terminal is True + assert session.closed is True + session.upstream.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_terminal_capacity_retry_stops_after_downstream_detach_without_double_counting( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + request_state = _accepted_capacity_retry_state() + other_request = _accepted_capacity_retry_state( + request_id="req-other", + response_id="resp-other", + request_text='{"type":"response.create","input":"other"}', + previous_response_id=None, + ) + + async def fake_acquire( + state: proxy_service._WebSocketRequestState, + *, + response_create_gate: asyncio.Semaphore, + **kwargs: object, + ) -> None: + del kwargs + await response_create_gate.acquire() + state.response_create_gate = response_create_gate + state.response_create_gate_acquired = True + state.awaiting_response_created = True + + class _DetachDuringSleepAsyncio: + def __getattr__(self, name: str) -> object: + return getattr(asyncio, name) + + async def sleep(self, delay: float) -> None: + assert delay == 0.25 + request_state.draining_until_terminal = True + session.queued_request_count = max(0, session.queued_request_count - 1) + await proxy_service._release_websocket_response_create_gate( + request_state, + session.response_create_gate, + ) + async with session.pending_lock: + session.pending_requests.append(other_request) + session.queued_request_count += 1 + + reconnect = AsyncMock() + monkeypatch.setattr(service, "_acquire_request_state_response_create_admission", fake_acquire) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + monkeypatch.setattr(http_bridge_request_submit_module, "backoff_seconds", lambda attempt: 0.25) + monkeypatch.setattr(http_bridge_request_submit_module, "asyncio", _DetachDuringSleepAsyncio()) + + retried = await service._retry_http_bridge_terminal_capacity_request( + session, + request_state, + error_code="server_is_overloaded", + ) + + assert retried is False + reconnect.assert_not_awaited() + assert list(session.pending_requests) == [other_request] + assert session.queued_request_count == 1 + + +@pytest.mark.asyncio +async def test_terminal_capacity_retry_leaves_pending_cleanup_to_waiting_detach( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = _accepted_capacity_retry_state() + session = _make_bridge_session( + pending_requests=deque([request_state]), + queued_request_count=1, + ) + close_upstream = cast(AsyncMock, session.upstream.close) + detach_task: asyncio.Task[bool] | None = None + + async def fake_acquire( + state: proxy_service._WebSocketRequestState, + *, + response_create_gate: asyncio.Semaphore, + **kwargs: object, + ) -> None: + del kwargs + await response_create_gate.acquire() + state.response_create_gate = response_create_gate + state.response_create_gate_acquired = True + state.awaiting_response_created = True + + async def fail_reconnect(*args: object, **kwargs: object) -> None: + nonlocal detach_task + del args, kwargs + detach_task = asyncio.create_task(service._detach_http_bridge_request(session, request_state=request_state)) + await asyncio.sleep(0) + assert request_state.event_queue is None + assert detach_task.done() is False + raise RuntimeError("retry reconnect failed") + + monkeypatch.setattr(service, "_acquire_request_state_response_create_admission", fake_acquire) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", fail_reconnect) + monkeypatch.setattr(http_bridge_request_submit_module, "backoff_seconds", lambda attempt: 0.0) + + retried = await service._retry_http_bridge_terminal_capacity_request( + session, + request_state, + error_code="server_is_overloaded", + preserve_for_reader_failure=True, + ) + + assert retried is False + assert detach_task is not None + assert await detach_task is True + assert request_state.draining_until_terminal is True + assert list(session.pending_requests) == [] + assert session.queued_request_count == 0 + assert request_state.response_create_gate_acquired is False + assert session.closed is True + close_upstream.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_terminal_capacity_retry_stops_when_backoff_exhausts_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = _accepted_capacity_retry_state(bridge_request_deadline=1.0) + session = _make_bridge_session() + clock = {"now": 0.9} + + async def fake_acquire( + state: proxy_service._WebSocketRequestState, + *, + response_create_gate: asyncio.Semaphore, + **kwargs: object, + ) -> None: + del kwargs + await response_create_gate.acquire() + state.response_create_gate = response_create_gate + state.response_create_gate_acquired = True + state.awaiting_response_created = True + + async def fake_sleep(delay: float) -> None: + assert delay == pytest.approx(0.1) + clock["now"] += delay + + reconnect = AsyncMock() + monkeypatch.setattr(service, "_acquire_request_state_response_create_admission", fake_acquire) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + monkeypatch.setattr(http_bridge_request_submit_module, "backoff_seconds", lambda _attempt: 0.25) + monkeypatch.setattr(http_bridge_request_submit_module.asyncio, "sleep", fake_sleep) + monkeypatch.setattr( + http_bridge_request_submit_module, + "_service_time", + lambda: SimpleNamespace(monotonic=lambda: clock["now"]), + ) + + retried = await service._retry_http_bridge_terminal_capacity_request( + session, + request_state, + error_code="server_is_overloaded", + ) + + assert retried is False + reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_terminal_capacity_retry_does_not_reconnect_over_competing_pending_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + request_state = _accepted_capacity_retry_state() + competing_request = _accepted_capacity_retry_state( + request_id="req-competing", + response_id="resp-competing", + request_text='{"type":"response.create","input":"competing"}', + previous_response_id=None, + ) + + async def fake_acquire( + state: proxy_service._WebSocketRequestState, + *, + response_create_gate: asyncio.Semaphore, + **kwargs: object, + ) -> None: + del kwargs + await response_create_gate.acquire() + state.response_create_gate = response_create_gate + state.response_create_gate_acquired = True + state.awaiting_response_created = True + async with session.pending_lock: + session.pending_requests.append(competing_request) + session.queued_request_count += 1 + + reconnect = AsyncMock() + monkeypatch.setattr(service, "_acquire_request_state_response_create_admission", fake_acquire) + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + + retried = await service._retry_http_bridge_terminal_capacity_request( + session, + request_state, + error_code="server_is_overloaded", + ) + + assert retried is False + reconnect.assert_not_awaited() + assert list(session.pending_requests) == [competing_request] + assert session.queued_request_count == 1 + + +@pytest.mark.asyncio +async def test_terminal_capacity_retry_does_not_overtake_admission_waiter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + session = _make_bridge_session() + session.queued_request_count = 1 + session.admission_waiter_count = 1 + request_state = _accepted_capacity_retry_state() + reconnect = AsyncMock() + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + + retried = await service._retry_http_bridge_terminal_capacity_request( + session, + request_state, + error_code="server_is_overloaded", + ) + + assert retried is False + reconnect.assert_not_awaited() + assert list(session.pending_requests) == [] + assert session.queued_request_count == 1 + + +@pytest.mark.asyncio +async def test_terminal_capacity_close_retry_keeps_pending_request_for_terminal_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = _accepted_capacity_retry_state(bridge_request_deadline=0.0) + session = _make_bridge_session( + pending_requests=deque([request_state]), + queued_request_count=1, + ) + + async def fake_acquire( + state: proxy_service._WebSocketRequestState, + *, + response_create_gate: asyncio.Semaphore, + **kwargs: object, + ) -> None: + del kwargs + await response_create_gate.acquire() + state.response_create_gate = response_create_gate + state.response_create_gate_acquired = True + state.awaiting_response_created = True + + monkeypatch.setattr(service, "_acquire_request_state_response_create_admission", fake_acquire) + + retried = await service._retry_http_bridge_terminal_capacity_request( + session, + request_state, + error_code="server_is_overloaded", + preserve_for_reader_failure=True, + ) + + assert retried is False + assert list(session.pending_requests) == [request_state] + assert session.queued_request_count == 1 + assert request_state.response_create_gate_acquired is True + + await service._fail_http_bridge_reader_and_maybe_retire( + session, + error_code="stream_incomplete", + error_message="HTTP bridge closed before the capacity retry could be sent", + ) + + assert request_state.event_queue is not None + terminal_block = await request_state.event_queue.get() + assert terminal_block is not None + terminal_event = json.loads(terminal_block.split("data: ", 1)[1]) + assert terminal_event["type"] == "response.failed" + assert terminal_event["response"]["error"]["code"] == "stream_incomplete" + + +@pytest.mark.asyncio +async def test_expired_terminal_capacity_request_is_not_generic_precreated_replay( + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = proxy_service.ProxyService(cast(Any, nullcontext())) + request_state = _accepted_capacity_retry_state(bridge_request_deadline=0.0) + session = _make_bridge_session(pending_requests=deque([request_state]), queued_request_count=1) + reconnect = AsyncMock() + monkeypatch.setattr(service, "_reconnect_http_bridge_session", reconnect) + + retried = await service._retry_http_bridge_precreated_request(session) + + assert retried is False + reconnect.assert_not_awaited() + assert list(session.pending_requests) == [request_state] + + @pytest.mark.asyncio async def test_legacy_forward_anchor_lookup_accepts_registered_turn_state_alias() -> None: key = proxy_service._HTTPBridgeSessionKey("session_header", "sid-123", None) @@ -16662,13 +17192,14 @@ async def test_http_bridge_replays_proxy_verified_full_resend_after_owner_quota( '{"type":"response.create","model":"gpt-5.6-sol",' '"input":[{"role":"user","content":[{"type":"input_text","text":"full resend"}]}]}' ) + started_at = time.monotonic() request_state = proxy_service._WebSocketRequestState( request_id="req-verified-owner-limit", model="gpt-5.6-sol", service_tier="priority", reasoning_effort="high", api_key_reservation=None, - started_at=1.0, + started_at=started_at, previous_response_id="resp_verified_owner", preferred_account_id="acc-limited", proxy_injected_previous_response_id=True, @@ -16684,17 +17215,17 @@ async def test_http_bridge_replays_proxy_verified_full_resend_after_owner_quota( skip_request_log=True, affinity_policy=proxy_service._AffinityPolicy( key="http_turn_verified_owner_limit", - kind=proxy_service.StickySessionKind.CODEX_SESSION, + kind=proxy_service.StickySessionKind.PROMPT_CACHE, ), ) account = cast(Any, SimpleNamespace(id="acc-limited", status=AccountStatus.ACTIVE)) session = proxy_service._HTTPBridgeSession( key=proxy_service._HTTPBridgeSessionKey( - "turn_state_header", + "prompt_cache", "http_turn_verified_owner_limit", None, ), - headers={"x-codex-turn-state": "http_turn_verified_owner_limit"}, + headers={"x-codex-prompt-cache-key": "http_turn_verified_owner_limit"}, affinity=request_state.affinity_policy, request_model="gpt-5.6-sol", account=account, @@ -16704,25 +17235,35 @@ async def test_http_bridge_replays_proxy_verified_full_resend_after_owner_quota( pending_lock=anyio.Lock(), response_create_gate=asyncio.Semaphore(1), queued_request_count=1, - last_used_at=1.0, + last_used_at=started_at, idle_ttl_seconds=120.0, upstream_turn_state="turn-old-account", downstream_turn_state="turn-client-alias", ) handle_stream_error = AsyncMock() release_create_lease = AsyncMock() + settings = SimpleNamespace( + prefer_earlier_reset_accounts=False, + prefer_earlier_reset_window="primary", + routing_strategy="usage_weighted", + ) + fallback_account = cast(Any, SimpleNamespace(id="acc-fallback", status=AccountStatus.ACTIVE)) + fallback_upstream = cast( + UpstreamWebSocket, + SimpleNamespace(response_header=lambda _name: None, close=AsyncMock(), send_text=AsyncMock()), + ) - async def retry_precreated(retry_session): - assert retry_session is session - assert session.upstream_turn_state is None - assert session.downstream_turn_state is None - assert request_state.previous_response_id is None - assert request_state.preferred_account_id is None - assert request_state.request_text == fresh_text - assert request_state.excluded_account_ids == {account.id} - assert request_state.affinity_policy.reallocate_sticky is True - assert list(session.pending_requests) == [request_state] - return True + async def select_account(_deadline: float, **_: object) -> proxy_service.AccountSelection: + return proxy_service.AccountSelection(account=fallback_account, error_message=None, error_code=None) + + async def ensure_fresh(account: object, **_: object) -> object: + return account + + async def open_upstream(_account: object, _headers: dict[str, str], **_: object) -> UpstreamWebSocket: + return fallback_upstream + + response_create_lease = object() + acquire_create_lease = AsyncMock(return_value=response_create_lease) monkeypatch.setattr(service, "_handle_stream_error", handle_stream_error) monkeypatch.setattr( @@ -16730,7 +17271,19 @@ async def retry_precreated(retry_session): "_release_request_state_account_response_create_lease", release_create_lease, ) - monkeypatch.setattr(service, "_retry_http_bridge_precreated_request", retry_precreated) + monkeypatch.setattr( + proxy_service, + "get_settings_cache", + lambda: SimpleNamespace(get=AsyncMock(return_value=settings)), + ) + monkeypatch.setattr(service, "_select_account_with_budget_for_stream", select_account) + monkeypatch.setattr(service, "_ensure_fresh_with_budget", ensure_fresh) + monkeypatch.setattr(service, "_open_upstream_websocket_with_budget", open_upstream) + monkeypatch.setattr( + service, + "_acquire_account_response_create_lease_or_overload", + acquire_create_lease, + ) await service._process_http_bridge_upstream_text( session, @@ -16749,6 +17302,19 @@ async def retry_precreated(retry_session): handle_stream_error.assert_awaited_once() release_create_lease.assert_awaited_once_with(request_state) + acquire_create_lease.assert_awaited_once() + cast(Any, fallback_upstream).send_text.assert_awaited_once_with(fresh_text) + assert session.account is fallback_account + assert session.upstream is fallback_upstream + assert session.upstream_turn_state is None + assert session.downstream_turn_state is None + assert request_state.previous_response_id is None + assert request_state.preferred_account_id is None + assert request_state.request_text == fresh_text + assert request_state.excluded_account_ids == {account.id} + assert request_state.affinity_policy.reallocate_sticky is True + assert request_state.account_response_create_lease is response_create_lease + assert list(session.pending_requests) == [request_state] assert request_state.event_queue is not None assert request_state.event_queue.empty() @@ -21159,7 +21725,7 @@ async def test_http_bridge_clean_close_before_response_does_not_penalize_account @pytest.mark.asyncio @pytest.mark.parametrize("routed", [False, True], ids=["direct-close", "routed-receive-error"]) -async def test_http_bridge_reader_maps_ordinary_websocket_receive_failure_to_stream_incomplete( +async def test_http_bridge_reader_maps_ordinary_websocket_receive_failure_to_account_neutral_stream_incomplete( monkeypatch: pytest.MonkeyPatch, routed: bool, ) -> None: @@ -21205,7 +21771,7 @@ async def fail_reader( assert session.last_upstream_close_code == (None if routed else 1011) assert len(failure_calls) == 1 assert failure_calls[0]["error_code"] == "stream_incomplete" - assert failure_calls[0]["penalize_account"] is True + assert failure_calls[0]["penalize_account"] is False assert failure_calls[0]["response_events_seen"] == 0 if routed: assert failure_calls[0]["upstream_close_code"] is None diff --git a/tests/unit/test_proxy_utils.py b/tests/unit/test_proxy_utils.py index bc84f388bf..6dbaf87e82 100644 --- a/tests/unit/test_proxy_utils.py +++ b/tests/unit/test_proxy_utils.py @@ -5226,6 +5226,22 @@ def test_http_downstream_always_websocket_policy_keeps_websocket_without_sticky_ ) +@pytest.mark.parametrize( + ("remaining_budget", "backoff", "expected_delay"), + [(1.0, 0.25, 0.25), (0.25, 0.25, None), (0.1, 0.25, None), (0.1, 0.0, 0.0)], +) +def test_transient_stream_retry_reserves_budget_for_next_attempt( + monkeypatch: pytest.MonkeyPatch, + remaining_budget: float, + backoff: float, + expected_delay: float | None, +) -> None: + monkeypatch.setattr(proxy_service, "_remaining_budget_seconds", lambda _deadline: remaining_budget) + monkeypatch.setattr(streaming_retry_module, "backoff_seconds", lambda _attempt: backoff) + + assert streaming_retry_module._transient_stream_retry_delay(123.0, 1) == expected_delay + + async def _capture_stream_retry_transport( monkeypatch: pytest.MonkeyPatch, *, @@ -15015,6 +15031,52 @@ async def fake_stream(*_args: object, **_kwargs: object): assert stream_attempts == 2 +@pytest.mark.asyncio +async def test_stream_responses_http_500_exhaustion_preserves_error_when_no_failover_candidate(monkeypatch): + settings = _make_proxy_settings() + service = proxy_service.ProxyService(_repo_factory(_RequestLogsRecorder())) + account = _make_account("acc_http_500_only") + seen_excluded_account_ids: list[set[str]] = [] + call_count = 0 + + 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_account_ids = set(cast(set[str] | None, kwargs.get("exclude_account_ids")) or set()) + seen_excluded_account_ids.append(excluded_account_ids) + if not excluded_account_ids: + return AccountSelection(account=account, error_message=None) + return AccountSelection(account=None, error_message="No active accounts available", error_code="no_accounts") + + async def fake_stream(*args: object, **kwargs: object): + nonlocal call_count + del args, kwargs + call_count += 1 + raise proxy_module.ProxyResponseError( + 500, + openai_error("server_error", "Upstream failed after retries"), + failure_phase="status", + ) + yield # pragma: no cover - keep this async generator-shaped + + monkeypatch.setattr(service._load_balancer, "select_account", select_account) + monkeypatch.setattr(service._load_balancer, "record_errors", AsyncMock()) + monkeypatch.setattr(service, "_ensure_fresh", AsyncMock(return_value=account)) + monkeypatch.setattr(service, "_handle_stream_error", AsyncMock()) + 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-stream"})] + + event = json.loads(chunks[-1].split("data: ", 1)[1]) + assert call_count == 3 + assert seen_excluded_account_ids == [set(), {account.id}] + assert event["type"] == "response.failed" + assert event["response"]["error"]["code"] == "server_error" + assert event["response"]["error"]["message"] == "Upstream failed after retries" + + @pytest.mark.asyncio async def test_stream_responses_empty_upstream_emits_terminal_failure(monkeypatch): settings = _make_proxy_settings()