fix(proxy): retry accepted capacity bridge failures

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] = []

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

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",

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 = (
@@ -2195,6 +2342,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