codexLB-up 1.25.0b10__py3-none-any.whl

app/core/clients/usage.py

@@ -0,0 +1,467 @@
+from __future__ import annotations
+
+import asyncio
+import codecs
+import contextlib
+import json
+import logging
+from typing import Final, Literal, cast
+
+import aiohttp
+from aiohttp_retry import ExponentialRetry, RetryClient
+from pydantic import BaseModel, ConfigDict, ValidationError
+
+from app.core.clients.codex import (
+    CodexClient,
+    CodexTransportError,
+    create_codex_session,
+    require_route_or_direct_egress_opt_in,
+)
+from app.core.clients.headers import build_chatgpt_auth_headers
+from app.core.clients.http import _safe_json, lease_retry_client
+from app.core.clients.native_egress import (
+    NativeEgressClient,
+    NativeEgressError,
+    NativeEgressRequest,
+    NativeEgressResponse,
+    NativeEgressTransportError,
+    NativeEgressUnavailable,
+    discover_native_egress_client,
+)
+from app.core.clients.proxy import _codex_response_status
+from app.core.config.settings import get_settings
+from app.core.types import JsonObject
+from app.core.upstream_proxy import ResolvedUpstreamRoute
+from app.core.usage.models import UsagePayload
+from app.core.utils.proxy_env import resolve_http_proxy_from_env
+from app.core.utils.request_id import get_request_id
+
+RETRYABLE_STATUS = {408, 429, 500, 502, 503, 504}
+RETRY_START_TIMEOUT = 0.5
+RETRY_MAX_TIMEOUT = 2.0
+# Per-call total timeout and retry budget of the usage / reset-credits fetches
+# (fixed; issue #1340 / PRINCIPLES.md P2). Callers may still pass explicit
+# ``timeout_seconds`` / ``max_retries`` overrides.
+USAGE_FETCH_TIMEOUT_SECONDS: Final[float] = 10.0
+USAGE_FETCH_MAX_RETRIES: Final[int] = 2
+
+logger = logging.getLogger(__name__)
+
+
+class UsageErrorDetail(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    code: str | None = None
+    message: str | None = None
+    error_description: str | None = None
+
+
+class UsageErrorEnvelope(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    error: UsageErrorDetail | str | None = None
+    error_description: str | None = None
+    message: str | None = None
+
+
+class UsageFetchError(Exception):
+    def __init__(self, status_code: int, message: str, code: str | None = None) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+        self.message = message
+        self.code = code
+
+
+class ConsumeRateLimitResetCreditResponse(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    code: Literal["reset", "nothing_to_reset", "no_credit", "already_redeemed"]
+    windows_reset: int = 0
+
+
+async def fetch_usage(
+    *,
+    access_token: str,
+    account_id: str | None,
+    base_url: str | None = None,
+    timeout_seconds: float | None = None,
+    max_retries: int | None = None,
+    client: RetryClient | None = None,
+    route: ResolvedUpstreamRoute | None = None,
+    codex_client: CodexClient | None = None,
+    allow_direct_egress: bool = False,
+) -> UsagePayload:
+    settings = get_settings()
+    usage_base = base_url or settings.upstream_base_url
+    url = _usage_url(usage_base)
+    timeout = aiohttp.ClientTimeout(total=timeout_seconds or USAGE_FETCH_TIMEOUT_SECONDS)
+    retries = max_retries if max_retries is not None else USAGE_FETCH_MAX_RETRIES
+    headers = _usage_headers(access_token, account_id)
+    retry_options = _retry_options(retries + 1)
+    require_route_or_direct_egress_opt_in(
+        route=route,
+        allow_direct_egress=allow_direct_egress,
+        operation="usage fetch",
+    )
+
+    try:
+        if route is not None:
+            return await _fetch_usage_via_codex(
+                url=url,
+                route=route,
+                headers=headers,
+                timeout_seconds=timeout_seconds or USAGE_FETCH_TIMEOUT_SECONDS,
+                retries=retries,
+                codex_client=codex_client,
+            )
+        native_client = discover_native_egress_client() if client is None else None
+        if native_client is not None:
+            try:
+                return await _fetch_usage_via_native(
+                    client=native_client,
+                    url=url,
+                    headers=headers,
+                    timeout_seconds=timeout_seconds or USAGE_FETCH_TIMEOUT_SECONDS,
+                    retries=retries,
+                )
+            except NativeEgressUnavailable:
+                pass
+        async with lease_retry_client(client) as retry_client:
+            async with retry_client.request(
+                "GET",
+                url,
+                headers=headers,
+                timeout=timeout,
+                retry_options=retry_options,
+            ) as resp:
+                data = await _safe_json(resp)
+                return _usage_payload_or_raise(data, resp.status)
+    except (aiohttp.ClientError, asyncio.TimeoutError, CodexTransportError, NativeEgressError) as exc:
+        logger.warning(
+            "Usage fetch error request_id=%s error=%s",
+            get_request_id(),
+            exc,
+        )
+        raise UsageFetchError(0, f"Usage fetch failed: {exc}") from exc
+
+
+async def _fetch_usage_via_native(
+    *,
+    client: NativeEgressClient,
+    url: str,
+    headers: dict[str, str],
+    timeout_seconds: float,
+    retries: int,
+) -> UsagePayload:
+    attempts = max(1, retries + 1)
+    retry_options = _retry_options(attempts)
+    request = NativeEgressRequest(
+        method="GET",
+        url=url,
+        # The helper enables response decompression when negotiation is explicit.
+        headers={
+            "Accept-Encoding": aiohttp.ClientRequest.DEFAULT_HEADERS[aiohttp.hdrs.ACCEPT_ENCODING],
+            **headers,
+        },
+        timeout_seconds=timeout_seconds,
+        proxy_url=resolve_http_proxy_from_env(url),
+    )
+    for attempt in range(attempts):
+        try:
+            response = await client.request(request)
+        except NativeEgressUnavailable as exc:
+            if attempt == 0:
+                raise
+            # Do not restart the budget through Python after a native attempt.
+            raise NativeEgressError("native helper became unavailable during usage retries") from exc
+        except NativeEgressTransportError:
+            if attempt < attempts - 1:
+                await asyncio.sleep(retry_options.get_timeout(attempt + 1))
+                continue
+            raise
+        async with response:
+            status = response.status
+            if status not in RETRYABLE_STATUS or attempt == attempts - 1:
+                data = await _native_usage_json(response)
+                return _usage_payload_or_raise(data, status)
+        await asyncio.sleep(retry_options.get_timeout(attempt + 1))
+    raise RuntimeError("unreachable native usage retry state")
+
+
+async def _native_usage_json(response: NativeEgressResponse) -> JsonObject:
+    # Match the default aiohttp session's charset and empty-body behavior,
+    # keeping body/IPC failures outside the JSON-syntax exception handler.
+    body = await response.read()
+    encoding = "utf-8"
+    content_type = aiohttp.helpers.parse_mimetype(response.headers.get("Content-Type", "").lower())
+    charset = content_type.parameters.get("charset")
+    if charset:
+        with contextlib.suppress(LookupError, ValueError):
+            encoding = codecs.lookup(charset).name
+    try:
+        value = json.loads(body.strip().decode(encoding)) if body.strip() else None
+    except ValueError:
+        return {"error": {"message": body.decode(encoding).strip()}}
+    return cast(JsonObject, value) if isinstance(value, dict) else {"error": {"message": str(value)}}
+
+
+async def consume_rate_limit_reset_credit(
+    *,
+    access_token: str,
+    account_id: str,
+    redeem_request_id: str,
+    base_url: str | None = None,
+    timeout_seconds: float | None = None,
+    max_retries: int | None = None,
+    client: RetryClient | None = None,
+    route: ResolvedUpstreamRoute | None = None,
+    codex_client: CodexClient | None = None,
+    allow_direct_egress: bool = False,
+) -> ConsumeRateLimitResetCreditResponse:
+    settings = get_settings()
+    usage_base = base_url or settings.upstream_base_url
+    url = _rate_limit_reset_url(usage_base)
+    timeout = aiohttp.ClientTimeout(total=timeout_seconds or USAGE_FETCH_TIMEOUT_SECONDS)
+    retries = max_retries if max_retries is not None else USAGE_FETCH_MAX_RETRIES
+    headers = _usage_headers(access_token, account_id)
+    payload = {"redeem_request_id": redeem_request_id}
+    retry_options = _retry_options(retries + 1)
+    require_route_or_direct_egress_opt_in(
+        route=route,
+        allow_direct_egress=allow_direct_egress,
+        operation="usage limit reset consume",
+    )
+
+    try:
+        if route is not None:
+            return await _consume_rate_limit_reset_via_codex(
+                url=url,
+                route=route,
+                headers=headers,
+                payload=payload,
+                timeout_seconds=timeout_seconds or USAGE_FETCH_TIMEOUT_SECONDS,
+                retries=retries,
+                codex_client=codex_client,
+            )
+        async with lease_retry_client(client) as retry_client:
+            async with retry_client.request(
+                "POST",
+                url,
+                headers=headers,
+                json=payload,
+                timeout=timeout,
+                retry_options=retry_options,
+            ) as resp:
+                data = await _safe_json(resp)
+                return _consume_rate_limit_reset_response_or_raise(data, resp.status)
+    except (aiohttp.ClientError, asyncio.TimeoutError, CodexTransportError) as exc:
+        logger.warning(
+            "Usage limit reset consume error request_id=%s error=%s",
+            get_request_id(),
+            exc,
+        )
+        raise UsageFetchError(0, f"Usage limit reset consume failed: {exc}") from exc
+
+
+async def _fetch_usage_via_codex(
+    *,
+    url: str,
+    route: ResolvedUpstreamRoute,
+    headers: dict[str, str],
+    timeout_seconds: float,
+    retries: int,
+    codex_client: CodexClient | None,
+) -> UsagePayload:
+    attempts = max(1, retries + 1)
+    owns_codex_client = codex_client is None
+    active_codex_client = codex_client or CodexClient(create_codex_session())
+    try:
+        for attempt in range(attempts):
+            try:
+                resp = await active_codex_client.request(
+                    "GET",
+                    url,
+                    route=route,
+                    headers=headers,
+                    timeout=timeout_seconds,
+                )
+            except CodexTransportError:
+                if attempt < attempts - 1:
+                    await asyncio.sleep(_retry_delay_seconds(attempt))
+                    continue
+                raise
+
+            data = await _safe_codex_json(resp)
+            status = _codex_response_status(resp)
+            if status in RETRYABLE_STATUS and attempt < attempts - 1:
+                await asyncio.sleep(_retry_delay_seconds(attempt))
+                continue
+            return _usage_payload_or_raise(data, status)
+    finally:
+        if owns_codex_client:
+            close = getattr(active_codex_client, "close", None)
+            if callable(close):
+                await close()
+    raise RuntimeError("unreachable usage retry state")
+
+
+async def _consume_rate_limit_reset_via_codex(
+    *,
+    url: str,
+    route: ResolvedUpstreamRoute,
+    headers: dict[str, str],
+    payload: dict[str, str],
+    timeout_seconds: float,
+    retries: int,
+    codex_client: CodexClient | None,
+) -> ConsumeRateLimitResetCreditResponse:
+    attempts = max(1, retries + 1)
+    owns_codex_client = codex_client is None
+    active_codex_client = codex_client or CodexClient(create_codex_session())
+    try:
+        for attempt in range(attempts):
+            try:
+                resp = await active_codex_client.request(
+                    "POST",
+                    url,
+                    route=route,
+                    headers=headers,
+                    json=payload,
+                    timeout=timeout_seconds,
+                )
+            except CodexTransportError:
+                if attempt < attempts - 1:
+                    await asyncio.sleep(_retry_delay_seconds(attempt))
+                    continue
+                raise
+
+            data = await _safe_codex_json(resp)
+            status = _codex_response_status(resp)
+            if status in RETRYABLE_STATUS and attempt < attempts - 1:
+                await asyncio.sleep(_retry_delay_seconds(attempt))
+                continue
+            return _consume_rate_limit_reset_response_or_raise(data, status)
+    finally:
+        if owns_codex_client:
+            close = getattr(active_codex_client, "close", None)
+            if callable(close):
+                await close()
+    raise RuntimeError("unreachable usage limit reset retry state")
+
+
+def _usage_payload_or_raise(data: JsonObject, status: int) -> UsagePayload:
+    if status >= 400:
+        code = _extract_error_code(data)
+        message = _extract_error_message(data) or f"Usage fetch failed ({status})"
+        logger.warning(
+            "Usage fetch failed request_id=%s status=%s code=%s message=%s",
+            get_request_id(),
+            status,
+            code,
+            message,
+        )
+        raise UsageFetchError(status, message, code=code)
+    try:
+        return UsagePayload.model_validate(data)
+    except ValidationError as exc:
+        logger.warning(
+            "Usage fetch invalid payload request_id=%s",
+            get_request_id(),
+        )
+        raise UsageFetchError(502, "Invalid usage payload") from exc
+
+
+def _consume_rate_limit_reset_response_or_raise(
+    data: JsonObject,
+    status: int,
+) -> ConsumeRateLimitResetCreditResponse:
+    if status >= 400:
+        code = _extract_error_code(data)
+        message = _extract_error_message(data) or f"Usage limit reset consume failed ({status})"
+        logger.warning(
+            "Usage limit reset consume failed request_id=%s status=%s code=%s message=%s",
+            get_request_id(),
+            status,
+            code,
+            message,
+        )
+        raise UsageFetchError(status, message, code=code)
+    try:
+        return ConsumeRateLimitResetCreditResponse.model_validate(data)
+    except ValidationError as exc:
+        logger.warning(
+            "Usage limit reset consume invalid payload request_id=%s",
+            get_request_id(),
+        )
+        raise UsageFetchError(502, "Invalid usage limit reset payload") from exc
+
+
+async def _safe_codex_json(response: object) -> JsonObject:
+    try:
+        json_method = getattr(response, "json", None)
+        if callable(json_method):
+            data = json_method()
+            if asyncio.iscoroutine(data):
+                data = await data
+            return data if isinstance(data, dict) else {"error": {"message": str(data)}}
+    except Exception:
+        pass
+    content = getattr(response, "content", None)
+    if isinstance(content, bytes):
+        return {"error": {"message": content.decode("utf-8", errors="replace").strip()}}
+    if isinstance(content, str):
+        return {"error": {"message": content.strip()}}
+    return {"error": {"message": ""}}
+
+
+def _usage_url(base_url: str) -> str:
+    normalized = base_url.rstrip("/")
+    if "/backend-api" not in normalized:
+        normalized = f"{normalized}/backend-api"
+    return f"{normalized}/wham/usage"
+
+
+def _rate_limit_reset_url(base_url: str) -> str:
+    normalized = base_url.rstrip("/")
+    if "/backend-api" not in normalized:
+        normalized = f"{normalized}/backend-api"
+    return f"{normalized}/wham/rate-limit-reset-credits/consume"
+
+
+def _usage_headers(access_token: str, account_id: str | None) -> dict[str, str]:
+    return build_chatgpt_auth_headers(access_token, account_id)
+
+
+def _extract_error_message(payload: JsonObject) -> str | None:
+    envelope = UsageErrorEnvelope.model_validate(payload)
+    error = envelope.error
+    if isinstance(error, UsageErrorDetail):
+        return error.message or error.error_description
+    if isinstance(error, str):
+        return envelope.error_description or error
+    return envelope.message
+
+
+def _extract_error_code(payload: JsonObject) -> str | None:
+    envelope = UsageErrorEnvelope.model_validate(payload)
+    error = envelope.error
+    if isinstance(error, UsageErrorDetail) and isinstance(error.code, str):
+        normalized = error.code.strip().lower()
+        return normalized or None
+    return None
+
+
+def _retry_options(attempts: int) -> ExponentialRetry:
+    return ExponentialRetry(
+        attempts=attempts,
+        start_timeout=RETRY_START_TIMEOUT,
+        max_timeout=RETRY_MAX_TIMEOUT,
+        factor=2.0,
+        statuses=RETRYABLE_STATUS,
+        exceptions={aiohttp.ClientError, asyncio.TimeoutError},
+        retry_all_server_errors=False,
+    )
+
+
+def _retry_delay_seconds(attempt: int) -> float:
+    return min(RETRY_MAX_TIMEOUT, RETRY_START_TIMEOUT * (2.0**attempt))

app/core/clock.py

@@ -0,0 +1,170 @@
+"""Clock and scheduler seams for deterministic proxy lifecycle tests.
+
+Production code reads time and schedules timers/tasks through these two
+protocols so a test can substitute a virtual implementation
+(``tests/simulation/virtual_time.py``) and drive a whole proxy turn without
+wall-clock sleeps. The real adapters are *verbatim* passthroughs: every
+``RealScheduler`` method is the corresponding ``asyncio``/``anyio`` call with
+no task registry, no extra timeout and no wrapper task, so injecting the
+defaults changes nothing about production behavior. The scheduler's timing
+and spawn members *are* the ``asyncio``/``anyio`` functions (``staticmethod``
+aliases), so the per-event relay sites (``scheduler.create_task`` and
+``scheduler.wait`` per upstream message, ``scheduler.wait_for(queue.get(), ...)``
+per delivered block) pay neither a wrapper frame nor a bound-method allocation
+over the raw call.
+
+Usage rules (also enforced by ``scripts/check_proxy_timing_seams.py`` once it
+lands):
+
+* Objects that own the collaborators (``ProxyService``, ``LoadBalancer``,
+  ``WorkAdmissionController``) use ``self._scheduler`` / ``self._clock``.
+* Mixin methods and module functions that receive an owner read the seam with
+  ``scheduler_for(owner)`` / ``clock_for(owner)``; partial doubles without the
+  attribute keep the production default.
+* Owner-less turn-path functions take keyword-only ``scheduler`` / ``clock``
+  parameters that callers pass explicitly.
+* ``asyncio.sleep(0)`` is a yield point, not a timer, and stays raw.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import Awaitable, Coroutine, Iterable
+from contextlib import AbstractContextManager
+from datetime import datetime, timezone
+from typing import Any, Protocol, TypeVar
+
+import anyio
+
+T = TypeVar("T")
+
+
+class Clock(Protocol):
+    def monotonic(self) -> float: ...
+
+    def time(self) -> float: ...
+
+    def now(self) -> datetime: ...
+
+
+class Scheduler(Protocol):
+    """Timing and task-spawn seam.
+
+    The timing members are declared as plain functions returning a coroutine
+    (an ``async def`` implementation satisfies them too) so the real adapter
+    can hand back the ``asyncio`` coroutine object without a wrapper frame.
+    The awaitable/coroutine parameters are positional-only so the
+    ``asyncio`` functions themselves (``fut``, ``coro``) satisfy the protocol.
+    """
+
+    def sleep(self, delay: float, result: T | None = None) -> Coroutine[Any, Any, T | None]: ...
+
+    def wait_for(self, fut: Awaitable[T], /, timeout: float | None) -> Coroutine[Any, Any, T]:
+        """Bound ``fut`` by ``timeout``.
+
+        Real: ``asyncio.wait_for`` verbatim. Virtual: the same inline
+        ``asyncio.timeout`` shape on a virtual deadline (same task, same
+        contextvars); see ``tests/simulation/test_virtual_time.py``.
+        """
+        ...
+
+    def wait(
+        self,
+        fs: Iterable[asyncio.Future[Any]],
+        *,
+        timeout: float | None = None,
+        return_when: str = asyncio.ALL_COMPLETED,
+    ) -> Coroutine[Any, Any, tuple[set[asyncio.Future[Any]], set[asyncio.Future[Any]]]]:
+        """Mirror ``asyncio.wait``: never cancels ``fs``, reports done at the deadline."""
+        ...
+
+    def create_task(self, coro: Coroutine[Any, Any, T], /, *, name: str | None = None) -> asyncio.Task[T]: ...
+
+    def fail_after(self, delay: float) -> AbstractContextManager[Any]:
+        """Mirror ``anyio.fail_after``: cancel the enclosing scope after ``delay``."""
+        ...
+
+    def drain(self) -> Coroutine[Any, Any, None]: ...
+
+    async def cancel_owned_tasks(self) -> None: ...
+
+
+class RealClock:
+    """Wall-clock reads through one Python frame each.
+
+    Kept as ``def`` wrappers (unlike ``RealScheduler``'s aliases) on purpose:
+    the unit suite steers production time in well over a hundred places with
+    ``monkeypatch.setattr(time, "time" | "monotonic", ...)``, which an
+    import-time alias would silently bypass. The cost is one frame per read
+    (two reads per relayed upstream message) on top of the clock read itself.
+    """
+
+    def monotonic(self) -> float:
+        return time.monotonic()
+
+    def time(self) -> float:
+        return time.time()
+
+    def now(self) -> datetime:
+        return datetime.now(timezone.utc)
+
+
+class RealScheduler:
+    """Pure passthrough to ``asyncio``/``anyio``.
+
+    Deliberately stateless: a registry of every rerouted task would extend task
+    lifetimes and let ``cancel_owned_tasks`` on the process-wide singleton
+    cancel production work. Ownership tracking exists only in the virtual
+    scheduler used by tests.
+
+    The per-event members are the primitives themselves, bound as
+    ``staticmethod`` so ``scheduler.create_task(...)`` compiles to the raw
+    ``asyncio.create_task(...)`` call: no Python frame and no bound-method
+    object per seam call on the relay paths (``create_task`` + ``wait`` per
+    upstream message, ``wait_for`` per delivered block). The aliases bind at
+    import, so a test that needs to observe or replace a timed wait must
+    subclass this class or inject a virtual scheduler (``tests/simulation``);
+    monkeypatching ``asyncio.wait`` afterwards does not reach them.
+    ``tests/unit/test_clock_real_parity.py`` pins the identities. ``sleep``
+    stays a one-frame wrapper: ``asyncio.sleep`` is overloaded on ``result``
+    in typeshed, which ty 0.0.73 cannot match against the protocol member,
+    and no per-event site sleeps (backoff and recovery waits only).
+    """
+
+    wait_for = staticmethod(asyncio.wait_for)
+    wait = staticmethod(asyncio.wait)
+    create_task = staticmethod(asyncio.create_task)
+    fail_after = staticmethod(anyio.fail_after)
+
+    def sleep(self, delay: float, result: T | None = None) -> Coroutine[Any, Any, T | None]:
+        return asyncio.sleep(delay, result=result)
+
+    def drain(self) -> Coroutine[Any, Any, None]:
+        return asyncio.sleep(0)
+
+    async def cancel_owned_tasks(self) -> None:
+        # The real scheduler owns nothing; production tasks are never cancelled here.
+        return None
+
+
+REAL_CLOCK = RealClock()
+REAL_SCHEDULER = RealScheduler()
+
+
+def scheduler_for(owner: object) -> Scheduler:
+    """Return the scheduler collaborator of ``owner``.
+
+    Proxy lifecycle behavior is spread over mixins that are also reached through
+    partial collaborators, so the seam reads the collaborator instead of
+    requiring every holder to carry one. Anything without an explicit scheduler
+    gets the real ``asyncio`` one, which keeps the production default intact.
+    """
+
+    return getattr(owner, "_scheduler", REAL_SCHEDULER)
+
+
+def clock_for(owner: object) -> Clock:
+    """Return the clock collaborator of ``owner``, defaulting to real time."""
+
+    return getattr(owner, "_clock", REAL_CLOCK)