codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/model_sources/service.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import json
+import uuid
+from urllib.parse import urlparse
+
+from app.core.crypto import TokenEncryptor
+from app.db.models import ModelSource, ModelSourceModel
+from app.modules.model_sources.repository import ModelSourcesRepository
+from app.modules.model_sources.schemas import (
+    ModelSourceCreateRequest,
+    ModelSourceModelInput,
+    ModelSourceModelResponse,
+    ModelSourceResponse,
+    ModelSourceUpdateRequest,
+)
+
+MODEL_SOURCE_KIND_OPENAI_COMPATIBLE = "openai_compatible"
+MODEL_SOURCE_HEALTH_UNKNOWN = "unknown"
+
+
+class ModelSourceNotFoundError(ValueError):
+    pass
+
+
+class ModelSourceValidationError(ValueError):
+    pass
+
+
+class ModelSourcesService:
+    def __init__(
+        self,
+        repository: ModelSourcesRepository,
+        *,
+        encryptor: TokenEncryptor | None = None,
+    ) -> None:
+        self._repository = repository
+        self._encryptor = encryptor or TokenEncryptor()
+
+    async def list_sources(self) -> list[ModelSourceResponse]:
+        rows = await self._repository.list_sources()
+        return [_to_response(row) for row in rows]
+
+    async def list_enabled_sources(self) -> list[ModelSource]:
+        return await self._repository.list_enabled_sources()
+
+    async def create_source(self, payload: ModelSourceCreateRequest) -> ModelSourceResponse:
+        model_rows = _model_inputs_to_rows(payload.models)
+        row = ModelSource(
+            id=f"src_{uuid.uuid4().hex}",
+            name=_normalize_name(payload.name),
+            kind=MODEL_SOURCE_KIND_OPENAI_COMPATIBLE,
+            base_url=_normalize_base_url(payload.base_url),
+            api_key_encrypted=_encrypt_optional(self._encryptor, payload.api_key),
+            is_enabled=True,
+            health_status=MODEL_SOURCE_HEALTH_UNKNOWN,
+            supports_chat_completions=payload.supports_chat_completions,
+            supports_responses=payload.supports_responses,
+            supports_audio_transcriptions=payload.supports_audio_transcriptions,
+            supports_embeddings=payload.supports_embeddings,
+            timeout_seconds=payload.timeout_seconds,
+            max_concurrency=payload.max_concurrency,
+            models=model_rows,
+        )
+        try:
+            created = await self._repository.create(row, commit=True)
+        except Exception:
+            await self._repository.rollback()
+            raise
+        return _to_response(created)
+
+    async def update_source(self, source_id: str, payload: ModelSourceUpdateRequest) -> ModelSourceResponse:
+        row = await self._repository.get_by_id(source_id)
+        if row is None:
+            raise ModelSourceNotFoundError(f"Model source not found: {source_id}")
+
+        fields = payload.model_fields_set
+        if "name" in fields and payload.name is not None:
+            row.name = _normalize_name(payload.name)
+        if "base_url" in fields and payload.base_url is not None:
+            row.base_url = _normalize_base_url(payload.base_url)
+        if "api_key" in fields:
+            row.api_key_encrypted = _encrypt_optional(self._encryptor, payload.api_key)
+        if "is_enabled" in fields and payload.is_enabled is not None:
+            row.is_enabled = payload.is_enabled
+        if "supports_chat_completions" in fields and payload.supports_chat_completions is not None:
+            row.supports_chat_completions = payload.supports_chat_completions
+        if "supports_responses" in fields and payload.supports_responses is not None:
+            row.supports_responses = payload.supports_responses
+        if "supports_audio_transcriptions" in fields and payload.supports_audio_transcriptions is not None:
+            row.supports_audio_transcriptions = payload.supports_audio_transcriptions
+        if "supports_embeddings" in fields and payload.supports_embeddings is not None:
+            row.supports_embeddings = payload.supports_embeddings
+        if "timeout_seconds" in fields:
+            row.timeout_seconds = payload.timeout_seconds
+        if "max_concurrency" in fields:
+            row.max_concurrency = payload.max_concurrency
+
+        models_replaced = False
+        try:
+            if "models" in fields and payload.models is not None:
+                await self._repository.replace_models(row, _model_inputs_to_rows(payload.models), commit=False)
+                models_replaced = True
+            await self._repository.commit()
+        except Exception:
+            await self._repository.rollback()
+            raise
+
+        if models_replaced:
+            # ``replace_models`` bulk-deletes and re-inserts child rows without
+            # touching the identity-mapped parent's already-loaded ``models``
+            # collection, so the post-commit read would return the stale
+            # pre-update list without an explicit refresh.
+            await self._repository.refresh_models(row)
+        refreshed = await self._repository.get_by_id(source_id)
+        if refreshed is None:
+            raise ModelSourceNotFoundError(f"Model source not found: {source_id}")
+        return _to_response(refreshed)
+
+    async def delete_source(self, source_id: str) -> None:
+        deleted = await self._repository.delete(source_id)
+        if not deleted:
+            raise ModelSourceNotFoundError(f"Model source not found: {source_id}")
+
+
+def _normalize_name(value: str) -> str:
+    name = value.strip()
+    if not name:
+        raise ModelSourceValidationError("Model source name is required")
+    return name
+
+
+def _normalize_base_url(value: str) -> str:
+    url = value.strip().rstrip("/")
+    parsed = urlparse(url)
+    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+        raise ModelSourceValidationError("Model source base_url must be an absolute HTTP(S) URL")
+    return url
+
+
+def _normalize_model_slug(value: str) -> str:
+    model = value.strip()
+    if not model:
+        raise ModelSourceValidationError("Model source model name is required")
+    return model
+
+
+def _validate_raw_metadata_json(value: str | None) -> str | None:
+    if value is None:
+        return None
+    try:
+        parsed = json.loads(value)
+    except json.JSONDecodeError as exc:
+        raise ModelSourceValidationError("raw_metadata_json must be valid JSON") from exc
+    if not isinstance(parsed, dict):
+        raise ModelSourceValidationError("raw_metadata_json must be a JSON object")
+    return json.dumps(parsed, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
+
+
+def _model_inputs_to_rows(models: list[ModelSourceModelInput]) -> list[ModelSourceModel]:
+    seen: set[str] = set()
+    rows: list[ModelSourceModel] = []
+    for item in models:
+        model = _normalize_model_slug(item.model)
+        if model in seen:
+            raise ModelSourceValidationError(f"Duplicate model source model: {model}")
+        seen.add(model)
+        rows.append(
+            ModelSourceModel(
+                model=model,
+                display_name=item.display_name.strip() if item.display_name else None,
+                context_window=item.context_window,
+                max_output_tokens=item.max_output_tokens,
+                supports_streaming=item.supports_streaming,
+                supports_tools=item.supports_tools,
+                supports_vision=item.supports_vision,
+                input_per_1m=item.input_per_1m,
+                cached_input_per_1m=item.cached_input_per_1m,
+                output_per_1m=item.output_per_1m,
+                audio_per_minute=item.audio_per_minute,
+                raw_metadata_json=_validate_raw_metadata_json(item.raw_metadata_json),
+                is_enabled=item.is_enabled,
+            )
+        )
+    return rows
+
+
+def _encrypt_optional(encryptor: TokenEncryptor, value: str | None) -> bytes | None:
+    if value is None:
+        return None
+    secret = value.strip()
+    if not secret:
+        return None
+    return encryptor.encrypt(secret)
+
+
+def _to_model_response(row: ModelSourceModel) -> ModelSourceModelResponse:
+    return ModelSourceModelResponse(
+        id=row.id,
+        source_id=row.source_id,
+        model=row.model,
+        display_name=row.display_name,
+        context_window=row.context_window,
+        max_output_tokens=row.max_output_tokens,
+        supports_streaming=row.supports_streaming,
+        supports_tools=row.supports_tools,
+        supports_vision=row.supports_vision,
+        input_per_1m=row.input_per_1m,
+        cached_input_per_1m=row.cached_input_per_1m,
+        output_per_1m=row.output_per_1m,
+        audio_per_minute=row.audio_per_minute,
+        raw_metadata_json=row.raw_metadata_json,
+        is_enabled=row.is_enabled,
+        created_at=row.created_at,
+        updated_at=row.updated_at,
+    )
+
+
+def _to_response(row: ModelSource) -> ModelSourceResponse:
+    return ModelSourceResponse(
+        id=row.id,
+        name=row.name,
+        kind=row.kind,
+        base_url=row.base_url,
+        is_enabled=row.is_enabled,
+        health_status=row.health_status,
+        supports_chat_completions=row.supports_chat_completions,
+        supports_responses=row.supports_responses,
+        supports_audio_transcriptions=row.supports_audio_transcriptions,
+        supports_embeddings=row.supports_embeddings,
+        timeout_seconds=row.timeout_seconds,
+        max_concurrency=row.max_concurrency,
+        created_at=row.created_at,
+        updated_at=row.updated_at,
+        models=[_to_model_response(model) for model in row.models],
+    )

app/modules/oauth/api.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import logging
+
+from fastapi import APIRouter, Body, Depends, Query
+from fastapi.responses import JSONResponse
+
+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.clients.oauth import OAuthError
+from app.core.errors import dashboard_error
+from app.dependencies import OauthContext, get_oauth_context
+from app.modules.oauth.schemas import (
+    ManualCallbackRequest,
+    ManualCallbackResponse,
+    OauthCompleteRequest,
+    OauthCompleteResponse,
+    OauthStartRequest,
+    OauthStartResponse,
+    OauthStatusResponse,
+)
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(
+    prefix="/api/oauth",
+    tags=["dashboard"],
+    dependencies=[Depends(validate_dashboard_session), Depends(set_dashboard_error_format)],
+)
+
+
+@router.post("/start", response_model=OauthStartResponse)
+async def start_oauth(
+    request: OauthStartRequest,
+    _write_access=Depends(require_dashboard_permission(Permission.ACCOUNTS_WRITE)),
+    context: OauthContext = Depends(get_oauth_context),
+) -> OauthStartResponse | JSONResponse:
+    try:
+        return await context.service.start_oauth(request)
+    except OAuthError as exc:
+        return JSONResponse(
+            status_code=502,
+            content=dashboard_error(exc.code, exc.message),
+        )
+    except NotImplementedError:
+        return JSONResponse(
+            status_code=501,
+            content=dashboard_error("not_implemented", "OAuth start is not implemented"),
+        )
+
+
+@router.get(
+    "/status",
+    response_model=OauthStatusResponse,
+    dependencies=[Depends(require_dashboard_permission(Permission.ACCOUNTS_WRITE))],
+)
+async def oauth_status(
+    flow_id: str | None = Query(default=None, alias="flowId"),
+    context: OauthContext = Depends(get_oauth_context),
+) -> OauthStatusResponse | JSONResponse:
+    return await context.service.oauth_status(flow_id=flow_id)
+
+
+@router.post("/complete", response_model=OauthCompleteResponse)
+async def complete_oauth(
+    request: OauthCompleteRequest | None = Body(default=None),
+    _write_access=Depends(require_dashboard_permission(Permission.ACCOUNTS_WRITE)),
+    context: OauthContext = Depends(get_oauth_context),
+) -> OauthCompleteResponse | JSONResponse:
+    try:
+        return await context.service.complete_oauth(request)
+    except NotImplementedError:
+        return JSONResponse(
+            status_code=501,
+            content=dashboard_error("not_implemented", "OAuth complete is not implemented"),
+        )
+
+
+@router.post("/manual-callback", response_model=ManualCallbackResponse)
+async def manual_callback(
+    request: ManualCallbackRequest,
+    _write_access=Depends(require_dashboard_permission(Permission.ACCOUNTS_WRITE)),
+    context: OauthContext = Depends(get_oauth_context),
+) -> ManualCallbackResponse | JSONResponse:
+    try:
+        return await context.service.manual_callback(request.callback_url, flow_id=request.flow_id)
+    except OAuthError as exc:
+        return JSONResponse(
+            status_code=502,
+            content=dashboard_error(exc.code, exc.message),
+        )
+    except Exception:
+        logger.exception("manual_callback failed")
+        return JSONResponse(
+            status_code=500,
+            content=dashboard_error("manual_callback_failed", "An internal error occurred."),
+        )

app/modules/oauth/repository.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, cast
+
+from sqlalchemy import CursorResult, delete, select, update
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.core.crypto import TokenEncryptor
+from app.core.utils.time import to_utc_naive, utcnow
+from app.db.models import OAuthDeviceFlowSlot, OAuthFlowState
+
+_TERMINAL_OAUTH_STATUSES = {"error", "success"}
+
+# Fixed key for the single-active dashboard device-code OAuth flow slot. The
+# dashboard runs at most one device flow at a time (a new device ``start``
+# supersedes any prior one), so a single global slot models "which device flow
+# is current" across replicas.
+DEVICE_FLOW_SLOT_KEY = "dashboard"
+
+
+def epoch_to_naive_utc(value: float | None) -> datetime | None:
+    """Convert an epoch-seconds float (as used by the in-memory store) to the
+    naive-UTC datetime the rest of the schema stores."""
+
+    if value is None:
+        return None
+    return datetime.fromtimestamp(value, tz=timezone.utc).replace(tzinfo=None)
+
+
+@dataclass(slots=True)
+class OAuthFlowRecord:
+    """Durable representation of a dashboard OAuth flow.
+
+    ``code_verifier`` is held in plaintext in memory only; the repository
+    encrypts it at rest and decrypts it on read.
+    """
+
+    flow_id: str
+    method: str
+    status: str
+    state_token: str | None = None
+    error_message: str | None = None
+    intended_account_id: str | None = None
+    code_verifier: str | None = None
+    device_auth_id: str | None = None
+    user_code: str | None = None
+    interval_seconds: int | None = None
+    expires_at: datetime | None = None
+    finished_at: datetime | None = None
+
+
+class OAuthFlowRepository:
+    """DB access for the shared ``oauth_flow_states`` table."""
+
+    def __init__(self, session: AsyncSession, encryptor: TokenEncryptor) -> None:
+        self._session = session
+        self._encryptor = encryptor
+
+    def _to_record(self, row: OAuthFlowState) -> OAuthFlowRecord:
+        verifier: str | None = None
+        if row.code_verifier_encrypted is not None:
+            verifier = self._encryptor.decrypt(row.code_verifier_encrypted)
+        return OAuthFlowRecord(
+            flow_id=row.flow_id,
+            method=row.method,
+            status=row.status,
+            state_token=row.state_token,
+            error_message=row.error_message,
+            intended_account_id=row.intended_account_id,
+            code_verifier=verifier,
+            device_auth_id=row.device_auth_id,
+            user_code=row.user_code,
+            interval_seconds=row.interval_seconds,
+            expires_at=row.expires_at,
+            finished_at=row.finished_at,
+        )
+
+    @staticmethod
+    def _is_expired_pending(row: OAuthFlowState, now: datetime) -> bool:
+        if row.status != "pending" or row.expires_at is None:
+            return False
+        # ``expires_at`` is ``DateTime(timezone=True)``: on PostgreSQL asyncpg
+        # returns an offset-AWARE datetime, while ``now`` (``utcnow``) is naive
+        # UTC. Normalize the row value to naive UTC before comparing so the
+        # comparison never raises ``TypeError`` for offset-naive vs -aware.
+        return to_utc_naive(row.expires_at) <= now
+
+    async def create(self, record: OAuthFlowRecord) -> None:
+        encrypted = None
+        if record.code_verifier is not None:
+            encrypted = self._encryptor.encrypt(record.code_verifier)
+        row = OAuthFlowState(
+            flow_id=record.flow_id,
+            state_token=record.state_token,
+            method=record.method,
+            status=record.status,
+            error_message=record.error_message,
+            intended_account_id=record.intended_account_id,
+            code_verifier_encrypted=encrypted,
+            device_auth_id=record.device_auth_id,
+            user_code=record.user_code,
+            interval_seconds=record.interval_seconds,
+            expires_at=record.expires_at,
+            created_at=utcnow(),
+            finished_at=record.finished_at,
+        )
+        self._session.add(row)
+        await self._session.commit()
+
+    async def get_by_flow_id(self, flow_id: str) -> OAuthFlowRecord | None:
+        row = await self._session.get(OAuthFlowState, flow_id)
+        if row is None or self._is_expired_pending(row, utcnow()):
+            return None
+        return self._to_record(row)
+
+    async def get_by_state_token(self, state_token: str) -> OAuthFlowRecord | None:
+        result = await self._session.execute(
+            select(OAuthFlowState).where(OAuthFlowState.state_token == state_token).limit(1)
+        )
+        row = result.scalar_one_or_none()
+        if row is None or self._is_expired_pending(row, utcnow()):
+            return None
+        return self._to_record(row)
+
+    async def set_status(
+        self,
+        flow_id: str,
+        *,
+        status: str,
+        error_message: str | None,
+    ) -> bool:
+        # Terminal writes are monotonic and enforced ATOMICALLY in SQL so the
+        # guard holds under real cross-session/cross-replica concurrency. A
+        # durable ``success`` is sticky: only a ``success`` write may land on an
+        # already-``success`` row (idempotent), and any non-success write is
+        # rejected by the ``WHERE`` predicate. This closes the read-then-write
+        # TOCTOU race where two pollers both load a still-``pending`` row and the
+        # later ``error`` writer would otherwise clobber a committed ``success``
+        # (e.g. a losing/duplicate device poller receiving an OAuth error for the
+        # now-consumed device code). ``pending -> terminal`` and
+        # ``error -> success`` remain allowed.
+        finished_at = utcnow() if status in _TERMINAL_OAUTH_STATUSES else None
+        statement = update(OAuthFlowState).where(OAuthFlowState.flow_id == flow_id)
+        if status != "success":
+            statement = statement.where(OAuthFlowState.status != "success")
+        statement = statement.values(
+            status=status, error_message=error_message, finished_at=finished_at
+        ).execution_options(synchronize_session=False)
+        result = cast(CursorResult[Any], await self._session.execute(statement))
+        await self._session.commit()
+        return int(result.rowcount or 0) > 0
+
+    async def purge_expired(self, *, terminal_keep: int) -> None:
+        now = utcnow()
+        await self._session.execute(
+            delete(OAuthFlowState).where(
+                OAuthFlowState.status == "pending",
+                OAuthFlowState.expires_at.is_not(None),
+                OAuthFlowState.expires_at <= now,
+            )
+        )
+        result = await self._session.execute(
+            select(OAuthFlowState.flow_id)
+            .where(OAuthFlowState.status.in_(tuple(_TERMINAL_OAUTH_STATUSES)))
+            .order_by(OAuthFlowState.finished_at.desc())
+            .offset(terminal_keep)
+        )
+        stale_terminal = [row[0] for row in result.all()]
+        if stale_terminal:
+            await self._session.execute(delete(OAuthFlowState).where(OAuthFlowState.flow_id.in_(stale_terminal)))
+        await self._session.commit()
+
+    # ------------------------------------------------------------------
+    # Single-active device-flow slot (atomic cross-replica coordination)
+    # ------------------------------------------------------------------
+
+    async def claim_device_slot(self, flow_id: str) -> None:
+        """Atomically make ``flow_id`` the single current device flow.
+
+        A single conditional UPSERT on the fixed slot key: two replicas starting
+        device OAuth simultaneously serialize on the slot row (PostgreSQL row
+        lock / SQLite single-writer lock), so the slot ends up naming exactly one
+        ``flow_id`` -- not two orphaned pending rows. ``generation`` is bumped on
+        every replacement so the slot is never silently left unchanged.
+        """
+
+        now = utcnow()
+        dialect = self._session.get_bind().dialect.name
+        if dialect == "postgresql":
+            insert_stmt = pg_insert(OAuthDeviceFlowSlot)
+        elif dialect == "sqlite":
+            insert_stmt = sqlite_insert(OAuthDeviceFlowSlot)
+        else:  # pragma: no cover - only sqlite/postgres are supported backends
+            raise RuntimeError(f"device-flow slot unsupported for dialect={dialect!r}")
+        statement = insert_stmt.values(
+            slot_key=DEVICE_FLOW_SLOT_KEY, flow_id=flow_id, generation=1, updated_at=now
+        ).on_conflict_do_update(
+            index_elements=[OAuthDeviceFlowSlot.slot_key],
+            set_={
+                "flow_id": flow_id,
+                "generation": OAuthDeviceFlowSlot.generation + 1,
+                "updated_at": now,
+            },
+        )
+        await self._session.execute(statement)
+        await self._session.commit()
+
+    async def consume_device_slot(self, flow_id: str) -> bool:
+        """Atomically consume the slot iff ``flow_id`` still holds it.
+
+        This is the poller's point of no return before persisting tokens: the
+        conditional delete removes the slot only when it still names this flow,
+        so a superseded poller (the slot was UPSERTed to a newer ``flow_id``)
+        matches zero rows and MUST abort without persisting an account. Atomic on
+        both backends (single conditional DELETE).
+        """
+
+        result = cast(
+            CursorResult[Any],
+            await self._session.execute(
+                delete(OAuthDeviceFlowSlot).where(
+                    OAuthDeviceFlowSlot.slot_key == DEVICE_FLOW_SLOT_KEY,
+                    OAuthDeviceFlowSlot.flow_id == flow_id,
+                )
+            ),
+        )
+        await self._session.commit()
+        return int(result.rowcount or 0) > 0
+
+    async def current_device_slot_flow_id(self) -> str | None:
+        result = await self._session.execute(
+            select(OAuthDeviceFlowSlot.flow_id).where(OAuthDeviceFlowSlot.slot_key == DEVICE_FLOW_SLOT_KEY)
+        )
+        return result.scalar_one_or_none()

app/modules/oauth/schemas.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from app.modules.shared.schemas import DashboardModel
+
+
+class OauthStartRequest(DashboardModel):
+    force_method: str | None = None
+    # When set, this OAuth flow is a TARGETED reauthentication of an existing
+    # local account row. The returned login's seat identity is verified against
+    # this row before any tokens are written, so a wrong browser identity cannot
+    # overwrite a different seat that shares the same Team/Business workspace.
+    account_id: str | None = None
+
+
+class OauthStartResponse(DashboardModel):
+    flow_id: str | None = None
+    method: str
+    authorization_url: str | None = None
+    callback_url: str | None = None
+    verification_url: str | None = None
+    user_code: str | None = None
+    device_auth_id: str | None = None
+    interval_seconds: int | None = None
+    expires_in_seconds: int | None = None
+
+
+class OauthStatusResponse(DashboardModel):
+    status: str
+    error_message: str | None = None
+
+
+class OauthCompleteRequest(DashboardModel):
+    flow_id: str | None = None
+    device_auth_id: str | None = None
+    user_code: str | None = None
+
+
+class OauthCompleteResponse(DashboardModel):
+    status: str
+
+
+class ManualCallbackRequest(DashboardModel):
+    callback_url: str
+    flow_id: str | None = None
+
+
+class ManualCallbackResponse(DashboardModel):
+    status: str
+    error_message: str | None = None