codexLB-up 1.25.0b10__py3-none-any.whl

app/core/auth/external_identity.py

@@ -0,0 +1,78 @@
+"""Request-time resolution of a provider-asserted identity to an account.
+
+Shared by the session dependency (which turns the outcome into a principal or
+a 401) and the session endpoint (which turns it into the session response, so
+a refused identity can render the "account not ready" screen). The outcome is
+memoised on the request so both read the same answer.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from fastapi import Request
+
+from app.core.auth.dashboard_mode import DashboardAuthMode, DashboardRequestAuth
+from app.core.auth.dashboard_users_cache import get_dashboard_users_cache
+from app.core.auth.providers import TrustedHeaderProvider
+from app.core.auth.providers.registry import get_auth_provider_registry
+from app.db.models import AuthProviderKind, DashboardAuthProvider, DashboardUser, DashboardUserStatus
+from app.modules.dashboard_users.identity_resolver import (
+    DenialCode,
+    Denied,
+    Resolution,
+    get_identity_resolution_cache,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ExternalResolution:
+    """Who the provider-asserted identity is on this install.
+
+    ``user`` is the active account (``denial`` is ``None``); otherwise
+    ``denial`` names why the request is refused: ``identity_not_provisioned``
+    (no account, none created) or ``account_disabled``.
+    """
+
+    provider: DashboardAuthProvider
+    user: DashboardUser | None
+    denial: DenialCode | None
+
+
+async def resolve_trusted_header_request(
+    request: Request, request_auth: DashboardRequestAuth
+) -> ExternalResolution | None:
+    """Resolve the trusted-header identity of ``request``; ``None`` when the provider is not active.
+
+    Cached on ``request.state`` so the dependency and the session endpoint agree.
+    """
+
+    cached = getattr(request.state, "external_resolution", None)
+    if isinstance(cached, ExternalResolution):
+        return cached
+    if request_auth.mode != DashboardAuthMode.TRUSTED_HEADER or not request_auth.actor:
+        return None
+    active = await get_auth_provider_registry().get_active(
+        AuthProviderKind.TRUSTED_HEADER, TrustedHeaderProvider.provider_key, request_auth.mode
+    )
+    if active is None:
+        return None
+    identity = TrustedHeaderProvider().identity_from_subject(request_auth.actor, groups=request_auth.groups)
+    if identity is None:
+        return None
+    actor_ip = request.client.host if request.client else None
+    result = await get_identity_resolution_cache().resolve(identity, active.row, actor_ip=actor_ip)
+    resolution = await _from_result(active.row, result)
+    request.state.external_resolution = resolution
+    return resolution
+
+
+async def _from_result(provider: DashboardAuthProvider, result: Resolution) -> ExternalResolution:
+    if isinstance(result, Denied):
+        return ExternalResolution(provider=provider, user=None, denial=result.code)
+    user = await get_dashboard_users_cache().get_user(result.user_id)
+    if user is None:
+        return ExternalResolution(provider=provider, user=None, denial="identity_not_provisioned")
+    if user.status != DashboardUserStatus.ACTIVE.value:
+        return ExternalResolution(provider=provider, user=None, denial="account_disabled")
+    return ExternalResolution(provider=provider, user=user, denial=None)

app/core/auth/guardian.py

@@ -0,0 +1,374 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import random
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from typing import Protocol, cast
+
+from app.core.auth.refresh import RefreshError
+from app.core.config.background_jobs import auth_guardian_blocked_by_topology, background_job_enabled
+from app.core.scheduling.leader_election_handle import (
+    LeaderElectionLike as _LeaderElectionLike,
+)
+from app.core.scheduling.leader_election_handle import (
+    get_leader_election as _get_leader_election,
+)
+from app.core.utils.time import to_utc_naive, utcnow
+from app.db.models import Account, AccountStatus
+from app.db.session import get_background_session
+from app.modules.accounts.auth_manager import AuthManager
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.proxy.account_cache import get_account_selection_cache
+
+logger = logging.getLogger(__name__)
+
+# Guardian cadence and backoff tuning (fixed; issue #1340 / PRINCIPLES.md P2).
+# The guardian is a background refresher; these values are implementation
+# details, not operator contract. ``AuthGuardianScheduler`` keeps them as
+# constructor fields so tests can exercise the behavior.
+_INTERVAL_SECONDS = 21600
+_MAX_REFRESH_AGE_SECONDS = 43200
+_BATCH_SIZE = 100
+_CONCURRENCY = 3
+_JITTER_SECONDS = 300.0
+_FAILURE_BACKOFF_BASE_SECONDS = 300.0
+_FAILURE_BACKOFF_MAX_SECONDS = 3600.0
+
+# RATE_LIMITED and QUOTA_EXCEEDED recover to ACTIVE through usage-refresh
+# reconciliation, then become guardian-eligible within the max-age window.
+_AUTH_GUARDIAN_ELIGIBLE_STATUSES = frozenset({AccountStatus.ACTIVE, AccountStatus.PAUSED})
+
+
+class _AccountsRepositoryLike(Protocol):
+    async def list_accounts(self, *, refresh_existing: bool = False) -> list[Account]: ...
+
+    async def get_by_id(self, account_id: str) -> Account | None: ...
+
+
+class _AuthManagerLike(Protocol):
+    async def ensure_fresh(self, account: Account, *, force: bool = False) -> Account: ...
+
+
+_RepoFactory = Callable[[], AbstractAsyncContextManager[_AccountsRepositoryLike]]
+_AuthManagerFactory = Callable[[_AccountsRepositoryLike], _AuthManagerLike]
+_LeaderElectionFactory = Callable[[], _LeaderElectionLike]
+_Sleep = Callable[[float], Awaitable[None]]
+
+
+@dataclass(slots=True)
+class _FailureBackoff:
+    attempts: int
+    retry_after_monotonic: float
+
+
+@dataclass(slots=True)
+class AuthGuardianScheduler:
+    interval_seconds: int
+    enabled: bool
+    max_age_seconds: int
+    batch_size: int
+    concurrency: int
+    jitter_seconds: float
+    failure_backoff_base_seconds: float = 300.0
+    failure_backoff_max_seconds: float = 3600.0
+    # When leader election is disabled, each refresh pass first counts live
+    # bridge-ring replicas and skips the pass if more than one is registered
+    # (the static instance ring is empty in Helm/compose deployments, so the
+    # build-time guard alone cannot see dynamically registered replicas).
+    # Defaults to True so directly constructed schedulers skip the dynamic
+    # check; build_auth_guardian_scheduler wires the real setting.
+    leader_election_enabled: bool = True
+    # M2 background jobs: the loop always runs; each pass reads the effective
+    # ``auth_guardian_enabled`` toggle (dashboard column, else the deprecated
+    # env alias, else the default) from the settings cache and skips while it is
+    # False, so a dashboard change applies on the next tick without a restart.
+    # ``topology_blocked`` is the static multi-replica-without-election gate
+    # computed at build time; the dashboard toggle cannot override it.
+    dashboard_enabled: Callable[[], Awaitable[bool]] = field(default_factory=lambda: _dashboard_guardian_enabled)
+    topology_blocked: bool = False
+    live_replica_count: Callable[[], Awaitable[int]] = field(default_factory=lambda: _count_live_bridge_ring_members)
+    leader_election_factory: _LeaderElectionFactory = field(default_factory=lambda: _get_leader_election)
+    repo_factory: _RepoFactory = field(default_factory=lambda: _default_accounts_repo_factory)
+    auth_manager_factory: _AuthManagerFactory = field(default_factory=lambda: _default_auth_manager_factory)
+    sleep: _Sleep = field(default_factory=lambda: asyncio.sleep)
+    now: Callable[[], datetime] = field(default_factory=lambda: utcnow)
+    _task: asyncio.Task[None] | None = None
+    _stop: asyncio.Event = field(default_factory=asyncio.Event)
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    _failures: dict[str, _FailureBackoff] = field(default_factory=dict)
+
+    async def start(self) -> None:
+        if not self.enabled:
+            return
+        if self._task and not self._task.done():
+            return
+        self._stop.clear()
+        self._task = asyncio.create_task(self._run_loop())
+
+    async def stop(self) -> None:
+        self._stop.set()
+        if self._task is not None:
+            self._task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
+
+    async def _run_loop(self) -> None:
+        while not self._stop.is_set():
+            jitter = _jitter_delay(self.jitter_seconds)
+            if jitter > 0:
+                try:
+                    await asyncio.wait_for(self._stop.wait(), timeout=jitter)
+                    break
+                except asyncio.TimeoutError:
+                    pass
+            try:
+                await self._refresh_once()
+            except Exception:
+                logger.exception("Auth Guardian refresh pass failed")
+            try:
+                await asyncio.wait_for(self._stop.wait(), timeout=self.interval_seconds)
+            except asyncio.TimeoutError:
+                continue
+
+    async def _refresh_once(self) -> None:
+        if not await self.dashboard_enabled():
+            logger.debug("Auth Guardian skipped refresh pass: disabled in the dashboard settings")
+            return
+        if self.topology_blocked:
+            # The builder already logged this once at WARNING; repeating it
+            # every pass for a static condition would only be noise.
+            logger.debug(
+                "Auth Guardian skipped refresh pass: multi-replica deployment without leader election; "
+                "set CODEX_LB_LEADER_ELECTION_ENABLED=true to run it leader-gated"
+            )
+            return
+        if not self.leader_election_enabled:
+            live_replicas = await self.live_replica_count()
+            if live_replicas > 1:
+                logger.warning(
+                    "Auth Guardian skipped refresh pass: %d live replicas registered in the bridge ring "
+                    "while leader election is disabled; set CODEX_LB_LEADER_ELECTION_ENABLED=true so a "
+                    "single elected replica performs proactive refresh work",
+                    live_replicas,
+                )
+                return
+        await self.leader_election_factory().run_if_leader(self._refresh_as_leader)
+
+    async def _refresh_as_leader(self) -> None:
+        async with self._lock:
+            async with self.repo_factory() as repo:
+                accounts = await repo.list_accounts(refresh_existing=True)
+                candidates = select_auth_guardian_candidates(
+                    accounts,
+                    now=self.now(),
+                    max_age_seconds=self.max_age_seconds,
+                    limit=len(accounts),
+                )
+                candidate_ids = [account.id for account in candidates if not self._in_backoff(account.id)]
+                candidate_ids = candidate_ids[: max(0, self.batch_size)]
+            if not candidate_ids:
+                return
+            semaphore = asyncio.Semaphore(max(1, self.concurrency))
+            await asyncio.gather(*(self._refresh_candidate(account_id, semaphore) for account_id in candidate_ids))
+
+    async def _refresh_candidate(self, account_id: str, semaphore: asyncio.Semaphore) -> None:
+        if self._in_backoff(account_id):
+            return
+        async with semaphore:
+            async with self.repo_factory() as repo:
+                account = await repo.get_by_id(account_id)
+                if account is None:
+                    self._failures.pop(account_id, None)
+                    return
+                source_status = account.status.value
+                if not _auth_guardian_account_is_stale_eligible(
+                    account,
+                    now=self.now(),
+                    max_age_seconds=self.max_age_seconds,
+                ):
+                    return
+                manager = self.auth_manager_factory(repo)
+                try:
+                    refresh_task = asyncio.create_task(manager.ensure_fresh(account, force=True))
+                    try:
+                        await asyncio.shield(refresh_task)
+                    except asyncio.CancelledError:
+                        with contextlib.suppress(Exception):
+                            await refresh_task
+                        raise
+                except RefreshError as exc:
+                    self._record_failure(account_id)
+                    if exc.is_permanent:
+                        get_account_selection_cache().invalidate()
+                    logger.warning(
+                        "Auth Guardian refresh failed account_id=%s account_alias=%s status=%s code=%s permanent=%s "
+                        "transport=%s",
+                        account.id,
+                        _safe_account_alias(account),
+                        source_status,
+                        exc.code,
+                        exc.is_permanent,
+                        exc.transport_error,
+                    )
+                    return
+                except Exception as exc:
+                    self._record_failure(account_id)
+                    logger.warning(
+                        "Auth Guardian refresh failed account_id=%s account_alias=%s status=%s error_type=%s",
+                        account.id,
+                        _safe_account_alias(account),
+                        source_status,
+                        exc.__class__.__name__,
+                        exc_info=True,
+                    )
+                    return
+                self._failures.pop(account_id, None)
+                get_account_selection_cache().invalidate()
+                logger.info(
+                    "Auth Guardian refreshed account_id=%s account_alias=%s status=%s",
+                    account.id,
+                    _safe_account_alias(account),
+                    source_status,
+                )
+
+    def _in_backoff(self, account_id: str) -> bool:
+        failure = self._failures.get(account_id)
+        if failure is None:
+            return False
+        if failure.retry_after_monotonic > time.monotonic():
+            return True
+        return False
+
+    def _record_failure(self, account_id: str) -> None:
+        previous = self._failures.get(account_id)
+        attempts = 1 if previous is None else previous.attempts + 1
+        base = max(0.0, float(self.failure_backoff_base_seconds))
+        cap = max(base, float(self.failure_backoff_max_seconds))
+        delay = min(cap, base * (2 ** min(attempts - 1, 6)))
+        delay += _jitter_delay(self.jitter_seconds)
+        self._failures[account_id] = _FailureBackoff(
+            attempts=attempts,
+            retry_after_monotonic=time.monotonic() + delay,
+        )
+
+
+def select_auth_guardian_candidates(
+    accounts: list[Account],
+    *,
+    now: datetime,
+    max_age_seconds: int,
+    limit: int,
+) -> list[Account]:
+    candidates = [
+        account
+        for account in accounts
+        if _auth_guardian_account_is_stale_eligible(
+            account,
+            now=now,
+            max_age_seconds=max_age_seconds,
+        )
+    ]
+    candidates.sort(key=lambda account: to_utc_naive(account.last_refresh))
+    return candidates[: max(0, limit)]
+
+
+def build_auth_guardian_scheduler() -> AuthGuardianScheduler:
+    from app.core.config.settings import get_settings
+
+    settings = get_settings()
+    # Deliberate exception to the "disabled election means every replica is
+    # leader" escape hatch: concurrent force token refreshes across replicas
+    # can invalidate rotated refresh tokens, so without election the guardian
+    # must not run in a multi-replica ring at all. The loop still starts (M2:
+    # the dashboard toggle is read at every pass), but every pass skips.
+    topology_blocked = auth_guardian_blocked_by_topology(settings)
+    if topology_blocked:
+        logger.warning(
+            "Auth Guardian disabled: multi-replica deployment without leader election; "
+            "set CODEX_LB_LEADER_ELECTION_ENABLED=true to run it leader-gated"
+        )
+    return AuthGuardianScheduler(
+        interval_seconds=_INTERVAL_SECONDS,
+        enabled=True,
+        topology_blocked=topology_blocked,
+        max_age_seconds=_MAX_REFRESH_AGE_SECONDS,
+        batch_size=_BATCH_SIZE,
+        concurrency=_CONCURRENCY,
+        jitter_seconds=_JITTER_SECONDS,
+        failure_backoff_base_seconds=_FAILURE_BACKOFF_BASE_SECONDS,
+        failure_backoff_max_seconds=_FAILURE_BACKOFF_MAX_SECONDS,
+        leader_election_enabled=settings.leader_election_enabled,
+    )
+
+
+def _auth_guardian_account_is_stale_eligible(
+    account: Account,
+    *,
+    now: datetime,
+    max_age_seconds: int,
+) -> bool:
+    if account.status not in _AUTH_GUARDIAN_ELIGIBLE_STATUSES:
+        return False
+    age = to_utc_naive(now) - to_utc_naive(account.last_refresh)
+    return age > timedelta(seconds=max_age_seconds)
+
+
+async def _dashboard_guardian_enabled() -> bool:
+    return await background_job_enabled("auth_guardian_enabled")
+
+
+async def _count_live_bridge_ring_members() -> int:
+    from sqlalchemy import func, select
+
+    from app.db.models import BridgeRingMember
+    from app.modules.proxy.ring_membership import RING_STALE_THRESHOLD_SECONDS
+
+    cutoff = utcnow() - timedelta(seconds=RING_STALE_THRESHOLD_SECONDS)
+    async with get_background_session() as session:
+        result = await session.execute(
+            select(func.count()).select_from(BridgeRingMember).where(BridgeRingMember.last_heartbeat_at >= cutoff)
+        )
+        return int(result.scalar_one() or 0)
+
+
+@asynccontextmanager
+async def _default_accounts_repo_factory() -> AsyncIterator[AccountsRepository]:
+    async with get_background_session() as session:
+        yield AccountsRepository(session)
+
+
+def _default_auth_manager_factory(repo: _AccountsRepositoryLike) -> _AuthManagerLike:
+    return AuthManager(cast(AccountsRepository, repo), refresh_repo_factory=_default_accounts_repo_factory)
+
+
+def _jitter_delay(max_seconds: float) -> float:
+    if max_seconds <= 0:
+        return 0.0
+    return random.uniform(0.0, max_seconds)
+
+
+def _safe_account_alias(account: Account) -> str:
+    alias = (account.alias or "").strip()
+    if alias:
+        return alias[:64]
+    return _mask_email(account.email)
+
+
+def _mask_email(email: str) -> str:
+    if "@" not in email:
+        return email[:2] + "***" if email else ""
+    local, domain = email.split("@", 1)
+    if not local:
+        return f"***@{domain}"
+    if len(local) == 1:
+        masked_local = f"{local}***"
+    else:
+        masked_local = f"{local[0]}***{local[-1]}"
+    return f"{masked_local}@{domain}"

app/core/auth/models.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from pydantic import AliasChoices, BaseModel, ConfigDict, Field, StrictStr, field_validator
+
+from app.core.types import JsonObject
+
+
+class OAuthTokenPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    access_token: StrictStr | None = None
+    refresh_token: StrictStr | None = None
+    id_token: StrictStr | None = None
+    authorization_code: StrictStr | None = None
+    code_verifier: StrictStr | None = None
+    error: JsonObject | StrictStr | None = None
+    error_description: StrictStr | None = None
+    message: StrictStr | None = None
+    error_code: StrictStr | None = None
+    code: StrictStr | None = None
+    status: StrictStr | None = None
+
+
+class DeviceCodePayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    device_auth_id: StrictStr | None = None
+    user_code: StrictStr | None = Field(
+        default=None,
+        validation_alias=AliasChoices("user_code", "usercode"),
+    )
+    interval: int | None = None
+    expires_in: int | None = None
+    expires_at: StrictStr | None = None
+
+    @field_validator("interval", mode="before")
+    @classmethod
+    def _parse_interval(cls, value: int | str | None) -> int | None:
+        if value is None:
+            return None
+        if isinstance(value, int):
+            return value
+        if isinstance(value, str):
+            stripped = value.strip()
+            if not stripped:
+                return None
+            if stripped.isdigit():
+                return int(stripped)
+        raise ValueError("Invalid interval")

app/core/auth/providers/__init__.py

@@ -0,0 +1,129 @@
+"""Sign-in providers: how a request proves who is calling.
+
+A provider either recognises the caller on every request from evidence the
+request carries (``resolve_identity``: the reverse-proxy header today, a
+signed assertion tomorrow) or runs a redirect dance (``begin_login`` /
+``complete_login``: OIDC, Phase 3). Every provider hands the same
+:class:`ExternalIdentity` to the shared identity resolver, which maps it to a
+dashboard account; providers never touch accounts themselves.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from fastapi import Request
+
+from app.core.auth.dashboard_mode import DashboardAuthMode, get_dashboard_request_auth
+from app.db.models import AuthProviderKind
+from app.modules.auth_providers.seed import DEFAULT_PROVIDER_KEY
+from app.modules.dashboard_users.repository import is_valid_email, normalize_email
+
+#: Column widths of ``dashboard_identities.subject`` / ``display_name``; the API validates the same bounds.
+MAX_SUBJECT_LENGTH = 512
+MAX_DISPLAY_NAME_LENGTH = 128
+
+__all__ = [
+    "DEFAULT_PROVIDER_KEY",
+    "MAX_DISPLAY_NAME_LENGTH",
+    "MAX_SUBJECT_LENGTH",
+    "AuthProvider",
+    "ExternalIdentity",
+    "PasswordProvider",
+    "TrustedHeaderProvider",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class ExternalIdentity:
+    """What a provider learned about the caller, before any account lookup.
+
+    ``subject`` is the stable identifier the provider guarantees (the header
+    value, an OIDC ``sub``); ``email`` and ``display_name`` are display and
+    linking hints; ``groups`` is the IdP's group snapshot (case-folded), which
+    role mappings match on.
+    """
+
+    provider: str
+    provider_key: str
+    subject: str
+    email: str | None = None
+    display_name: str | None = None
+    groups: tuple[str, ...] = ()
+
+
+@runtime_checkable
+class AuthProvider(Protocol):
+    """One way of signing in.
+
+    Header-style providers implement :meth:`resolve_identity`. Redirect-style
+    providers (OIDC) add ``begin_login(request) -> RedirectResponse`` and
+    ``complete_login(request) -> ExternalIdentity`` when they arrive; both
+    still end in the shared identity resolver. ``kind`` matches
+    ``dashboard_auth_providers.kind`` and ``provider_key`` the row's key.
+    """
+
+    @property
+    def kind(self) -> AuthProviderKind: ...
+
+    @property
+    def provider_key(self) -> str: ...
+
+    def resolve_identity(self, request: Request) -> ExternalIdentity | None:
+        """The caller's identity when the request carries the provider's evidence, else ``None``."""
+        ...
+
+
+class PasswordProvider:
+    """Local username/password sign-in. The login flow itself lives in
+    ``dashboard_auth``; this marker exists so the provider list and the
+    providers table have one entry per way of signing in."""
+
+    kind = AuthProviderKind.PASSWORD
+    provider_key = DEFAULT_PROVIDER_KEY
+
+    def resolve_identity(self, request: Request) -> ExternalIdentity | None:
+        del request
+        return None
+
+
+class TrustedHeaderProvider:
+    """The reverse-proxy identity header (``CODEX_LB_DASHBOARD_AUTH_PROXY_HEADER``).
+
+    Reuses the request-auth resolution that already checks the trusted-proxy
+    peer and the singular header. The subject is case-folded: two proxies
+    spelling the same person differently must not become two accounts. The
+    raw value is kept as the display name and, when it parses as an e-mail,
+    as the identity's e-mail. Groups arrive on the same request-auth value,
+    read from ``CODEX_LB_DASHBOARD_AUTH_PROXY_GROUPS_HEADER``.
+    """
+
+    kind = AuthProviderKind.TRUSTED_HEADER
+    provider_key = DEFAULT_PROVIDER_KEY
+
+    def resolve_identity(self, request: Request) -> ExternalIdentity | None:
+        request_auth = get_dashboard_request_auth(request)
+        if request_auth is None or request_auth.mode != DashboardAuthMode.TRUSTED_HEADER or not request_auth.actor:
+            return None
+        return self.identity_from_subject(request_auth.actor, groups=request_auth.groups)
+
+    def identity_from_subject(self, raw_subject: str, *, groups: tuple[str, ...] = ()) -> ExternalIdentity | None:
+        """``None`` when the value cannot be an identity here (empty or wider than the subject column).
+
+        ``groups`` defaults to none so the paths without a request (an invite's
+        expected identity) can still name an identity.
+        """
+
+        subject = raw_subject.strip()
+        if not subject or len(subject) > MAX_SUBJECT_LENGTH:
+            return None
+        email = normalize_email(subject)
+        return ExternalIdentity(
+            provider=self.kind.value,
+            provider_key=self.provider_key,
+            subject=subject.casefold(),
+            email=email if email is not None and is_valid_email(email) else None,
+            display_name=subject[:MAX_DISPLAY_NAME_LENGTH],
+            groups=groups,
+        )