codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/firewall/__init__.py

@@ -0,0 +1 @@
+"""Firewall dashboard module."""

app/modules/firewall/api.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Body, Depends
+
+from app.core.auth.dashboard_access import Permission
+from app.core.auth.dependencies import (
+    require_dashboard_permission,
+    set_dashboard_error_format,
+    validate_dashboard_session,
+)
+from app.core.cache.invalidation import NAMESPACE_FIREWALL, get_cache_invalidation_poller
+from app.core.exceptions import DashboardBadRequestError, DashboardConflictError, DashboardNotFoundError
+from app.core.middleware.firewall_cache import get_firewall_ip_cache
+from app.dependencies import FirewallContext, get_firewall_context
+from app.modules.firewall.schemas import (
+    FirewallDeleteResponse,
+    FirewallIpCreateRequest,
+    FirewallIpEntry,
+    FirewallIpsResponse,
+)
+from app.modules.firewall.service import FirewallIpAlreadyExistsError, FirewallValidationError
+
+router = APIRouter(
+    prefix="/api/firewall",
+    tags=["dashboard"],
+    dependencies=[Depends(validate_dashboard_session), Depends(set_dashboard_error_format)],
+)
+
+
+@router.get("/ips", response_model=FirewallIpsResponse)
+async def list_firewall_ips(
+    context: FirewallContext = Depends(get_firewall_context),
+) -> FirewallIpsResponse:
+    payload = await context.service.list_ips()
+    return FirewallIpsResponse(
+        mode=payload.mode,
+        entries=[
+            FirewallIpEntry(ip_address=entry.ip_address, created_at=entry.created_at) for entry in payload.entries
+        ],
+    )
+
+
+@router.post("/ips", response_model=FirewallIpEntry)
+async def add_firewall_ip(
+    payload: FirewallIpCreateRequest = Body(...),
+    _security_access=Depends(require_dashboard_permission(Permission.SECURITY_WRITE)),
+    context: FirewallContext = Depends(get_firewall_context),
+) -> FirewallIpEntry:
+    try:
+        created = await context.service.add_ip(payload.ip_address)
+    except FirewallValidationError as exc:
+        raise DashboardBadRequestError(str(exc), code="invalid_ip") from exc
+    except FirewallIpAlreadyExistsError as exc:
+        raise DashboardConflictError(str(exc), code="ip_exists") from exc
+    get_firewall_ip_cache().invalidate_all()
+    poller = get_cache_invalidation_poller()
+    if poller is not None:
+        await poller.bump(NAMESPACE_FIREWALL)
+    return FirewallIpEntry(ip_address=created.ip_address, created_at=created.created_at)
+
+
+@router.delete("/ips/{ip_address}", response_model=FirewallDeleteResponse)
+async def delete_firewall_ip(
+    ip_address: str,
+    _security_access=Depends(require_dashboard_permission(Permission.SECURITY_WRITE)),
+    context: FirewallContext = Depends(get_firewall_context),
+) -> FirewallDeleteResponse:
+    try:
+        deleted = await context.service.remove_ip(ip_address)
+    except FirewallValidationError as exc:
+        raise DashboardBadRequestError(str(exc), code="invalid_ip") from exc
+    if not deleted:
+        raise DashboardNotFoundError("IP address not found", code="ip_not_found")
+    get_firewall_ip_cache().invalidate_all()
+    poller = get_cache_invalidation_poller()
+    if poller is not None:
+        await poller.bump(NAMESPACE_FIREWALL)
+    return FirewallDeleteResponse(status="deleted")

app/modules/firewall/repository.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.models import ApiFirewallAllowlist
+
+
+class FirewallRepositoryConflictError(ValueError):
+    pass
+
+
+class FirewallRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def list_entries(self) -> Sequence[ApiFirewallAllowlist]:
+        result = await self._session.execute(
+            select(ApiFirewallAllowlist).order_by(ApiFirewallAllowlist.created_at, ApiFirewallAllowlist.ip_address)
+        )
+        return list(result.scalars().all())
+
+    async def list_ip_addresses(self) -> set[str]:
+        result = await self._session.execute(select(ApiFirewallAllowlist.ip_address))
+        return {row[0] for row in result.all()}
+
+    async def exists(self, ip_address: str) -> bool:
+        row = await self._session.get(ApiFirewallAllowlist, ip_address)
+        return row is not None
+
+    async def add(self, ip_address: str) -> ApiFirewallAllowlist:
+        row = ApiFirewallAllowlist(ip_address=ip_address)
+        self._session.add(row)
+        try:
+            await self._session.commit()
+        except IntegrityError as exc:
+            await self._session.rollback()
+            raise FirewallRepositoryConflictError("IP address already exists") from exc
+        await self._session.refresh(row)
+        return row
+
+    async def delete(self, ip_address: str) -> bool:
+        row = await self._session.get(ApiFirewallAllowlist, ip_address)
+        if row is None:
+            return False
+        await self._session.delete(row)
+        await self._session.commit()
+        return True

app/modules/firewall/schemas.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Literal
+
+from pydantic import Field
+
+from app.modules.shared.schemas import DashboardModel
+
+FirewallMode = Literal["allow_all", "allowlist_active"]
+
+
+class FirewallIpEntry(DashboardModel):
+    ip_address: str
+    created_at: datetime
+
+
+class FirewallIpsResponse(DashboardModel):
+    mode: FirewallMode
+    entries: list[FirewallIpEntry] = Field(default_factory=list)
+
+
+class FirewallIpCreateRequest(DashboardModel):
+    ip_address: str
+
+
+class FirewallDeleteResponse(DashboardModel):
+    status: str

app/modules/firewall/service.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import ipaddress
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Literal, Protocol
+
+from app.modules.firewall.repository import FirewallRepositoryConflictError
+
+
+class FirewallEntryLike(Protocol):
+    ip_address: str
+    created_at: datetime
+
+
+class FirewallRepositoryPort(Protocol):
+    async def list_entries(self) -> Sequence[FirewallEntryLike]: ...
+
+    async def list_ip_addresses(self) -> set[str]: ...
+
+    async def exists(self, ip_address: str) -> bool: ...
+
+    async def add(self, ip_address: str) -> FirewallEntryLike: ...
+
+    async def delete(self, ip_address: str) -> bool: ...
+
+
+class FirewallValidationError(ValueError):
+    pass
+
+
+class FirewallIpAlreadyExistsError(ValueError):
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class FirewallIpEntryData:
+    ip_address: str
+    created_at: datetime
+
+
+@dataclass(frozen=True, slots=True)
+class FirewallListData:
+    mode: Literal["allow_all", "allowlist_active"]
+    entries: list[FirewallIpEntryData]
+
+
+class FirewallService:
+    def __init__(self, repository: FirewallRepositoryPort) -> None:
+        self._repository = repository
+
+    async def list_ips(self) -> FirewallListData:
+        rows = await self._repository.list_entries()
+        entries = [FirewallIpEntryData(ip_address=row.ip_address, created_at=row.created_at) for row in rows]
+        mode: Literal["allow_all", "allowlist_active"] = "allow_all" if not entries else "allowlist_active"
+        return FirewallListData(mode=mode, entries=entries)
+
+    async def add_ip(self, ip_address: str) -> FirewallIpEntryData:
+        normalized = normalize_ip_address(ip_address)
+        if await self._repository.exists(normalized):
+            raise FirewallIpAlreadyExistsError("IP address already exists")
+        try:
+            row = await self._repository.add(normalized)
+        except FirewallRepositoryConflictError as exc:
+            raise FirewallIpAlreadyExistsError("IP address already exists") from exc
+        return FirewallIpEntryData(ip_address=row.ip_address, created_at=row.created_at)
+
+    async def remove_ip(self, ip_address: str) -> bool:
+        normalized = normalize_ip_address(ip_address)
+        return await self._repository.delete(normalized)
+
+    async def is_ip_allowed(self, ip_address: str | None) -> bool:
+        allowlist = await self._repository.list_ip_addresses()
+        if not allowlist:
+            return True
+        if ip_address is None:
+            return False
+        try:
+            normalized = normalize_ip_address(ip_address)
+        except FirewallValidationError:
+            return False
+        return normalized in allowlist
+
+
+def normalize_ip_address(value: str) -> str:
+    raw = value.strip()
+    if not raw:
+        raise FirewallValidationError("IP address is required")
+    try:
+        return str(ipaddress.ip_address(raw))
+    except ValueError as exc:
+        raise FirewallValidationError("Invalid IP address") from exc

app/modules/fleet/__init__.py

@@ -0,0 +1 @@
+"""Fleet summary endpoints for external local dashboards."""

app/modules/fleet/api.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+
+from fastapi import APIRouter, Depends, Security
+
+from app.core.auth.dependencies import set_dashboard_error_format, validate_usage_api_key
+from app.core.config.settings_cache import get_settings_cache
+from app.core.exceptions import DashboardServiceUnavailableError
+from app.core.shutdown import is_control_plane_task_admission_open, wait_for_tasks_to_drain
+from app.core.utils.time import utcnow
+from app.db.models import AccountStatus
+from app.db.session import get_background_session
+from app.dependencies import AccountsContext, get_accounts_context
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.api_keys.service import ApiKeyData
+from app.modules.fleet.mappers import build_fleet_account_summaries
+from app.modules.fleet.observability import build_fleet_observability
+from app.modules.fleet.schemas import FleetObservabilityResponse, FleetRefreshResponse, FleetSummaryResponse
+from app.modules.proxy.account_cache import get_account_selection_cache
+from app.modules.proxy.rate_limit_cache import get_rate_limit_headers_cache
+from app.modules.usage.repository import AdditionalUsageRepository, UsageRepository
+from app.modules.usage.updater import UsageUpdater
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(
+    prefix="/api/fleet",
+    tags=["fleet"],
+    dependencies=[Depends(set_dashboard_error_format)],
+)
+
+_REFRESH_SKIP_STATUSES = {AccountStatus.PAUSED, AccountStatus.REAUTH_REQUIRED, AccountStatus.DEACTIVATED}
+# Every route-owned refresh is registered at creation. Caller cancellation only
+# changes who observes its outcome; it does not establish task ownership.
+_BACKGROUND_REFRESH_TASKS: set[asyncio.Task[FleetRefreshResponse]] = set()
+
+
+def _visible_account_ids(api_key: ApiKeyData) -> list[str] | None:
+    return list(api_key.assigned_account_ids) if api_key.account_assignment_scope_enabled else None
+
+
+def _usage_sections(raw: str) -> set[str]:
+    if not raw or not raw.strip():
+        return set()
+    return {section.strip() for section in raw.split(",") if section.strip()}
+
+
+async def _can_view_fleet_usage(api_key: ApiKeyData) -> bool:
+    sections = _usage_sections(api_key.usage_sections)
+    if "account_pool_usage" not in sections or "upstream_limits" not in sections:
+        return False
+    settings = await get_settings_cache().get()
+    return not bool(getattr(settings, "hide_upstream_quota_from_api_keys", False))
+
+
+@router.get("/summary", response_model=FleetSummaryResponse)
+async def get_fleet_summary(
+    context: AccountsContext = Depends(get_accounts_context),
+    api_key: ApiKeyData = Security(validate_usage_api_key),
+) -> FleetSummaryResponse:
+    """Read-only, minimal per-account capacity summary for fleet consumers."""
+
+    visible_account_ids = _visible_account_ids(api_key)
+    include_usage = await _can_view_fleet_usage(api_key)
+    accounts = await context.service.list_accounts(account_ids=visible_account_ids)
+    persisted_status_by_account_id: dict[str, str] | None = None
+    if not include_usage:
+        persisted_accounts = (
+            await context.repository.list_accounts_by_ids(visible_account_ids, refresh_existing=True)
+            if visible_account_ids is not None
+            else await context.repository.list_accounts(refresh_existing=True)
+        )
+        persisted_status_by_account_id = {account.id: account.status.value for account in persisted_accounts}
+    return FleetSummaryResponse(
+        accounts=build_fleet_account_summaries(
+            accounts,
+            include_usage=include_usage,
+            persisted_status_by_account_id=persisted_status_by_account_id,
+        )
+    )
+
+
+@router.get("/observability", response_model=FleetObservabilityResponse)
+async def get_fleet_observability(
+    context: AccountsContext = Depends(get_accounts_context),
+    api_key: ApiKeyData = Security(validate_usage_api_key),
+) -> FleetObservabilityResponse:
+    """Read-only pressure and sticky-session summary for fleet consumers."""
+
+    return await build_fleet_observability(
+        context.session,
+        visible_account_ids=_visible_account_ids(api_key),
+        include_usage=await _can_view_fleet_usage(api_key),
+    )
+
+
+@router.post("/refresh", response_model=FleetRefreshResponse)
+async def refresh_fleet_usage(
+    api_key: ApiKeyData = Security(validate_usage_api_key),
+) -> FleetRefreshResponse:
+    """Request a bounded usage refresh using codex-lb's normal refresh rules."""
+
+    if not is_control_plane_task_admission_open():
+        raise DashboardServiceUnavailableError("Server is draining")
+    task = asyncio.create_task(
+        _refresh_fleet_usage_with_owned_session(_visible_account_ids(api_key)),
+        name="fleet-usage-refresh",
+    )
+    _BACKGROUND_REFRESH_TASKS.add(task)
+    task.add_done_callback(_discard_refresh_task)
+    try:
+        return await asyncio.shield(task)
+    except asyncio.CancelledError:
+        task.add_done_callback(_handle_cancelled_refresh_task_done)
+        raise
+
+
+async def _refresh_fleet_usage_with_owned_session(visible_account_ids: list[str] | None) -> FleetRefreshResponse:
+    async with get_background_session() as session:
+        accounts_repo = AccountsRepository(session)
+        usage_repo = UsageRepository(session)
+        additional_usage_repo = AdditionalUsageRepository(session)
+        accounts = (
+            await accounts_repo.list_accounts_by_ids(visible_account_ids, refresh_existing=True)
+            if visible_account_ids is not None
+            else await accounts_repo.list_accounts(refresh_existing=True)
+        )
+        eligible_accounts = [account for account in accounts if account.status not in _REFRESH_SKIP_STATUSES]
+        latest_primary = await usage_repo.latest_by_account(window="primary", account_ids=visible_account_ids)
+        usage_written = await UsageUpdater(
+            usage_repo,
+            accounts_repo,
+            additional_usage_repo,
+        ).refresh_accounts(eligible_accounts, latest_primary, own_singleflight_sessions=True)
+        if usage_written:
+            await get_rate_limit_headers_cache().invalidate()
+            get_account_selection_cache().invalidate()
+        return FleetRefreshResponse(
+            usage_written=usage_written,
+            account_count=len(accounts),
+            attempted_count=len(eligible_accounts),
+            generated_at=utcnow(),
+        )
+
+
+async def drain_background_refresh_tasks(timeout_seconds: float) -> bool:
+    pending = await wait_for_tasks_to_drain(_BACKGROUND_REFRESH_TASKS, timeout_seconds)
+    for task in sorted(pending, key=lambda pending_task: pending_task.get_name()):
+        logger.warning("Fleet refresh task did not drain before shutdown: %s", task.get_name())
+    return not pending
+
+
+def _handle_cancelled_refresh_task_done(task: asyncio.Task[FleetRefreshResponse]) -> None:
+    try:
+        _log_cancelled_refresh_task_exception(task)
+    finally:
+        _BACKGROUND_REFRESH_TASKS.discard(task)
+
+
+def _discard_refresh_task(task: asyncio.Task[FleetRefreshResponse]) -> None:
+    _BACKGROUND_REFRESH_TASKS.discard(task)
+
+
+def _log_cancelled_refresh_task_exception(task: asyncio.Task[FleetRefreshResponse]) -> None:
+    if task.cancelled():
+        return
+    try:
+        exc = task.exception()
+    except asyncio.CancelledError:
+        return
+    if exc is not None:
+        logger.error(
+            "Fleet usage refresh failed after request cancellation",
+            exc_info=(type(exc), exc, exc.__traceback__),
+        )

app/modules/fleet/mappers.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from app.modules.accounts.schemas import AccountSummary
+from app.modules.fleet.schemas import FleetAccountSummary, FleetWindowSummary
+
+
+def fleet_account_summary_from_account(
+    account: AccountSummary,
+    *,
+    include_usage: bool = True,
+    persisted_status_by_account_id: dict[str, str] | None = None,
+) -> FleetAccountSummary:
+    """Project a dashboard account into the minimal fleet payload."""
+
+    usage = account.usage
+    if include_usage:
+        status = account.status
+    elif persisted_status_by_account_id is None:
+        status = "unknown"
+    else:
+        status = persisted_status_by_account_id.get(account.account_id, "unknown")
+    return FleetAccountSummary(
+        account_id=account.account_id,
+        display_name=account.display_name,
+        email=account.email,
+        status=status,
+        plan_type=account.plan_type,
+        primary=FleetWindowSummary(
+            remaining_percent=usage.primary_remaining_percent if include_usage and usage is not None else None,
+            reset_at=account.reset_at_primary if include_usage else None,
+            window_minutes=account.window_minutes_primary if include_usage else None,
+        ),
+        secondary=FleetWindowSummary(
+            remaining_percent=usage.secondary_remaining_percent if include_usage and usage is not None else None,
+            reset_at=account.reset_at_secondary if include_usage else None,
+            window_minutes=account.window_minutes_secondary if include_usage else None,
+        ),
+        last_refresh_at=account.last_refresh_at if include_usage else None,
+        usage_refreshed_at=account.usage_refreshed_at if include_usage else None,
+    )
+
+
+def build_fleet_account_summaries(
+    accounts: list[AccountSummary],
+    *,
+    include_usage: bool = True,
+    persisted_status_by_account_id: dict[str, str] | None = None,
+) -> list[FleetAccountSummary]:
+    return [
+        fleet_account_summary_from_account(
+            account,
+            include_usage=include_usage,
+            persisted_status_by_account_id=persisted_status_by_account_id,
+        )
+        for account in accounts
+    ]