codexLB-up 1.25.0b10__py3-none-any.whl

app/core/usage/live_snapshots.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from typing import Any, Mapping
+
+_HEADER_PREFIX = "x-codex-"
+_EVENT_TYPE = "codex.rate_limits"
+# Cheap containment probe so stream hot paths can skip JSON parsing for the
+# overwhelming majority of event blocks.
+EVENT_MARKER = '"codex.rate_limits"'
+
+
+@dataclass(frozen=True, slots=True)
+class LiveUsageWindow:
+    used_percent: float
+    window_minutes: int | None
+    reset_at: int | None
+
+
+@dataclass(frozen=True, slots=True)
+class LiveRateLimitSnapshot:
+    primary: LiveUsageWindow | None
+    secondary: LiveUsageWindow | None
+    credits_has: bool | None
+    credits_unlimited: bool | None
+    credits_balance: float | None
+
+    @property
+    def has_windows(self) -> bool:
+        return self.primary is not None or self.secondary is not None
+
+
+def _parse_float(value: object) -> float | None:
+    if value is None or isinstance(value, bool):
+        return None
+    if not isinstance(value, (int, float, str)):
+        return None
+    try:
+        parsed = float(value)
+    except ValueError:
+        return None
+    # Live signals run synchronously on the serving path; non-finite values
+    # (NaN/Infinity strings) must degrade to "unparseable", not raise later
+    # in int() coercion.
+    if not math.isfinite(parsed):
+        return None
+    return parsed
+
+
+def _parse_int(value: object) -> int | None:
+    parsed = _parse_float(value)
+    if parsed is None:
+        return None
+    return int(parsed)
+
+
+def _parse_bool(value: object) -> bool | None:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in {"true", "1"}:
+            return True
+        if lowered in {"false", "0"}:
+            return False
+    return None
+
+
+def _header_window(headers: Mapping[str, str], slot: str) -> LiveUsageWindow | None:
+    used = _parse_float(headers.get(f"{_HEADER_PREFIX}{slot}-used-percent"))
+    if used is None:
+        return None
+    return LiveUsageWindow(
+        used_percent=used,
+        window_minutes=_parse_int(headers.get(f"{_HEADER_PREFIX}{slot}-window-minutes")),
+        reset_at=_parse_int(headers.get(f"{_HEADER_PREFIX}{slot}-reset-at")),
+    )
+
+
+def parse_rate_limit_headers(headers: Mapping[str, str] | None) -> LiveRateLimitSnapshot | None:
+    """Parse upstream ``x-codex-*`` rate-limit response headers.
+
+    Only the default ``codex`` limit family is consumed; additional
+    ``x-<limit-id>-*`` families are out of scope for live ingestion.
+    """
+    if not headers:
+        return None
+    lowered = {key.lower(): value for key, value in headers.items()}
+    primary = _header_window(lowered, "primary")
+    secondary = _header_window(lowered, "secondary")
+    if primary is None and secondary is None:
+        return None
+    return LiveRateLimitSnapshot(
+        primary=primary,
+        secondary=secondary,
+        credits_has=_parse_bool(lowered.get(f"{_HEADER_PREFIX}credits-has-credits")),
+        credits_unlimited=_parse_bool(lowered.get(f"{_HEADER_PREFIX}credits-unlimited")),
+        credits_balance=_parse_float(lowered.get(f"{_HEADER_PREFIX}credits-balance")),
+    )
+
+
+def _event_window(payload: object) -> LiveUsageWindow | None:
+    if not isinstance(payload, dict):
+        return None
+    used = _parse_float(payload.get("used_percent"))
+    if used is None:
+        return None
+    return LiveUsageWindow(
+        used_percent=used,
+        window_minutes=_parse_int(payload.get("window_minutes")),
+        reset_at=_parse_int(
+            payload.get("reset_at") if payload.get("reset_at") is not None else payload.get("resets_at")
+        ),
+    )
+
+
+def parse_rate_limit_event(payload: Mapping[str, Any]) -> LiveRateLimitSnapshot | None:
+    """Parse a ``codex.rate_limits`` stream event payload."""
+    if payload.get("type") != _EVENT_TYPE:
+        return None
+    # Only the default ``codex`` limit family maps onto the account's main
+    # usage rows; events for model-specific or individual limit buckets
+    # (discriminated by limit_id / metered_limit_name / limit_name) must not
+    # corrupt global selection state.
+    for discriminator in ("limit_id", "metered_limit_name", "limit_name"):
+        value = payload.get(discriminator)
+        if value is not None and value != "codex":
+            return None
+    rate_limits = payload.get("rate_limits")
+    if not isinstance(rate_limits, dict):
+        return None
+    primary = _event_window(rate_limits.get("primary"))
+    secondary = _event_window(rate_limits.get("secondary"))
+    if primary is None and secondary is None:
+        return None
+    credits = payload.get("credits")
+    if not isinstance(credits, dict):
+        credits = rate_limits.get("credits")
+    if not isinstance(credits, dict):
+        credits = {}
+    return LiveRateLimitSnapshot(
+        primary=primary,
+        secondary=secondary,
+        credits_has=_parse_bool(credits.get("has_credits")),
+        credits_unlimited=_parse_bool(credits.get("unlimited")),
+        credits_balance=_parse_float(credits.get("balance")),
+    )
+
+
+def parse_rate_limit_event_text(text: str) -> LiveRateLimitSnapshot | None:
+    """Parse a raw stream chunk that may contain a rate-limit event.
+
+    Accepts either a bare JSON object or an SSE ``data:`` block. Callers
+    should gate on :data:`EVENT_MARKER` first to keep hot paths cheap.
+    """
+    if EVENT_MARKER not in text:
+        return None
+    for line in text.splitlines():
+        candidate = line.strip()
+        if candidate.startswith("data:"):
+            candidate = candidate[5:].strip()
+        if not candidate.startswith("{") or EVENT_MARKER not in candidate:
+            continue
+        try:
+            payload = json.loads(candidate)
+        except json.JSONDecodeError:
+            continue
+        if isinstance(payload, dict):
+            snapshot = parse_rate_limit_event(payload)
+            if snapshot is not None:
+                return snapshot
+    return None

app/core/usage/logs.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+from app.core.usage.pricing import (
+    UsageCostBreakdown,
+    UsageTokens,
+    calculate_cost_breakdown_from_usage,
+    calculate_cost_from_usage,
+    get_pricing_for_model,
+)
+
+# Request-log status classification shared by every error-metric surface
+# (usage builders, time rollups, reports, fleet): `cancelled` is a normal
+# client-side terminal (written when the downstream client disconnects before
+# the final event lands — routing health already treats it as non-penalizing),
+# so only statuses outside this tuple count as errors.
+SUCCESS_STATUS = "success"
+CANCELLED_STATUS = "cancelled"
+NON_ERROR_STATUSES: tuple[str, ...] = (SUCCESS_STATUS, CANCELLED_STATUS)
+# The error code cancelled rows carry; excluded read-side from historical
+# error-satellite rollup rows that were folded under the legacy
+# `status != 'success'` filter.
+CLIENT_DISCONNECT_ERROR_CODE = "client_disconnected"
+
+
+class RequestLogLike(Protocol):
+    @property
+    def model(self) -> str | None: ...
+
+    @property
+    def service_tier(self) -> str | None: ...
+
+    @property
+    def input_tokens(self) -> int | None: ...
+
+    @property
+    def output_tokens(self) -> int | None: ...
+
+    @property
+    def cached_input_tokens(self) -> int | None: ...
+
+    @property
+    def reasoning_tokens(self) -> int | None: ...
+
+    @property
+    def cost_usd(self) -> float | None: ...
+
+
+def cached_input_tokens_from_log(log: RequestLogLike) -> int | None:
+    cached_tokens = log.cached_input_tokens
+    if cached_tokens is None:
+        return None
+    cached_tokens = max(0, int(cached_tokens))
+    input_tokens = log.input_tokens
+    if input_tokens is not None:
+        cached_tokens = min(cached_tokens, int(input_tokens))
+    return cached_tokens
+
+
+def usage_tokens_from_log(log: RequestLogLike) -> UsageTokens | None:
+    input_tokens = log.input_tokens
+    if input_tokens is None:
+        return None
+    output_tokens = output_tokens_from_log(log)
+    if output_tokens is None:
+        return None
+    cached_tokens = cached_input_tokens_from_log(log) or 0
+    return UsageTokens(
+        input_tokens=float(input_tokens),
+        output_tokens=float(output_tokens),
+        cached_input_tokens=float(cached_tokens),
+    )
+
+
+def output_tokens_from_log(log: RequestLogLike) -> int | None:
+    output_tokens = log.output_tokens
+    if output_tokens is not None:
+        return int(output_tokens)
+    reasoning_tokens = log.reasoning_tokens
+    if reasoning_tokens is None:
+        return None
+    return int(reasoning_tokens)
+
+
+def calculated_cost_from_log(log: RequestLogLike, *, precision: int | None = None) -> float | None:
+    if not log.model:
+        return None
+    usage = usage_tokens_from_log(log)
+    if not usage:
+        return None
+    resolved = get_pricing_for_model(log.model, None, None)
+    if not resolved:
+        return None
+    _, price = resolved
+    cost = calculate_cost_from_usage(usage, price, service_tier=log.service_tier)
+    if cost is None:
+        return None
+    if precision is None:
+        return cost
+    return round(cost, precision)
+
+
+def cost_from_log(log: RequestLogLike, *, precision: int | None = None) -> float | None:
+    cost = log.cost_usd
+    if cost is None:
+        return None
+    if precision is None:
+        return float(cost)
+    return round(float(cost), precision)
+
+
+def _totals_match(left: float | None, right: float | None, *, precision: int | None) -> bool:
+    if left is None or right is None:
+        return False
+    if precision is None:
+        return left == right
+    return abs(left - right) < (10 ** (-precision)) / 2
+
+
+def cost_breakdown_from_log(log: RequestLogLike, *, precision: int | None = None) -> UsageCostBreakdown:
+    full_breakdown: UsageCostBreakdown | None = None
+    input_usd: float | None = None
+    cached_input_usd: float | None = None
+    output_usd: float | None = None
+    raw_total_usd: float | None = None
+    total_usd: float | None = None
+    if log.model:
+        resolved = get_pricing_for_model(log.model, None, None)
+        if resolved is not None:
+            _, price = resolved
+            input_tokens = log.input_tokens
+            cached_tokens = cached_input_tokens_from_log(log)
+            output_tokens = output_tokens_from_log(log)
+            usage = usage_tokens_from_log(log)
+            if usage is not None:
+                raw_full_breakdown = calculate_cost_breakdown_from_usage(usage, price, service_tier=log.service_tier)
+                if raw_full_breakdown is not None:
+                    raw_total_usd = raw_full_breakdown.total_usd
+                full_breakdown = calculate_cost_breakdown_from_usage(
+                    usage,
+                    price,
+                    service_tier=log.service_tier,
+                    precision=precision,
+                )
+                if full_breakdown is not None:
+                    total_usd = full_breakdown.total_usd
+            if input_tokens is not None and cached_tokens is not None:
+                input_breakdown = calculate_cost_breakdown_from_usage(
+                    UsageTokens(
+                        input_tokens=float(input_tokens),
+                        output_tokens=0.0,
+                        cached_input_tokens=float(cached_tokens),
+                    ),
+                    price,
+                    service_tier=log.service_tier,
+                    precision=precision,
+                )
+                if input_breakdown is not None:
+                    input_usd = input_breakdown.input_usd
+                    cached_input_usd = input_breakdown.cached_input_usd
+            if output_tokens is not None:
+                output_breakdown = calculate_cost_breakdown_from_usage(
+                    UsageTokens(
+                        input_tokens=float(input_tokens or 0),
+                        output_tokens=float(output_tokens),
+                        cached_input_tokens=float(cached_tokens or 0),
+                    ),
+                    price,
+                    service_tier=log.service_tier,
+                    precision=precision,
+                )
+                if output_breakdown is not None:
+                    output_usd = output_breakdown.output_usd
+
+    persisted_cost = cost_from_log(log, precision=precision)
+    if persisted_cost is not None:
+        persisted_raw_cost = cost_from_log(log)
+        if not _totals_match(persisted_raw_cost, raw_total_usd, precision=precision):
+            return UsageCostBreakdown(
+                input_usd=None,
+                cached_input_usd=None,
+                output_usd=None,
+                total_usd=persisted_cost,
+            )
+        return UsageCostBreakdown(
+            input_usd=input_usd,
+            cached_input_usd=cached_input_usd,
+            output_usd=output_usd,
+            total_usd=persisted_cost,
+        )
+    if full_breakdown is not None:
+        return UsageCostBreakdown(
+            input_usd=input_usd,
+            cached_input_usd=cached_input_usd,
+            output_usd=output_usd,
+            total_usd=total_usd,
+        )
+    return UsageCostBreakdown(
+        input_usd=input_usd,
+        cached_input_usd=cached_input_usd,
+        output_usd=output_usd,
+        total_usd=None,
+    )
+
+
+def total_tokens_from_log(log: RequestLogLike) -> int | None:
+    input_tokens = log.input_tokens
+    output_tokens = output_tokens_from_log(log)
+    if input_tokens is None and output_tokens is None:
+        return None
+    return (input_tokens or 0) + (output_tokens or 0)

app/core/usage/metadata_scheduler.py

@@ -0,0 +1,118 @@
+"""Own upstream metadata refresh and bounded historical cost repair."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+
+from app.core.clients.codex_version import get_codex_version_cache
+from app.core.config.settings import get_settings
+from app.core.metadata_store import write_metadata
+from app.core.scheduling.leader_election_handle import get_leader_election
+from app.core.usage.pricing_catalog import (
+    BUNDLE_PATH,
+    decode_snapshot,
+    encode_snapshot,
+    fetch_catalogs,
+    get_active_prices,
+    install_prices,
+    merge_catalogs,
+    snapshot_updated_at,
+)
+from app.db.session import get_background_session
+from app.modules.request_logs.cost_backfill import backfill_missing_costs
+
+logger = logging.getLogger(__name__)
+_REFRESH_SECONDS = 3600
+_FAILURE_RETRY_SECONDS = 300
+
+
+@dataclass
+class MetadataRefreshScheduler:
+    interval_seconds: float = 5.0
+    _task: asyncio.Task[None] | None = None
+    _stop: asyncio.Event = field(default_factory=asyncio.Event)
+    _refresh_at: float = 0.0
+    _backfill_at: float = 0.0
+    _cursor: int = 0
+
+    async def start(self) -> None:
+        if self._task is not None and not self._task.done():
+            return
+        path = get_settings().data_dir / "pricing-cache.json"
+        try:
+            cached = json.loads(await asyncio.to_thread(path.read_text))
+            bundled = json.loads(BUNDLE_PATH.read_text())
+            prices = decode_snapshot(cached)
+            if snapshot_updated_at(cached) < snapshot_updated_at(bundled):
+                prices = merge_catalogs(decode_snapshot(bundled), prices)
+            install_prices(prices)
+        except (OSError, ValueError):
+            get_active_prices()
+        await get_codex_version_cache().restore(get_settings().data_dir / "codex-version-cache.json")
+        self._stop.clear()
+        self._task = asyncio.create_task(self._run_loop())
+
+    async def stop(self) -> None:
+        self._stop.set()
+        if self._task is not None:
+            self._task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
+
+    async def _refresh(self) -> None:
+        await get_codex_version_cache().get_version()
+        try:
+            prices = await fetch_catalogs()
+            previous = get_active_prices()
+            install_prices(prices)
+            if get_active_prices() != previous:
+                self._cursor = 0
+                self._backfill_at = 0.0
+            await asyncio.to_thread(
+                write_metadata, get_settings().data_dir / "pricing-cache.json", encode_snapshot(get_active_prices())
+            )
+            logger.info("Refreshed OpenAI pricing metadata: models=%d", len(prices))
+            self._refresh_at = time.monotonic() + _REFRESH_SECONDS
+        except (OSError, ValueError):
+            self._refresh_at = time.monotonic() + _FAILURE_RETRY_SECONDS
+            logger.warning("Metadata refresh failed; retaining last-good prices", exc_info=True)
+
+    async def _backfill(self) -> bool:
+        async with get_background_session() as session:
+            batch = await backfill_missing_costs(session, after_id=self._cursor)
+        self._cursor = batch.last_id
+        if batch.scanned < 200:
+            self._cursor = 0
+            self._backfill_at = time.monotonic() + _REFRESH_SECONDS
+        if batch.updated:
+            logger.info("Backfilled missing request costs: updated=%d last_id=%d", batch.updated, batch.last_id)
+        return True
+
+    async def _run_loop(self) -> None:
+        while not self._stop.is_set():
+            if time.monotonic() >= self._refresh_at:
+                try:
+                    await self._refresh()
+                except Exception:
+                    self._refresh_at = time.monotonic() + _FAILURE_RETRY_SECONDS
+                    logger.exception("Upstream metadata refresh failed")
+            if time.monotonic() >= self._backfill_at:
+                try:
+                    await get_leader_election().run_if_leader(self._backfill)
+                except Exception:
+                    self._backfill_at = time.monotonic() + _FAILURE_RETRY_SECONDS
+                    logger.exception("Missing-cost backfill failed")
+            try:
+                await asyncio.wait_for(self._stop.wait(), timeout=self.interval_seconds)
+            except TimeoutError:
+                pass
+
+
+def build_metadata_refresh_scheduler() -> MetadataRefreshScheduler:
+    return MetadataRefreshScheduler()

app/core/usage/models.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict
+
+
+class UsageWindow(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    used_percent: float | None = None
+    reset_at: int | None = None
+    limit_window_seconds: int | None = None
+    reset_after_seconds: int | None = None
+
+
+class RateLimitPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    primary_window: UsageWindow | None = None
+    secondary_window: UsageWindow | None = None
+
+
+class CreditsPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    has_credits: bool | None = None
+    unlimited: bool | None = None
+    balance: str | None = None
+
+
+class RateLimitResetCreditsPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    available_count: int | None = None
+
+
+class AdditionalRateLimitPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    limit_name: str
+    metered_feature: str
+    rate_limit: RateLimitPayload | None = None
+
+
+class UsagePayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    plan_type: str | None = None
+    workspace_id: str | None = None
+    workspace_label: str | None = None
+    seat_type: str | None = None
+    rate_limit: RateLimitPayload | None = None
+    credits: CreditsPayload | None = None
+    rate_limit_reset_credits: RateLimitResetCreditsPayload | None = None
+    additional_rate_limits: list[AdditionalRateLimitPayload] | None = None