codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/rate_limit_reset_credits/redeem_coordination.py

@@ -0,0 +1,274 @@
+"""Cross-replica coordination primitives for reset-credit redemption.
+
+Two shared-database primitives back the redeem path:
+
+- A durable idempotency ledger (``reset_credit_redeem_requests``) that pins the
+  (account_id, redeem_request_id) pair to the credit selected on the first
+  attempt, so a retry served by ANY replica retargets the same credit instead
+  of burning a second one.
+- A per-account claim row (``reset_credit_redeem_claims``) that serializes
+  redemption across processes sharing one SQLite file via a single atomic
+  conditional upsert with a lease the holder renews on a heartbeat while the
+  redeem section runs. PostgreSQL keeps ``pg_advisory_xact_lock``.
+
+All statements run on dedicated short-lived sessions committed immediately;
+they never join the caller's transaction.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from datetime import UTC, datetime, timedelta
+
+from sqlalchemy import delete, select, update
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+from app.db.models import ResetCreditRedeemClaim, ResetCreditRedeemRequest
+from app.db.session import SessionLocal, close_session
+
+logger = logging.getLogger(__name__)
+
+REDEEM_CLAIM_LEASE_SECONDS = 30.0
+REDEEM_CLAIM_RENEW_INTERVAL_SECONDS = 10.0
+REDEEM_CLAIM_RETRY_INTERVAL_SECONDS = 0.1
+REDEEM_CLAIM_TIMEOUT_SECONDS = 15.0
+REDEEM_REQUEST_TTL = timedelta(hours=24)
+
+
+class RedeemClaimTimeoutError(Exception):
+    """The per-account redeem claim stayed held past the acquisition timeout."""
+
+
+def new_redeem_claim_holder_id() -> str:
+    return uuid.uuid4().hex
+
+
+async def try_acquire_redeem_claim(
+    account_id: str,
+    holder_id: str,
+    *,
+    lease_seconds: float = REDEEM_CLAIM_LEASE_SECONDS,
+) -> bool:
+    """Attempt to claim the per-account redeem slot; True when claimed.
+
+    Uses one atomic ``INSERT ... ON CONFLICT(account_id) DO UPDATE ... WHERE
+    expires_at < now`` so only a missing or lease-expired claim can be taken —
+    the same conditional-upsert shape as the scheduler-leader lease.
+    """
+    now = datetime.now(UTC)
+    session = SessionLocal()
+    try:
+        insert_stmt = sqlite_insert(ResetCreditRedeemClaim).values(
+            account_id=account_id,
+            holder_id=holder_id,
+            expires_at=now + timedelta(seconds=lease_seconds),
+        )
+        stmt = insert_stmt.on_conflict_do_update(
+            index_elements=[ResetCreditRedeemClaim.account_id],
+            set_={
+                "holder_id": insert_stmt.excluded.holder_id,
+                "expires_at": insert_stmt.excluded.expires_at,
+            },
+            where=ResetCreditRedeemClaim.expires_at < now,
+        )
+        result = await session.execute(stmt.returning(ResetCreditRedeemClaim.account_id))
+        await session.commit()
+        return result.scalar_one_or_none() is not None
+    finally:
+        await close_session(session)
+
+
+async def acquire_redeem_claim(
+    account_id: str,
+    holder_id: str,
+    *,
+    lease_seconds: float = REDEEM_CLAIM_LEASE_SECONDS,
+    retry_interval_seconds: float = REDEEM_CLAIM_RETRY_INTERVAL_SECONDS,
+    timeout_seconds: float = REDEEM_CLAIM_TIMEOUT_SECONDS,
+) -> None:
+    """Acquire the per-account redeem claim, retrying until the timeout."""
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout_seconds
+    while True:
+        if await try_acquire_redeem_claim(account_id, holder_id, lease_seconds=lease_seconds):
+            return
+        remaining = deadline - loop.time()
+        if remaining <= 0:
+            raise RedeemClaimTimeoutError(
+                f"reset-credit redeem claim for account {account_id} not acquired within {timeout_seconds}s"
+            )
+        await asyncio.sleep(min(retry_interval_seconds, remaining))
+
+
+async def renew_redeem_claim(
+    account_id: str,
+    holder_id: str,
+    *,
+    lease_seconds: float = REDEEM_CLAIM_LEASE_SECONDS,
+) -> bool:
+    """Extend the holder's lease; False when the row is no longer held by this holder."""
+    now = datetime.now(UTC)
+    session = SessionLocal()
+    try:
+        result = await session.execute(
+            update(ResetCreditRedeemClaim)
+            .where(
+                ResetCreditRedeemClaim.account_id == account_id,
+                ResetCreditRedeemClaim.holder_id == holder_id,
+            )
+            .values(expires_at=now + timedelta(seconds=lease_seconds))
+            .returning(ResetCreditRedeemClaim.account_id)
+        )
+        await session.commit()
+        return result.scalar_one_or_none() is not None
+    finally:
+        await close_session(session)
+
+
+async def renew_redeem_claim_periodically(
+    account_id: str,
+    holder_id: str,
+    *,
+    lease_seconds: float = REDEEM_CLAIM_LEASE_SECONDS,
+    renew_interval_seconds: float = REDEEM_CLAIM_RENEW_INTERVAL_SECONDS,
+) -> None:
+    """Heartbeat that keeps a held claim's lease alive while the redeem section runs.
+
+    Mirrors the scheduler-leader renew pattern: the holder extends
+    ``expires_at`` every ``renew_interval_seconds`` (a fraction of the lease)
+    so a legitimately slow redemption — usage fetch retries plus the upstream
+    consume can exceed one lease — is not taken over by a second process.
+    Transient renewal errors are logged and retried on the next tick; a renew
+    reporting the row gone means an expired claim was already taken over, so
+    the loop stops (exclusivity is lost and lease takeover semantics apply).
+    The caller cancels this task before releasing the claim.
+    """
+    while True:
+        await asyncio.sleep(renew_interval_seconds)
+        try:
+            renewed = await renew_redeem_claim(account_id, holder_id, lease_seconds=lease_seconds)
+        except Exception:
+            logger.warning(
+                "reset-credit redeem claim renewal failed account_id=%s (retrying next tick)",
+                account_id,
+                exc_info=True,
+            )
+            continue
+        if not renewed:
+            logger.warning(
+                "reset-credit redeem claim lost before renewal account_id=%s holder_id=%s",
+                account_id,
+                holder_id,
+            )
+            return
+
+
+async def release_redeem_claim(account_id: str, holder_id: str) -> None:
+    """Release the claim if this holder still owns it (lease expiry is the backstop)."""
+    session = SessionLocal()
+    try:
+        await session.execute(
+            delete(ResetCreditRedeemClaim).where(
+                ResetCreditRedeemClaim.account_id == account_id,
+                ResetCreditRedeemClaim.holder_id == holder_id,
+            )
+        )
+        await session.commit()
+    except Exception:
+        logger.warning(
+            "reset-credit redeem claim release failed account_id=%s (lease expiry will recover)",
+            account_id,
+            exc_info=True,
+        )
+    finally:
+        await close_session(session)
+
+
+async def get_pinned_redeem_credit_id(account_id: str, redeem_request_id: str) -> str | None:
+    """Read the credit pinned to this redeem_request_id by any replica.
+
+    Rows older than the 24h TTL are ignored (the read TTL matches the purge
+    TTL applied by ``pin_redeem_request``), so an expired pin reads as absent.
+    The caller then re-selects a fresh credit and re-pins it via
+    ``pin_redeem_request`` instead of forwarding the redemption for a stale
+    credit id that the purge path would otherwise drop.
+    """
+    now = datetime.now(UTC)
+    session = SessionLocal()
+    try:
+        return await session.scalar(
+            select(ResetCreditRedeemRequest.credit_id).where(
+                ResetCreditRedeemRequest.account_id == account_id,
+                ResetCreditRedeemRequest.redeem_request_id == redeem_request_id,
+                ResetCreditRedeemRequest.created_at >= now - REDEEM_REQUEST_TTL,
+            )
+        )
+    finally:
+        await close_session(session)
+
+
+async def pin_redeem_request(account_id: str, redeem_request_id: str, credit_id: str) -> str:
+    """Durably pin the selected credit to this redeem request; first writer wins.
+
+    Returns the authoritative credit id (the previously pinned one on
+    conflict). Rows older than the 24h TTL for this account are purged in the
+    same transaction BEFORE the insert, so a ``redeem_request_id`` reused after
+    its prior row has aged past the TTL is re-pinned to the new credit instead
+    of colliding with the stale row via ``ON CONFLICT DO NOTHING`` and losing
+    the new pin.
+    """
+    now = datetime.now(UTC)
+    session = SessionLocal()
+    try:
+        values = {
+            "account_id": account_id,
+            "redeem_request_id": redeem_request_id,
+            "credit_id": credit_id,
+            "created_at": now,
+        }
+        # Purge expired rows first: an expired row for the SAME
+        # (account_id, redeem_request_id) would otherwise absorb the insert
+        # below via ON CONFLICT DO NOTHING and then be deleted, leaving the new
+        # attempt with no durable pin.
+        await session.execute(
+            delete(ResetCreditRedeemRequest).where(
+                ResetCreditRedeemRequest.account_id == account_id,
+                ResetCreditRedeemRequest.created_at < now - REDEEM_REQUEST_TTL,
+            )
+        )
+        dialect = session.get_bind().dialect.name
+        if dialect == "postgresql":
+            await session.execute(
+                pg_insert(ResetCreditRedeemRequest)
+                .values(**values)
+                .on_conflict_do_nothing(
+                    index_elements=[
+                        ResetCreditRedeemRequest.account_id,
+                        ResetCreditRedeemRequest.redeem_request_id,
+                    ]
+                )
+            )
+        else:
+            await session.execute(
+                sqlite_insert(ResetCreditRedeemRequest)
+                .values(**values)
+                .on_conflict_do_nothing(
+                    index_elements=[
+                        ResetCreditRedeemRequest.account_id,
+                        ResetCreditRedeemRequest.redeem_request_id,
+                    ]
+                )
+            )
+        await session.commit()
+        stored = await session.scalar(
+            select(ResetCreditRedeemRequest.credit_id).where(
+                ResetCreditRedeemRequest.account_id == account_id,
+                ResetCreditRedeemRequest.redeem_request_id == redeem_request_id,
+            )
+        )
+        return stored if stored is not None else credit_id
+    finally:
+        await close_session(session)

app/modules/rate_limit_reset_credits/store.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+import anyio
+
+from app.core.clients.rate_limit_reset_credits import (
+    RateLimitResetCreditsSnapshot,
+    ResetCreditItem,
+    _nearest_available_expires_at,
+)
+
+
+class RateLimitResetCreditsStore:
+    """In-memory cache of the most recent reset-credits snapshot per account.
+
+    Mirrors the lock-guarded shape of :class:`RateLimitHeadersCache` /
+    :class:`AccountSelectionCache`. Snapshots are keyed by account id and are
+    repopulated by each replica's refresh scheduler on every tick; reads from
+    the dashboard (GET + the AccountSummary mapper) never hit upstream.
+    """
+
+    def __init__(self) -> None:
+        self._snapshots: dict[str, RateLimitResetCreditsSnapshot] = {}
+        self._lock = anyio.Lock()
+        self._clear_generation = 0
+        self._account_generations: dict[str, int] = {}
+
+    async def set(self, account_id: str, snapshot: RateLimitResetCreditsSnapshot) -> None:
+        async with self._lock:
+            self._snapshots[account_id] = snapshot
+            self._bump_account_generation(account_id)
+
+    def generation(self, account_id: str) -> int:
+        return self._generation_for(account_id)
+
+    async def set_if_generation(
+        self,
+        account_id: str,
+        snapshot: RateLimitResetCreditsSnapshot,
+        expected_generation: int,
+    ) -> bool:
+        async with self._lock:
+            if self._generation_for(account_id) != expected_generation:
+                return False
+            self._snapshots[account_id] = snapshot
+            self._bump_account_generation(account_id)
+            return True
+
+    async def mark_credit_redeemed(
+        self,
+        account_id: str,
+        credit_id: str,
+        *,
+        redeemed_at: datetime | None,
+    ) -> None:
+        async with self._lock:
+            snapshot = self._snapshots.get(account_id)
+            if snapshot is None:
+                return
+            updated_credits, matched = _mark_credit_redeemed(snapshot.credits, credit_id, redeemed_at=redeemed_at)
+            if not matched:
+                return
+            available_count = sum(1 for credit in updated_credits if credit.status == "available")
+            self._snapshots[account_id] = snapshot.model_copy(
+                update={
+                    "available_count": available_count,
+                    "nearest_expires_at": _nearest_available_expires_at(updated_credits),
+                    "credits": updated_credits,
+                }
+            )
+            self._bump_account_generation(account_id)
+
+    def get(self, account_id: str) -> RateLimitResetCreditsSnapshot | None:
+        return self._snapshots.get(account_id)
+
+    async def invalidate(self, account_id: str | None = None) -> None:
+        async with self._lock:
+            if account_id is None:
+                self._snapshots.clear()
+                self._clear_generation += 1
+                return
+            self._snapshots.pop(account_id, None)
+            self._bump_account_generation(account_id)
+
+    def _generation_for(self, account_id: str) -> int:
+        return self._clear_generation + self._account_generations.get(account_id, 0)
+
+    def _bump_account_generation(self, account_id: str) -> None:
+        self._account_generations[account_id] = self._account_generations.get(account_id, 0) + 1
+
+
+_rate_limit_reset_credits_store = RateLimitResetCreditsStore()
+
+
+def get_rate_limit_reset_credits_store() -> RateLimitResetCreditsStore:
+    return _rate_limit_reset_credits_store
+
+
+def _mark_credit_redeemed(
+    credits: list[ResetCreditItem],
+    credit_id: str,
+    *,
+    redeemed_at: datetime | None,
+) -> tuple[list[ResetCreditItem], bool]:
+    matched = False
+    updated: list[ResetCreditItem] = []
+    for credit in credits:
+        if credit.id != credit_id:
+            updated.append(credit)
+            continue
+        matched = True
+        updated.append(credit.model_copy(update={"status": "redeemed", "redeemed_at": redeemed_at}))
+    return updated, matched

app/modules/reports/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

app/modules/reports/api.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from datetime import date
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, HTTPException, Query, status
+
+from app.core.auth.dependencies import (
+    set_dashboard_error_format,
+    validate_dashboard_session,
+)
+from app.core.exceptions import DashboardBadRequestError
+from app.dependencies import ReportsContext, get_reports_caches, get_reports_context
+from app.modules.reports.cache import ReportCacheKey, ReportsCaches
+from app.modules.reports.repository import DailyReportRangeTooLargeError
+from app.modules.reports.schemas import ReportsOptionsResponse, ReportsResponse, ThreadIdentityResponse
+from app.modules.reports.service import InvalidReportDateRangeError, resolve_report_range
+
+router = APIRouter(
+    prefix="/api/reports",
+    tags=["dashboard"],
+    dependencies=[Depends(validate_dashboard_session), Depends(set_dashboard_error_format)],
+)
+
+
+@router.get("", response_model=ReportsResponse)
+async def get_reports(
+    context: ReportsContext = Depends(get_reports_context),
+    caches: ReportsCaches = Depends(get_reports_caches),
+    start_date: Annotated[date | None, Query()] = None,
+    end_date: Annotated[date | None, Query()] = None,
+    report_timezone: Annotated[str | None, Query(alias="timezone")] = None,
+    account_id: Annotated[list[str] | None, Query()] = None,
+    api_key_id: Annotated[list[str] | None, Query()] = None,
+    model: Annotated[str | None, Query()] = None,
+    useragent_group: Annotated[str | None, Query()] = None,
+) -> ReportsResponse:
+    try:
+        key = _cache_key(start_date, end_date, report_timezone, account_id, api_key_id, model, useragent_group)
+        return await caches.reports.get(
+            key,
+            lambda: context.service.get_reports(
+                start_date=key.start,
+                end_date=key.end,
+                report_timezone=key.timezone,
+                account_ids=list(key.accounts) or None,
+                api_key_ids=list(key.api_keys) or None,
+                model=key.model or None,
+                useragent_group=key.useragent or None,
+            ),
+        )
+    except InvalidReportDateRangeError as exc:
+        raise DashboardBadRequestError(str(exc), code="invalid_report_date_range") from exc
+    except DailyReportRangeTooLargeError as exc:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc
+
+
+@router.get("/thread-identity", response_model=ThreadIdentityResponse)
+async def get_thread_identity(
+    context: ReportsContext = Depends(get_reports_context),
+    caches: ReportsCaches = Depends(get_reports_caches),
+    start_date: Annotated[date | None, Query()] = None,
+    end_date: Annotated[date | None, Query()] = None,
+    report_timezone: Annotated[str | None, Query(alias="timezone")] = None,
+) -> ThreadIdentityResponse:
+    """Thread identity and cache locality over the selected range, pool-wide.
+
+    The Reports filters are intentionally not accepted here: filtering by
+    account would pin every conversation to one account and make the
+    accounts-per-conversation factor meaningless.
+    """
+    try:
+        key = _cache_key(start_date, end_date, report_timezone, None, None)
+        return await caches.thread_identity.get(
+            key,
+            lambda: context.service.get_thread_identity(
+                start_date=key.start,
+                end_date=key.end,
+                report_timezone=key.timezone,
+            ),
+        )
+    except InvalidReportDateRangeError as exc:
+        raise DashboardBadRequestError(str(exc), code="invalid_report_date_range") from exc
+    except DailyReportRangeTooLargeError as exc:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc
+
+
+@router.get("/options", response_model=ReportsOptionsResponse)
+async def get_report_options(
+    context: ReportsContext = Depends(get_reports_context),
+    caches: ReportsCaches = Depends(get_reports_caches),
+    start_date: Annotated[date | None, Query()] = None,
+    end_date: Annotated[date | None, Query()] = None,
+    report_timezone: Annotated[str | None, Query(alias="timezone")] = None,
+    account_id: Annotated[list[str] | None, Query()] = None,
+    api_key_id: Annotated[list[str] | None, Query()] = None,
+) -> ReportsOptionsResponse:
+    try:
+        key = _cache_key(start_date, end_date, report_timezone, account_id, api_key_id)
+        return await caches.options.get(
+            key,
+            lambda: context.service.get_options(
+                start_date=key.start,
+                end_date=key.end,
+                report_timezone=key.timezone,
+                account_ids=list(key.accounts) or None,
+                api_key_ids=list(key.api_keys) or None,
+            ),
+        )
+    except InvalidReportDateRangeError as exc:
+        raise DashboardBadRequestError(str(exc), code="invalid_report_date_range") from exc
+    except DailyReportRangeTooLargeError as exc:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc
+
+
+def _cache_key(
+    start: date | None,
+    end: date | None,
+    timezone: str | None,
+    accounts: list[str] | None,
+    api_keys: list[str] | None,
+    model: str | None = None,
+    useragent: str | None = None,
+) -> ReportCacheKey:
+    start, end, tz = resolve_report_range(start, end, timezone)
+    return ReportCacheKey(
+        start,
+        end,
+        str(tz),
+        tuple(sorted(set(accounts or []))),
+        tuple(sorted(set(api_keys or []))),
+        model or "",
+        useragent or "",
+    )

app/modules/reports/cache.py

@@ -0,0 +1,74 @@
+"""Small process-local report caches; one owned computation per cache at a time."""
+
+from __future__ import annotations
+
+import asyncio
+from collections import OrderedDict
+from collections.abc import Awaitable, Callable, Hashable
+from dataclasses import dataclass
+from datetime import date
+from time import monotonic
+
+from app.modules.reports.schemas import ReportsOptionsResponse, ReportsResponse, ThreadIdentityResponse
+
+
+@dataclass(frozen=True)
+class ReportCacheKey:
+    start: date
+    end: date
+    timezone: str
+    accounts: tuple[str, ...]
+    api_keys: tuple[str, ...]
+    model: str
+    useragent: str
+
+
+class ReportCache[K: Hashable, T]:
+    def __init__(self, *, ttl_seconds: float = 60, max_entries: int = 64) -> None:
+        self._ttl = ttl_seconds
+        self._capacity = max_entries
+        self._entries: OrderedDict[K, tuple[float, T]] = OrderedDict()
+        self._compute_slot = asyncio.Semaphore(1)
+
+    def _cached(self, key: K) -> tuple[float, T] | None:
+        # All state access runs synchronously on the application's event loop.
+        # No await occurs while inspecting, pruning or publishing entries.
+        now = monotonic()
+        for expired in [key for key, (expires, _) in self._entries.items() if expires <= now]:
+            del self._entries[expired]
+        if key in self._entries:
+            self._entries.move_to_end(key)
+            return self._entries[key]
+        return None
+
+    async def get(self, key: K, compute: Callable[[], Awaitable[T]]) -> T:
+        cached = self._cached(key)
+        if cached is not None:
+            return cached[1]
+        # Cache hits bypass the compute limit. Misses recheck after admission
+        # so identical requests reuse the completed result. The caller owns
+        # computation and its AsyncSession; cancellation releases the slot.
+        async with self._compute_slot:
+            cached = self._cached(key)
+            if cached is not None:
+                return cached[1]
+            value = await compute()
+            self._entries[key] = (monotonic() + self._ttl, value)
+            while len(self._entries) > self._capacity:
+                self._entries.popitem(last=False)
+            return value
+
+
+# Thread identity scans raw request logs instead of the hourly rollup, so a
+# miss is seconds rather than milliseconds. The figures move slowly (they are
+# whole-window aggregates), so they tolerate a longer TTL than the cost report.
+_THREAD_IDENTITY_TTL_SECONDS = 300.0
+
+
+class ReportsCaches:
+    def __init__(self) -> None:
+        self.reports = ReportCache[ReportCacheKey, ReportsResponse]()
+        self.options = ReportCache[ReportCacheKey, ReportsOptionsResponse]()
+        self.thread_identity = ReportCache[ReportCacheKey, ThreadIdentityResponse](
+            ttl_seconds=_THREAD_IDENTITY_TTL_SECONDS
+        )

app/modules/reports/filters.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from sqlalchemy import and_, or_
+
+from app.db.models import RequestLog
+
+_INTERNAL_LIMIT_WARMUP_SOURCE = "limit_warmup"
+_INTERNAL_WARMUP_REQUEST_KINDS = ("warmup", "limit_warmup")
+MISSING_USERAGENT_GROUP = "Missing User-Agent"
+
+
+def _useragent_group_filter_clause(useragent_group: str | None):
+    if not useragent_group:
+        return None
+    if useragent_group == MISSING_USERAGENT_GROUP:
+        return RequestLog.useragent_group.is_(None)
+    return RequestLog.useragent_group == useragent_group
+
+
+def _normal_traffic_clause():
+    return and_(
+        or_(RequestLog.source.is_(None), RequestLog.source != _INTERNAL_LIMIT_WARMUP_SOURCE),
+        or_(
+            RequestLog.request_kind.is_(None),
+            RequestLog.request_kind.not_in(_INTERNAL_WARMUP_REQUEST_KINDS),
+        ),
+    )