diff --git a/app/core/balancer/__init__.py b/app/core/balancer/__init__.py index 2e29cad9db..c8d4dc525b 100644 --- a/app/core/balancer/__init__.py +++ b/app/core/balancer/__init__.py @@ -1,4 +1,5 @@ from app.core.balancer.logic import ( + ERROR_BACKOFF_THRESHOLD, HEALTH_TIER_DRAINING, HEALTH_TIER_HEALTHY, HEALTH_TIER_PROBING, @@ -36,6 +37,7 @@ "HEALTH_TIER_DRAINING", "HEALTH_TIER_HEALTHY", "HEALTH_TIER_PROBING", + "ERROR_BACKOFF_THRESHOLD", "REAUTH_REQUIRED_FAILURE_CODES", "AccountState", "RoutingCost", diff --git a/app/core/balancer/logic.py b/app/core/balancer/logic.py index a3caf85afa..96d9dc4a8f 100644 --- a/app/core/balancer/logic.py +++ b/app/core/balancer/logic.py @@ -91,6 +91,7 @@ DRAIN_SECONDARY_THRESHOLD_PCT = 90.0 DRAIN_ERROR_WINDOW_SECONDS = 60.0 DRAIN_ERROR_COUNT_THRESHOLD = 2 +ERROR_BACKOFF_THRESHOLD = 3 PROBE_QUIET_SECONDS = 60.0 PROBE_SUCCESS_STREAK_REQUIRED = 3 ROUTING_POLICY_NORMAL = "normal" @@ -481,8 +482,8 @@ def select_account( state.error_count = 0 if state.cooldown_until and current < state.cooldown_until: continue - if state.error_count >= 3: - backoff = min(300, 30 * (2 ** (state.error_count - 3))) + if state.error_count >= ERROR_BACKOFF_THRESHOLD: + backoff = min(300, 30 * (2 ** (state.error_count - ERROR_BACKOFF_THRESHOLD))) if state.last_error_at and current - state.last_error_at < backoff: in_error_backoff.append(state) continue @@ -519,7 +520,7 @@ def select_account( if allow_backoff_fallback and (len(in_error_backoff) > 1 or (in_error_backoff and hard_blocked_exists)): def _backoff_expires_at(s: AccountState) -> float: - backoff = min(300, 30 * (2 ** (s.error_count - 3))) + backoff = min(300, 30 * (2 ** (s.error_count - ERROR_BACKOFF_THRESHOLD))) return (s.last_error_at or 0.0) + backoff available.append(min(in_error_backoff, key=_backoff_expires_at)) diff --git a/app/core/clients/codex.py b/app/core/clients/codex.py index 8bae151434..3d160958bc 100644 --- a/app/core/clients/codex.py +++ b/app/core/clients/codex.py @@ -212,11 +212,20 @@ async def request_with_route_metadata( retryable_same_contract=False, ) from None return CodexRequestResult(response, candidate, index > 0) - except CodexTransportError: - if index == len(endpoints) - 1 or not allow_fallback: + except CodexTransportError as exc: + # A confirmed pre-dispatch connect failure proves the request + # never left for upstream, so trying the next endpoint in the + # same resolved pool is safe even for a non-idempotent POST. + # TLS verification failures are stable endpoint configuration + # errors rather than transient connect losses; they keep the + # idempotent-only rule. + if index == len(endpoints) - 1 or not ( + allow_fallback or (exc.retryable_same_contract and not exc.is_tls_verification_failure) + ): raise except Exception as exc: - if index == len(endpoints) - 1 or not allow_fallback: + pre_dispatch = is_pre_dispatch_connection_failure(exc) and not isinstance(exc, aiohttp.ClientSSLError) + if index == len(endpoints) - 1 or not (allow_fallback or pre_dispatch): raise _transport_error( "request", endpoint.id, diff --git a/app/core/clients/proxy.py b/app/core/clients/proxy.py index 11105dc0e3..717448360d 100644 --- a/app/core/clients/proxy.py +++ b/app/core/clients/proxy.py @@ -92,6 +92,7 @@ CODEX_INSTALLATION_ID_HEADER = "x-codex-installation-id" CODEX_TURN_METADATA_HEADER = "x-codex-turn-metadata" +CODEX_LB_REQUIRED_CAPABILITY_HEADER = "x-codex-lb-required-capability" CODEX_RESPONSES_LITE_HEADER = "x-openai-internal-codex-responses-lite" CODEX_RESPONSES_LITE_WEBSOCKET_METADATA_KEY = "ws_request_header_x_openai_internal_codex_responses_lite" @@ -103,6 +104,7 @@ "forwarded", "x-real-ip", CODEX_INSTALLATION_ID_HEADER, + CODEX_LB_REQUIRED_CAPABILITY_HEADER, "true-client-ip", } INTERNAL_OPENAI_UPSTREAM_HEADERS = frozenset( @@ -488,6 +490,24 @@ def __init__( self.failed_session = failed_session +def is_confirmed_pre_dispatch_transport_error(exc: ProxyResponseError) -> bool: + """Return whether the transport proved the upstream request never dispatched. + + Only this provenance authorizes replaying a movable request on another + account: a typed connector failure while reaching the account's routed + proxy endpoint, before any request bytes could leave for upstream. + Host-wide network loss (``proxy_network_unavailable``) stays on its + account-neutral process recovery path instead of penalizing the selected + account, and ambiguous dispatch outcomes remain non-replayable. + """ + + if not (exc.retryable_same_contract and exc.failure_phase == "connect"): + return False + error = exc.payload.get("error") + error_code = error.get("code") if isinstance(error, dict) else None + return error_code != PROCESS_NETWORK_UNAVAILABLE_CODE + + def _process_network_failure_error( message: str, exc: Exception, diff --git a/app/core/clients/proxy_websocket.py b/app/core/clients/proxy_websocket.py index 06eb2d5fbb..aa6b47e7c2 100644 --- a/app/core/clients/proxy_websocket.py +++ b/app/core/clients/proxy_websocket.py @@ -769,9 +769,22 @@ async def _connect_upstream_websocket( status_code if policy.preserve_handshake_status else 502, openai_error(error_code, message, error_type="server_error"), failure_phase="connect", + # Carry the client's dispatch provenance across the sanitizing + # boundary: a typed connector failure against the routed proxy + # proves no ``response.create`` frame could have reached + # upstream, so service-level failover may replay the request + # on another account. TLS verification failures are stable + # endpoint configuration errors and stay non-replayable. retryable_same_contract=( - policy.retry_routed_network_errors and error_code == PROCESS_NETWORK_UNAVAILABLE_CODE + (policy.retry_routed_network_errors and error_code == PROCESS_NETWORK_UNAVAILABLE_CODE) + or (exc.retryable_same_contract and not exc.is_tls_verification_failure) ), + failure_detail=( + "proxy_connect_pre_dispatch" + if exc.retryable_same_contract and not exc.is_tls_verification_failure + else "transport_error" + ), + failure_exception_type=type(exc).__name__, ) from exc except Exception: if owns_codex_client: diff --git a/app/core/usage/pricing.py b/app/core/usage/pricing.py index d2865b2251..5da4a9a145 100644 --- a/app/core/usage/pricing.py +++ b/app/core/usage/pricing.py @@ -104,34 +104,34 @@ def _normalize_usage(usage: UsageTokens | ResponseUsage | None) -> UsageTokens | long_context_output_per_1m=45.0, ), "gpt-5.6-terra": ModelPrice( - input_per_1m=2.5, - cached_input_per_1m=0.25, - output_per_1m=15.0, - priority_input_per_1m=5.0, - priority_cached_input_per_1m=0.5, - priority_output_per_1m=30.0, - flex_input_per_1m=1.25, - flex_cached_input_per_1m=0.125, - flex_output_per_1m=7.5, + input_per_1m=2.0, + cached_input_per_1m=0.2, + output_per_1m=12.0, + priority_input_per_1m=4.0, + priority_cached_input_per_1m=0.4, + priority_output_per_1m=24.0, + flex_input_per_1m=1.0, + flex_cached_input_per_1m=0.1, + flex_output_per_1m=6.0, long_context_threshold_tokens=272_000, - long_context_input_per_1m=5.0, - long_context_cached_input_per_1m=0.5, - long_context_output_per_1m=22.5, + long_context_input_per_1m=4.0, + long_context_cached_input_per_1m=0.4, + long_context_output_per_1m=18.0, ), "gpt-5.6-luna": ModelPrice( - input_per_1m=1.0, - cached_input_per_1m=0.1, - output_per_1m=6.0, - priority_input_per_1m=2.0, - priority_cached_input_per_1m=0.2, - priority_output_per_1m=12.0, - flex_input_per_1m=0.5, - flex_cached_input_per_1m=0.05, - flex_output_per_1m=3.0, + input_per_1m=0.2, + cached_input_per_1m=0.02, + output_per_1m=1.2, + priority_input_per_1m=0.4, + priority_cached_input_per_1m=0.04, + priority_output_per_1m=2.4, + flex_input_per_1m=0.1, + flex_cached_input_per_1m=0.01, + flex_output_per_1m=0.6, long_context_threshold_tokens=272_000, - long_context_input_per_1m=2.0, - long_context_cached_input_per_1m=0.2, - long_context_output_per_1m=9.0, + long_context_input_per_1m=0.4, + long_context_cached_input_per_1m=0.04, + long_context_output_per_1m=1.8, ), "gpt-5.5": ModelPrice( input_per_1m=5.0, diff --git a/app/db/alembic/versions/20260731_000000_add_capability_lineage_markers.py b/app/db/alembic/versions/20260731_000000_add_capability_lineage_markers.py new file mode 100644 index 0000000000..15e81b7e06 --- /dev/null +++ b/app/db/alembic/versions/20260731_000000_add_capability_lineage_markers.py @@ -0,0 +1,38 @@ +"""add durable capability lineage markers + +Revision ID: 20260731_000000_add_capability_lineage_markers +Revises: 20260725_000000_add_http_bridge_pending_tool_calls +Create Date: 2026-07-31 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "20260731_000000_add_capability_lineage_markers" +down_revision = "20260725_000000_add_http_bridge_pending_tool_calls" +branch_labels = None +depends_on = None + +_TABLE = "capability_lineage_markers" + + +def upgrade() -> None: + bind = op.get_bind() + if sa.inspect(bind).has_table(_TABLE): + return + op.create_table( + _TABLE, + sa.Column("marker_hash", sa.String(length=64), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False), + sa.Column("last_seen_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False), + sa.PrimaryKeyConstraint("marker_hash"), + ) + + +def downgrade() -> None: + bind = op.get_bind() + if not sa.inspect(bind).has_table(_TABLE): + return + op.drop_table(_TABLE) diff --git a/app/db/models.py b/app/db/models.py index 6f0cb0df9b..521a47d02b 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -719,6 +719,23 @@ class StickySession(Base): ) +class CapabilityLineageMarker(Base): + __tablename__ = "capability_lineage_markers" + + marker_hash: Mapped[str] = mapped_column(String(64), primary_key=True) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + server_default=func.now(), + nullable=False, + ) + last_seen_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + server_default=func.now(), + onupdate=func.now(), + nullable=False, + ) + + class DashboardSettings(Base): __tablename__ = "dashboard_settings" diff --git a/app/dependencies.py b/app/dependencies.py index 72783ea3f1..8fd8131e07 100644 --- a/app/dependencies.py +++ b/app/dependencies.py @@ -32,6 +32,7 @@ from app.modules.model_sources.repository import ModelSourcesRepository from app.modules.model_sources.service import ModelSourcesService from app.modules.oauth.service import OauthService +from app.modules.proxy.capability_lineage_repository import CapabilityLineageRepository from app.modules.proxy.repo_bundle import ProxyRepositories from app.modules.proxy.service import ProxyService from app.modules.proxy.sticky_repository import StickySessionsRepository @@ -220,6 +221,7 @@ async def _proxy_repo_context() -> AsyncIterator[ProxyRepositories]: api_keys=ApiKeysRepository(session), additional_usage=AdditionalUsageRepository(session), quota_planner=QuotaPlannerRepository(session), + capability_lineage=CapabilityLineageRepository(session), ) diff --git a/app/main.py b/app/main.py index df28262a46..4b59dbe295 100644 --- a/app/main.py +++ b/app/main.py @@ -687,6 +687,7 @@ def create_app() -> FastAPI: app.include_router(dashboard_api.router) app.include_router(usage_api.router) app.include_router(request_logs_api.router) + app.include_router(request_logs_api.conversations_router) app.include_router(quota_planner_api.router) app.include_router(reports_api.router) app.include_router(conversation_archive_api.router) diff --git a/app/modules/dashboard/timeframes.py b/app/modules/dashboard/timeframes.py new file mode 100644 index 0000000000..5552571e58 --- /dev/null +++ b/app/modules/dashboard/timeframes.py @@ -0,0 +1,15 @@ +from datetime import datetime, timedelta +from typing import cast + +from app.core.utils.time import utcnow +from app.modules.dashboard.builders import _OVERVIEW_TIMEFRAME_CONFIGS +from app.modules.dashboard.schemas import DashboardOverviewTimeframeKey + +CONVERSATION_TIMEFRAME_KEYS: frozenset[str] = frozenset({"1d", "7d", "30d"}) + + +def resolve_conversation_timeframe(key: str) -> tuple[int, datetime]: + """Return the window duration and rolling start for a conversation timeframe.""" + _key = cast(DashboardOverviewTimeframeKey, key) + timeframe = _OVERVIEW_TIMEFRAME_CONFIGS[_key] + return timeframe.window_minutes, utcnow() - timedelta(minutes=timeframe.window_minutes) diff --git a/app/modules/proxy/_service/api_key_usage.py b/app/modules/proxy/_service/api_key_usage.py index 2159e3802d..76345f406c 100644 --- a/app/modules/proxy/_service/api_key_usage.py +++ b/app/modules/proxy/_service/api_key_usage.py @@ -9,9 +9,12 @@ import anyio +from app.core.clients.proxy import ProxyResponseError +from app.core.errors import openai_error from app.core.exceptions import ProxyAuthError, ProxyRateLimitError from app.core.openai.models import CompactResponsePayload from app.core.utils.request_id import get_request_id +from app.db.models import Account from app.modules.api_keys.service import ( ApiKeyData, ApiKeyInvalidError, @@ -58,6 +61,7 @@ def _api_key_reservation_heartbeat_seconds() -> float: class _ApiKeyUsageServiceProtocol(Protocol): _repo_factory: ProxyRepoFactory _background_cleanup_tasks: set[asyncio.Task[None]] + _load_balancer: Any def _normalize_service_tier_value(value: Any) -> str | None: @@ -129,6 +133,30 @@ async def _release_websocket_request_state_reservation( ) -> None: self._cancel_request_state_api_key_reservation_heartbeat(request_state) await self._release_websocket_reservation(request_state.api_key_reservation) + request_state.api_key_reservation = None + lifecycle = request_state.deferred_account_backoff_lifecycle + if lifecycle is not None: + lifecycle.settlement_confirmed = True + pending_backoffs = ( + lifecycle.pending_backoffs if lifecycle is not None else request_state.deferred_account_error_backoffs + ) + if pending_backoffs: + await self._drain_deferred_account_error_backoffs(pending_backoffs) + + async def _drain_deferred_account_error_backoffs( + self, + pending_backoffs: dict[str, Account], + ) -> None: + if not pending_backoffs: + return + proxy = cast(_ApiKeyUsageServiceProtocol, self) + while pending_backoffs: + account_id, account = pending_backoffs.popitem() + try: + await proxy._load_balancer.record_error_backoff(account) + except BaseException: + pending_backoffs.setdefault(account_id, account) + raise async def _maybe_touch_api_key_reservation( self, @@ -289,13 +317,36 @@ async def _settle_compact_api_key_usage( ) else: await api_keys_service.release_usage_reservation(reservation_id) - except Exception: + except Exception as exc: logger.warning( "Failed to settle compact API key reservation key_id=%s request_id=%s", api_key.id, get_request_id(), exc_info=True, ) + try: + async with proxy._repo_factory() as repos: + api_keys_service = _service_api_keys_service()(repos.api_keys) + await api_keys_service.release_usage_reservation(reservation_id) + except Exception: + logger.warning( + "Failed to release compact API key reservation after settlement failure " + "key_id=%s request_id=%s", + api_key.id, + get_request_id(), + exc_info=True, + ) + raise ProxyResponseError( + 502, + openai_error( + "usage_settlement_failed", + "Compact API key usage could not be settled", + error_type="server_error", + ), + failure_phase="usage_settlement", + failure_detail="compact_api_key_usage_persistence_failed", + failure_exception_type=type(exc).__name__, + ) from exc async def _settle_stream_api_key_usage( self, @@ -343,9 +394,10 @@ async def _settle_once() -> bool: ) return False - # Detach unconditionally instead of shield-awaiting: the tracking - # callback already schedules a release when settlement fails or is - # cancelled, the caller's finally-net skips via + # Detach unconditionally instead of shield-awaiting: failed or cancelled + # settlements release through the tracking callback (detached callers) + # or the synchronous fallback below (ordering-sensitive callers), while + # the caller's finally-net skips via # usage_settlement_transferred, and reservations keep counting toward # limits until finalized/released, so a briefly-lagging settlement can # only over-restrict, never over-admit. Awaiting the ~5+2N-statement @@ -358,16 +410,36 @@ async def _settle_once() -> bool: api_key=api_key, api_key_reservation=api_key_reservation, request_id=request_id, + release_on_failure=not wait_for_settlement, ) if wait_for_settlement: # Ordering-sensitive callers (the websocket error path) must # commit the settlement before load-balancer health writes; they # opt into waiting while everything else stays detached. + settlement_committed = False with anyio.CancelScope(shield=True): - try: - await asyncio.shield(task) - except Exception: # failures release via the tracking callback - pass + while True: + try: + settlement_committed = await asyncio.shield(task) + break + except asyncio.CancelledError: + # Shield keeps caller cancellation from cancelling the + # settlement task. Wait until that task is actually done + # before a fallback release can touch the same reservation. + if task.cancelled(): + break + except Exception: + break + if not settlement_committed: + # Ordering-sensitive callers cannot let the subsequent + # health write race a background fallback. Await the release + # here so the reservation is no longer live on return. + return await self._release_unsettled_stream_api_key_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + request_id=request_id, + ) + return True return True def _track_stream_usage_settlement_task( @@ -377,10 +449,18 @@ def _track_stream_usage_settlement_task( api_key: ApiKeyData, api_key_reservation: ApiKeyUsageReservationData, request_id: str, + release_on_failure: bool = True, ) -> None: proxy = cast(_ApiKeyUsageServiceProtocol, self) proxy._background_cleanup_tasks.add(cast(asyncio.Task[None], task)) + async def _release_after_failed_settlement() -> None: + await self._release_unsettled_stream_api_key_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + request_id=request_id, + ) + def _settlement_done(done_task: asyncio.Task[bool]) -> None: proxy._background_cleanup_tasks.discard(cast(asyncio.Task[None], done_task)) try: @@ -391,16 +471,12 @@ def _settlement_done(done_task: asyncio.Task[bool]) -> None: api_key.id, request_id, ) - release_coro = self._release_unsettled_stream_api_key_usage( - api_key=api_key, - api_key_reservation=api_key_reservation, - request_id=request_id, - ) - self._schedule_cancel_safe_cleanup( - release_coro, - action="release_stream_api_key_reservation_after_cancelled_settlement", - request_id=request_id, - ) + if release_on_failure: + self._schedule_cancel_safe_cleanup( + _release_after_failed_settlement(), + action="release_stream_api_key_reservation_after_cancelled_settlement", + request_id=request_id, + ) except Exception as exc: logger.warning( "Stream API key settlement task failed key_id=%s request_id=%s", @@ -409,14 +485,9 @@ def _settlement_done(done_task: asyncio.Task[bool]) -> None: exc_info=(type(exc), exc, exc.__traceback__), ) else: - if not settled: - release_coro = self._release_unsettled_stream_api_key_usage( - api_key=api_key, - api_key_reservation=api_key_reservation, - request_id=request_id, - ) + if not settled and release_on_failure: self._schedule_cancel_safe_cleanup( - release_coro, + _release_after_failed_settlement(), action="release_stream_api_key_reservation_after_failed_settlement", request_id=request_id, ) @@ -456,7 +527,7 @@ async def _release_unsettled_stream_api_key_usage( api_key: ApiKeyData, api_key_reservation: ApiKeyUsageReservationData, request_id: str, - ) -> None: + ) -> bool: proxy = cast(_ApiKeyUsageServiceProtocol, self) with anyio.CancelScope(shield=True): try: @@ -465,6 +536,7 @@ async def _release_unsettled_stream_api_key_usage( await api_keys_service.release_usage_reservation( api_key_reservation.reservation_id, ) + return True except Exception: logger.warning( "Failed to release stream API key reservation key_id=%s request_id=%s", @@ -472,3 +544,4 @@ async def _release_unsettled_stream_api_key_usage( request_id, exc_info=True, ) + return False diff --git a/app/modules/proxy/_service/compact.py b/app/modules/proxy/_service/compact.py index 4623f61c46..f8eea368f6 100644 --- a/app/modules/proxy/_service/compact.py +++ b/app/modules/proxy/_service/compact.py @@ -1133,6 +1133,8 @@ async def _call_compact( log_status = "success" return response except ProxyResponseError as exc: + if exc.failure_phase == "usage_settlement": + raise compact_continuity_error = _compact_previous_response_not_found_error(exc) if compact_continuity_error is not None: await proxy._settle_compact_api_key_usage( diff --git a/app/modules/proxy/_service/http_bridge/mixin.py b/app/modules/proxy/_service/http_bridge/mixin.py index d8dfeb0c51..eaef79c319 100644 --- a/app/modules/proxy/_service/http_bridge/mixin.py +++ b/app/modules/proxy/_service/http_bridge/mixin.py @@ -5,6 +5,7 @@ import logging from collections import deque from collections.abc import Collection +from dataclasses import replace from typing import Any, Literal, TypeVar, overload from uuid import uuid4 @@ -120,6 +121,7 @@ ) from app.modules.proxy._service.http_bridge.owner_forwarding import _HTTPBridgeOwnerForwardingMixin from app.modules.proxy._service.http_bridge.protocol import _HTTPBridgeServiceProtocol +from app.modules.proxy._service.http_bridge.proxy_failover import _HTTPBridgePreDispatchFailover from app.modules.proxy._service.http_bridge.request_submit import _HTTPBridgeRequestSubmitMixin from app.modules.proxy._service.http_bridge.service_stubs import ( _await_cancelled_task, @@ -164,6 +166,7 @@ _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS, # noqa: F401 _clear_websocket_precreated_replay_fallback, _copy_websocket_route_metadata_to_session, + _DeferredAccountBackoffLifecycle, _HTTPBridgeOwnerForward, _HTTPBridgeSession, _HTTPBridgeSessionKey, @@ -361,6 +364,8 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession": ... @overload @@ -393,6 +398,8 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession | _HTTPBridgeOwnerForward": ... async def _get_or_create_http_bridge_session( @@ -424,6 +431,8 @@ async def _get_or_create_http_bridge_session( request_deadline: float | None = None, session_header_fallback_key: "_HTTPBridgeSessionKey | None" = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession | _HTTPBridgeOwnerForward": settings = _service_get_settings() request_scope_id = ensure_request_scope_id() @@ -1484,6 +1493,8 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: "request_usage_budget": request_usage_budget, "request_deadline": request_deadline, "exclude_account_ids": exclude_account_ids, + "deferred_account_backoff_lifecycle": deferred_account_backoff_lifecycle, + "defer_account_health_writes": defer_account_health_writes, } try: create_signature = inspect.signature(create_session) @@ -1500,6 +1511,8 @@ def bind_account_neutral_recovery_owner(session: _HTTPBridgeSession) -> None: "request_deadline", "exclude_account_ids", "preferred_account_is_continuity_owner", + "deferred_account_backoff_lifecycle", + "defer_account_health_writes", ): if optional_kwarg not in create_signature.parameters: create_kwargs.pop(optional_kwarg, None) @@ -1721,6 +1734,8 @@ async def _create_http_bridge_session( request_usage_budget: ApiKeyRequestUsageBudget | None = None, request_deadline: float | None = None, exclude_account_ids: Collection[str] | None = None, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_writes: bool = False, ) -> "_HTTPBridgeSession": request_state = _WebSocketRequestState( request_id=f"http_bridge_connect_{uuid4().hex}", @@ -1744,7 +1759,11 @@ async def _create_http_bridge_session( if require_preferred_account: fallback_on_preferred_account_unavailable = False retry_same_account_once = preferred_account_id is not None - preferred_candidate_id = preferred_account_id + proxy_connect_failover = _HTTPBridgePreDispatchFailover( + excluded_account_ids, + preferred_account_id, + affinity.reallocate_sticky, + ) selected_account_lease: AccountLease | None = None while True: select_kwargs = { @@ -1752,14 +1771,18 @@ async def _create_http_bridge_session( "kind": "http_bridge", "request_stage": request_stage, "api_key": api_key, - "affinity_policy": affinity, + "affinity_policy": ( + replace(affinity, reallocate_sticky=True) + if proxy_connect_failover.reallocate_sticky and not affinity.reallocate_sticky + else affinity + ), "prefer_earlier_reset_accounts": settings.prefer_earlier_reset_accounts, "prefer_earlier_reset_window": _prefer_earlier_reset_window(settings), "routing_strategy": _routing_strategy(settings), "model": request_model, "service_tier": request_service_tier, "exclude_account_ids": excluded_account_ids, - "preferred_account_id": preferred_candidate_id, + "preferred_account_id": proxy_connect_failover.preferred_account_id, "preferred_account_is_continuity_owner": preferred_account_is_continuity_owner, "lease_kind": "stream", "estimated_lease_tokens": _estimated_lease_tokens_from_request_usage_budget(request_usage_budget), @@ -1775,6 +1798,11 @@ async def _create_http_bridge_session( preferred_account_id=preferred_account_id, selected_account_id=None, ) + if proxy_connect_failover.last_error is not None: + # No eligible replacement exists after a confirmed + # pre-dispatch route failure: preserve the original + # sanitized failure instead of generating ``no_accounts``. + raise proxy_connect_failover.last_error is_local_account_cap = _is_local_account_cap_code(selection.error_code) if ( require_preferred_account @@ -1832,6 +1860,17 @@ async def _create_http_bridge_session( ) break except ProxyResponseError as exc: + if await proxy_connect_failover.handle( + self, + account, + selected_account_lease, + exc, + required_account=require_preferred_account and selected_is_preferred, + deferred_account_backoff_lifecycle=deferred_account_backoff_lifecycle, + defer_account_health_write=defer_account_health_writes, + ): + selected_account_lease = None + continue if exc.status_code != 401 or _remaining_budget_seconds(deadline) <= 0: await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None @@ -1857,6 +1896,17 @@ async def _create_http_bridge_session( ) break except ProxyResponseError as retry_exc: + if await proxy_connect_failover.handle( + self, + account, + selected_account_lease, + retry_exc, + required_account=require_preferred_account and selected_is_preferred, + deferred_account_backoff_lifecycle=deferred_account_backoff_lifecycle, + defer_account_health_write=defer_account_health_writes, + ): + selected_account_lease = None + continue if retry_exc.status_code != 401: await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None @@ -1867,7 +1917,7 @@ async def _create_http_bridge_session( selected_account_lease = None raise excluded_account_ids.add(account.id) - preferred_candidate_id = None + proxy_connect_failover.preferred_account_id = None await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None continue @@ -1879,7 +1929,7 @@ async def _create_http_bridge_session( selected_account_lease = None raise excluded_account_ids.add(account.id) - preferred_candidate_id = None + proxy_connect_failover.preferred_account_id = None await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None continue @@ -1904,7 +1954,7 @@ async def _create_http_bridge_session( ), ) from exc excluded_account_ids.add(account.id) - preferred_candidate_id = None + proxy_connect_failover.preferred_account_id = None await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None continue @@ -1943,7 +1993,7 @@ async def _create_http_bridge_session( ), ) from exc excluded_account_ids.add(account.id) - preferred_candidate_id = None + proxy_connect_failover.preferred_account_id = None await self._load_balancer.release_account_lease(selected_account_lease) selected_account_lease = None continue diff --git a/app/modules/proxy/_service/http_bridge/protocol.py b/app/modules/proxy/_service/http_bridge/protocol.py index b88edce03e..841d02c014 100644 --- a/app/modules/proxy/_service/http_bridge/protocol.py +++ b/app/modules/proxy/_service/http_bridge/protocol.py @@ -102,4 +102,5 @@ async def _maybe_touch_request_state_api_key_reservation(self, *args: Any, **kwa async def _reserve_websocket_api_key_usage(self, *args: Any, **kwargs: Any) -> Any: ... async def _release_websocket_reservation(self, *args: Any, **kwargs: Any) -> None: ... async def _release_websocket_request_state_reservation(self, *args: Any, **kwargs: Any) -> None: ... + async def _drain_deferred_account_error_backoffs(self, *args: Any, **kwargs: Any) -> None: ... def _schedule_cancel_safe_cleanup(self, *args: Any, **kwargs: Any) -> None: ... diff --git a/app/modules/proxy/_service/http_bridge/proxy_failover.py b/app/modules/proxy/_service/http_bridge/proxy_failover.py new file mode 100644 index 0000000000..c8a474302d --- /dev/null +++ b/app/modules/proxy/_service/http_bridge/proxy_failover.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from app.core.clients.proxy import ProxyResponseError, is_confirmed_pre_dispatch_transport_error +from app.db.models import Account +from app.modules.proxy._service.http_bridge.protocol import _HTTPBridgeServiceProtocol +from app.modules.proxy._service.support import _DeferredAccountBackoffLifecycle +from app.modules.proxy.load_balancer import AccountLease + + +@dataclass +class _HTTPBridgePreDispatchFailover: + """Bridge-session startup failover for confirmed dead account proxy routes. + + Only a transport failure that proves the upstream request never dispatched + may move a session to another account. The failed account's stream lease + is released before the bounded transient backoff floor is recorded; keyed + requests defer that write until their singular reservation settles. A + hard-required account fails closed on the original sanitized failure. The + preserved ``last_error`` keeps that failure authoritative when selection + cannot produce a replacement, instead of a generated ``no_accounts``. + """ + + excluded_account_ids: set[str] + preferred_account_id: str | None + reallocate_sticky: bool + last_error: ProxyResponseError | None = None + + async def handle( + self, + service: _HTTPBridgeServiceProtocol, + account: Account, + lease: AccountLease | None, + exc: ProxyResponseError, + *, + required_account: bool, + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None, + defer_account_health_write: bool = False, + ) -> bool: + if not is_confirmed_pre_dispatch_transport_error(exc): + return False + await service._load_balancer.release_account_lease(lease) + if ( + defer_account_health_write + and deferred_account_backoff_lifecycle is not None + and not deferred_account_backoff_lifecycle.settlement_confirmed + ): + deferred_account_backoff_lifecycle.pending_backoffs.setdefault(account.id, account) + else: + await service._load_balancer.record_error_backoff(account) + if required_account: + raise exc + self.last_error = exc + self.excluded_account_ids.add(account.id) + self.preferred_account_id = None + self.reallocate_sticky = True + return True diff --git a/app/modules/proxy/_service/http_bridge/streaming.py b/app/modules/proxy/_service/http_bridge/streaming.py index 1fce2706ea..2503978997 100644 --- a/app/modules/proxy/_service/http_bridge/streaming.py +++ b/app/modules/proxy/_service/http_bridge/streaming.py @@ -48,9 +48,7 @@ from app.core.types import JsonValue from app.core.utils.request_id import ensure_request_id, ensure_request_scope_id from app.core.utils.sse import format_sse_event, parse_sse_data_json -from app.db.models import ( - StickySessionKind, -) +from app.db.models import StickySessionKind from app.modules.api_keys.service import ( ApiKeyData, ApiKeyUsageReservationData, @@ -144,6 +142,8 @@ _WEBSOCKET_FULL_REPLAY_WAIT_POLL_SECONDS, # noqa: F401 _account_capacity_wait_payload, _account_selection_recovery_sleep_seconds_from_message, + _DeferredAccountBackoffLifecycle, + _DeferredAccountBackoffTracker, _event_type_from_payload, _HTTPBridgeOwnerForward, _HTTPBridgeSession, @@ -658,6 +658,7 @@ async def _stream_http_bridge_or_retry( return request_scope_id = ensure_request_scope_id() + deferred_account_backoff_tracker = _DeferredAccountBackoffTracker() try: async for line in self._stream_via_http_bridge( payload, @@ -685,14 +686,40 @@ async def _stream_http_bridge_or_retry( enforce_openai_sdk_contract=enforce_openai_sdk_contract, capacity_startup_wait_event=capacity_startup_wait_event, capacity_startup_ready_event=capacity_startup_ready_event, + deferred_account_backoff_tracker=deferred_account_backoff_tracker, ): yield line finally: with anyio.CancelScope(shield=True): - await _release_http_bridge_unanchored_handoffs_for_request( - self, - request_scope_id=request_scope_id, - ) + try: + lifecycle = deferred_account_backoff_tracker.current_lifecycle + if lifecycle is not None: + pending_backoffs = lifecycle.pending_backoffs + if lifecycle.settlement_confirmed: + await self._drain_deferred_account_error_backoffs(pending_backoffs) + elif not lifecycle.settlement_owned and ( + pending_backoffs or lifecycle.reservation != api_key_reservation + ): + # Session creation can fail before the request is + # submitted. Until submit returns, this wrapper owns + # the current lifecycle and may release exactly that + # reservation. Once ownership transfers, the request + # finalizer is the only safe settlement owner. + try: + await self._release_websocket_reservation(lifecycle.reservation) + except Exception: + logger.warning( + "Failed to release HTTP bridge API key reservation before deferred backoff", + exc_info=True, + ) + else: + lifecycle.settlement_confirmed = True + await self._drain_deferred_account_error_backoffs(pending_backoffs) + finally: + await _release_http_bridge_unanchored_handoffs_for_request( + self, + request_scope_id=request_scope_id, + ) async def _stream_via_http_bridge( self: Any, @@ -722,11 +749,14 @@ async def _stream_via_http_bridge( enforce_openai_sdk_contract: bool = True, capacity_startup_wait_event: asyncio.Event | None = None, capacity_startup_ready_event: asyncio.Event | None = None, + deferred_account_backoff_tracker: _DeferredAccountBackoffTracker | None = None, ) -> AsyncIterator[str]: del suppress_text_done_events request_id = ensure_request_id() dashboard_settings = await _service_get_settings_cache().get() runtime_config = _http_bridge_runtime_config(dashboard_settings, _service_get_settings()) + if deferred_account_backoff_tracker is None: + deferred_account_backoff_tracker = _DeferredAccountBackoffTracker() bridge_payload = payload.to_payload() bridge_client_metadata = _response_create_client_metadata( bridge_payload, @@ -739,6 +769,33 @@ async def _stream_via_http_bridge( if bridge_client_metadata is not None or "client_metadata" in bridge_payload: payload = payload.model_copy(update={"client_metadata": bridge_client_metadata}) + def begin_bridge_lifecycle( + reservation: ApiKeyUsageReservationData | None, + ) -> _DeferredAccountBackoffLifecycle: + previous_lifecycle = deferred_account_backoff_tracker.current_lifecycle + same_reservation = bool( + previous_lifecycle is not None + and ( + previous_lifecycle.reservation is reservation + or ( + previous_lifecycle.reservation is not None + and reservation is not None + and previous_lifecycle.reservation.reservation_id == reservation.reservation_id + ) + ) + ) + pending_backoffs = ( + previous_lifecycle.pending_backoffs + if previous_lifecycle is not None and not previous_lifecycle.settlement_owned and same_reservation + else {} + ) + lifecycle = _DeferredAccountBackoffLifecycle( + reservation=reservation, + pending_backoffs=pending_backoffs, + ) + deferred_account_backoff_tracker.current_lifecycle = lifecycle + return lifecycle + def prepare_bridge_request( request_payload: ResponsesRequest, *, @@ -765,8 +822,25 @@ def prepare_bridge_request( ) request_state.capacity_startup_wait_event = capacity_startup_wait_event request_state.capacity_startup_ready_event = capacity_startup_ready_event + lifecycle = begin_bridge_lifecycle(request_state.api_key_reservation) + request_state.deferred_account_error_backoffs = lifecycle.pending_backoffs + request_state.deferred_account_backoff_tracker = deferred_account_backoff_tracker + request_state.deferred_account_backoff_lifecycle = lifecycle return request_state, text_data + async def release_unowned_bridge_lifecycle( + lifecycle: _DeferredAccountBackoffLifecycle | None, + request_state: _WebSocketRequestState | None, + ) -> None: + if lifecycle is None or lifecycle.settlement_owned: + return + if request_state is not None: + await self._release_websocket_request_state_reservation(request_state) + return + await self._release_websocket_reservation(lifecycle.reservation) + lifecycle.settlement_confirmed = True + await self._drain_deferred_account_error_backoffs(lifecycle.pending_backoffs) + incoming_turn_state_header = _sticky_key_from_turn_state_header(headers) if not forwarded_request else None incoming_session_header = _sticky_key_from_session_header(headers) if not forwarded_request else None explicit_prompt_cache_key = _prompt_cache_key_from_request_model(payload) @@ -1344,6 +1418,8 @@ def switch_to_account_neutral_replay() -> None: request_deadline=request_deadline, session_header_fallback_key=session_header_fallback_key, exclude_account_ids=fresh_replay_excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as exc: if not owner_unavailable_allows_account_neutral_replay(exc): @@ -1609,6 +1685,8 @@ def switch_to_account_neutral_replay() -> None: session_header_fallback_key=session_header_fallback_key, request_deadline=request_deadline, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: if owner_unavailable_allows_account_neutral_replay(capacity_exc): @@ -1716,6 +1794,7 @@ def switch_to_account_neutral_replay() -> None: request_scope_id=owner_recovery_scope_id, ) retry_request_state: _WebSocketRequestState | None = None + retry_unowned_lifecycle: _DeferredAccountBackoffLifecycle | None = None try: retry_api_key_reservation = api_key_reservation retry_reservation_reacquired = False @@ -1729,6 +1808,7 @@ def switch_to_account_neutral_replay() -> None: request_usage_budget=estimate_api_key_request_usage(recovery_payload), ) retry_reservation_reacquired = True + retry_unowned_lifecycle = begin_bridge_lifecycle(retry_api_key_reservation) retry_request_state, retry_text_data = prepare_bridge_request( recovery_payload, @@ -1762,7 +1842,18 @@ def switch_to_account_neutral_replay() -> None: yield event_block except BaseException: if retry_reservation_reacquired and retry_api_key_reservation is not None: - await self._release_websocket_reservation(retry_api_key_reservation) + retry_lifecycle = ( + retry_request_state.deferred_account_backoff_lifecycle + if retry_request_state is not None + else retry_unowned_lifecycle + ) + try: + await release_unowned_bridge_lifecycle(retry_lifecycle, retry_request_state) + except Exception: + logger.warning( + "Failed to release owner-recovery HTTP bridge reservation", + exc_info=True, + ) raise finally: if owner_recovery_scope_id is not None: @@ -2042,6 +2133,8 @@ def switch_to_account_neutral_replay() -> None: request_deadline=request_deadline, session_header_fallback_key=session_header_fallback_key, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: wait_plan = _http_bridge_capacity_wait_plan(capacity_exc, request_deadline=request_deadline) @@ -2146,6 +2239,8 @@ def switch_to_account_neutral_replay() -> None: request_usage_budget=request_state.request_usage_budget, request_deadline=request_deadline, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: wait_plan = _http_bridge_capacity_wait_plan(capacity_exc, request_deadline=request_deadline) @@ -2335,6 +2430,8 @@ def switch_to_account_neutral_replay() -> None: request_usage_budget=estimate_api_key_request_usage(retry_payload), request_deadline=request_deadline, exclude_account_ids=request_state.excluded_account_ids or None, + deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle, + defer_account_health_writes=request_state.api_key_reservation is not None, ) except ProxyResponseError as capacity_exc: wait_plan = _http_bridge_capacity_wait_plan(capacity_exc, request_deadline=request_deadline) @@ -2371,6 +2468,8 @@ def switch_to_account_neutral_replay() -> None: session, request_scope_id=local_recovery_scope_id, ) + retry_request_state: _WebSocketRequestState | None = None + retry_unowned_lifecycle: _DeferredAccountBackoffLifecycle | None = None try: retry_api_key_reservation = api_key_reservation retry_reservation_reacquired = False @@ -2384,6 +2483,7 @@ def switch_to_account_neutral_replay() -> None: request_usage_budget=estimate_api_key_request_usage(retry_payload), ) retry_reservation_reacquired = True + retry_unowned_lifecycle = begin_bridge_lifecycle(retry_api_key_reservation) retry_request_state, retry_text_data = prepare_bridge_request( retry_payload, @@ -2416,7 +2516,18 @@ def switch_to_account_neutral_replay() -> None: pass except BaseException: if retry_reservation_reacquired and retry_api_key_reservation is not None: - await self._release_websocket_reservation(retry_api_key_reservation) + retry_lifecycle = ( + retry_request_state.deferred_account_backoff_lifecycle + if retry_request_state is not None + else retry_unowned_lifecycle + ) + try: + await release_unowned_bridge_lifecycle(retry_lifecycle, retry_request_state) + except Exception: + logger.warning( + "Failed to release local-recovery HTTP bridge reservation", + exc_info=True, + ) raise finally: if local_recovery_scope_id is not None: @@ -2495,6 +2606,9 @@ async def _stream_http_bridge_session_events( text_data=text_data, queue_limit=queue_limit, ) + lifecycle = request_state.deferred_account_backoff_lifecycle + if lifecycle is not None: + lifecycle.settlement_owned = True except ProxyResponseError as exc: if request_state.bridge_soft_capacity_reroute_allowed: raise diff --git a/app/modules/proxy/_service/streaming/protocol.py b/app/modules/proxy/_service/streaming/protocol.py index 8e3d17a6b8..279a59390a 100644 --- a/app/modules/proxy/_service/streaming/protocol.py +++ b/app/modules/proxy/_service/streaming/protocol.py @@ -6,6 +6,7 @@ class _StreamingServiceProtocol(Protocol): _acquire_account_response_create_lease_or_overload: Any _cancel_api_key_reservation_heartbeat_task: Any + _drain_deferred_account_error_backoffs: Any _encryptor: Any _ensure_fresh_with_budget: Any _get_work_admission: Any diff --git a/app/modules/proxy/_service/streaming/retry.py b/app/modules/proxy/_service/streaming/retry.py index 35e12de908..202bef4630 100644 --- a/app/modules/proxy/_service/streaming/retry.py +++ b/app/modules/proxy/_service/streaming/retry.py @@ -14,7 +14,12 @@ from app.core.auth.refresh import RefreshError, is_transient_refresh_contention, refresh_contention_kind from app.core.balancer import failover_decision from app.core.balancer.types import UpstreamError -from app.core.clients.proxy import ProxyResponseError, _resolve_stream_transport, pop_stream_timeout_overrides +from app.core.clients.proxy import ( + ProxyResponseError, + _resolve_stream_transport, + is_confirmed_pre_dispatch_transport_error, + pop_stream_timeout_overrides, +) from app.core.errors import openai_error, response_failed_event from app.core.openai.requests import ResponsesRequest, extract_input_file_ids from app.core.resilience.network_recovery import ( @@ -351,6 +356,7 @@ async def _stream_with_retry( network_recovery = ProcessNetworkRecovery(transport="stream", request_id=request_id) settlement = _StreamSettlement() last_transient_exc: ProxyResponseError | None = None + last_pre_dispatch_transport_error: ProxyResponseError | None = None last_account_model_rejection: ProxyResponseError | None = None last_account_model_rejection_account_id: str | None = None account_model_replacement_account_id: str | None = None @@ -367,6 +373,7 @@ async def _stream_with_retry( require_preferred_account = False last_retryable_stream_error: _RetryableStreamError | None = None pending_post_refresh_transient_penalties: list[tuple[Account, UpstreamError, str, int, int]] = [] + deferred_account_error_backoffs: dict[str, Account] = {} post_refresh_transient_replacement_selected = False require_security_work_authorized = False account_leases: list[AccountLease] = [] @@ -389,20 +396,48 @@ async def _release_tracked_stream_lease(lease: AccountLease | None) -> None: pass await proxy._load_balancer.release_account_lease(lease) + def _render_dispatch_transport_error(exc: ProxyResponseError) -> str: + # Terminal render of the preserved sanitized transport failure: + # the client sees the original upstream-unavailable error instead + # of a misleading generated ``no_accounts`` response. + error = _parse_openai_error(exc.payload) + error_code = ( + _normalize_error_code( + error.code if error else None, + error.type if error else None, + ) + or "upstream_unavailable" + ) + error_message = error.message if error and error.message else "Upstream transport failed" + event = response_failed_event( + error_code, + error_message, + error_type=(error.type if error else None) or "server_error", + response_id=request_id, + error_param=error.param if error else None, + ) + _apply_error_metadata(event["response"]["error"], error) + return format_sse_event(event) + async def _settle_stream_usage_before_pending_penalty( current_settlement: _StreamSettlement, ) -> bool: apply_pending_penalty = post_refresh_transient_replacement_selected and bool( pending_post_refresh_transient_penalties ) + wait_for_health_write = apply_pending_penalty or bool(deferred_account_error_backoffs) + settle_kwargs = {"wait_for_settlement": True} if wait_for_health_write else {} + settled_result = await proxy._settle_stream_api_key_usage( + api_key, + api_key_reservation, + current_settlement, + request_id, + **settle_kwargs, + ) + if not settled_result: + return False + await proxy._drain_deferred_account_error_backoffs(deferred_account_error_backoffs) if apply_pending_penalty: - settled_result = await proxy._settle_stream_api_key_usage( - api_key, - api_key_reservation, - current_settlement, - request_id, - wait_for_settlement=True, - ) pending_penalties = list(pending_post_refresh_transient_penalties) pending_post_refresh_transient_penalties.clear() for pending_penalty in pending_penalties: @@ -421,24 +456,27 @@ async def _settle_stream_usage_before_pending_penalty( ) if transient_retry_count > 1: await proxy._load_balancer.record_errors(failed_account, transient_retry_count - 1) - return settled_result - return await proxy._settle_stream_api_key_usage( - api_key, - api_key_reservation, - current_settlement, - request_id, - ) + return settled_result + + async def _record_or_defer_confirmed_route_backoff(account: Account) -> None: + if api_key is not None and api_key_reservation is not None: + deferred_account_error_backoffs.setdefault(account.id, account) + return + await proxy._load_balancer.record_error_backoff(account) async def _drain_pending_post_refresh_penalty_on_terminal( current_settlement: _StreamSettlement, - ) -> None: + ) -> bool: nonlocal post_refresh_transient_replacement_selected, settled - if pending_post_refresh_transient_penalties: + if pending_post_refresh_transient_penalties or deferred_account_error_backoffs: # A failed replacement selection still ends the request. Mark # it as terminal so the deferred failure is settled and # recorded before this path returns or re-raises. - post_refresh_transient_replacement_selected = True + if pending_post_refresh_transient_penalties: + post_refresh_transient_replacement_selected = True settled = await _settle_stream_usage_before_pending_penalty(current_settlement) + return settled + return True async def _wait_for_process_network_recovery( account: Account, @@ -547,6 +585,12 @@ async def _iter_stream_once() -> AsyncIterator[str]: ): yield line except ProxyResponseError as exc: + if is_confirmed_pre_dispatch_transport_error(exc): + # Keep dispatch provenance intact for the outer account + # failover handler. Converting this into the generic + # transient wrapper would authorize same-account replay + # and lose the confirmed dead-route backoff semantics. + raise error = _parse_openai_error(exc.payload) error_code = _normalize_error_code( error.code if error else None, @@ -1071,6 +1115,13 @@ async def _retry_account_model_rejection( selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES or not (propagate_http_errors and last_transient_exc is not None) ) + and ( + selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES + # A preserved confirmed pre-dispatch failure is + # terminal for this request: waiting for capacity + # recovery cannot resurrect the dead proxy route. + or last_pre_dispatch_transport_error is None + ) and ( selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES or (last_retryable_stream_error is None and last_security_work_retry_error is None) @@ -1106,6 +1157,15 @@ async def _retry_account_model_rejection( account_id=last_account_model_rejection_account_id, ) return + if last_pre_dispatch_transport_error is not None: + # No eligible replacement exists: preserve the original + # sanitized upstream-unavailable failure instead of + # generating a misleading ``no_accounts`` response. + await _drain_pending_post_refresh_penalty_on_terminal(settlement) + if propagate_http_errors: + raise last_pre_dispatch_transport_error + yield _render_dispatch_transport_error(last_pre_dispatch_transport_error) + return if selection.error_code in _LOCAL_ACCOUNT_CAP_ERROR_CODES: await _drain_pending_post_refresh_penalty_on_terminal(settlement) no_accounts_msg = selection.error_message or "Local account capacity is exhausted" @@ -1300,6 +1360,13 @@ async def _retry_account_model_rejection( post_refresh_transient_replacement_selected = True account_id_value = account.id + if last_pre_dispatch_transport_error is not None: + # The preserved connect failure is only authoritative when + # replacement selection is empty. Once another account is + # actually attempted, its terminal outcome takes precedence. + if last_transient_exc is last_pre_dispatch_transport_error: + last_transient_exc = None + last_pre_dispatch_transport_error = None if last_account_model_rejection is not None and account.id != last_account_model_rejection_account_id: # The original 400 is only the fallback when account # selection cannot produce a replacement. Once this @@ -1771,7 +1838,7 @@ async def _retry_account_model_rejection( settlement.error = tex.error settlement.account_health_error = _facade()._should_penalize_stream_error(error_code) settled = await _settle_stream_usage_before_pending_penalty(settlement) - if settlement.account_health_error: + if settled and settlement.account_health_error: await proxy._handle_stream_error( account, _stream_settlement_error_payload(settlement), @@ -1885,6 +1952,43 @@ async def _retry_account_model_rejection( _facade()._raise_proxy_budget_exhausted() if _facade()._is_account_neutral_error_code(code): raise + if is_confirmed_pre_dispatch_transport_error(tex): + # The transport proved the request never + # dispatched: this account's proxy route is + # dead. Release the account's stream lease + # before recording health so its slot never + # outlives the failed route, then jump + # straight to the bounded transient backoff + # floor so independent requests stop + # rediscovering the dead route one generic + # error at a time. + await _release_tracked_stream_lease(current_account_lease) + current_account_lease = None + await _record_or_defer_confirmed_route_backoff(account) + can_try_other_account = ( + not require_preferred_account + and account.id != file_preferred_account_id + and attempt < max_attempts - 1 + ) + if not can_try_other_account: + # Hard account ownership or exhausted + # attempts: fail closed on the original + # sanitized failure without crossing + # accounts. + raise + last_transient_exc = tex + last_pre_dispatch_transport_error = tex + transient_failed_account_id = account.id + excluded_account_ids.add(account.id) + affinity = replace(affinity, reallocate_sticky=True) + _facade().logger.info( + "Retrying stream after confirmed pre-dispatch proxy connect failure " + "request_id=%s account_id=%s attempt=%d", + request_id, + account.id, + attempt + 1, + ) + break classified = await proxy._handle_stream_error( account, _upstream_error_from_openai(error), @@ -1991,13 +2095,13 @@ async def _retry_account_model_rejection( finally: pop_stream_timeout_overrides(stream_timeout_tokens) settled = await _settle_stream_usage_before_pending_penalty(settlement) - if settlement.account_health_error: + if settled and settlement.account_health_error: await proxy._handle_stream_error( account, _stream_settlement_error_payload(settlement), settlement.error_code or "upstream_error", ) - elif settlement.record_success: + elif settled and settlement.record_success: await proxy._load_balancer.record_success(account) network_recovery.log_recovered() upstream_transport_metric_status = settlement.status @@ -2051,8 +2155,8 @@ async def _retry_account_model_rejection( ) continue except _TerminalStreamError as exc: - await _drain_pending_post_refresh_penalty_on_terminal(settlement) - if _facade()._should_penalize_stream_error(exc.code): + health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement) + if health_write_allowed and _facade()._should_penalize_stream_error(exc.code): await proxy._handle_stream_error(account, exc.error, exc.code) return except ProxyResponseError as exc: @@ -2349,7 +2453,7 @@ async def _retry_account_model_rejection( settlement.error = _upstream_error_from_openai(error) settlement.account_health_error = _facade()._should_penalize_stream_error(error_code) settled = await _settle_stream_usage_before_pending_penalty(settlement) - if settlement.account_health_error: + if settled and settlement.account_health_error: await proxy._handle_stream_error( account, _stream_settlement_error_payload(settlement), @@ -2407,6 +2511,56 @@ async def _retry_account_model_rejection( if _facade()._is_account_neutral_error_code(error_code): await _drain_pending_post_refresh_penalty_on_terminal(settlement) raise + if is_confirmed_pre_dispatch_transport_error(retry_exc): + # This retry still failed before dispatch. Handle + # the proven dead route before the generic + # failover policy: that policy does not know + # about hard account ownership and may otherwise + # cross a previous-response, turn-state, file, or + # single-account boundary. + await _release_tracked_stream_lease(current_account_lease) + current_account_lease = None + await _record_or_defer_confirmed_route_backoff(account) + last_transient_exc = retry_exc + last_pre_dispatch_transport_error = retry_exc + + verified_owner_replay_moved = False + if ( + attempt < max_attempts - 1 + and routing_strategy != "single_account" + and file_preferred_account_id is None + and turn_state_owner_account_id is None + ): + verified_owner_replay_moved = _move_verified_fresh_replay_from_owner( + account_id=account.id, + outcome="owner_post_refresh_proxy_connect_failure", + ) + + can_try_other_account = bool( + attempt < max_attempts - 1 + and routing_strategy != "single_account" + and file_preferred_account_id is None + and turn_state_owner_account_id is None + and not require_preferred_account + ) + if can_try_other_account: + excluded_account_ids.add(account.id) + if not verified_owner_replay_moved: + affinity = replace(affinity, reallocate_sticky=True) + _facade().logger.info( + "Retrying post-refresh stream after confirmed pre-dispatch proxy " + "connect failure request_id=%s account_id=%s attempt=%d", + request_id, + account.id, + attempt + 1, + ) + continue + + # Hard ownership or an exhausted attempt budget: + # stop here. The shared terminal path settles the + # reservation, drains the deferred backoff floor, + # and preserves this sanitized failure. + break current_error_payload = _upstream_error_from_openai(error) current_error_code = error_code or "upstream_error" classified = classify_upstream_failure( @@ -2451,13 +2605,14 @@ async def _retry_account_model_rejection( ) excluded_account_ids.add(account.id) continue - await _drain_pending_post_refresh_penalty_on_terminal(settlement) - await proxy._handle_stream_error( - account, - current_error_payload, - current_error_code, - http_status=retry_exc.status_code, - ) + health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement) + if health_write_allowed: + await proxy._handle_stream_error( + account, + current_error_payload, + current_error_code, + http_status=retry_exc.status_code, + ) if propagate_http_errors: raise error_message = error.message if error else None @@ -2475,17 +2630,25 @@ async def _retry_account_model_rejection( failed_account is account for failed_account, *_rest in pending_post_refresh_transient_penalties ) - if pending_post_refresh_transient_penalties: - await _drain_pending_post_refresh_penalty_on_terminal(settlement) - if settlement.account_health_error and not current_account_penalty_queued: + ordered_settlement_required = bool( + pending_post_refresh_transient_penalties or deferred_account_error_backoffs + ) + health_write_allowed = True + if ordered_settlement_required: + health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement) + if ( + health_write_allowed + and settlement.account_health_error + and not current_account_penalty_queued + ): await proxy._handle_stream_error( account, _stream_settlement_error_payload(settlement), settlement.error_code or "upstream_error", ) - elif settlement.record_success: + elif health_write_allowed and settlement.record_success: await proxy._load_balancer.record_success(account) - if not settled: + if not settled and not ordered_settlement_required: settled = await _settle_stream_usage_before_pending_penalty(settlement) upstream_transport_metric_status = settlement.status _record_upstream_transport_metric_once(settlement.status) @@ -2521,8 +2684,8 @@ async def _retry_account_model_rejection( excluded_account_ids.add(account.id) require_security_work_authorized = True continue - await _drain_pending_post_refresh_penalty_on_terminal(settlement) - if _facade()._should_penalize_stream_error(error_code): + health_write_allowed = await _drain_pending_post_refresh_penalty_on_terminal(settlement) + if health_write_allowed and _facade()._should_penalize_stream_error(error_code): await proxy._handle_stream_error( account, _upstream_error_from_openai(error), @@ -2572,6 +2735,12 @@ async def _retry_account_model_rejection( return if propagate_http_errors and last_transient_exc is not None: raise last_transient_exc + if last_pre_dispatch_transport_error is not None: + # Attempt budget exhausted after confirmed pre-dispatch route + # failures: surface the original sanitized failure rather than + # a generated ``no_accounts`` response. + yield _render_dispatch_transport_error(last_pre_dispatch_transport_error) + return if last_retryable_stream_error is not None: retries_exhausted_msg = str(last_retryable_stream_error.error.get("message") or "Upstream error") event = response_failed_event( @@ -2666,11 +2835,17 @@ async def _retry_account_model_rejection( and api_key is not None and api_key_reservation is not None ): - release_coro = proxy._release_unsettled_stream_api_key_usage( - api_key=api_key, - api_key_reservation=api_key_reservation, - request_id=request_id, - ) + + async def _release_reservation_then_drain_backoffs() -> None: + released = await proxy._release_unsettled_stream_api_key_usage( + api_key=api_key, + api_key_reservation=api_key_reservation, + request_id=request_id, + ) + if released: + await proxy._drain_deferred_account_error_backoffs(deferred_account_error_backoffs) + + release_coro = _release_reservation_then_drain_backoffs() current_task = asyncio.current_task() if current_task is not None and current_task.cancelling(): proxy._schedule_cancel_safe_cleanup( @@ -2680,3 +2855,5 @@ async def _retry_account_model_rejection( ) else: await release_coro + elif settled and deferred_account_error_backoffs: + await proxy._drain_deferred_account_error_backoffs(deferred_account_error_backoffs) diff --git a/app/modules/proxy/_service/support.py b/app/modules/proxy/_service/support.py index 0d842e78f9..2e94184e09 100644 --- a/app/modules/proxy/_service/support.py +++ b/app/modules/proxy/_service/support.py @@ -739,6 +739,19 @@ class _RequestLogFailureMetadata: bridge_stage: str | None = None +@dataclass(slots=True) +class _DeferredAccountBackoffLifecycle: + reservation: ApiKeyUsageReservationData | None + pending_backoffs: dict[str, Account] = field(default_factory=dict) + settlement_owned: bool = False + settlement_confirmed: bool = False + + +@dataclass(slots=True) +class _DeferredAccountBackoffTracker: + current_lifecycle: _DeferredAccountBackoffLifecycle | None = None + + @dataclass class _WebSocketRequestState: request_id: str @@ -820,6 +833,7 @@ class _WebSocketRequestState: request_stage: str = "first_turn" preferred_account_id: str | None = None require_security_work_authorized: bool = False + durable_capability_lineage_required: bool = False file_required_preferred_account: bool = False bridge_soft_capacity_reroute_allowed: bool = False error_code_override: str | None = None @@ -868,6 +882,13 @@ class _WebSocketRequestState: client_ip: str | None = None downstream_visible: bool = False last_downstream_sequence_number: int | None = None + # Confirmed pre-dispatch account-route failures must not mutate account + # health while this request's API-key reservation is still live. The + # account objects are keyed by id so repeated connect attempts cannot + # stack the same backoff floor more than once before settlement. + deferred_account_error_backoffs: dict[str, Account] = field(default_factory=dict) + deferred_account_backoff_tracker: _DeferredAccountBackoffTracker | None = None + deferred_account_backoff_lifecycle: _DeferredAccountBackoffLifecycle | None = None deferred_reasoning_downstream_texts: list[str] = field(default_factory=list) suppress_next_created_downstream: bool = False replay_downstream_response_id: str | None = None diff --git a/app/modules/proxy/_service/websocket/helpers.py b/app/modules/proxy/_service/websocket/helpers.py index 7201fa13ec..e3a3c2e8db 100644 --- a/app/modules/proxy/_service/websocket/helpers.py +++ b/app/modules/proxy/_service/websocket/helpers.py @@ -14,8 +14,8 @@ from app.core.clients.files import create_file as core_create_file # noqa: F401 from app.core.clients.files import finalize_file as core_finalize_file # noqa: F401 from app.core.clients.http import lease_http_session as lease_http_session # noqa: F401 -from app.core.clients.proxy import CodexControlResponse as CodexControlResponse from app.core.clients.proxy import ( # noqa: F401 # noqa: F401 + CODEX_LB_REQUIRED_CAPABILITY_HEADER, ImageFetchSession, ProxyResponseError, UpstreamProxyRouteTrace, @@ -31,6 +31,7 @@ push_stream_timeout_overrides, push_transcribe_timeout_overrides, ) +from app.core.clients.proxy import CodexControlResponse as CodexControlResponse from app.core.clients.proxy import codex_control_request as core_codex_control_request # noqa: F401 from app.core.clients.proxy import compact_responses as core_compact_responses # noqa: F401 from app.core.clients.proxy import transcribe_audio as core_transcribe_audio # noqa: F401 @@ -1696,9 +1697,15 @@ def _websocket_receive_timeout_for_pending_requests( ) +class _WebSocketJsonObject(dict[str, JsonValue]): + def __init__(self, pairs: list[tuple[str, JsonValue]]) -> None: + super().__init__(pairs) + self.raw_pairs = tuple(pairs) + + def _parse_websocket_payload(text: str) -> dict[str, JsonValue] | None: try: - payload = json.loads(text) + payload = json.loads(text, object_pairs_hook=_WebSocketJsonObject) except json.JSONDecodeError: return None if not isinstance(payload, dict): @@ -1706,6 +1713,31 @@ def _parse_websocket_payload(text: str) -> dict[str, JsonValue] | None: return payload +def _websocket_capability_metadata_values(payload: dict[str, JsonValue]) -> tuple[JsonValue, ...] | None: + normalized_name = CODEX_LB_REQUIRED_CAPABILITY_HEADER.lower() + payload_pairs = payload.raw_pairs if isinstance(payload, _WebSocketJsonObject) else tuple(payload.items()) + misplaced_values = [value for key, value in payload_pairs if key.lower() == normalized_name] + if misplaced_values: + # The reserved per-frame carrier is valid only inside + # ``client_metadata``. Surface a deliberately ambiguous carrier set so + # the shared parser rejects any top-level placement before selection. + return (misplaced_values[0], misplaced_values[0]) + if not isinstance(payload, _WebSocketJsonObject): + return None + metadata_objects = [value for key, value in payload.raw_pairs if key == "client_metadata"] + values: list[JsonValue] = [] + for metadata in metadata_objects: + if not isinstance(metadata, _WebSocketJsonObject): + continue + values.extend(value for key, value in metadata.raw_pairs if key.lower() == normalized_name) + if len(metadata_objects) > 1 and values: + # A duplicate top-level metadata container can otherwise erase the + # sole marker through last-key-wins JSON decoding. Treat the carrier + # as ambiguous so routing fails before selection. + values.append(values[0]) + return tuple(values) + + def _is_websocket_response_create(payload: dict[str, JsonValue]) -> bool: payload_type = payload.get("type") return isinstance(payload_type, str) and payload_type == "response.create" diff --git a/app/modules/proxy/_service/websocket/mixin.py b/app/modules/proxy/_service/websocket/mixin.py index 4e292e0046..8d215f85c4 100644 --- a/app/modules/proxy/_service/websocket/mixin.py +++ b/app/modules/proxy/_service/websocket/mixin.py @@ -44,6 +44,7 @@ apply_codex_installation_headers, apply_codex_installation_metadata, filter_inbound_headers, + is_confirmed_pre_dispatch_transport_error, pop_compact_timeout_overrides, pop_stream_timeout_overrides, pop_transcribe_timeout_overrides, @@ -401,6 +402,7 @@ _trim_websocket_previous_response_input_items, _upstream_websocket_disconnect_message, _websocket_auth_failure_requires_reauth, + _websocket_capability_metadata_values, _websocket_client_previous_response_full_resend_is_retry_safe, _websocket_connect_deadline, _websocket_continuity_anchor_for_payload, @@ -433,6 +435,17 @@ _sticky_key_from_turn_state_header, ) from app.modules.proxy.api_key_usage import estimate_api_key_request_usage +from app.modules.proxy.capability_routing import ( + CAPABILITY_ROUTING_UNAVAILABLE_CODE, + CAPABILITY_ROUTING_UNAVAILABLE_MESSAGE, + RoutingCapability, + RoutingIntent, + _capability_lineage_unavailable_error, + capability_lineage_aliases, + parse_routing_intent, + reject_capability_signal_outside_response_create, + strip_capability_metadata, +) from app.modules.proxy.continuity import resolve_required_account_id from app.modules.proxy.durable_bridge_coordinator import ( DurableBridgeLookup as DurableBridgeLookup, @@ -475,12 +488,23 @@ def _facade() -> Any: logger = logging.getLogger(__name__) _WEBSOCKET_PINNED_REFRESH_UNAVAILABLE_MESSAGE = "Account refresh is temporarily unavailable; retry later." +_CAPABILITY_REQUIRED_NO_AUTHORIZED_ACCOUNTS_MESSAGE = ( + "This request requires Trusted Access for Cyber, but no eligible account is marked as " + "security-work-authorized. codex-lb did not fall back to an ordinary account." +) +_CAPABILITY_REQUIRED_NO_AUTHORIZED_ACCOUNTS_ACTION = "fail_closed_capability_routing" class _WebSocketReplaySequenceRegression(Exception): pass +class _CapabilityLineagePropagationError(Exception): + def __init__(self, error: ProxyResponseError) -> None: + super().__init__("Capability lineage propagation failed") + self.error = error + + def _log_websocket_persist_conflict(context: str, exc: RefreshError, account_id: str) -> None: """Surface a post-exchange persist/status CAS conflict distinctly in logs. @@ -529,6 +553,40 @@ async def _reject_websocket_owner_switch_blocked( await _release_websocket_response_create_gate(request_state, response_create_gate) +async def _reject_websocket_capability_switch_blocked( + proxy: Any, + websocket: WebSocket, + *, + client_send_lock: anyio.Lock, + request_state: _WebSocketRequestState, + account: Account, + api_key: ApiKeyData | None, + response_create_gate: asyncio.Semaphore, + downstream_activity: _DownstreamWebSocketActivity, +) -> None: + error_message = ( + "Required capability cannot switch accounts while another response is still streaming; " + "retry after the terminal frame." + ) + await proxy._release_websocket_request_state_reservation(request_state) + await proxy._write_websocket_connect_failure( + account_id=account.id, + api_key=api_key, + request_state=request_state, + error_code="continuity_owner_conflict", + error_message=error_message, + ) + await proxy._emit_websocket_terminal_error( + websocket, + client_send_lock=client_send_lock, + request_state=request_state, + error_code="continuity_owner_conflict", + error_message=error_message, + downstream_activity=downstream_activity, + ) + await _release_websocket_response_create_gate(request_state, response_create_gate) + + async def _close_downstream_after_sequenced_replay_refusal( websocket: WebSocket, downstream_activity: _DownstreamWebSocketActivity, @@ -748,6 +806,7 @@ async def proxy_responses_websocket( api_key: ApiKeyData | None, client_ip: str | None = None, synthesized_turn_state: str | None = None, + capability_header_values: tuple[str, ...] | None = None, ) -> None: proxy = cast(_WebSocketServiceProtocol, self) _ = proxy @@ -775,6 +834,7 @@ async def proxy_responses_websocket( ) account: Account | None = None account_lease: AccountLease | None = None + upstream_requires_security_work_authorized: bool | None = None upstream_turn_state: str | None = _sticky_key_from_turn_state_header(headers) client_turn_state_header: str | None = _sticky_key_from_turn_state_header(filtered_headers) upstream_account_id: str | None = None @@ -788,6 +848,7 @@ async def release_current_account_lease() -> None: async def retire_current_upstream() -> None: nonlocal account, upstream, upstream_control, upstream_reader + nonlocal upstream_requires_security_work_authorized if upstream_control is not None: upstream_control.reconnect_requested = True if upstream_reader is not None: @@ -805,6 +866,7 @@ async def retire_current_upstream() -> None: upstream = None await release_current_account_lease() account = None + upstream_requires_security_work_authorized = None try: while True: @@ -934,9 +996,26 @@ async def retire_current_upstream() -> None: text_data = message.get("text") bytes_data = message.get("bytes") + if bytes_data is not None: + async with client_send_lock: + await websocket.send_text( + _serialize_websocket_error_event( + _wrapped_websocket_error_event(400, openai_invalid_payload_error()) + ) + ) + continue + if text_data is not None: payload = _parse_websocket_payload(text_data) - if payload is not None and _is_websocket_response_create(payload): + if payload is None: + async with client_send_lock: + await websocket.send_text( + _serialize_websocket_error_event( + _wrapped_websocket_error_event(400, openai_invalid_payload_error()) + ) + ) + continue + if _is_websocket_response_create(payload): try: prepared_request = await proxy._prepare_websocket_response_create_request( payload, @@ -953,6 +1032,7 @@ async def retire_current_upstream() -> None: conversation_id=conversation_id, client_ip=client_ip, synthesized_turn_state=synthesized_turn_state, + capability_header_values=capability_header_values, ) if await _websocket_full_replay_should_wait_for_continuity( prepared_request.request_state, @@ -991,6 +1071,7 @@ async def retire_current_upstream() -> None: conversation_id=conversation_id, client_ip=client_ip, synthesized_turn_state=synthesized_turn_state, + capability_header_values=capability_header_values, ) request_state = prepared_request.request_state request_affinity = prepared_request.affinity_policy @@ -1043,6 +1124,21 @@ async def retire_current_upstream() -> None: ) ) continue + elif payload is not None: + try: + reject_capability_signal_outside_response_create( + api_key=api_key, + client_metadata=payload.get("client_metadata"), + client_metadata_values=_websocket_capability_metadata_values(payload), + ) + except ProxyResponseError as exc: + async with client_send_lock: + await websocket.send_text( + _serialize_websocket_error_event( + _wrapped_websocket_error_event(exc.status_code, exc.payload) + ) + ) + continue if upstream_reader is not None and upstream_reader.done(): try: @@ -1175,6 +1271,117 @@ async def retire_current_upstream() -> None: payload = None continue + if ( + request_state is not None + and upstream is not None + and account is not None + and request_state.require_security_work_authorized + ): + capability_account_reusable = False + if upstream_requires_security_work_authorized: + try: + ( + revalidated_account, + _error_code, + _error_message, + ) = await proxy._revalidate_open_websocket_account( + account, + request_state=request_state, + api_key=request_state.api_key or api_key, + ) + except ProxyResponseError as exc: + error = _parse_openai_error(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 and error.message else "Upstream error" + await proxy._release_websocket_request_state_reservation(request_state) + await proxy._write_websocket_connect_failure( + account_id=account.id, + api_key=api_key, + request_state=request_state, + error_code=error_code or "upstream_error", + error_message=error_message, + ) + await proxy._emit_websocket_terminal_error( + websocket, + client_send_lock=client_send_lock, + request_state=request_state, + error_code=error_code or "upstream_error", + error_message=error_message, + error_type=error.type if error and error.type else "server_error", + error_param=error.param if error else None, + downstream_activity=downstream_activity, + ) + request_state = None + text_data = None + payload = None + continue + except BaseException as exc: + await proxy._release_websocket_request_state_reservation(request_state) + if not isinstance(exc, Exception): + raise + _facade().logger.exception( + "Capability account revalidation failed request_id=%s account_id=%s", + request_state.request_log_id or request_state.request_id, + account.id, + ) + await proxy._write_websocket_connect_failure( + account_id=account.id, + api_key=api_key, + request_state=request_state, + error_code=CAPABILITY_ROUTING_UNAVAILABLE_CODE, + error_message=CAPABILITY_ROUTING_UNAVAILABLE_MESSAGE, + ) + await proxy._emit_websocket_terminal_error( + websocket, + client_send_lock=client_send_lock, + request_state=request_state, + error_code=CAPABILITY_ROUTING_UNAVAILABLE_CODE, + error_message=CAPABILITY_ROUTING_UNAVAILABLE_MESSAGE, + error_type="server_error", + downstream_activity=downstream_activity, + ) + request_state = None + text_data = None + payload = None + continue + if revalidated_account is not None: + account = revalidated_account + capability_account_reusable = True + + if not capability_account_reusable: + async with pending_lock: + capability_switch_blocked = any( + pending_request is not request_state for pending_request in pending_requests + ) + if capability_switch_blocked and request_state in pending_requests: + pending_requests.remove(request_state) + if capability_switch_blocked: + await _reject_websocket_capability_switch_blocked( + proxy, + websocket, + client_send_lock=client_send_lock, + request_state=request_state, + account=account, + api_key=api_key, + response_create_gate=response_create_gate, + downstream_activity=downstream_activity, + ) + request_state = None + text_data = None + payload = None + continue + await retire_current_upstream() + upstream_turn_state = None + if synthesized_turn_state is not None: + filtered_headers = { + key: value + for key, value in filtered_headers.items() + if key.lower() != "x-codex-turn-state" + } + if ( request_state is not None and upstream is not None @@ -1193,6 +1400,7 @@ async def retire_current_upstream() -> None: affinity_policy=request_state.affinity_policy, model=request_state.model, preferred_account_id=account.id, + require_security_work_authorized=request_state.require_security_work_authorized, fallback_on_preferred_account_unavailable=False, ) if ownership_selection.account is None: @@ -1321,14 +1529,6 @@ async def retire_current_upstream() -> None: } if upstream is None: - if text_data is not None and payload is None: - async with client_send_lock: - await websocket.send_text( - _serialize_websocket_error_event( - _wrapped_websocket_error_event(400, openai_invalid_payload_error()) - ) - ) - continue if request_state is None: async with client_send_lock: await websocket.send_text( @@ -1402,6 +1602,7 @@ async def retire_current_upstream() -> None: # owner when a transparent replay reconnects. upstream_turn_state = None upstream_account_id = account.id + upstream_requires_security_work_authorized = request_state.require_security_work_authorized upstream_turn_state = _facade()._upstream_turn_state_from_socket(upstream) or upstream_turn_state upstream_control = _WebSocketUpstreamControl() upstream_reader = asyncio.create_task( @@ -1471,10 +1672,6 @@ async def retire_current_upstream() -> None: archive_request_id = None if request_state is None else request_state.archive_request_id with _websocket_archive_request_context(archive_request_id): await upstream.send_text(text_data) - elif bytes_data is not None: - archive_request_id = None if request_state is None else request_state.archive_request_id - with _websocket_archive_request_context(archive_request_id): - await upstream.send_bytes(bytes_data) except ProxyResponseError as exc: error = _parse_openai_error(exc.payload) error_code = _normalize_error_code(error.code if error else None, error.type if error else None) @@ -1659,10 +1856,19 @@ async def _prepare_websocket_response_create_request( conversation_id: str | None = None, client_ip: str | None = None, synthesized_turn_state: str | None = None, + capability_header_values: tuple[str, ...] | None = None, ) -> _PreparedWebSocketRequest: proxy = cast(_WebSocketServiceProtocol, self) _ = proxy refreshed_api_key = await proxy._refresh_websocket_api_key_policy(api_key) + raw_client_metadata = payload.get("client_metadata") + capability_intent = parse_routing_intent( + headers, + api_key=refreshed_api_key, + client_metadata=raw_client_metadata, + header_values=capability_header_values, + client_metadata_values=_websocket_capability_metadata_values(payload), + ) responses_payload = normalize_responses_request_payload( payload, openai_compat=openai_cache_affinity, @@ -1677,6 +1883,10 @@ async def _prepare_websocket_response_create_request( service_tier_was_enforced=service_tier_was_enforced, ) normalized_payload = responses_payload.to_payload() + stripped_client_metadata = strip_capability_metadata(normalized_payload.get("client_metadata")) + if stripped_client_metadata is not normalized_payload.get("client_metadata"): + responses_payload = responses_payload.model_copy(update={"client_metadata": stripped_client_metadata}) + normalized_payload = responses_payload.to_payload() body_uses_responses_lite = _payload_uses_responses_lite(normalized_payload) trusted_incremental_responses_lite = bool( not body_uses_responses_lite @@ -1785,6 +1995,21 @@ async def _prepare_websocket_response_create_request( responses_payload.previous_response_id, len(missing_call_ids), ) + session_id = _owner_lookup_session_id_from_headers( + headers, + synthesized_turn_state=synthesized_turn_state, + ) + capability_route = await proxy._capability_router.route( + capability_intent, + api_key_id=refreshed_api_key.id if refreshed_api_key is not None else None, + aliases=capability_lineage_aliases( + headers, + session_id=_sticky_key_from_session_header(headers), + turn_state=_sticky_key_from_turn_state_header(headers) or synthesized_turn_state, + previous_response_ids=(responses_payload.previous_response_id,), + client_metadata=client_metadata, + ), + ) reservation = await proxy._reserve_websocket_api_key_usage( refreshed_api_key, request_model=responses_payload.model, @@ -1794,7 +2019,6 @@ async def _prepare_websocket_response_create_request( request_usage_budget=estimate_api_key_request_usage(responses_payload), ) try: - session_id = _owner_lookup_session_id_from_headers(headers, synthesized_turn_state=synthesized_turn_state) request_state, text_data = proxy._prepare_response_bridge_request_state( responses_payload, api_key=refreshed_api_key, @@ -1815,6 +2039,8 @@ async def _prepare_websocket_response_create_request( request_state.client_ip = client_ip request_state.responses_lite_model = next_responses_lite_model request_state.expose_stale_previous_response_classifier = codex_session_affinity + request_state.require_security_work_authorized = capability_route.require_security_work_authorized + request_state.durable_capability_lineage_required = capability_route.require_security_work_authorized original_full_resend_input: JsonValue | None = None if session_anchor is not None: request_state.proxy_injected_previous_response_id = True @@ -2015,6 +2241,13 @@ async def _connect_proxy_websocket( ) -> tuple[Account | None, UpstreamWebSocket | None]: proxy = cast(_WebSocketServiceProtocol, self) _ = proxy + + async def _record_or_defer_confirmed_route_backoff(account: Account) -> None: + if request_state.api_key_reservation is not None: + request_state.deferred_account_error_backoffs.setdefault(account.id, account) + return + await proxy._load_balancer.record_error_backoff(account) + if ( request_state.useragent is None and request_state.useragent_group is None @@ -2175,6 +2408,7 @@ async def _connect_proxy_websocket( last_failover_account = account continue except ProxyResponseError as exc: + confirmed_pre_dispatch = is_confirmed_pre_dispatch_transport_error(exc) if selected_account_model_replacement: # The account/model retry budget selected this replacement; # its connection failure must be surfaced rather than @@ -2188,9 +2422,16 @@ async def _connect_proxy_websocket( attempt=attempt + 1, max_attempts=max_attempts, deterministic_failover_enabled=getattr(base_settings, "deterministic_failover_enabled", True), + require_preferred_account=require_preferred_account, ) if action == "failover_next": + # Release the dead route's stream lease before recording + # the backoff so its concurrency slot never outlives the + # failed connection attempt. await proxy._load_balancer.release_account_lease(selected_stream_lease) + selected_stream_lease = None + if confirmed_pre_dispatch: + await _record_or_defer_confirmed_route_backoff(account) last_failover_exc = exc last_failover_account = account excluded_account_ids.add(account.id) @@ -2200,6 +2441,8 @@ async def _connect_proxy_websocket( error_message = error.message if error else None await proxy._load_balancer.release_account_lease(selected_stream_lease) selected_stream_lease = None + if confirmed_pre_dispatch: + await _record_or_defer_confirmed_route_backoff(account) await proxy._emit_websocket_connect_failure( websocket, client_send_lock=client_send_lock, @@ -2378,7 +2621,17 @@ async def _heartbeat(remaining_seconds: float) -> None: if account: request_state.websocket_stream_lease = selection.lease return account - if defer_no_account_error and not _facade()._is_local_account_cap_code(selection.error_code): + durable_capability_pool_missing = bool( + require_security_work_authorized + and request_state.durable_capability_lineage_required + and not require_preferred_account + and not _facade()._is_local_account_cap_code(selection.error_code) + ) + if ( + defer_no_account_error + and not durable_capability_pool_missing + and not _facade()._is_local_account_cap_code(selection.error_code) + ): _facade().logger.warning( "Websocket account selection deferred no-account error request_id=%s model=%s " "preferred_account_id=%s require_preferred=%s error_code=%s error=%s excluded_count=%s", @@ -2393,7 +2646,9 @@ async def _heartbeat(remaining_seconds: float) -> None: return None error_code = selection.error_code or "no_accounts" error_message = selection.error_message or "No active accounts available" - if require_security_work_authorized and error_code == _facade()._NO_SECURITY_WORK_AUTHORIZED_ACCOUNTS_CODE: + if durable_capability_pool_missing or ( + require_security_work_authorized and error_code == _facade()._NO_SECURITY_WORK_AUTHORIZED_ACCOUNTS_CODE + ): await proxy._emit_websocket_security_work_missing_pool( websocket, client_send_lock=client_send_lock, @@ -2485,14 +2740,35 @@ async def _emit_websocket_security_work_missing_pool( ) -> None: proxy = cast(_WebSocketServiceProtocol, self) _ = proxy + if request_state.durable_capability_lineage_required: + _clear_websocket_precreated_replay_fallback(request_state) + error_code = _facade()._NO_SECURITY_WORK_AUTHORIZED_ACCOUNTS_CODE + error_message = _CAPABILITY_REQUIRED_NO_AUTHORIZED_ACCOUNTS_MESSAGE + error_type = "server_error" + status_code = 503 + advisory_message = _CAPABILITY_REQUIRED_NO_AUTHORIZED_ACCOUNTS_MESSAGE + advisory_action = _CAPABILITY_REQUIRED_NO_AUTHORIZED_ACCOUNTS_ACTION + else: + error_code = request_state.error_code_override or _facade()._NO_SECURITY_WORK_AUTHORIZED_ACCOUNTS_CODE + error_message = ( + request_state.error_message_override or _facade()._SECURITY_WORK_NO_AUTHORIZED_ACCOUNTS_MESSAGE + ) + error_type = request_state.error_type_override or ( + "invalid_request_error" if request_state.error_code_override is not None else "server_error" + ) + status_code = request_state.error_http_status_override or ( + 400 if request_state.error_code_override is not None else 503 + ) + advisory_message = _facade()._SECURITY_WORK_NO_AUTHORIZED_ACCOUNTS_MESSAGE + advisory_action = "forward_original_security_work_error" async with client_send_lock: await websocket.send_text( json.dumps( _facade()._security_work_advisory_event( code=_facade()._NO_SECURITY_WORK_AUTHORIZED_ACCOUNTS_CODE, - message=_facade()._SECURITY_WORK_NO_AUTHORIZED_ACCOUNTS_MESSAGE, + message=advisory_message, request_id=request_state.request_log_id or request_state.request_id, - action="forward_original_security_work_error", + action=advisory_action, ), ensure_ascii=True, separators=(",", ":"), @@ -2504,14 +2780,14 @@ async def _emit_websocket_security_work_missing_pool( account_id=account_id, api_key=api_key, request_state=request_state, - status_code=400, + status_code=status_code, payload=openai_error( - request_state.error_code_override or _facade()._SECURITY_WORK_AUTHORIZATION_REQUIRED_CODE, - request_state.error_message_override or "Security work authorization is required", - error_type=request_state.error_type_override or "invalid_request_error", + error_code, + error_message, + error_type=error_type, ), - error_code=request_state.error_code_override or _facade()._SECURITY_WORK_AUTHORIZATION_REQUIRED_CODE, - error_message=request_state.error_message_override or "Security work authorization is required", + error_code=error_code, + error_message=error_message, ) async def _try_open_websocket_connect_attempt( @@ -2880,13 +3156,25 @@ async def _decide_websocket_failover_action( attempt: int, max_attempts: int, deterministic_failover_enabled: bool, + require_preferred_account: bool = False, ) -> str: proxy = cast(_WebSocketServiceProtocol, self) _ = proxy - classified = await proxy._handle_websocket_connect_error(account, exc) - failure_class = classified["failure_class"] if isinstance(classified, dict) else "non_retryable" + confirmed_pre_dispatch = is_confirmed_pre_dispatch_transport_error(exc) + if confirmed_pre_dispatch: + # A proven pre-dispatch proxy connect failure is account-local + # transient evidence. The caller applies the bounded transient + # backoff floor once the failed lease is released, so the generic + # single-error health write is skipped here. Hard account + # ownership fails closed on the original sanitized failure. + failure_class = "retryable_transient" + else: + classified = await proxy._handle_websocket_connect_error(account, exc) + failure_class = classified["failure_class"] if isinstance(classified, dict) else "non_retryable" candidates_remaining = max_attempts - attempt - if exc.status_code == 401 and candidates_remaining > 0: + if confirmed_pre_dispatch: + action = "surface" if require_preferred_account or candidates_remaining <= 0 else "failover_next" + elif exc.status_code == 401 and candidates_remaining > 0: action = "failover_next" elif deterministic_failover_enabled: action = failover_decision( @@ -3532,6 +3820,35 @@ async def _relay_upstream_websocket_messages( break except asyncio.CancelledError: raise + except _CapabilityLineagePropagationError as exc: + error = _parse_openai_error(exc.error.payload) + error_code = _normalize_error_code( + error.code if error else None, + error.type if error else None, + ) + error_message = error.message if error and error.message else "Required capability lineage is unavailable" + await proxy._fail_pending_websocket_requests( + account=account, + account_id_value=account_id_value, + pending_requests=pending_requests, + pending_lock=pending_lock, + error_code=error_code or "capability_lineage_unavailable", + error_message=error_message, + api_key=api_key, + websocket=websocket, + client_send_lock=client_send_lock, + response_create_gate=response_create_gate, + downstream_activity=downstream_activity, + penalize_account=False, + ) + upstream_control.reconnect_requested = True + try: + await upstream.close() + except Exception: + _facade().logger.debug( + "Failed to retire upstream websocket after capability lineage propagation failure", + exc_info=True, + ) except _WebSocketReplaySequenceRegression as exc: _facade().logger.warning( "Refusing websocket replay after non-advancing sequence account_id=%s detail=%s", @@ -3655,11 +3972,6 @@ async def _process_upstream_websocket_text( request_state = _assign_websocket_response_id(pending_requests, response_id) created_request_state = request_state release_create_gate = request_state is not None - if request_state is not None and continuity_state is not None: - _record_websocket_responses_lite_acceptance( - continuity_state, - request_state=request_state, - ) elif response_id is not None: request_state = _find_websocket_request_state_by_response_id(pending_requests, response_id) release_create_gate = False @@ -3813,6 +4125,38 @@ async def _process_upstream_websocket_text( else: request_state = None + if ( + event_type == "response.created" + and response_id is not None + and created_request_state is not None + and created_request_state.durable_capability_lineage_required + ): + capability_api_key = created_request_state.api_key + try: + if capability_api_key is None: + raise _capability_lineage_unavailable_error() + await proxy._capability_router.route( + RoutingIntent.requiring(RoutingCapability.TRUSTED_CYBER), + api_key_id=capability_api_key.id, + aliases=capability_lineage_aliases( + {}, + previous_response_ids=_websocket_continuity_response_ids( + created_request_state, + response_id, + ), + ), + ) + except ProxyResponseError as exc: + async with pending_lock: + created_request_state.response_id = None + raise _CapabilityLineagePropagationError(exc) from exc + + if event_type == "response.created" and created_request_state is not None and continuity_state is not None: + _record_websocket_responses_lite_acceptance( + continuity_state, + request_state=created_request_state, + ) + if event_type == "response.created" and release_create_gate and created_request_state is not None: await _release_websocket_response_create_gate(created_request_state, response_create_gate) @@ -4351,8 +4695,7 @@ async def _finalize_websocket_request_state( if request_state.draining_until_terminal: await _release_websocket_response_create_gate(request_state, response_create_gate) - await proxy._release_websocket_reservation(request_state.api_key_reservation) - request_state.api_key_reservation = None + await proxy._release_websocket_request_state_reservation(request_state) return if request_state.latency_first_token_ms is None: @@ -4432,25 +4775,41 @@ async def _finalize_websocket_request_state( settlement.account_health_error = False proxy._cancel_request_state_api_key_reservation_heartbeat(request_state) await _release_websocket_response_create_gate(request_state, response_create_gate) - await proxy._settle_stream_api_key_usage( + lifecycle = request_state.deferred_account_backoff_lifecycle + settlement_confirmed = await proxy._settle_stream_api_key_usage( api_key, request_state.api_key_reservation, settlement, response_id, # The reservation must be settled before the load-balancer # health write below (settlement-ordering invariant). - wait_for_settlement=settlement.account_health_error, + wait_for_settlement=( + lifecycle is not None + or settlement.account_health_error + or bool(request_state.deferred_account_error_backoffs) + ), ) - if settlement.account_health_error: - await proxy._handle_stream_error( - account, - _stream_settlement_error_payload(settlement), - settlement.error_code or "upstream_error", + if settlement_confirmed: + request_state.api_key_reservation = None + if lifecycle is not None: + lifecycle.settlement_confirmed = True + pending_backoffs = ( + lifecycle.pending_backoffs if lifecycle is not None else request_state.deferred_account_error_backoffs ) + if pending_backoffs: + await proxy._drain_deferred_account_error_backoffs(pending_backoffs) + if settlement.account_health_error: + if settlement_confirmed: + await proxy._handle_stream_error( + account, + _stream_settlement_error_payload(settlement), + settlement.error_code or "upstream_error", + ) upstream_control.reconnect_requested = True upstream_control.retire_after_drain = True elif settlement.record_success: - await proxy._load_balancer.record_success(account) + if settlement_confirmed: + await proxy._load_balancer.record_success(account) for remembered_response_id in _websocket_continuity_response_ids(request_state, response_id): proxy._remember_websocket_previous_response_owner( previous_response_id=remembered_response_id, @@ -4698,23 +5057,6 @@ async def _fail_pending_websocket_requests( penalty_message = request_state.error_message_override or error_message break - if ( - remaining - and penalize_account - and account is not None - and isinstance(account, Account) - and penalty_code is not None - ): - try: - await proxy._handle_stream_error(account, {"message": penalty_message or error_message}, penalty_code) - except Exception: - _facade().logger.warning( - "Failed to record websocket pending-request health penalty account_id=%s error_code=%s", - account_id_value, - penalty_code, - exc_info=True, - ) - last_index = len(remaining) - 1 for index, request_state in enumerate(remaining): proxy._cancel_request_state_api_key_reservation_heartbeat(request_state) @@ -4788,7 +5130,10 @@ async def _fail_pending_websocket_requests( ) await proxy._write_request_log( account_id=account_id_value, - api_key=api_key, + # HTTP-bridge callers fan a shared session failure out to + # requests from multiple API keys, so they pass api_key=None; + # each request_state carries its own authenticated key. + api_key=request_state.api_key or api_key, request_id=request_state.response_id or request_state.request_log_id or request_state.request_id, archive_request_id=request_state.archive_request_id, model=request_state.model or "", @@ -4833,6 +5178,23 @@ async def _fail_pending_websocket_requests( status=status, ) + if ( + remaining + and penalize_account + and account is not None + and isinstance(account, Account) + and penalty_code is not None + ): + try: + await proxy._handle_stream_error(account, {"message": penalty_message or error_message}, penalty_code) + except Exception: + _facade().logger.warning( + "Failed to record websocket pending-request health penalty account_id=%s error_code=%s", + account_id_value, + penalty_code, + exc_info=True, + ) + async def _emit_websocket_terminal_error( self, websocket: WebSocket, diff --git a/app/modules/proxy/_service/websocket/protocol.py b/app/modules/proxy/_service/websocket/protocol.py index 322a377096..61a05edb43 100644 --- a/app/modules/proxy/_service/websocket/protocol.py +++ b/app/modules/proxy/_service/websocket/protocol.py @@ -4,12 +4,14 @@ class _WebSocketServiceProtocol(Protocol): + _capability_router: Any _acquire_account_response_create_lease_or_overload: Any _acquire_request_state_response_create_admission: Any _cancel_request_state_api_key_reservation_heartbeat: Any _connect_proxy_websocket: Any _decide_websocket_failover_action: Any _downstream_websocket_is_idle: Any + _drain_deferred_account_error_backoffs: Any _emit_pending_websocket_keepalive: Any _emit_websocket_connect_failure: Any _emit_websocket_connect_timeout: Any diff --git a/app/modules/proxy/api.py b/app/modules/proxy/api.py index fcbc012a8a..296979a67d 100644 --- a/app/modules/proxy/api.py +++ b/app/modules/proxy/api.py @@ -44,6 +44,7 @@ from app.core.cache.invalidation import NAMESPACE_RESET_CREDITS, bump_cache_invalidation_local from app.core.clients.files import FileProxyError from app.core.clients.proxy import ( + CODEX_LB_REQUIRED_CAPABILITY_HEADER, CodexControlRequestPrivacyPolicy, CodexControlResponse, ProxyResponseError, @@ -1106,6 +1107,7 @@ async def responses_websocket( api_key=api_key, client_ip=resolve_request_client_host(websocket), synthesized_turn_state=turn_state if client_turn_state is None else None, + capability_header_values=tuple(websocket.headers.getlist(CODEX_LB_REQUIRED_CAPABILITY_HEADER)), ) @@ -1410,6 +1412,7 @@ async def v1_responses_websocket( api_key=api_key, client_ip=resolve_request_client_host(websocket), synthesized_turn_state=turn_state if client_turn_state is None else None, + capability_header_values=tuple(websocket.headers.getlist(CODEX_LB_REQUIRED_CAPABILITY_HEADER)), ) diff --git a/app/modules/proxy/capability_lineage.py b/app/modules/proxy/capability_lineage.py new file mode 100644 index 0000000000..338e754917 --- /dev/null +++ b/app/modules/proxy/capability_lineage.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from collections.abc import Collection +from dataclasses import dataclass +from hashlib import sha256 + +_MARKER_DOMAIN = "capability-lineage/v1" + + +@dataclass(frozen=True, slots=True) +class CapabilityLineageAlias: + kind: str + value: str + + +def normalize_capability_lineage_aliases( + aliases: Collection[CapabilityLineageAlias], +) -> tuple[CapabilityLineageAlias, ...]: + normalized: dict[tuple[str, str], CapabilityLineageAlias] = {} + for alias in aliases: + kind = alias.kind.strip() + value = alias.value.strip() + if kind and value: + normalized[(kind, value)] = CapabilityLineageAlias(kind=kind, value=value) + return tuple(normalized.values()) + + +def capability_lineage_marker_hash( + *, + capability: str, + api_key_scope: str, + alias: CapabilityLineageAlias, +) -> str: + normalized_capability = capability.strip() + normalized_scope = api_key_scope.strip() + normalized_aliases = normalize_capability_lineage_aliases((alias,)) + if not normalized_capability or not normalized_scope or not normalized_aliases: + raise ValueError("capability, API-key scope, and lineage alias must be non-empty") + normalized_alias = normalized_aliases[0] + payload = "\0".join( + ( + _MARKER_DOMAIN, + normalized_capability, + normalized_scope, + normalized_alias.kind, + normalized_alias.value, + ) + ) + return sha256(payload.encode("utf-8")).hexdigest() + + +def capability_lineage_marker_hashes( + *, + capability: str, + api_key_scope: str, + aliases: Collection[CapabilityLineageAlias], +) -> tuple[str, ...]: + return tuple( + capability_lineage_marker_hash( + capability=capability, + api_key_scope=api_key_scope, + alias=alias, + ) + for alias in normalize_capability_lineage_aliases(aliases) + ) diff --git a/app/modules/proxy/capability_lineage_repository.py b/app/modules/proxy/capability_lineage_repository.py new file mode 100644 index 0000000000..f09ba0f078 --- /dev/null +++ b/app/modules/proxy/capability_lineage_repository.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from collections.abc import Collection + +from sqlalchemy import func, select +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.dialects.sqlite import insert as sqlite_insert +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.sql import Insert + +from app.db.models import CapabilityLineageMarker +from app.db.session import sqlite_writer_section +from app.modules.proxy.capability_lineage import ( + CapabilityLineageAlias, + capability_lineage_marker_hashes, +) + + +class CapabilityLineageRepository: + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def is_required( + self, + *, + capability: str, + api_key_scope: str, + aliases: Collection[CapabilityLineageAlias], + ) -> bool: + marker_hashes = capability_lineage_marker_hashes( + capability=capability, + api_key_scope=api_key_scope, + aliases=aliases, + ) + if not marker_hashes: + return False + statement = ( + select(CapabilityLineageMarker.marker_hash) + .where(CapabilityLineageMarker.marker_hash.in_(marker_hashes)) + .limit(1) + ) + required = (await self._session.execute(statement)).scalar_one_or_none() is not None + # End SQLite's read snapshot before a caller upgrades this dedicated + # repository session into a writer while another connection commits. + await self._session.commit() + return required + + async def require( + self, + *, + capability: str, + api_key_scope: str, + aliases: Collection[CapabilityLineageAlias], + ) -> tuple[str, ...]: + marker_hashes = capability_lineage_marker_hashes( + capability=capability, + api_key_scope=api_key_scope, + aliases=aliases, + ) + if not marker_hashes: + return () + async with sqlite_writer_section(): + # Overlapping alias sets must lock rows in one global order on + # PostgreSQL so concurrent reconnects cannot deadlock A->B/B->A. + for marker_hash in sorted(marker_hashes): + await self._session.execute(self._upsert_statement(marker_hash)) + await self._session.commit() + return marker_hashes + + def _upsert_statement(self, marker_hash: str) -> Insert: + dialect = self._session.get_bind().dialect.name + if dialect == "postgresql": + insert_fn = pg_insert + elif dialect == "sqlite": + insert_fn = sqlite_insert + else: + raise RuntimeError(f"Capability lineage persistence unsupported for dialect={dialect!r}") + statement = insert_fn(CapabilityLineageMarker).values(marker_hash=marker_hash) + return statement.on_conflict_do_update( + index_elements=[CapabilityLineageMarker.marker_hash], + set_={"last_seen_at": func.now()}, + ) diff --git a/app/modules/proxy/capability_routing.py b/app/modules/proxy/capability_routing.py new file mode 100644 index 0000000000..6e6c0238c6 --- /dev/null +++ b/app/modules/proxy/capability_routing.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import re +from collections.abc import Collection, Mapping +from dataclasses import dataclass +from enum import StrEnum + +from app.core.clients.proxy import CODEX_LB_REQUIRED_CAPABILITY_HEADER, ProxyResponseError +from app.core.errors import openai_error +from app.core.types import JsonValue +from app.core.utils.json_guards import is_json_mapping +from app.modules.api_keys.service import ApiKeyData +from app.modules.proxy.capability_lineage import ( + CapabilityLineageAlias, + normalize_capability_lineage_aliases, +) +from app.modules.proxy.repo_bundle import ProxyRepoFactory + +REQUIRED_CAPABILITY_HEADER = CODEX_LB_REQUIRED_CAPABILITY_HEADER +CODEX_PARENT_THREAD_ID_HEADER = "x-codex-parent-thread-id" +CODEX_WINDOW_ID_HEADER = "x-codex-window-id" +CAPABILITY_SIGNAL_UNTRUSTED_CODE = "capability_signal_untrusted" +UNSUPPORTED_REQUIRED_CAPABILITY_CODE = "unsupported_required_capability" +CAPABILITY_LINEAGE_UNAVAILABLE_CODE = "capability_lineage_unavailable" +CAPABILITY_ROUTING_UNAVAILABLE_CODE = "capability_routing_unavailable" +CAPABILITY_ROUTING_UNAVAILABLE_MESSAGE = "Required capability routing is unavailable; retry later." + +_TERMINAL_WINDOW_SLOT_RE = re.compile(r":\d+$") + + +class RoutingCapability(StrEnum): + TRUSTED_CYBER = "trusted_cyber" + + +@dataclass(frozen=True, slots=True) +class RoutingIntent: + required_capabilities: frozenset[RoutingCapability] = frozenset() + + @classmethod + def empty(cls) -> RoutingIntent: + return cls() + + @classmethod + def requiring(cls, capability: RoutingCapability) -> RoutingIntent: + return cls(required_capabilities=frozenset({capability})) + + @property + def requires_trusted_cyber(self) -> bool: + return RoutingCapability.TRUSTED_CYBER in self.required_capabilities + + +@dataclass(frozen=True, slots=True) +class CapabilityRoute: + require_security_work_authorized: bool + aliases: tuple[CapabilityLineageAlias, ...] + + +class CapabilityRouter: + def __init__(self, repo_factory: ProxyRepoFactory) -> None: + self._repo_factory = repo_factory + + async def route( + self, + intent: RoutingIntent, + *, + api_key_id: str | None, + aliases: Collection[CapabilityLineageAlias], + ) -> CapabilityRoute: + normalized_aliases = normalize_capability_lineage_aliases(aliases) + explicitly_required = intent.requires_trusted_cyber + if api_key_id is None or not normalized_aliases: + return CapabilityRoute( + require_security_work_authorized=explicitly_required, + aliases=normalized_aliases, + ) + + try: + async with self._repo_factory() as repositories: + repository = repositories.capability_lineage + if repository is None: + raise RuntimeError("capability lineage repository is unavailable") + inherited_requirement = False + if not explicitly_required: + inherited_requirement = await repository.is_required( + capability=RoutingCapability.TRUSTED_CYBER.value, + api_key_scope=api_key_id, + aliases=normalized_aliases, + ) + required = explicitly_required or inherited_requirement + if required: + marker_hashes = await repository.require( + capability=RoutingCapability.TRUSTED_CYBER.value, + api_key_scope=api_key_id, + aliases=normalized_aliases, + ) + if not marker_hashes: + raise RuntimeError("capability lineage marker was not persisted") + except Exception as exc: + raise _capability_lineage_unavailable_error() from exc + + return CapabilityRoute( + require_security_work_authorized=required, + aliases=normalized_aliases, + ) + + +def capability_lineage_aliases( + headers: Mapping[str, str], + *, + session_id: str | None = None, + turn_state: str | None = None, + previous_response_ids: Collection[str | None] = (), + client_metadata: JsonValue | None = None, +) -> tuple[CapabilityLineageAlias, ...]: + aliases: list[CapabilityLineageAlias] = [] + _append_alias(aliases, "session_header", session_id) + _append_alias(aliases, "turn_state", turn_state) + for previous_response_id in previous_response_ids: + _append_alias(aliases, "previous_response", previous_response_id) + + parent_thread_ids = ( + *_header_values(headers, CODEX_PARENT_THREAD_ID_HEADER), + *_metadata_string_values(client_metadata, CODEX_PARENT_THREAD_ID_HEADER), + ) + for parent_thread_id in parent_thread_ids: + _append_alias(aliases, "codex_task", parent_thread_id) + + window_ids = ( + *_header_values(headers, CODEX_WINDOW_ID_HEADER), + *_metadata_string_values(client_metadata, CODEX_WINDOW_ID_HEADER), + ) + for window_id in window_ids: + _append_alias(aliases, "codex_window", window_id) + stable_task_id = _TERMINAL_WINDOW_SLOT_RE.sub("", window_id.strip()) + _append_alias(aliases, "codex_task", stable_task_id) + return normalize_capability_lineage_aliases(aliases) + + +def parse_routing_intent( + headers: Mapping[str, str], + *, + api_key: ApiKeyData | None, + client_metadata: JsonValue | None = None, + header_values: Collection[str] | None = None, + client_metadata_values: Collection[JsonValue] | None = None, +) -> RoutingIntent: + values: tuple[JsonValue, ...] = ( + *(tuple(header_values) if header_values is not None else _header_values(headers, REQUIRED_CAPABILITY_HEADER)), + *( + tuple(client_metadata_values) + if client_metadata_values is not None + else _metadata_values(client_metadata, REQUIRED_CAPABILITY_HEADER) + ), + ) + if not values: + return RoutingIntent.empty() + if api_key is None: + raise ProxyResponseError( + 403, + openai_error( + CAPABILITY_SIGNAL_UNTRUSTED_CODE, + "Required capability signal requires an authenticated proxy API key.", + error_type="permission_error", + ), + ) + if len(values) != 1: + raise _unsupported_capability_error() + raw_capability = values[0] + if not isinstance(raw_capability, str): + raise _unsupported_capability_error() + try: + capability = RoutingCapability(raw_capability) + except ValueError as exc: + raise _unsupported_capability_error() from exc + return RoutingIntent.requiring(capability) + + +def reject_capability_signal_outside_response_create( + *, + api_key: ApiKeyData | None, + client_metadata: JsonValue | None, + client_metadata_values: Collection[JsonValue] | None = None, +) -> None: + intent = parse_routing_intent( + {}, + api_key=api_key, + client_metadata=client_metadata, + client_metadata_values=client_metadata_values, + ) + if intent.required_capabilities: + raise _unsupported_capability_error() + + +def strip_capability_metadata(client_metadata: JsonValue | None) -> JsonValue | None: + if not is_json_mapping(client_metadata): + return client_metadata + normalized_name = REQUIRED_CAPABILITY_HEADER.lower() + return { + key: value + for key, value in client_metadata.items() + if not isinstance(key, str) or key.lower() != normalized_name + } + + +def _header_values(headers: Mapping[str, str], name: str) -> tuple[str, ...]: + normalized_name = name.lower() + return tuple(value for header_name, value in headers.items() if header_name.lower() == normalized_name) + + +def _metadata_values(client_metadata: JsonValue | None, name: str) -> tuple[JsonValue, ...]: + if not is_json_mapping(client_metadata): + return () + normalized_name = name.lower() + return tuple( + value for key, value in client_metadata.items() if isinstance(key, str) and key.lower() == normalized_name + ) + + +def _metadata_string_values(client_metadata: JsonValue | None, name: str) -> tuple[str, ...]: + return tuple(value for value in _metadata_values(client_metadata, name) if isinstance(value, str)) + + +def _append_alias(aliases: list[CapabilityLineageAlias], kind: str, value: str | None) -> None: + if value is not None: + aliases.append(CapabilityLineageAlias(kind=kind, value=value)) + + +def _unsupported_capability_error() -> ProxyResponseError: + return ProxyResponseError( + 400, + openai_error( + UNSUPPORTED_REQUIRED_CAPABILITY_CODE, + "Required routing capability is unsupported.", + error_type="invalid_request_error", + ), + ) + + +def _capability_lineage_unavailable_error() -> ProxyResponseError: + return ProxyResponseError( + 503, + openai_error( + CAPABILITY_LINEAGE_UNAVAILABLE_CODE, + "Required capability lineage is unavailable; retry later.", + error_type="server_error", + ), + ) diff --git a/app/modules/proxy/load_balancer.py b/app/modules/proxy/load_balancer.py index 27244b522a..57b73c6261 100644 --- a/app/modules/proxy/load_balancer.py +++ b/app/modules/proxy/load_balancer.py @@ -13,6 +13,7 @@ from app.core import usage as usage_core from app.core.balancer import ( + ERROR_BACKOFF_THRESHOLD, HEALTH_TIER_DRAINING, HEALTH_TIER_HEALTHY, HEALTH_TIER_PROBING, @@ -1510,7 +1511,17 @@ async def mark_permanent_failure(self, account: Account, error_code: str) -> boo async def record_error(self, account: Account) -> None: await self.record_errors(account, 1) - async def record_errors(self, account: Account, count: int) -> None: + async def record_error_backoff(self, account: Account) -> None: + """Record one error and immediately enter the bounded transient backoff.""" + await self.record_errors(account, 1, minimum_error_count=ERROR_BACKOFF_THRESHOLD) + + async def record_errors( + self, + account: Account, + count: int, + *, + minimum_error_count: int = 0, + ) -> None: """Record *count* transient errors in a single lock acquisition.""" if count < 1: return @@ -1518,7 +1529,7 @@ async def record_errors(self, account: Account, count: int) -> None: async with lock: account_snapshot = _clone_account(account) state = self._state_for(account) - state.error_count += count + state.error_count = max(state.error_count + count, minimum_error_count) state.last_error_at = time.time() self._sync_runtime_state(account, state) runtime = self._runtime.get(account.id) diff --git a/app/modules/proxy/repo_bundle.py b/app/modules/proxy/repo_bundle.py index 5f95d281bc..596a471f29 100644 --- a/app/modules/proxy/repo_bundle.py +++ b/app/modules/proxy/repo_bundle.py @@ -6,6 +6,7 @@ from app.modules.accounts.repository import AccountsRepository from app.modules.api_keys.repository import ApiKeysRepository +from app.modules.proxy.capability_lineage_repository import CapabilityLineageRepository from app.modules.proxy.sticky_repository import StickySessionsRepository from app.modules.quota_planner.repository import QuotaPlannerRepository from app.modules.request_logs.repository import RequestLogsRepository @@ -21,6 +22,7 @@ class ProxyRepositories: api_keys: ApiKeysRepository additional_usage: AdditionalUsageRepository quota_planner: QuotaPlannerRepository | None = None + capability_lineage: CapabilityLineageRepository | None = None ProxyRepoFactory = Callable[[], AsyncContextManager[ProxyRepositories]] diff --git a/app/modules/proxy/service.py b/app/modules/proxy/service.py index 3b6e75b4f1..b077a794fd 100644 --- a/app/modules/proxy/service.py +++ b/app/modules/proxy/service.py @@ -703,6 +703,7 @@ from app.modules.proxy.affinity import ( _sticky_key_for_responses_request as _sticky_key_for_responses_request, ) +from app.modules.proxy.capability_routing import CapabilityRouter from app.modules.proxy.durable_bridge_coordinator import ( DurableBridgeLookup as DurableBridgeLookup, ) @@ -934,6 +935,7 @@ def __init__( self._repo_factory = repo_factory self._encryptor = TokenEncryptor() self._load_balancer = LoadBalancer(repo_factory) + self._capability_router = CapabilityRouter(repo_factory) self._live_websocket_connector = live_websocket_connector self._ring_membership = RingMembershipService(SessionLocal) self._durable_bridge = DurableBridgeSessionCoordinator(SessionLocal) @@ -1407,10 +1409,16 @@ async def _select_account_with_budget_compatible( require_unambiguous_account=affinity_policy.require_unambiguous_account, sticky_max_age_seconds=affinity_policy.max_age_seconds, ) + required_capability_kwargs = {} + if kwargs.get("require_security_work_authorized") is True: + required_capability_kwargs["require_security_work_authorized"] = kwargs.pop( + "require_security_work_authorized" + ) return await _call_with_supported_optional_kwargs( self._select_account_with_budget, deadline, optional_kwargs=kwargs, + **required_capability_kwargs, ) @asynccontextmanager diff --git a/app/modules/request_logs/api.py b/app/modules/request_logs/api.py index 4e43599aab..fab29112c4 100644 --- a/app/modules/request_logs/api.py +++ b/app/modules/request_logs/api.py @@ -1,17 +1,22 @@ from __future__ import annotations -from datetime import datetime +from datetime import datetime, timedelta -from fastapi import APIRouter, Depends, Query +from fastapi import APIRouter, Depends, HTTPException, Query from app.core.auth.dashboard_access import DashboardPrincipal, DashboardRole from app.core.auth.dependencies import ( ensure_dashboard_admin_access, + require_dashboard_admin_access, set_dashboard_error_format, validate_dashboard_session, ) +from app.core.utils.time import to_utc_naive, utcnow from app.dependencies import RequestLogsContext, get_request_logs_context +from app.modules.dashboard.timeframes import resolve_conversation_timeframe from app.modules.request_logs.schemas import ( + ConversationDetailsResponse, + ConversationsResponse, RequestLogApiKeyOption, RequestLogFilterOptionsResponse, RequestLogModelOption, @@ -25,7 +30,14 @@ dependencies=[Depends(validate_dashboard_session), Depends(set_dashboard_error_format)], ) +conversations_router = APIRouter( + prefix="/api/conversations", + tags=["dashboard"], + dependencies=[Depends(require_dashboard_admin_access), Depends(set_dashboard_error_format)], +) + _MODEL_OPTION_DELIMITER = ":::" +_CONVERSATION_MAX_LOOKBACK = timedelta(days=30) def _parse_model_option(value: str) -> ServiceRequestLogModelOption | None: @@ -127,3 +139,57 @@ async def list_request_log_filter_options( ], statuses=options.statuses, ) + + +@conversations_router.get("/", response_model=ConversationsResponse, include_in_schema=False) +@conversations_router.get("", response_model=ConversationsResponse) +async def list_conversations( + limit: int = Query(50, ge=1, le=1000), + offset: int = Query(0, ge=0), + search: str | None = Query(default=None), + since: datetime | None = Query(default=None), + timeframe: str | None = Query(default=None, pattern="^(1d|7d|30d)$"), + context: RequestLogsContext = Depends(get_request_logs_context), +) -> ConversationsResponse: + if timeframe is not None and since is not None: + raise HTTPException(status_code=422, detail="timeframe and since cannot be supplied together") + + if timeframe is not None: + _, effective_since = resolve_conversation_timeframe(timeframe) + else: + cutoff = utcnow() - _CONVERSATION_MAX_LOOKBACK + effective_since = to_utc_naive(since) if since is not None else cutoff + if effective_since < cutoff: + effective_since = cutoff + page = await context.service.list_conversations( + limit=limit, + offset=offset, + search=search, + since=effective_since, + cache_mode="timeframe" if timeframe else "since", + timeframe=timeframe, + ) + return ConversationsResponse( + conversations=page.conversations, + total=page.total, + has_more=page.has_more, + ) + + +@conversations_router.get("/{conversation_id:path}", response_model=ConversationDetailsResponse) +async def get_conversation_details( + conversation_id: str, + context: RequestLogsContext = Depends(get_request_logs_context), +) -> ConversationDetailsResponse: + details = await context.service.get_conversation_details(conversation_id) + if details is None: + raise HTTPException(status_code=404, detail="Conversation not found") + return ConversationDetailsResponse( + conversation_id=details.conversation_id, + start=details.start, + latest=details.latest, + account_count=details.account_count, + total_elapsed_time=details.total_elapsed_time, + dominant_useragent_group=details.dominant_useragent_group, + model_stats=details.model_stats, + ) diff --git a/app/modules/request_logs/repository.py b/app/modules/request_logs/repository.py index e8dd8fbb24..50732755bc 100644 --- a/app/modules/request_logs/repository.py +++ b/app/modules/request_logs/repository.py @@ -64,6 +64,17 @@ class RequestLogsResult: _recent_count_cache: dict[tuple, tuple[int, float]] = {} +def _escape_like(value: str) -> str: + return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + + +def _normalize_conversation_id(value: str | None) -> str | None: + if value is None: + return None + normalized = value.strip(_CONVERSATION_WHITESPACE) + return normalized or None + + def _clear_recent_count_cache() -> None: _recent_count_cache.clear() @@ -93,6 +104,57 @@ class PreviousResponseOwnerRecord: session_id: str | None +@dataclass(frozen=True, slots=True) +class ConversationListSummary: + conversation_id: str + first_requested_at: datetime + last_requested_at: datetime + request_count: int + account_count: int + total_tokens: int + cached_input_tokens: int | None + cost_usd: float + + +@dataclass(frozen=True, slots=True) +class ConversationFacet: + conversation_id: str + value: str + request_count: int + + +@dataclass(frozen=True, slots=True) +class ConversationListResult: + summaries: list[ConversationListSummary] + account_facets: list[ConversationFacet] + api_key_facets: list[ConversationFacet] + model_facets: list[ConversationFacet] + total: int + + +@dataclass(frozen=True, slots=True) +class ConversationModelStatRow: + model: str + reasoning_effort: str | None + request_count: int + total_elapsed_ms: int + input_tokens: int + cached_input_tokens: int | None + output_tokens: int + cost_usd: float + + +@dataclass(frozen=True, slots=True) +class ConversationDetailsResult: + conversation_id: str + started_at: datetime + last_requested_at: datetime + account_count: int + total_elapsed_ms: int + useragent_group: str | None + model_stats: list[ConversationModelStatRow] + + class RequestLogsRepository: def __init__(self, session: AsyncSession) -> None: self._session = session @@ -109,6 +171,272 @@ def _conversation_id_expr() -> ColumnElement: ) return func.nullif(trimmed, "") + def _conversation_output_expr(self) -> ColumnElement: + return func.coalesce(RequestLog.output_tokens, RequestLog.reasoning_tokens, 0) + + def _conversation_cached_expr(self) -> ColumnElement: + dialect = self._session.get_bind().dialect.name + least = func.least if dialect == "postgresql" else func.min + greatest = func.greatest if dialect == "postgresql" else func.max + return case( + (RequestLog.cached_input_tokens.is_(None), None), + (RequestLog.input_tokens.is_(None), greatest(0, RequestLog.cached_input_tokens)), + else_=greatest(0, least(RequestLog.cached_input_tokens, RequestLog.input_tokens)), + ) + + def _eligible_conversation_row_conditions(self) -> list[ColumnElement[bool]]: + return [ + RequestLog.deleted_at.is_(None), + self._exclude_warmup_clause(), + ] + + def _conversation_conditions(self) -> list[ColumnElement[bool]]: + return [ + *self._eligible_conversation_row_conditions(), + RequestLog.conversation_id.is_not(None), + RequestLog.conversation_id != "", + ] + + def _reasoning_effort_sort_key(self) -> list[ColumnElement]: + rank = case( + (RequestLog.reasoning_effort.is_(None), 0), + (RequestLog.reasoning_effort == "", 1), + else_=2, + ) + return [rank, func.coalesce(RequestLog.reasoning_effort, "")] + + async def list_conversations( + self, + *, + limit: int = 50, + offset: int = 0, + search: str | None = None, + since: datetime | None = None, + cache_mode: str = "since", + timeframe: str | None = None, + ) -> ConversationListResult: + conversation_id = RequestLog.conversation_id + base_conditions = self._conversation_conditions() + search_value = search.strip() if search and search.strip() else None + candidate_conditions = [*base_conditions] + if since is not None: + candidate_conditions.append(RequestLog.requested_at >= since) + if search_value is not None: + pattern = f"%{_escape_like(search_value)}%" + candidate_conditions.append( + or_( + conversation_id.ilike(pattern, escape="\\"), + RequestLog.useragent_group.ilike(pattern, escape="\\"), + ) + ) + if since is not None or search_value is not None: + candidate_ids = ( + select(conversation_id.label("conversation_id")).where(*candidate_conditions).distinct().subquery() + ) + conditions = [*base_conditions, conversation_id.in_(select(candidate_ids.c.conversation_id))] + else: + conditions = base_conditions + + output = self._conversation_output_expr() + cached = self._conversation_cached_expr() + summary_conditions = [*conditions] + summary_stmt = ( + select( + conversation_id.label("conversation_id"), + func.min(RequestLog.requested_at).label("first_requested_at"), + func.max(RequestLog.requested_at).label("last_requested_at"), + func.count().label("request_count"), + func.count(func.distinct(RequestLog.account_id)).label("account_count"), + func.coalesce(func.sum(func.coalesce(RequestLog.input_tokens, 0) + output), 0).label("total_tokens"), + func.sum(cached).label("cached_input_tokens"), + func.coalesce(func.sum(RequestLog.cost_usd), 0.0).label("cost_usd"), + ) + .where(*summary_conditions) + .group_by(conversation_id) + ) + summary_subquery = summary_stmt.subquery() + filtered_summary_stmt = select(summary_subquery) + ttl_seconds = _COUNT_CACHE_TTL_SECONDS + if ttl_seconds <= 0: + total = int( + ( + await self._session.execute(select(func.count()).select_from(filtered_summary_stmt.subquery())) + ).scalar_one() + ) + else: + normalized_search = (search.strip() or None) if search else None + if cache_mode == "timeframe": + mode_token = ("timeframe", timeframe) + else: + mode_token = ("since", since) + cache_key = ("conversation-count", normalized_search, mode_token) + total = _cached_recent_count(cache_key) + if total is None: + total = int( + ( + await self._session.execute(select(func.count()).select_from(filtered_summary_stmt.subquery())) + ).scalar_one() + ) + _store_recent_count(cache_key, total, ttl_seconds) + page_rows = ( + await self._session.execute( + filtered_summary_stmt.order_by( + summary_subquery.c.last_requested_at.desc(), summary_subquery.c.conversation_id.asc() + ) + .offset(offset) + .limit(limit) + ) + ).all() + summaries = [ + ConversationListSummary( + conversation_id=row.conversation_id, + first_requested_at=row.first_requested_at, + last_requested_at=row.last_requested_at, + request_count=int(row.request_count), + account_count=int(row.account_count), + total_tokens=int(row.total_tokens), + cached_input_tokens=(int(row.cached_input_tokens) if row.cached_input_tokens is not None else None), + cost_usd=float(row.cost_usd or 0.0), + ) + for row in page_rows + ] + + page_ids = [summary.conversation_id for summary in summaries] + account_facets: list[ConversationFacet] = [] + api_key_facets: list[ConversationFacet] = [] + model_facets: list[ConversationFacet] = [] + if page_ids: + # Candidate-ID membership already selected conversations with + # activity in the window. Facets must describe the same full + # eligible conversation rows as that summary, including history + # before `since`. + facet_conditions = [*conditions] + account_facets = await self._conversation_facets(facet_conditions, page_ids, RequestLog.account_id) + api_key_facets = await self._conversation_facets(facet_conditions, page_ids, RequestLog.api_key_id) + model_facets = await self._conversation_facets(facet_conditions, page_ids, RequestLog.model) + return ConversationListResult( + summaries=summaries, + account_facets=account_facets, + api_key_facets=api_key_facets, + model_facets=model_facets, + total=total, + ) + + async def _conversation_facets( + self, + conditions: list[ColumnElement[bool]], + page_ids: list[str], + value_column: InstrumentedAttribute[str | None] | InstrumentedAttribute[str], + ) -> list[ConversationFacet]: + conversation_id = RequestLog.conversation_id + facet_conditions = [*conditions, conversation_id.in_(page_ids), value_column.is_not(None)] + if getattr(value_column, "key", None) == RequestLog.api_key_id.key: + facet_conditions.append(value_column.in_(select(ApiKey.id))) + stmt = ( + select( + conversation_id.label("conversation_id"), + value_column.label("value"), + func.count().label("request_count"), + ) + .where(*facet_conditions) + .group_by(conversation_id, value_column) + .order_by( + conversation_id.asc(), + func.count().desc(), + func.max(RequestLog.requested_at).desc(), + value_column.asc(), + ) + ) + rows = (await self._session.execute(stmt)).all() + return [ + ConversationFacet( + conversation_id=row.conversation_id, + value=row.value, + request_count=int(row.request_count), + ) + for row in rows + ] + + async def get_conversation_details(self, conversation_id: str) -> ConversationDetailsResult | None: + target = _normalize_conversation_id(conversation_id) + if not target: + return None + normalized_id = RequestLog.conversation_id + conditions = [*self._conversation_conditions(), normalized_id == target] + # Details intentionally remain window-agnostic; only list membership uses since. + summary = ( + await self._session.execute( + select( + func.min(RequestLog.requested_at).label("started_at"), + func.max(RequestLog.requested_at).label("last_requested_at"), + func.count(func.distinct(RequestLog.account_id)).label("account_count"), + func.coalesce(func.sum(func.coalesce(RequestLog.latency_ms, 0)), 0).label("total_elapsed_ms"), + ).where(*conditions) + ) + ).one() + if summary.started_at is None: + return None + + dominant = ( + await self._session.execute( + select(RequestLog.useragent_group) + .where(*conditions, RequestLog.useragent_group.is_not(None)) + .group_by(RequestLog.useragent_group) + .order_by( + func.count().desc(), + func.max(RequestLog.requested_at).desc(), + RequestLog.useragent_group.asc(), + ) + .limit(1) + ) + ).scalar_one_or_none() + + output = self._conversation_output_expr() + cached = self._conversation_cached_expr() + model_rows = ( + await self._session.execute( + select( + RequestLog.model.label("model"), + RequestLog.reasoning_effort.label("reasoning_effort"), + func.count().label("request_count"), + func.coalesce(func.sum(func.coalesce(RequestLog.latency_ms, 0)), 0).label("total_elapsed_ms"), + func.coalesce(func.sum(func.coalesce(RequestLog.input_tokens, 0)), 0).label("input_tokens"), + func.sum(cached).label("cached_input_tokens"), + func.coalesce(func.sum(output), 0).label("output_tokens"), + func.coalesce(func.sum(RequestLog.cost_usd), 0.0).label("cost_usd"), + ) + .where(*conditions) + .group_by(RequestLog.model, RequestLog.reasoning_effort) + .order_by( + func.count().desc(), + func.max(RequestLog.requested_at).desc(), + RequestLog.model.asc(), + *self._reasoning_effort_sort_key(), + ) + ) + ).all() + return ConversationDetailsResult( + conversation_id=target, + started_at=summary.started_at, + last_requested_at=summary.last_requested_at, + account_count=int(summary.account_count), + total_elapsed_ms=int(summary.total_elapsed_ms), + useragent_group=dominant, + model_stats=[ + ConversationModelStatRow( + model=row.model, + reasoning_effort=row.reasoning_effort, + request_count=int(row.request_count), + total_elapsed_ms=int(row.total_elapsed_ms), + input_tokens=int(row.input_tokens), + cached_input_tokens=(int(row.cached_input_tokens) if row.cached_input_tokens is not None else None), + output_tokens=int(row.output_tokens), + cost_usd=float(row.cost_usd or 0.0), + ) + for row in model_rows + ], + ) + def _bucket_epoch_expr(self, bucket_seconds: int) -> ColumnElement: bind = self._session.get_bind() dialect = bind.dialect.name if bind else "sqlite" @@ -300,7 +628,7 @@ async def aggregate_conversations_by_bucket( ) .where( RequestLog.requested_at >= since, - self._exclude_warmup_clause(), + *self._eligible_conversation_row_conditions(), conversation_id.is_not(None), ) .group_by(bucket_col) @@ -371,7 +699,7 @@ async def _aggregate_activity(self, since: datetime, until: datetime | None) -> func.count(self._conversation_id_expr()).label("conversation_request_count"), ).where( RequestLog.requested_at >= since, - self._exclude_warmup_clause(), + *self._eligible_conversation_row_conditions(), ) if until is not None: conversation_stmt = conversation_stmt.where(RequestLog.requested_at < until) @@ -585,7 +913,7 @@ async def add_log( resolved_useragent_group = ( useragent_group if not isinstance(useragent_group, str) or useragent_group.strip() else None ) - resolved_conversation_id = (conversation_id or "").strip() or None + resolved_conversation_id = _normalize_conversation_id(conversation_id) resolved_client_ip = client_ip if not isinstance(client_ip, str) or client_ip.strip() else None log = RequestLog( account_id=account_id, diff --git a/app/modules/request_logs/schemas.py b/app/modules/request_logs/schemas.py index da88c94146..7073b436c4 100644 --- a/app/modules/request_logs/schemas.py +++ b/app/modules/request_logs/schemas.py @@ -87,3 +87,50 @@ class RequestLogFilterOptionsResponse(DashboardModel): model_options: list[RequestLogModelOption] = Field(default_factory=list) api_keys: list[RequestLogApiKeyOption] = Field(default_factory=list) statuses: list[str] = Field(default_factory=list) + + +class ConversationModelEffort(DashboardModel): + model: str + reasoning_effort: str | None = None + + +class ConversationModelStat(DashboardModel): + model_effort: ConversationModelEffort + reqs: int + total_elapsed_time: int + total_input_tokens: int + cached_input_tokens: int | None + total_output_tokens: int + total_cost_usd: float + + +class ConversationEntry(DashboardModel): + conversation_id: str + first_request: datetime + last_request: datetime + request_count: int + representative_account: str | None = None + remaining_account_count: int + api_key_id: str | None = None + api_key_name: str | None = None + representative_model: str | None = None + remaining_model_count: int + total_tokens: int + cached_input_tokens: int | None + total_cost_usd: float + + +class ConversationsResponse(DashboardModel): + conversations: list[ConversationEntry] = Field(default_factory=list) + total: int + has_more: bool + + +class ConversationDetailsResponse(DashboardModel): + conversation_id: str + start: datetime + latest: datetime + account_count: int + total_elapsed_time: int + dominant_useragent_group: str | None = None + model_stats: list[ConversationModelStat] = Field(default_factory=list) diff --git a/app/modules/request_logs/service.py b/app/modules/request_logs/service.py index d4bd19732c..39f33b0ca8 100644 --- a/app/modules/request_logs/service.py +++ b/app/modules/request_logs/service.py @@ -9,8 +9,19 @@ normalize_log_status, to_request_log_entry, ) -from app.modules.request_logs.repository import RequestLogsRepository -from app.modules.request_logs.schemas import RequestLogConversation, RequestLogEntry +from app.modules.request_logs.repository import ( + ConversationDetailsResult, + ConversationFacet, + ConversationListResult, + RequestLogsRepository, +) +from app.modules.request_logs.schemas import ( + ConversationEntry, + ConversationModelEffort, + ConversationModelStat, + RequestLogConversation, + RequestLogEntry, +) @dataclass(frozen=True, slots=True) @@ -50,6 +61,24 @@ class RequestLogsPage: conversation: RequestLogConversation | None +@dataclass(frozen=True, slots=True) +class ConversationsPage: + conversations: list[ConversationEntry] + total: int + has_more: bool + + +@dataclass(frozen=True, slots=True) +class ConversationDetails: + conversation_id: str + start: datetime + latest: datetime + account_count: int + total_elapsed_time: int + dominant_useragent_group: str | None + model_stats: list[ConversationModelStat] + + class RequestLogsService: def __init__(self, repo: RequestLogsRepository) -> None: self._repo = repo @@ -166,6 +195,36 @@ async def list_filter_options( statuses=_normalize_status_values(status_values), ) + async def list_conversations( + self, + *, + limit: int = 50, + offset: int = 0, + search: str | None = None, + since: datetime | None = None, + cache_mode: str = "since", + timeframe: str | None = None, + ) -> ConversationsPage: + result = await self._repo.list_conversations( + limit=limit, + offset=offset, + search=search, + since=since, + cache_mode=cache_mode, + timeframe=timeframe, + ) + api_key_ids = [facet.value for facet in result.api_key_facets] + api_key_names = await self._repo.get_api_key_names_by_ids(api_key_ids) + return ConversationsPage( + conversations=_to_conversations(result, api_key_names), + total=result.total, + has_more=offset + limit < result.total, + ) + + async def get_conversation_details(self, conversation_id: str) -> ConversationDetails | None: + result = await self._repo.get_conversation_details(conversation_id) + return _to_conversation_details(result) if result is not None else None + def _map_status_filter(status: list[str] | None) -> RequestLogStatusFilter: if not status: @@ -207,3 +266,68 @@ def _normalize_status_values(values: list[tuple[str, str | None]]) -> list[str]: normalized = {normalize_log_status(status, error_code) for status, error_code in values} ordered = ["ok", "rate_limit", "quota", "error"] return [status for status in ordered if status in normalized] + + +def _first_facets(facets: list[ConversationFacet]) -> dict[str, ConversationFacet]: + first: dict[str, ConversationFacet] = {} + for facet in facets: + first.setdefault(facet.conversation_id, facet) + return first + + +def _to_conversations(result: ConversationListResult, api_key_names: dict[str, str]) -> list[ConversationEntry]: + accounts = _first_facets(result.account_facets) + api_keys = _first_facets(result.api_key_facets) + models = _first_facets(result.model_facets) + account_counts = {summary.conversation_id: summary.account_count for summary in result.summaries} + model_counts: dict[str, int] = {} + for facet in result.model_facets: + model_counts[facet.conversation_id] = model_counts.get(facet.conversation_id, 0) + 1 + + entries: list[ConversationEntry] = [] + for summary in result.summaries: + account = accounts.get(summary.conversation_id) + api_key = api_keys.get(summary.conversation_id) + model = models.get(summary.conversation_id) + api_key_name = api_key_names.get(api_key.value) if api_key is not None else None + entries.append( + ConversationEntry( + conversation_id=summary.conversation_id, + first_request=summary.first_requested_at, + last_request=summary.last_requested_at, + request_count=summary.request_count, + representative_account=account.value if account else None, + remaining_account_count=max(0, account_counts[summary.conversation_id] - 1), + api_key_id=api_key.value if api_key is not None and api_key_name is not None else None, + api_key_name=api_key_name, + representative_model=model.value if model else None, + remaining_model_count=max(0, model_counts.get(summary.conversation_id, 0) - 1), + total_tokens=summary.total_tokens, + cached_input_tokens=summary.cached_input_tokens, + total_cost_usd=summary.cost_usd, + ) + ) + return entries + + +def _to_conversation_details(result: ConversationDetailsResult) -> ConversationDetails: + return ConversationDetails( + conversation_id=result.conversation_id, + start=result.started_at, + latest=result.last_requested_at, + account_count=result.account_count, + total_elapsed_time=result.total_elapsed_ms, + dominant_useragent_group=result.useragent_group, + model_stats=[ + ConversationModelStat( + model_effort=ConversationModelEffort(model=row.model, reasoning_effort=row.reasoning_effort), + reqs=row.request_count, + total_elapsed_time=row.total_elapsed_ms, + total_input_tokens=row.input_tokens, + cached_input_tokens=row.cached_input_tokens, + total_output_tokens=row.output_tokens, + total_cost_usd=row.cost_usd, + ) + for row in result.model_stats + ], + ) diff --git a/deploy/helm/codex-lb/templates/grafana-dashboard.yaml b/deploy/helm/codex-lb/templates/grafana-dashboard.yaml index f82ac67e2e..9bb9ba9a54 100644 --- a/deploy/helm/codex-lb/templates/grafana-dashboard.yaml +++ b/deploy/helm/codex-lb/templates/grafana-dashboard.yaml @@ -11,7 +11,15 @@ metadata: grafana_folder: {{ .Values.metrics.grafanaDashboard.folder | quote }} data: {{- range $path, $_ := .Files.Glob "dashboards/*.json" }} - {{ base $path }}: |- + {{- $filename := base $path }} + {{- if hasKey $.Values.metrics.grafanaDashboard.titles $filename }} + {{- $dashboard := $.Files.Get $path | fromJson }} + {{- $_ := set $dashboard "title" (index $.Values.metrics.grafanaDashboard.titles $filename) }} + {{ $filename }}: |- + {{ $dashboard | toPrettyJson | nindent 4 }} + {{- else }} + {{ $filename }}: |- {{ $.Files.Get $path | nindent 4 }} {{- end }} + {{- end }} {{- end }} diff --git a/deploy/helm/codex-lb/values.yaml b/deploy/helm/codex-lb/values.yaml index 171a10e329..404913ae1e 100644 --- a/deploy/helm/codex-lb/values.yaml +++ b/deploy/helm/codex-lb/values.yaml @@ -417,6 +417,8 @@ metrics: enabled: false # @param metrics.grafanaDashboard.folder Grafana dashboard folder label folder: codex-lb + # @param metrics.grafanaDashboard.titles Dashboard title overrides keyed by JSON filename + titles: {} # @section Tracing parameters tracing: diff --git a/docs/conversations.md b/docs/conversations.md new file mode 100644 index 0000000000..5777160d81 --- /dev/null +++ b/docs/conversations.md @@ -0,0 +1,126 @@ +# Conversations + +The dashboard's **Conversations** view groups request logs by their +`conversation_id`. It is a read-only, derived view: codex-lb does not store a +separate conversation entity. + +Conversation functionality turns request-log metadata into an operator view of +recent conversation activity. It identifies related requests, summarizes their +usage, and lets admins inspect model and reasoning-effort breakdowns without +displaying raw prompt or response content. + +## How IDs Are Extracted + +codex-lb extracts a conversation ID from the inbound request's `User-Agent` and +conversation headers while creating request-log metadata. The value is stored +as the nullable `conversation_id` field and is then used for dashboard and +report aggregation. The extraction does not modify or reject the proxied +request. + +| Client | Header lookup order | +| --- | --- | +| OpenCode | `x-parent-session-id`, `x-opencode-session`, `x-session-id`, `x-session-affinity` | +| Codex | `thread-id` | + +User-agent and header-name matching is case-insensitive. The first non-empty +header value is selected and surrounding whitespace is removed. Requests from +unsupported clients, or requests without a usable matching header, have no +conversation ID and are not grouped. Empty IDs are also excluded from the +conversation view. The same metadata is carried through normal HTTP, +WebSocket, and supported control or auxiliary request-log paths. + +## Dashboard Data + +The dashboard overview shows **Active Conversations** for the selected `1d`, +`7d`, or `30d` timeframe, the average requests per conversation, and a +conversation trend alongside the other request, token, and cost metrics. + +Open **Dashboard**, use the view selector next to the Requests heading, and +choose **Conversations**. The view supports: + +- Activity windows of `1d`, `7d`, or `30d` (the default is `7d`) +- Pagination + +The dashboard view does not expose a free-text search control. API clients may +search by conversation ID or user-agent family using the `search` parameter +documented below. + +The conversation table shows: + +- Last request time and total conversation duration +- Conversation ID +- Representative account and API key, with counts for additional accounts +- Representative model, with counts for additional models +- Request count, total and cached tokens, and total cost + +Selecting **Details** opens a dialog with the conversation ID, start and latest +timestamps, account count, total elapsed time, dominant user-agent group, and +per-model/reasoning-effort totals for requests, elapsed time, input tokens, +cached input tokens, output tokens, and cost. + +A conversation appears when it has at least one eligible request inside the +selected activity window. Its displayed start time and aggregate totals still +include the conversation's full eligible history. For example, a conversation +that started before the window but was active during it is listed, and its +`firstRequest` can be older than the selected window. + +## Reports Data + +The Reports page does not list conversation IDs or raw conversation content. It +shows aggregate counts for the selected date range, timezone, and account, +model, and user-agent filters: + +- The summary card shows **Active Conversations**, the number of distinct + non-empty conversation IDs in the report range. +- The daily breakdown includes a sortable **Conversations** column with the + distinct conversation count for each local report day. +- The daily breakdown CSV export includes the same Conversations value for + each day. + +The summary count is distinct across the whole report range, while a +conversation active on multiple days can appear in multiple daily counts. Do +not add the daily values to reproduce the summary total. + +## API + +Both endpoints require an authenticated dashboard **admin** principal. Guest +requests receive HTTP `403` with error code `admin_access_required`. + +### List conversations + +```http +GET /api/conversations?timeframe=7d&search=conv&limit=25&offset=0 +``` + +Query parameters: + +- `timeframe` — `1d`, `7d`, or `30d` +- `since` — an explicit ISO 8601 activity-window start; it cannot be used with + `timeframe` and is capped at a 30-day lookback +- `search` — optional conversation ID or user-agent family search text +- `limit` — page size from `1` to `1000` (the API default is `50`) +- `offset` — zero-based page offset + +The response contains `conversations`, `total`, and `hasMore`. Each list entry +includes the conversation ID, first and last request timestamps, request count, +representative account and model information, token totals, cached input +tokens, and total cost. + +### Get conversation details + +```http +GET /api/conversations/{conversation_id} +``` + +The detail response includes `conversationId`, `start`, `latest`, +`accountCount`, `totalElapsedTime`, `dominantUseragentGroup`, and +`modelStats`. Each model statistic includes the model and reasoning effort, +request count, elapsed time, input/output/cached token totals, and cost. + +Conversations are derived from eligible request-log rows. Empty conversation +IDs and rows excluded by the request-log retention filters do not create a +conversation entry. + +--- + +*Spec: [conversations-api](https://github.com/Soju06/codex-lb/tree/main/openspec/specs/conversations-api)* diff --git a/docs/deployment/kubernetes.md b/docs/deployment/kubernetes.md index 3184c69fe8..69f57b0d36 100644 --- a/docs/deployment/kubernetes.md +++ b/docs/deployment/kubernetes.md @@ -95,6 +95,26 @@ gatewayApi: Gateway defaults to a single HTTP listener on port 80; override `gatewayApi.gateway.listeners` for TLS or other ports. +## Grafana dashboard hierarchy + +The chart can assign concise titles to its packaged dashboards without copying +their JSON. When the Grafana sidecar maps annotation paths to filesystem-backed +nested folders, the following values produce `Applications / Codex LB / +Overview` and `Applications / Codex LB / TTFT Breakdown`: + +```yaml +metrics: + grafanaDashboard: + enabled: true + folder: Applications/Codex LB + titles: + codex-lb.json: Overview + ttft-breakdown.json: TTFT Breakdown +``` + +The title map is keyed by the JSON filenames packaged in the chart. Omitting it +preserves the default dashboard titles. + ## Full chart reference For external database, production config, ingress, observability, and more see the diff --git a/docs/index.md b/docs/index.md index 87e45adb9d..668ef7f154 100644 --- a/docs/index.md +++ b/docs/index.md @@ -20,6 +20,7 @@ Load balancer for ChatGPT accounts. Pool multiple accounts, track usage, manage - [Client Setup](client-setup.md) — Codex CLI, OpenCode, OpenClaw, Python SDK - [Configuration](configuration.md) — the few settings that matter - [Authentication](authentication.md) — dashboard auth modes +- [Conversations](conversations.md) — dashboard view and conversation APIs - [API Keys](api-keys.md) — protecting proxy routes - [Routing](routing.md) — routing strategy guide - [Database](database.md) — SQLite / PostgreSQL, data paths, Postgres upgrades diff --git a/frontend/screenshots/capture.spec.ts b/frontend/screenshots/capture.spec.ts index 4e87a8ccf8..480e48c5fc 100644 --- a/frontend/screenshots/capture.spec.ts +++ b/frontend/screenshots/capture.spec.ts @@ -17,7 +17,12 @@ import { upstreamProxyAdmin, unauthenticatedSession, } from "./fixtures"; -import { createAccountSummary } from "../src/test/mocks/factories"; +import { + createAccountSummary, + createConversationDetails, + createConversationEntry, + createConversationsResponse, +} from "../src/test/mocks/factories"; const __dirname = path.dirname(fileURLToPath(import.meta.url)); const SCREENSHOT_DIR = path.resolve(__dirname, "../../docs/screenshots"); @@ -66,6 +71,15 @@ async function interceptApi( const slice = requestLogs.slice(offset, offset + limit); return fulfill(route, createRequestLogsResponse(slice, requestLogs.length, offset + limit < requestLogs.length)); } + if (p === "/api/conversations") { + return fulfill( + route, + createConversationsResponse([createConversationEntry({ conversationId: "conv_abc" })], 1, false), + ); + } + if (p === "/api/conversations/conv_abc") { + return fulfill(route, createConversationDetails({ conversationId: "conv_abc" })); + } if (p === "/api/accounts") return fulfill(route, { accounts: accountList }); const trendsMatch = p.match(/^\/api\/accounts\/([^/]+)\/trends$/); if (trendsMatch) { @@ -119,6 +133,7 @@ async function capture( fullPage?: boolean; session?: SessionOverride; waitFor?: string; + beforeScreenshot?: (page: Page) => Promise; }, ) { await applyTheme(page, opts.theme); @@ -143,6 +158,10 @@ async function capture( // Short settle for JS-driven rendering (Recharts SVG mutations etc.) await page.waitForTimeout(SETTLE_MS); + if (opts.beforeScreenshot) { + await opts.beforeScreenshot(page); + } + // For fullPage captures, un-fix the sticky footer so it flows at the document bottom // instead of floating at the original viewport boundary. if (opts.fullPage) { @@ -173,6 +192,39 @@ test("dashboard — dark", async ({ page }) => { await capture(page, { file: "dashboard-dark.jpg", theme: "dark", route: "/dashboard" }); }); +test("dashboard conversations — desktop", async ({ page }) => { + await capture(page, { + file: "dashboard-conversations.jpg", + theme: "light", + route: "/dashboard?view=conversations", + waitFor: '[data-slot="table"]', + }); +}); + +test("dashboard conversations — narrow", async ({ page }) => { + await page.setViewportSize({ width: 390, height: 844 }); + await capture(page, { + file: "dashboard-conversations-narrow.jpg", + theme: "light", + route: "/dashboard?view=conversations", + waitFor: '[data-slot="table"]', + }); +}); + +test("dashboard conversation details dialog", async ({ page }) => { + await capture(page, { + file: "dashboard-conversation-details.jpg", + theme: "light", + route: "/dashboard?view=conversations", + waitFor: '[data-slot="table"]', + beforeScreenshot: async (currentPage) => { + await currentPage.getByRole("button", { name: /view details/i }).click(); + await currentPage.getByRole("dialog").waitFor(); + await currentPage.getByTestId("conversation-details-information").waitFor(); + }, + }); +}); + test("accounts — light", async ({ page }) => { await capture(page, { file: "accounts.jpg", theme: "light", route: "/accounts" }); }); diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 13d6bb8c08..53ae46bb49 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -1,8 +1,11 @@ -import { lazy, Suspense } from "react"; +import { lazy, Suspense, useState } from "react"; import { Navigate, Outlet, Route, Routes } from "react-router-dom"; import { AppHeader } from "@/components/layout/app-header"; -import { StatusBar } from "@/components/layout/status-bar"; +import { + STATUS_BAR_DEFAULT_HEIGHT_PX, + StatusBar, +} from "@/components/layout/status-bar"; import { Toaster } from "@/components/ui/sonner"; import { TooltipProvider } from "@/components/ui/tooltip"; import { AuthGate } from "@/features/auth/components/auth-gate"; @@ -36,9 +39,14 @@ function AppLayout() { const startAdminLogin = useAuthStore((state) => state.startAdminLogin); const timeFormat = useTimeFormatStore((state) => state.timeFormat); const isGuest = role === "guest"; + const [statusBarHeight, setStatusBarHeight] = useState(STATUS_BAR_DEFAULT_HEIGHT_PX); return ( -
+
{ void logout(); @@ -52,7 +60,7 @@ function AppLayout() { - +
); } diff --git a/frontend/src/__integration__/dashboard-flow.test.tsx b/frontend/src/__integration__/dashboard-flow.test.tsx index 45d2a1bdac..cc80db7acb 100644 --- a/frontend/src/__integration__/dashboard-flow.test.tsx +++ b/frontend/src/__integration__/dashboard-flow.test.tsx @@ -10,6 +10,8 @@ import { createAccountSummary, createDashboardOverview, createDashboardProjections, + createConversationEntry, + createConversationsResponse, createDefaultRequestLogs, createRequestLogEntry, createRequestLogFilterOptions, @@ -72,7 +74,7 @@ describe("dashboard flow integration", () => { renderWithProviders(); expect(await screen.findByRole("heading", { name: "Dashboard" })).toBeInTheDocument(); - expect(await screen.findByText("Request Logs")).toBeInTheDocument(); + expect(await screen.findByRole("heading", { name: "Request Logs" })).toBeInTheDocument(); await waitFor(() => { expect(overviewCalls).toBeGreaterThan(0); @@ -281,4 +283,108 @@ describe("dashboard flow integration", () => { expect(projectionsCalls).toBe(projectionsCallsBeforeRetry); expect(optionsCalls).toBe(optionsCallsBeforeRetry); }); + + it("switches to conversations without reinterpreting request-log URL state", async () => { + const user = userEvent.setup({ delay: null }); + server.use( + http.get("/api/conversations", () => + HttpResponse.json( + createConversationsResponse([ + createConversationEntry({ conversationId: "opencode_conversation" }), + ], 1, false), + ), + ), + ); + window.history.pushState( + {}, + "", + "/dashboard?search=requestlog&limit=10&offset=25&conversationSearch=opencode&conversationLimit=15&conversationOffset=7", + ); + + renderWithProviders(); + + expect(await screen.findByRole("heading", { name: "Request Logs" })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Request Logs" })); + await user.click(screen.getByRole("menuitemradio", { name: "Conversations" })); + + expect(await screen.findByText("opencode_conversation")).toBeInTheDocument(); + expect(window.location.search).toContain("view=conversations"); + expect(window.location.search).toContain("search=requestlog"); + expect(window.location.search).toContain("limit=10"); + expect(window.location.search).toContain("offset=25"); + expect(screen.queryByRole("searchbox")).not.toBeInTheDocument(); + expect(window.location.search).toContain("conversationSearch=opencode"); + expect(window.location.search).toContain("conversationLimit=15"); + expect(window.location.search).toContain("conversationOffset=7"); + + await user.click(screen.getByRole("button", { name: "Conversations" })); + await user.click(screen.getByRole("menuitemradio", { name: "Request Logs" })); + + await waitFor(() => expect(window.location.search).not.toContain("view=conversations")); + expect(window.location.search).toContain("search=requestlog"); + expect(window.location.search).toContain("limit=10"); + expect(window.location.search).toContain("offset=25"); + expect(window.location.search).toContain("conversationSearch=opencode"); + expect(window.location.search).toContain("conversationLimit=15"); + expect(window.location.search).toContain("conversationOffset=7"); + }); + + it("refetches the overview (stat boxes) when the conversation timeframe changes", async () => { + const user = userEvent.setup({ delay: null }); + + let overviewCalls = 0; + const overviewTimeframes: string[] = []; + + server.use( + http.get("/api/dashboard/overview", ({ request }) => { + overviewCalls += 1; + const timeframe = (new URL(request.url).searchParams.get("timeframe") ?? "7d") as string; + overviewTimeframes.push(timeframe); + return HttpResponse.json(createDashboardOverview({ + timeframe: + timeframe === "1d" + ? { key: "1d", windowMinutes: 1440, bucketSeconds: 3600, bucketCount: 24 } + : timeframe === "30d" + ? { key: "30d", windowMinutes: 43200, bucketSeconds: 86400, bucketCount: 30 } + : { key: "7d", windowMinutes: 10080, bucketSeconds: 21600, bucketCount: 28 }, + })); + }), + http.get("/api/conversations", () => + HttpResponse.json(createConversationsResponse([ + createConversationEntry({ conversationId: "opencode_conversation" }), + ], 1, false)), + ), + ); + + window.history.pushState({}, "", "/dashboard?view=conversations&overviewTimeframe=1d"); + renderWithProviders(); + + expect(await screen.findByRole("heading", { name: "Dashboard" })).toBeInTheDocument(); + await waitFor(() => expect(overviewCalls).toBeGreaterThan(0)); + expect(overviewTimeframes.at(-1)).toBe("7d"); + + const overviewAfterLoad = overviewCalls; + + // Change the date range from the conversations-mode selector (top right). + const timeframeSelect = screen.getByRole("combobox", { name: "Conversation timeframe" }); + await user.click(timeframeSelect); + await user.click(await screen.findByRole("option", { name: "30d" })); + + // Regression: the overview query MUST refetch with the new timeframe so the + // stat boxes (requests/tokens/cost/etc.) update alongside the conversation list. + await waitFor(() => { + expect(overviewCalls).toBeGreaterThan(overviewAfterLoad); + }); + expect(overviewTimeframes.at(-1)).toBe("30d"); + expect(window.location.search).toContain("conversationTimeframe=30d"); + expect(window.location.search).toContain("overviewTimeframe=1d"); + expect(window.location.search).not.toContain("overviewTimeframe=30d"); + + await user.click(screen.getByRole("button", { name: "Conversations" })); + await user.click(await screen.findByRole("menuitemradio", { name: "Request Logs" })); + + await waitFor(() => { + expect(overviewTimeframes.at(-1)).toBe("1d"); + }); + }); }); diff --git a/frontend/src/components/layout/status-bar.test.tsx b/frontend/src/components/layout/status-bar.test.tsx index 651c4b5b7d..5bc9abdf9b 100644 --- a/frontend/src/components/layout/status-bar.test.tsx +++ b/frontend/src/components/layout/status-bar.test.tsx @@ -1,14 +1,17 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render, screen } from "@testing-library/react"; +import { act, render, screen } from "@testing-library/react"; import { HttpResponse, http } from "msw"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; -import { StatusBar } from "@/components/layout/status-bar"; +import { StatusBar, type StatusBarProps } from "@/components/layout/status-bar"; import i18n from "@/i18n"; -import { createDashboardSettings } from "@/test/mocks/factories"; +import { + createDashboardOverview, + createDashboardSettings, +} from "@/test/mocks/factories"; import { server } from "@/test/mocks/server"; -function renderStatusBar() { +function renderStatusBar(props: StatusBarProps = {}) { const queryClient = new QueryClient({ defaultOptions: { queries: { @@ -19,7 +22,7 @@ function renderStatusBar() { return render( - + , ); } @@ -33,6 +36,117 @@ function mockSettings( } describe("StatusBar", () => { + it("shows ready service independently from stale usage", async () => { + server.use( + http.get("/health/ready", () => HttpResponse.json({ status: "ok" })), + http.get("/api/dashboard/overview", () => + HttpResponse.json( + createDashboardOverview({ + lastSyncAt: new Date(Date.now() - 120_000).toISOString(), + }), + ), + ), + ); + + renderStatusBar(); + + expect(await screen.findByText("Ready")).toBeInTheDocument(); + expect(await screen.findByText("Stale")).toBeInTheDocument(); + }); + + it("shows fresh usage independently from an unready service", async () => { + server.use( + http.get("/health/ready", () => + HttpResponse.json({ detail: "Service unavailable" }, { status: 503 }), + ), + http.get("/api/dashboard/overview", () => + HttpResponse.json( + createDashboardOverview({ + lastSyncAt: new Date().toISOString(), + }), + ), + ), + ); + + renderStatusBar(); + + expect(await screen.findByText("Not ready")).toBeInTheDocument(); + expect(await screen.findByText("Synced")).toBeInTheDocument(); + }); + + it("shows checking while readiness is pending without blocking usage status", async () => { + let resolveReadiness!: () => void; + const pendingReadiness = new Promise((resolve) => { + resolveReadiness = resolve; + }); + server.use( + http.get("/health/ready", async () => { + await pendingReadiness; + return HttpResponse.json({ status: "ok" }); + }), + http.get("/api/dashboard/overview", () => + HttpResponse.json( + createDashboardOverview({ + lastSyncAt: new Date().toISOString(), + }), + ), + ), + ); + + renderStatusBar(); + + expect(await screen.findByText("Synced")).toBeInTheDocument(); + expect(screen.getByText("Checking")).toBeInTheDocument(); + + await act(async () => { + resolveReadiness(); + }); + expect(await screen.findByText("Ready")).toBeInTheDocument(); + }); + + it("reports its resized height so wrapped rows remain clear of page content", () => { + let resizeCallback: ResizeObserverCallback | undefined; + const observe = vi.fn(); + const disconnect = vi.fn(); + const offsetHeight = vi + .spyOn(HTMLElement.prototype, "offsetHeight", "get") + .mockReturnValue(40); + const resizeObserver = vi + .spyOn(globalThis, "ResizeObserver") + .mockImplementation( + class { + constructor(callback: ResizeObserverCallback) { + resizeCallback = callback; + } + + observe = observe; + unobserve = vi.fn(); + disconnect = disconnect; + }, + ); + const onHeightChange = vi.fn(); + + try { + const { unmount } = renderStatusBar({ onHeightChange }); + + expect(onHeightChange).toHaveBeenLastCalledWith(40); + expect(observe).toHaveBeenCalledWith(screen.getByRole("contentinfo")); + + offsetHeight.mockReturnValue(72); + act(() => { + resizeCallback?.([], {} as ResizeObserver); + }); + + expect(onHeightChange).toHaveBeenLastCalledWith(72); + + unmount(); + expect(disconnect).toHaveBeenCalledOnce(); + } finally { + resizeObserver.mockRestore(); + offsetHeight.mockRestore(); + } + }); + it("links to the official GitHub repository", () => { renderStatusBar(); diff --git a/frontend/src/components/layout/status-bar.tsx b/frontend/src/components/layout/status-bar.tsx index ee1ed4110a..d1478c42d8 100644 --- a/frontend/src/components/layout/status-bar.tsx +++ b/frontend/src/components/layout/status-bar.tsx @@ -1,16 +1,24 @@ -import { useEffect, useState } from "react"; -import { Activity, ArrowRightLeft, ArrowUpCircle, Tag } from "lucide-react"; +import { useEffect, useLayoutEffect, useRef, useState } from "react"; +import { ArrowRightLeft, ArrowUpCircle, Tag } from "lucide-react"; import { useQuery } from "@tanstack/react-query"; import { useTranslation } from "react-i18next"; import type { TFunction } from "i18next"; import { getDashboardOverview } from "@/features/dashboard/api"; import { DEFAULT_OVERVIEW_TIMEFRAME } from "@/features/dashboard/schemas"; +import { getServiceReadiness } from "@/features/health/api"; import { getRuntimeVersion } from "@/features/runtime/api"; import { getSettings } from "@/features/settings/api"; import { formatTimeLong } from "@/utils/formatters"; const GITHUB_REPOSITORY_URL = "https://github.com/soju06/codex-lb"; +const STATUS_REFRESH_INTERVAL_MS = 60_000; +const USAGE_FRESHNESS_THRESHOLD_MS = 60_000; +export const STATUS_BAR_DEFAULT_HEIGHT_PX = 40; + +export interface StatusBarProps { + onHeightChange?: (height: number) => void; +} type RoutingStrategy = | "usage_weighted" @@ -80,12 +88,20 @@ function getRoutingLabel( return strategyLabel; } -export function StatusBar() { +export function StatusBar({ onHeightChange }: StatusBarProps = {}) { const { t } = useTranslation(); + const footerRef = useRef(null); + const readinessQuery = useQuery({ + queryKey: ["health", "ready"], + queryFn: getServiceReadiness, + refetchInterval: STATUS_REFRESH_INTERVAL_MS, + refetchIntervalInBackground: false, + retry: false, + }); const { data: lastSyncAt = null } = useQuery({ queryKey: ["dashboard", "overview", DEFAULT_OVERVIEW_TIMEFRAME], queryFn: () => getDashboardOverview({ timeframe: DEFAULT_OVERVIEW_TIMEFRAME }), - refetchInterval: 60_000, + refetchInterval: STATUS_REFRESH_INTERVAL_MS, refetchIntervalInBackground: false, select: (data) => data.lastSyncAt, }); @@ -101,15 +117,26 @@ export function StatusBar() { staleTime: 6 * 60 * 60 * 1000, }); const lastSync = formatTimeLong(lastSyncAt); - const [isLive, setIsLive] = useState(false); + const [isUsageSynced, setIsUsageSynced] = useState(false); useEffect(() => { function check() { - setIsLive(lastSyncAt ? Date.now() - new Date(lastSyncAt).getTime() < 60_000 : false); + setIsUsageSynced( + lastSyncAt + ? Date.now() - new Date(lastSyncAt).getTime() < USAGE_FRESHNESS_THRESHOLD_MS + : false, + ); } check(); const id = setInterval(check, 10_000); return () => clearInterval(id); }, [lastSyncAt]); + const serviceReadiness = readinessQuery.isPending + ? "checking" + : !readinessQuery.isError && readinessQuery.data?.status === "ok" + ? "ready" + : "notReady"; + const serviceStatusLabel = t(`statusBar.${serviceReadiness}`); + const usageStatusLabel = t(isUsageSynced ? "statusBar.synced" : "statusBar.stale"); const routingLabel = settings ? getRoutingLabel( @@ -127,17 +154,52 @@ export function StatusBar() { ? t("statusBar.updateAvailableWithVersion", { version: latestVersion }) : t("statusBar.updateAvailable"); + useLayoutEffect(() => { + const footer = footerRef.current; + if (!footer || !onHeightChange) { + return; + } + + const reportHeight = () => { + onHeightChange(Math.max(STATUS_BAR_DEFAULT_HEIGHT_PX, footer.offsetHeight)); + }; + reportHeight(); + + const observer = new ResizeObserver(reportHeight); + observer.observe(footer); + return () => observer.disconnect(); + }, [onHeightChange]); + return ( -