codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/sticky_sessions/schemas.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Literal
+
+from pydantic import Field, model_validator
+
+from app.db.models import StickySessionKind
+from app.modules.shared.schemas import DashboardModel
+
+StickySessionSortBy = Literal["updated_at", "created_at", "account", "key"]
+StickySessionSortDir = Literal["asc", "desc"]
+
+
+class StickySessionEntryResponse(DashboardModel):
+    key: str
+    display_name: str
+    kind: StickySessionKind
+    created_at: datetime
+    updated_at: datetime
+    expires_at: datetime | None = None
+    is_stale: bool
+
+
+class StickySessionsListResponse(DashboardModel):
+    entries: list[StickySessionEntryResponse] = Field(default_factory=list)
+    stale_prompt_cache_count: int = 0
+    total: int = 0
+    has_more: bool = False
+
+
+class StickySessionIdentifier(DashboardModel):
+    key: str = Field(min_length=1)
+    kind: StickySessionKind
+
+
+class StickySessionsDeleteRequest(DashboardModel):
+    sessions: list[StickySessionIdentifier] = Field(min_length=1, max_length=500)
+
+    @model_validator(mode="after")
+    def validate_unique_sessions(self) -> "StickySessionsDeleteRequest":
+        seen: set[tuple[str, StickySessionKind]] = set()
+        for session in self.sessions:
+            target = (session.key, session.kind)
+            if target in seen:
+                raise ValueError("duplicate sticky session targets are not allowed")
+            seen.add(target)
+        return self
+
+
+class StickySessionDeleteFailure(DashboardModel):
+    key: str
+    kind: StickySessionKind
+    reason: str
+
+
+class StickySessionsDeleteResponse(DashboardModel):
+    deleted_count: int
+    deleted: list[StickySessionIdentifier] = Field(default_factory=list)
+    failed: list[StickySessionDeleteFailure] = Field(default_factory=list)
+
+
+class StickySessionsDeleteFilteredRequest(DashboardModel):
+    stale_only: bool = False
+    account_query: str = ""
+    key_query: str = ""
+
+
+class StickySessionsDeleteFilteredResponse(DashboardModel):
+    deleted_count: int
+
+
+class StickySessionsPurgeRequest(DashboardModel):
+    stale_only: Literal[True] = True
+
+
+class StickySessionsPurgeResponse(DashboardModel):
+    deleted_count: int

app/modules/sticky_sessions/service.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+
+from app.core.utils.time import to_utc_naive, utcnow
+from app.db.models import StickySessionKind
+from app.modules.proxy.sticky_repository import (
+    StickySessionListEntryRecord,
+    StickySessionsRepository,
+    is_reserved_sticky_session_key,
+)
+from app.modules.settings.repository import SettingsRepository
+from app.modules.sticky_sessions.schemas import StickySessionSortBy, StickySessionSortDir
+
+
+@dataclass(frozen=True, slots=True)
+class StickySessionEntryData:
+    key: str
+    display_name: str
+    kind: StickySessionKind
+    created_at: datetime
+    updated_at: datetime
+    expires_at: datetime | None
+    is_stale: bool
+
+
+@dataclass(frozen=True, slots=True)
+class StickySessionListData:
+    entries: list[StickySessionEntryData]
+    stale_prompt_cache_count: int
+    total: int
+    has_more: bool
+
+
+@dataclass(frozen=True, slots=True)
+class StickySessionDeleteFailureData:
+    key: str
+    kind: StickySessionKind
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class StickySessionsDeleteData:
+    deleted: list[tuple[str, StickySessionKind]]
+    failed: list[StickySessionDeleteFailureData]
+
+    @property
+    def deleted_count(self) -> int:
+        return len(self.deleted)
+
+
+class StickySessionsService:
+    def __init__(
+        self,
+        repository: StickySessionsRepository,
+        settings_repository: SettingsRepository,
+    ) -> None:
+        self._repository = repository
+        self._settings_repository = settings_repository
+
+    async def list_entries(
+        self,
+        *,
+        kind: StickySessionKind | None = None,
+        stale_only: bool = False,
+        account_query: str | None = None,
+        key_query: str | None = None,
+        sort_by: StickySessionSortBy = "updated_at",
+        sort_dir: StickySessionSortDir = "desc",
+        offset: int = 0,
+        limit: int = 100,
+    ) -> StickySessionListData:
+        settings = await self._settings_repository.get_or_create()
+        ttl_seconds = settings.openai_cache_affinity_max_age_seconds
+        stale_cutoff = utcnow() - timedelta(seconds=ttl_seconds)
+        normalized_account_query = account_query.strip() if account_query else None
+        normalized_key_query = key_query.strip() if key_query else None
+        stale_prompt_cache_count = await self._count_stale_prompt_cache_entries(kind=kind, stale_cutoff=stale_cutoff)
+        if stale_only and kind not in (None, StickySessionKind.PROMPT_CACHE):
+            return StickySessionListData(
+                entries=[],
+                stale_prompt_cache_count=stale_prompt_cache_count,
+                total=0,
+                has_more=False,
+            )
+        effective_kind = StickySessionKind.PROMPT_CACHE if stale_only else kind
+        total = await self._repository.count_entries(
+            kind=effective_kind,
+            updated_before=stale_cutoff if stale_only else None,
+            account_query=normalized_account_query,
+            key_query=normalized_key_query,
+        )
+        rows = await self._repository.list_entries(
+            kind=effective_kind,
+            updated_before=stale_cutoff if stale_only else None,
+            account_query=normalized_account_query,
+            key_query=normalized_key_query,
+            sort_by=sort_by,
+            sort_dir=sort_dir,
+            offset=offset,
+            limit=limit,
+        )
+        entries = [self._to_entry(row, ttl_seconds=ttl_seconds) for row in rows]
+        return StickySessionListData(
+            entries=entries,
+            stale_prompt_cache_count=stale_prompt_cache_count,
+            total=total,
+            has_more=offset + len(entries) < total,
+        )
+
+    async def delete_entries(self, entries: Sequence[tuple[str, StickySessionKind]]) -> StickySessionsDeleteData:
+        failed: list[StickySessionDeleteFailureData] = []
+        seen: set[tuple[str, StickySessionKind]] = set()
+        targets: list[tuple[str, StickySessionKind]] = []
+
+        for key, kind in entries:
+            if not key:
+                continue
+            target = (key, kind)
+            if target in seen:
+                continue
+            seen.add(target)
+            if is_reserved_sticky_session_key(key):
+                failed.append(StickySessionDeleteFailureData(key=key, kind=kind, reason="not_found"))
+                continue
+            targets.append(target)
+
+        deleted = await self._repository.delete_entries(targets)
+        deleted_set = set(deleted)
+
+        for key, kind in targets:
+            if (key, kind) not in deleted_set:
+                failed.append(StickySessionDeleteFailureData(key=key, kind=kind, reason="not_found"))
+
+        return StickySessionsDeleteData(deleted=deleted, failed=failed)
+
+    async def delete_filtered_entries(
+        self,
+        *,
+        kind: StickySessionKind | None = None,
+        stale_only: bool = False,
+        account_query: str | None = None,
+        key_query: str | None = None,
+    ) -> int:
+        settings = await self._settings_repository.get_or_create()
+        stale_cutoff = utcnow() - timedelta(seconds=settings.openai_cache_affinity_max_age_seconds)
+        if stale_only and kind not in (None, StickySessionKind.PROMPT_CACHE):
+            return 0
+        effective_kind = StickySessionKind.PROMPT_CACHE if stale_only else kind
+        normalized_account_query = account_query.strip() if account_query else None
+        normalized_key_query = key_query.strip() if key_query else None
+        targets = await self._repository.list_entry_identifiers(
+            kind=effective_kind,
+            updated_before=stale_cutoff if stale_only else None,
+            account_query=normalized_account_query,
+            key_query=normalized_key_query,
+        )
+        deleted = await self._repository.delete_entries(targets)
+        return len(deleted)
+
+    async def purge_entries(self) -> int:
+        settings = await self._settings_repository.get_or_create()
+        cutoff = utcnow() - timedelta(seconds=settings.openai_cache_affinity_max_age_seconds)
+        return await self._repository.purge_prompt_cache_before(cutoff)
+
+    def _to_entry(self, row: StickySessionListEntryRecord, *, ttl_seconds: int) -> StickySessionEntryData:
+        sticky_session = row.sticky_session
+        expires_at: datetime | None = None
+        is_stale = False
+        if sticky_session.kind == StickySessionKind.PROMPT_CACHE:
+            expires_at = to_utc_naive(sticky_session.updated_at) + timedelta(seconds=ttl_seconds)
+            is_stale = expires_at <= utcnow()
+        return StickySessionEntryData(
+            key=sticky_session.key,
+            display_name=row.display_name,
+            kind=sticky_session.kind,
+            created_at=sticky_session.created_at,
+            updated_at=sticky_session.updated_at,
+            expires_at=expires_at,
+            is_stale=is_stale,
+        )
+
+    async def _count_stale_prompt_cache_entries(
+        self,
+        *,
+        kind: StickySessionKind | None,
+        stale_cutoff: datetime,
+    ) -> int:
+        if kind not in (None, StickySessionKind.PROMPT_CACHE):
+            return 0
+        return await self._repository.count_entries(
+            kind=StickySessionKind.PROMPT_CACHE,
+            updated_before=stale_cutoff,
+        )

app/modules/telemetry/__init__.py

@@ -0,0 +1,5 @@
+"""Anonymous, schema-allowlisted telemetry support."""
+
+from app.modules.telemetry.snapshot import TelemetrySnapshotBuilder
+
+__all__ = ["TelemetrySnapshotBuilder"]

app/modules/telemetry/api.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import platform
+
+from fastapi import APIRouter, Body, Depends, Query
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app import __version__
+from app.core.auth.dependencies import (
+    require_dashboard_write_access,
+    set_dashboard_error_format,
+    validate_dashboard_session,
+)
+from app.db.session import get_session
+from app.modules.telemetry.consent import ResolvedConsent, TelemetryConsentStore
+from app.modules.telemetry.schemas import (
+    TelemetryConsentResponse,
+    TelemetryConsentUpdate,
+    TelemetrySnapshotEnvelope,
+    build_snapshot_envelope,
+)
+from app.modules.telemetry.sender import TelemetrySender
+from app.modules.telemetry.snapshot import TelemetrySnapshotBuilder, deployment_method
+
+logger = logging.getLogger(__name__)
+
+_OPT_OUT_TASKS: set[asyncio.Task[None]] = set()
+
+router = APIRouter(
+    prefix="/api/settings",
+    tags=["dashboard"],
+    dependencies=[Depends(validate_dashboard_session), Depends(set_dashboard_error_format)],
+)
+
+
+@router.get("/telemetry", response_model=TelemetryConsentResponse)
+async def get_telemetry_consent(
+    include_preview: bool = Query(default=False),
+    session: AsyncSession = Depends(get_session),
+) -> TelemetryConsentResponse:
+    store = TelemetryConsentStore(session)
+    consent = await store.resolve()
+    return await _response(
+        session,
+        store,
+        consent,
+        include_preview=include_preview or (consent.state == "undecided" and consent.source == "default"),
+    )
+
+
+@router.put("/telemetry", response_model=TelemetryConsentResponse)
+async def update_telemetry_consent(
+    payload: TelemetryConsentUpdate = Body(...),
+    _write_access=Depends(require_dashboard_write_access),
+    session: AsyncSession = Depends(get_session),
+) -> TelemetryConsentResponse:
+    store = TelemetryConsentStore(session)
+    previous = await store.resolve()
+    consent = await store.set_decision(payload.enabled)
+    if previous.active and not consent.active:
+        try:
+            identity = await store.get_or_create_identity()
+            task = asyncio.create_task(
+                TelemetrySender().send_opt_out(
+                    identity,
+                    app_version=__version__,
+                    deployment_mode=deployment_method(),
+                    os_arch=f"{platform.system().lower()}/{platform.machine().lower()}",
+                ),
+                name="anonymous-telemetry-opt-out",
+            )
+            _OPT_OUT_TASKS.add(task)
+            task.add_done_callback(_handle_opt_out_task_done)
+        except Exception as exc:
+            logger.debug("Unable to schedule anonymous telemetry opt-out", exc_info=exc)
+    return await _response(session, store, consent, include_preview=False)
+
+
+async def _response(
+    session: AsyncSession,
+    store: TelemetryConsentStore,
+    consent: ResolvedConsent,
+    *,
+    include_preview: bool,
+) -> TelemetryConsentResponse:
+    preview: TelemetrySnapshotEnvelope | None = None
+    if include_preview:
+        identity = await store.get_or_create_identity()
+        snapshot_consent = "enabled" if consent.state == "disabled" else consent.state
+        snapshot = await TelemetrySnapshotBuilder(session).build(
+            identity.instance_id,
+            consent=snapshot_consent,
+        )
+        preview = build_snapshot_envelope(snapshot)
+    return TelemetryConsentResponse(
+        state=consent.state,
+        source=consent.source,
+        active=consent.active,
+        preview=preview,
+    )
+
+
+def _handle_opt_out_task_done(task: asyncio.Task[None]) -> None:
+    try:
+        if task.cancelled():
+            return
+        if exc := task.exception():
+            logger.debug(
+                "Anonymous telemetry opt-out background task failed",
+                exc_info=(type(exc), exc, exc.__traceback__),
+            )
+    finally:
+        _OPT_OUT_TASKS.discard(task)

app/modules/telemetry/clients.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+CLIENT_FAMILY_BY_RAW_GROUP: dict[str, str] = {
+    "codex_exec": "codex-cli",
+    "codex-tui": "codex-cli",
+    "codex_cli_rs": "codex-cli",
+    "Codex Desktop": "codex-desktop",
+    "codex_vscode": "codex-vscode",
+    "AsyncOpenAI": "openai-sdk-python",
+    "OpenAI": "openai-sdk-js",
+    "ai": "vercel-ai-sdk",
+    "ai-sdk": "vercel-ai-sdk",
+    "opencode": "opencode",
+    "Mozilla": "browser",
+    "curl": "script",
+    "undici": "script",
+    "node": "script",
+    "Python-urllib": "script",
+    "python-requests": "script",
+    "aiohttp": "script",
+}
+
+CANONICAL_CLIENT_FAMILIES = frozenset({*CLIENT_FAMILY_BY_RAW_GROUP.values(), "other"})
+
+
+@dataclass(frozen=True, slots=True)
+class ClientCount:
+    raw_group: str | None
+    requests: int
+
+
+def client_family(raw_group: str | None) -> str:
+    return CLIENT_FAMILY_BY_RAW_GROUP.get(raw_group or "", "other")
+
+
+def client_shares(rows: Iterable[ClientCount]) -> tuple[dict[str, float], float]:
+    counts: defaultdict[str, int] = defaultdict(int)
+    total = 0
+    for row in rows:
+        requests = max(0, row.requests)
+        family = client_family(row.raw_group)
+        if family not in CANONICAL_CLIENT_FAMILIES:
+            raise ValueError(f"non-canonical telemetry client family: {family}")
+        counts[family] += requests
+        total += requests
+    if total == 0:
+        return {}, 0.0
+    shares = {family: _ratio(count, total) for family, count in sorted(counts.items()) if count > 0}
+    return shares, shares.get("other", 0.0)
+
+
+def catalog_model_name(model: str | None, catalog: frozenset[str]) -> str:
+    normalized = (model or "").strip()
+    return normalized if normalized in catalog else "other"
+
+
+def _ratio(numerator: int | float, denominator: int | float) -> float:
+    return round(float(numerator) / float(denominator), 6) if denominator else 0.0

app/modules/telemetry/consent.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass
+from typing import Literal, cast
+from uuid import uuid4
+
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat, PublicFormat
+from sqlalchemy import or_, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.config.settings import Settings, get_settings
+from app.core.crypto import TokenEncryptor
+from app.db.models import DashboardSettings
+from app.modules.settings.repository import SettingsRepository
+
+ConsentState = Literal["undecided", "enabled", "disabled"]
+ConsentSource = Literal["env", "persisted", "default"]
+_VALID_STATES = frozenset({"undecided", "enabled", "disabled"})
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedConsent:
+    state: ConsentState
+    source: ConsentSource
+    active: bool
+
+
+@dataclass(frozen=True, slots=True)
+class TelemetryIdentity:
+    instance_id: str
+    private_key: Ed25519PrivateKey
+
+    @property
+    def public_key_hex(self) -> str:
+        public_bytes = self.private_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
+        return public_bytes.hex()
+
+
+def resolve_consent(telemetry_enabled: bool | None, persisted_state: str) -> ResolvedConsent:
+    """Resolve effective telemetry consent: dashboard decision > env > default.
+
+    A persisted dashboard decision is authoritative. ``CODEX_LB_TELEMETRY_ENABLED``
+    only decides while no decision has been saved (``undecided``), so a headless
+    opt-out set before first boot still works and a later dashboard decision is
+    never silently overridden by the environment.
+    """
+    if persisted_state not in _VALID_STATES:
+        raise ValueError(f"invalid telemetry consent state: {persisted_state}")
+    state = cast("ConsentState", persisted_state)
+    if state != "undecided":
+        return ResolvedConsent(state=state, source="persisted", active=state == "enabled")
+    if telemetry_enabled is not None:
+        env_state: ConsentState = "enabled" if telemetry_enabled else "disabled"
+        return ResolvedConsent(state=env_state, source="env", active=telemetry_enabled)
+    return ResolvedConsent(state="undecided", source="default", active=True)
+
+
+class TelemetryConsentStore:
+    def __init__(
+        self,
+        session: AsyncSession,
+        *,
+        settings: Settings | None = None,
+        encryptor: TokenEncryptor | None = None,
+    ) -> None:
+        self._session = session
+        self._settings = settings or get_settings()
+        self._encryptor = encryptor or TokenEncryptor()
+        self._repository = SettingsRepository(session)
+
+    async def resolve(self) -> ResolvedConsent:
+        row = await self._repository.get_or_create()
+        return resolve_consent(self._settings.telemetry_enabled, row.telemetry_consent)
+
+    async def set_decision(self, enabled: bool) -> ResolvedConsent:
+        row = await self._repository.get_or_create()
+        row.telemetry_consent = "enabled" if enabled else "disabled"
+        await self._repository.commit_refresh(row)
+        return resolve_consent(self._settings.telemetry_enabled, row.telemetry_consent)
+
+    async def get_or_create_identity(self) -> TelemetryIdentity:
+        row = await self._repository.get_or_create()
+        if row.telemetry_instance_id is None or row.telemetry_private_key_encrypted is None:
+            await self._mint_identity_if_missing()
+            self._session.expire_all()
+            row = await self._repository.get_or_create()
+        if row.telemetry_instance_id is None or row.telemetry_private_key_encrypted is None:
+            raise RuntimeError("telemetry identity could not be persisted")
+        raw_private_key = base64.b64decode(self._encryptor.decrypt(row.telemetry_private_key_encrypted))
+        private_key = Ed25519PrivateKey.from_private_bytes(raw_private_key)
+        return TelemetryIdentity(instance_id=row.telemetry_instance_id, private_key=private_key)
+
+    async def _mint_identity_if_missing(self) -> None:
+        private_key = Ed25519PrivateKey.generate()
+        raw_private_key = private_key.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
+        encrypted = self._encryptor.encrypt(base64.b64encode(raw_private_key).decode("ascii"))
+        await self._session.execute(
+            update(DashboardSettings)
+            .where(
+                DashboardSettings.id == 1,
+                or_(
+                    DashboardSettings.telemetry_instance_id.is_(None),
+                    DashboardSettings.telemetry_private_key_encrypted.is_(None),
+                ),
+            )
+            .values(
+                telemetry_instance_id=str(uuid4()),
+                telemetry_private_key_encrypted=encrypted,
+                version=DashboardSettings.version + 1,
+            )
+            .execution_options(synchronize_session=False)
+        )
+        await self._session.commit()

app/modules/telemetry/scheduler.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from dataclasses import dataclass, field
+from functools import partial
+
+from app.core.scheduling.leader_election_handle import get_leader_election as _get_leader_election
+from app.db.session import get_background_session
+from app.modules.telemetry.consent import TelemetryConsentStore
+from app.modules.telemetry.sender import TelemetrySender
+from app.modules.telemetry.snapshot import TelemetrySnapshotBuilder
+
+logger = logging.getLogger(__name__)
+
+TELEMETRY_INTERVAL_SECONDS = 24 * 60 * 60
+TELEMETRY_FIELDS_DOCUMENTATION = "https://soju06.github.io/codex-lb/telemetry/"
+
+
+@dataclass(slots=True)
+class TelemetryScheduler:
+    sender: TelemetrySender = field(default_factory=TelemetrySender)
+    interval_seconds: float = TELEMETRY_INTERVAL_SECONDS
+    _task: asyncio.Task[None] | None = None
+    _stop: asyncio.Event = field(default_factory=asyncio.Event)
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+    async def start(self) -> None:
+        if self._task and not self._task.done():
+            return
+        self._stop.clear()
+        self._task = asyncio.create_task(self._run_loop(), name="anonymous-telemetry-scheduler")
+
+    async def stop(self) -> None:
+        self._stop.set()
+        if self._task is None:
+            return
+        self._task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._task
+        self._task = None
+
+    async def _run_loop(self) -> None:
+        first_tick = True
+        while not self._stop.is_set():
+            await self._tick(log_undecided_notice=first_tick)
+            first_tick = False
+            try:
+                await asyncio.wait_for(self._stop.wait(), timeout=self.interval_seconds)
+            except asyncio.TimeoutError:
+                continue
+
+    async def _tick(self, *, log_undecided_notice: bool = False) -> None:
+        await _get_leader_election().run_if_leader(
+            partial(self._tick_as_leader, log_undecided_notice=log_undecided_notice)
+        )
+
+    async def _tick_as_leader(self, *, log_undecided_notice: bool = False) -> None:
+        async with self._lock:
+            try:
+                async with get_background_session() as session:
+                    store = TelemetryConsentStore(session)
+                    consent = await store.resolve()
+                    if log_undecided_notice and consent.state == "undecided" and consent.source == "default":
+                        logger.info(
+                            "Anonymous telemetry is active; collected fields: %s; disable with "
+                            "CODEX_LB_TELEMETRY_ENABLED=false",
+                            TELEMETRY_FIELDS_DOCUMENTATION,
+                        )
+                    if not consent.active:
+                        return
+                    assert consent.state != "disabled"
+                    identity = await store.get_or_create_identity()
+                    snapshot = await TelemetrySnapshotBuilder(session).build(
+                        identity.instance_id,
+                        consent=consent.state,
+                    )
+                await self.sender.send_snapshot(snapshot)
+            except Exception as exc:
+                logger.debug("Anonymous telemetry scheduler tick failed", exc_info=exc)
+
+
+def build_telemetry_scheduler() -> TelemetryScheduler:
+    return TelemetryScheduler()