codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/api_keys/api.py

@@ -0,0 +1,344 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Body, Depends, Request, Response
+
+from app.core.audit.service import AuditActor, AuditService, AuditTarget
+from app.core.auth.dashboard_access import DashboardPrincipal, Permission
+from app.core.auth.dependencies import (
+    require_dashboard_permission,
+    require_dashboard_write_access,
+    set_dashboard_error_format,
+    validate_dashboard_session,
+)
+from app.core.exceptions import DashboardBadRequestError, DashboardConflictError, DashboardNotFoundError
+from app.dependencies import ApiKeysContext, get_api_keys_context
+from app.modules.api_keys.repository import ApiKeyOwnerDisabledError
+from app.modules.api_keys.schemas import (
+    ApiKeyAccountCostResponse,
+    ApiKeyCreateRequest,
+    ApiKeyCreateResponse,
+    ApiKeyResponse,
+    ApiKeyTrendsResponse,
+    ApiKeyUpdateRequest,
+    ApiKeyUsage7DayResponse,
+    ApiKeyUsageSummaryResponse,
+    LimitRuleResponse,
+)
+from app.modules.api_keys.service import (
+    ApiKeyCreateData,
+    ApiKeyData,
+    ApiKeyNotFoundError,
+    ApiKeyUpdateData,
+    ApiKeyValidationError,
+    LimitRuleInput,
+)
+
+router = APIRouter(
+    prefix="/api/api-keys",
+    tags=["dashboard"],
+    dependencies=[
+        Depends(validate_dashboard_session),
+        # The key inventory (policies, assignments, limits, usage) is not a
+        # guest-safe read; every route here needs api_keys:read at minimum.
+        Depends(require_dashboard_permission(Permission.API_KEYS_READ)),
+        Depends(set_dashboard_error_format),
+    ],
+)
+
+
+def _to_response(row: ApiKeyData) -> ApiKeyResponse:
+    return ApiKeyResponse(
+        id=row.id,
+        name=row.name,
+        key_prefix=row.key_prefix,
+        allowed_models=row.allowed_models,
+        apply_to_codex_model=row.apply_to_codex_model,
+        enforced_model=row.enforced_model,
+        enforced_reasoning_effort=row.enforced_reasoning_effort,
+        allowed_reasoning_efforts=row.allowed_reasoning_efforts,
+        enforced_service_tier=row.enforced_service_tier,
+        traffic_class=row.traffic_class,
+        transport_policy_override=row.transport_policy_override,
+        thread_cache_identity_override=row.thread_cache_identity_override,
+        usage_sections=row.usage_sections,
+        expires_at=row.expires_at,
+        is_active=row.is_active,
+        account_assignment_scope_enabled=row.account_assignment_scope_enabled,
+        source_assignment_scope_enabled=row.source_assignment_scope_enabled,
+        assigned_account_ids=row.assigned_account_ids,
+        assigned_source_ids=row.assigned_source_ids,
+        created_at=row.created_at,
+        last_used_at=row.last_used_at,
+        limits=[
+            LimitRuleResponse(
+                id=li.id,
+                limit_type=li.limit_type,
+                limit_window=li.limit_window,
+                max_value=li.max_value,
+                current_value=li.current_value,
+                model_filter=li.model_filter,
+                reset_at=li.reset_at,
+            )
+            for li in row.limits
+        ],
+        usage_summary=(
+            ApiKeyUsageSummaryResponse(
+                request_count=row.usage_summary.request_count,
+                total_tokens=row.usage_summary.total_tokens,
+                cached_input_tokens=row.usage_summary.cached_input_tokens,
+                total_cost_usd=row.usage_summary.total_cost_usd,
+            )
+            if row.usage_summary is not None
+            else None
+        ),
+        pooled_remaining_percent_primary=(row.pooled_credits.remaining_percent_primary if row.pooled_credits else None),
+        pooled_remaining_percent_secondary=(
+            row.pooled_credits.remaining_percent_secondary if row.pooled_credits else None
+        ),
+        pooled_capacity_credits_primary=(row.pooled_credits.capacity_credits_primary if row.pooled_credits else 0.0),
+    )
+
+
+def _build_limit_inputs(payload: ApiKeyCreateRequest | ApiKeyUpdateRequest) -> list[LimitRuleInput]:
+    limit_inputs: list[LimitRuleInput] = []
+
+    if hasattr(payload, "limits") and payload.limits is not None:
+        for lr in payload.limits:
+            limit_inputs.append(
+                LimitRuleInput(
+                    limit_type=lr.limit_type,
+                    limit_window=lr.limit_window,
+                    max_value=lr.max_value,
+                    model_filter=lr.model_filter,
+                )
+            )
+    elif (
+        hasattr(payload, "weekly_token_limit")
+        and "weekly_token_limit" in payload.model_fields_set
+        and payload.weekly_token_limit is not None
+    ):
+        # Legacy: convert weeklyTokenLimit to a limit rule
+        limit_inputs.append(
+            LimitRuleInput(
+                limit_type="total_tokens",
+                limit_window="weekly",
+                max_value=payload.weekly_token_limit,
+            )
+        )
+
+    return limit_inputs
+
+
+@router.post("", response_model=ApiKeyCreateResponse, include_in_schema=False)
+@router.post("/", response_model=ApiKeyCreateResponse)
+async def create_api_key(
+    request: Request,
+    response: Response,
+    payload: ApiKeyCreateRequest = Body(...),
+    principal: DashboardPrincipal = Depends(require_dashboard_write_access),
+    context: ApiKeysContext = Depends(get_api_keys_context),
+) -> ApiKeyCreateResponse:
+    limit_inputs = _build_limit_inputs(payload)
+
+    try:
+        created = await context.service.create_key(
+            ApiKeyCreateData(
+                name=payload.name,
+                allowed_models=payload.allowed_models,
+                apply_to_codex_model=payload.apply_to_codex_model,
+                enforced_model=payload.enforced_model,
+                enforced_reasoning_effort=payload.enforced_reasoning_effort,
+                allowed_reasoning_efforts=payload.allowed_reasoning_efforts,
+                enforced_service_tier=payload.enforced_service_tier,
+                traffic_class=payload.traffic_class or "foreground",
+                transport_policy_override=payload.transport_policy_override,
+                thread_cache_identity_override=payload.thread_cache_identity_override,
+                usage_sections=(
+                    payload.usage_sections
+                    if payload.usage_sections is not None
+                    else "upstream_limits,account_pool_usage"
+                ),
+                expires_at=payload.expires_at,
+                assigned_account_ids=payload.assigned_account_ids,
+                assigned_source_ids=payload.assigned_source_ids,
+                limits=limit_inputs,
+            )
+        )
+    except ApiKeyValidationError as exc:
+        raise DashboardBadRequestError(str(exc), code="invalid_api_key_payload") from exc
+    resp = _to_response(created)
+    AuditService.log_async(
+        "api_key_created",
+        actor_ip=request.client.host if request.client else None,
+        actor=AuditActor.from_principal(principal),
+        target=AuditTarget("api_key", created.id),
+        details={"key_id": created.id},
+    )
+    response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, private"
+    response.headers["Pragma"] = "no-cache"
+    response.headers["Expires"] = "0"
+    return ApiKeyCreateResponse(
+        **resp.model_dump(),
+        key=created.key,
+    )
+
+
+@router.get("", response_model=list[ApiKeyResponse], include_in_schema=False)
+@router.get("/", response_model=list[ApiKeyResponse])
+async def list_api_keys(
+    context: ApiKeysContext = Depends(get_api_keys_context),
+) -> list[ApiKeyResponse]:
+    rows = await context.service.list_keys()
+    return [_to_response(row) for row in rows]
+
+
+@router.patch("/{key_id}", response_model=ApiKeyResponse)
+async def update_api_key(
+    request: Request,
+    key_id: str,
+    payload: ApiKeyUpdateRequest = Body(...),
+    principal: DashboardPrincipal = Depends(require_dashboard_write_access),
+    context: ApiKeysContext = Depends(get_api_keys_context),
+) -> ApiKeyResponse:
+    fields = payload.model_fields_set
+
+    limits_set = "limits" in fields or "weekly_token_limit" in fields
+    limit_inputs = _build_limit_inputs(payload) if limits_set else None
+
+    update = ApiKeyUpdateData(
+        name=payload.name,
+        name_set="name" in fields,
+        allowed_models=payload.allowed_models,
+        allowed_models_set="allowed_models" in fields,
+        apply_to_codex_model=payload.apply_to_codex_model,
+        apply_to_codex_model_set="apply_to_codex_model" in fields,
+        enforced_model=payload.enforced_model,
+        enforced_model_set="enforced_model" in fields,
+        enforced_reasoning_effort=payload.enforced_reasoning_effort,
+        enforced_reasoning_effort_set="enforced_reasoning_effort" in fields,
+        allowed_reasoning_efforts=payload.allowed_reasoning_efforts,
+        allowed_reasoning_efforts_set="allowed_reasoning_efforts" in fields,
+        enforced_service_tier=payload.enforced_service_tier,
+        enforced_service_tier_set="enforced_service_tier" in fields,
+        traffic_class=payload.traffic_class,
+        traffic_class_set="traffic_class" in fields,
+        transport_policy_override=payload.transport_policy_override,
+        transport_policy_override_set="transport_policy_override" in fields,
+        thread_cache_identity_override=payload.thread_cache_identity_override,
+        thread_cache_identity_override_set="thread_cache_identity_override" in fields,
+        usage_sections=payload.usage_sections,
+        usage_sections_set="usage_sections" in fields,
+        expires_at=payload.expires_at,
+        expires_at_set="expires_at" in fields,
+        is_active=payload.is_active,
+        is_active_set="is_active" in fields,
+        assigned_account_ids=payload.assigned_account_ids,
+        assigned_account_ids_set="assigned_account_ids" in fields,
+        assigned_source_ids=payload.assigned_source_ids,
+        assigned_source_ids_set="assigned_source_ids" in fields,
+        limits=limit_inputs,
+        limits_set=limits_set,
+        reset_usage=bool(payload.reset_usage),
+    )
+    try:
+        row = await context.service.update_key(key_id, update)
+    except ApiKeyNotFoundError as exc:
+        raise DashboardNotFoundError(str(exc)) from exc
+    except ApiKeyValidationError as exc:
+        raise DashboardBadRequestError(str(exc), code="invalid_api_key_payload") from exc
+    except ApiKeyOwnerDisabledError as exc:
+        raise DashboardConflictError(str(exc), code="owner_disabled") from exc
+    if "is_active" in fields and payload.is_active is False and row.is_active is False:
+        AuditService.log_async(
+            "api_key_revoked",
+            actor_ip=request.client.host if request.client else None,
+            actor=AuditActor.from_principal(principal),
+            target=AuditTarget("api_key", key_id),
+            details={"key_id": row.id},
+        )
+    return _to_response(row)
+
+
+@router.delete("/{key_id}")
+async def delete_api_key(
+    request: Request,
+    key_id: str,
+    principal: DashboardPrincipal = Depends(require_dashboard_write_access),
+    context: ApiKeysContext = Depends(get_api_keys_context),
+) -> Response:
+    try:
+        await context.service.delete_key(key_id)
+    except ApiKeyNotFoundError as exc:
+        raise DashboardNotFoundError(str(exc)) from exc
+    AuditService.log_async(
+        "api_key_revoked",
+        actor_ip=request.client.host if request.client else None,
+        actor=AuditActor.from_principal(principal),
+        target=AuditTarget("api_key", key_id),
+        details={"key_id": key_id},
+    )
+    return Response(status_code=204)
+
+
+@router.post("/{key_id}/regenerate", response_model=ApiKeyCreateResponse)
+async def regenerate_api_key(
+    key_id: str,
+    response: Response,
+    _write_access=Depends(require_dashboard_write_access),
+    context: ApiKeysContext = Depends(get_api_keys_context),
+) -> ApiKeyCreateResponse:
+    try:
+        row = await context.service.regenerate_key(key_id)
+    except ApiKeyNotFoundError as exc:
+        raise DashboardNotFoundError(str(exc)) from exc
+    resp = _to_response(row)
+    response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, private"
+    response.headers["Pragma"] = "no-cache"
+    response.headers["Expires"] = "0"
+    return ApiKeyCreateResponse(
+        **resp.model_dump(),
+        key=row.key,
+    )
+
+
+@router.get("/{key_id}/trends", response_model=ApiKeyTrendsResponse)
+async def get_api_key_trends(
+    key_id: str,
+    context: ApiKeysContext = Depends(get_api_keys_context),
+) -> ApiKeyTrendsResponse:
+    from app.modules.api_keys.schemas import ApiKeyTrendPoint
+
+    result = await context.service.get_key_trends(key_id)
+    if result is None:
+        raise DashboardNotFoundError(f"API key not found: {key_id}")
+    return ApiKeyTrendsResponse(
+        key_id=result.key_id,
+        cost=[ApiKeyTrendPoint(t=p.t, v=p.v) for p in result.cost],
+        tokens=[ApiKeyTrendPoint(t=p.t, v=p.v) for p in result.tokens],
+    )
+
+
+@router.get("/{key_id}/usage-7d", response_model=ApiKeyUsage7DayResponse)
+async def get_api_key_usage_7d(
+    key_id: str,
+    context: ApiKeysContext = Depends(get_api_keys_context),
+) -> ApiKeyUsage7DayResponse:
+    result = await context.service.get_key_usage_7d(key_id)
+    if result is None:
+        raise DashboardNotFoundError(f"API key not found: {key_id}")
+    return ApiKeyUsage7DayResponse(
+        key_id=result.key_id,
+        total_tokens=result.total_tokens,
+        total_cost_usd=result.total_cost_usd,
+        total_requests=result.total_requests,
+        cached_input_tokens=result.cached_input_tokens,
+        account_costs=[
+            ApiKeyAccountCostResponse(
+                account_id=ac.account_id,
+                email=ac.email,
+                cost_usd=ac.cost_usd,
+                is_deleted=ac.is_deleted,
+            )
+            for ac in result.account_costs
+        ],
+    )

app/modules/api_keys/last_used_coalescer.py

@@ -0,0 +1,216 @@
+"""Write-behind coalescing for ``api_keys.last_used_at``.
+
+Settlement paths record touches into a process-local map instead of writing
+the column inside every reservation-settlement transaction (production showed
+~616 such commits per 10 minutes for only 7 active keys). A replica-local
+periodic flusher folds the pending map into the database in one transaction —
+at most one guarded UPDATE per touched key per interval.
+
+The flush is monotonic (greatest-wins): ``WHERE last_used_at IS NULL OR
+last_used_at < :new`` never moves the stored value backwards, so concurrent
+replicas can flush out of order without leader election. On a hard crash at
+most one flush interval of ``last_used_at`` freshness is lost — accepted, the
+column is display-only (dashboard ``lastUsedAt``).
+
+Graceful shutdown is lossless: the scheduler's final flush retries transient
+failures a bounded number of times (logging the pending keys and timestamps
+at WARNING if every attempt fails), and the coalescer enters write-through
+mode before that final flush so a touch recorded by a settlement task that
+outlived the shutdown drain flushes itself immediately instead of parking in
+a pending map that no longer has a flusher.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from datetime import datetime
+
+from sqlalchemy import or_, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.models import ApiKey
+from app.db.session import get_background_session, sqlite_writer_section
+
+logger = logging.getLogger(__name__)
+
+FLUSH_INTERVAL_SECONDS = 30
+SHUTDOWN_FLUSH_ATTEMPTS = 3
+SHUTDOWN_FLUSH_RETRY_DELAY_SECONDS = 0.5
+
+
+class ApiKeyLastUsedCoalescer:
+    """Process-local pending map of the newest observed used-at per API key.
+
+    Only touched from the single asyncio event loop, so plain dict mutation
+    is safe; ``record`` keeps the per-key maximum and ``flush`` swaps the map
+    before writing so records arriving mid-flush land in the next interval.
+
+    In shutdown write-through mode (entered by the flush scheduler before its
+    final flush) every ``record`` immediately flushes the pending map itself:
+    the periodic loop is gone by then, so parking the touch would lose it at
+    process exit.
+    """
+
+    def __init__(self) -> None:
+        self._pending: dict[str, datetime] = {}
+        self._shutdown_write_through = False
+
+    async def record(self, key_id: str, used_at: datetime) -> None:
+        self._merge(key_id, used_at)
+        if self._shutdown_write_through:
+            await self.flush_with_retries()
+
+    def set_shutdown_write_through(self, enabled: bool) -> None:
+        self._shutdown_write_through = enabled
+
+    def pending_snapshot(self) -> dict[str, datetime]:
+        return dict(self._pending)
+
+    def clear(self) -> None:
+        self._pending.clear()
+        self._shutdown_write_through = False
+
+    async def flush(self) -> int:
+        """Write all pending touches in one transaction; returns keys flushed.
+
+        On failure the swapped batch is merged back into the pending map
+        (per-key maximum wins) so touches survive until a later flush. The
+        restore intentionally catches ``BaseException``: cancellation raises
+        ``asyncio.CancelledError`` (a ``BaseException`` since Python 3.8), and
+        a cancelled in-flight flush must not drop the swapped batch.
+        """
+        if not self._pending:
+            return 0
+        batch, self._pending = self._pending, {}
+        try:
+            async with sqlite_writer_section(), get_background_session() as session:
+                await self._apply_batch(session, batch)
+                await session.commit()
+        except BaseException:
+            for key_id, used_at in batch.items():
+                self._merge(key_id, used_at)
+            raise
+        return len(batch)
+
+    async def flush_with_retries(self) -> None:
+        """Shutdown-path flush: bounded retries, never raises (except cancel).
+
+        Used for the graceful-shutdown final flush and for write-through
+        records after it — in both cases no periodic tick remains to pick up
+        a retained batch, so a transient DB error must be retried here. If
+        every attempt fails the pending keys and timestamps are logged at
+        WARNING so operators can reconstruct the lost values.
+        """
+        for attempt in range(1, SHUTDOWN_FLUSH_ATTEMPTS + 1):
+            try:
+                await self.flush()
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                if attempt == SHUTDOWN_FLUSH_ATTEMPTS:
+                    logger.warning(
+                        "Final API key last_used_at flush failed after %d attempts; unflushed touches: %s",
+                        SHUTDOWN_FLUSH_ATTEMPTS,
+                        {key_id: used_at.isoformat() for key_id, used_at in sorted(self._pending.items())},
+                        exc_info=True,
+                    )
+                    return
+                logger.warning(
+                    "Final API key last_used_at flush failed (attempt %d/%d); retrying",
+                    attempt,
+                    SHUTDOWN_FLUSH_ATTEMPTS,
+                    exc_info=True,
+                )
+                await asyncio.sleep(SHUTDOWN_FLUSH_RETRY_DELAY_SECONDS)
+            else:
+                return
+
+    def _merge(self, key_id: str, used_at: datetime) -> None:
+        current = self._pending.get(key_id)
+        if current is None or current < used_at:
+            self._pending[key_id] = used_at
+
+    @staticmethod
+    async def _apply_batch(session: AsyncSession, batch: dict[str, datetime]) -> None:
+        for key_id, used_at in batch.items():
+            await session.execute(
+                update(ApiKey)
+                .where(ApiKey.id == key_id)
+                .where(or_(ApiKey.last_used_at.is_(None), ApiKey.last_used_at < used_at))
+                .values(last_used_at=used_at)
+            )
+
+
+class ApiKeyLastUsedFlushScheduler:
+    """Replica-local periodic flush loop for the last-used coalescer.
+
+    Every replica flushes its own process-local pending map, so this loop
+    MUST NOT be leader-gated; the monotonic guarded UPDATE keeps concurrent
+    replica flushes safe. ``stop()`` signals the loop to exit and awaits it —
+    it MUST NOT cancel the task, because cancelling an in-flight ``flush()``
+    would inject ``CancelledError`` into the awaited DB call after the pending
+    map was already swapped out — then switches the coalescer to shutdown
+    write-through mode and performs the final graceful-shutdown flush with
+    bounded retries. Write-through is entered BEFORE the final flush so there
+    is no window in which a late producer (e.g. a settlement task that
+    outlived the shutdown drain) can park a touch that nothing will flush.
+    """
+
+    def __init__(
+        self,
+        coalescer: ApiKeyLastUsedCoalescer,
+        *,
+        interval_seconds: float = FLUSH_INTERVAL_SECONDS,
+    ) -> None:
+        self._coalescer = coalescer
+        self._interval_seconds = interval_seconds
+        self._task: asyncio.Task[None] | None = None
+        self._stop = asyncio.Event()
+
+    async def start(self) -> None:
+        if self._task and not self._task.done():
+            return
+        self._coalescer.set_shutdown_write_through(False)
+        self._stop.clear()
+        self._task = asyncio.create_task(self._run_loop())
+
+    async def stop(self) -> None:
+        if self._task:
+            self._stop.set()
+            # No cancel(): the stop event wakes wait_for immediately and an
+            # in-flight flush runs to completion instead of losing its batch.
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
+        self._coalescer.set_shutdown_write_through(True)
+        await self._coalescer.flush_with_retries()
+
+    async def _run_loop(self) -> None:
+        while not await self._wait_or_stop(self._interval_seconds):
+            await self._flush_once()
+
+    async def _wait_or_stop(self, delay_seconds: float) -> bool:
+        try:
+            await asyncio.wait_for(self._stop.wait(), timeout=delay_seconds)
+        except asyncio.TimeoutError:
+            return False
+        return True
+
+    async def _flush_once(self) -> None:
+        try:
+            await self._coalescer.flush()
+        except Exception:
+            logger.exception("API key last_used_at flush failed; touches retained for the next flush")
+
+
+_api_key_last_used_coalescer = ApiKeyLastUsedCoalescer()
+
+
+def get_api_key_last_used_coalescer() -> ApiKeyLastUsedCoalescer:
+    return _api_key_last_used_coalescer
+
+
+def build_api_key_last_used_flush_scheduler() -> ApiKeyLastUsedFlushScheduler:
+    return ApiKeyLastUsedFlushScheduler(get_api_key_last_used_coalescer())

app/modules/api_keys/limit_windows.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+
+from app.db.models import LimitWindow
+
+
+def next_limit_reset(now: datetime, window: LimitWindow) -> datetime:
+    if window == LimitWindow.FIVE_HOURS:
+        return now + timedelta(hours=5)
+    if window == LimitWindow.SEVEN_DAYS:
+        return now + timedelta(days=7)
+    if window == LimitWindow.DAILY:
+        return now + timedelta(days=1)
+    if window == LimitWindow.WEEKLY:
+        return now + timedelta(days=7)
+    if window == LimitWindow.MONTHLY:
+        return now + timedelta(days=30)
+    return now + timedelta(days=7)
+
+
+def advance_limit_reset(reset_at: datetime, now: datetime, window: LimitWindow) -> datetime:
+    delta = limit_window_delta(window)
+    next_reset = reset_at
+    while next_reset <= now:
+        next_reset += delta
+    return next_reset
+
+
+def limit_window_delta(window: LimitWindow) -> timedelta:
+    if window == LimitWindow.FIVE_HOURS:
+        return timedelta(hours=5)
+    if window == LimitWindow.SEVEN_DAYS:
+        return timedelta(days=7)
+    if window == LimitWindow.DAILY:
+        return timedelta(days=1)
+    if window == LimitWindow.WEEKLY:
+        return timedelta(days=7)
+    if window == LimitWindow.MONTHLY:
+        return timedelta(days=30)
+    return timedelta(days=7)