codexLB-up 1.25.0b10__py3-none-any.whl

app/modules/model_sources/projection.py

@@ -0,0 +1,55 @@
+"""Responses body projection for OpenAI-compatible model sources.
+
+``strip_source_telemetry`` removes the Codex client telemetry a source must
+never see -- whole fields where the field itself is Codex-only, a single key
+inside the standard ``stream_options`` object -- and forwards everything else
+verbatim: direct source routing never fails closed on an unknown field.
+"""
+
+from __future__ import annotations
+
+from app.core.types import JsonValue
+
+type MutableJsonObject = dict[str, JsonValue]
+
+# Codex client telemetry that an OpenAI-compatible source rejects or must not see.
+# Whole top-level fields, because neither is a Responses API field:
+# ``client_metadata`` (installation/session/thread/turn ids and turn metadata)
+# and ``access_programs``.
+STRIPPED_TELEMETRY_FIELDS: frozenset[str] = frozenset({"client_metadata", "access_programs"})
+
+# ``stream_options`` is a standard Responses field (``include_obfuscation``)
+# that Codex extends with ``reasoning_summary_delivery`` (its reasoning-summary
+# delivery mode), so only that key is telemetry: it is removed from the object
+# and the object is dropped once the removal leaves it empty -- the shape every
+# Codex body has -- while an SDK client's ``include_obfuscation`` is forwarded
+# unchanged.
+STREAM_OPTIONS_FIELD = "stream_options"
+STRIPPED_STREAM_OPTIONS_KEYS: frozenset[str] = frozenset({"reasoning_summary_delivery"})
+
+
+def strip_source_telemetry(payload: MutableJsonObject) -> MutableJsonObject:
+    """Remove exactly the Codex telemetry in place; returns ``payload``.
+
+    ``STRIPPED_TELEMETRY_FIELDS`` are removed whole; from a ``stream_options``
+    object only ``STRIPPED_STREAM_OPTIONS_KEYS`` are removed, and the object
+    itself is dropped when that removal leaves nothing behind. Everything else
+    -- ``prompt_cache_key``, ``tools``, ``include``, ``max_output_tokens``,
+    ``prompt_cache_retention``, ``background``, ``max_tool_calls``, a
+    ``stream_options.include_obfuscation``, a ``stream_options`` that is not an
+    object, and any field this proxy has never seen -- is forwarded untouched,
+    so direct source routing never fails closed on an unknown field.
+    """
+
+    for field in STRIPPED_TELEMETRY_FIELDS:
+        payload.pop(field, None)
+    stream_options = payload.get(STREAM_OPTIONS_FIELD)
+    if isinstance(stream_options, dict):
+        removed = False
+        for key in STRIPPED_STREAM_OPTIONS_KEYS:
+            if key in stream_options:
+                del stream_options[key]
+                removed = True
+        if removed and not stream_options:
+            del payload[STREAM_OPTIONS_FIELD]
+    return payload

app/modules/model_sources/repository.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+from sqlalchemy import ColumnElement, and_, delete, or_, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import selectinload
+
+from app.db.models import ModelSource, ModelSourceModel
+
+
+def _enablement_filter(only_disabled: bool) -> ColumnElement[bool]:
+    """Enabled-state predicate for a per-capability source lookup.
+
+    ``only_disabled`` selects the exact complement of the routable set: rows
+    that match the model and the route shape but that an operator switched off,
+    at the source or at the individual model. Routing needs that complement to
+    tell "no source serves this model" apart from "the source that serves this
+    model is switched off" -- the latter must not fall through to a
+    subscription account, which rejects the model outright.
+    """
+    if only_disabled:
+        return or_(ModelSource.is_enabled.is_(False), ModelSourceModel.is_enabled.is_(False))
+    return and_(ModelSource.is_enabled.is_(True), ModelSourceModel.is_enabled.is_(True))
+
+
+class ModelSourcesRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def list_sources(self) -> list[ModelSource]:
+        result = await self._session.execute(
+            select(ModelSource).options(selectinload(ModelSource.models)).order_by(ModelSource.name)
+        )
+        return list(result.scalars().unique().all())
+
+    async def list_enabled_sources(self) -> list[ModelSource]:
+        result = await self._session.execute(
+            select(ModelSource)
+            .options(selectinload(ModelSource.models))
+            .where(ModelSource.is_enabled.is_(True))
+            .order_by(ModelSource.name)
+        )
+        return list(result.scalars().unique().all())
+
+    async def get_by_id(self, source_id: str) -> ModelSource | None:
+        result = await self._session.execute(
+            select(ModelSource).options(selectinload(ModelSource.models)).where(ModelSource.id == source_id)
+        )
+        return result.scalar_one_or_none()
+
+    async def find_chat_source_for_model(
+        self,
+        model: str,
+        *,
+        allowed_source_ids: set[str] | None = None,
+        require_streaming: bool = False,
+        only_disabled: bool = False,
+    ) -> ModelSource | None:
+        stmt = (
+            select(ModelSource)
+            .options(selectinload(ModelSource.models))
+            .join(ModelSourceModel, ModelSourceModel.source_id == ModelSource.id)
+            .where(ModelSource.kind == "openai_compatible")
+            .where(ModelSource.supports_chat_completions.is_(True))
+            .where(ModelSourceModel.model == model)
+            .where(_enablement_filter(only_disabled))
+            .order_by(ModelSource.name, ModelSource.id)
+            .limit(1)
+        )
+        if require_streaming:
+            stmt = stmt.where(ModelSourceModel.supports_streaming.is_(True))
+        if allowed_source_ids is not None:
+            if not allowed_source_ids:
+                return None
+            stmt = stmt.where(ModelSource.id.in_(allowed_source_ids))
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def find_responses_source_for_model(
+        self,
+        model: str,
+        *,
+        allowed_source_ids: set[str] | None = None,
+        require_streaming: bool = False,
+        only_disabled: bool = False,
+    ) -> ModelSource | None:
+        stmt = (
+            select(ModelSource)
+            .options(selectinload(ModelSource.models))
+            .join(ModelSourceModel, ModelSourceModel.source_id == ModelSource.id)
+            .where(ModelSource.kind == "openai_compatible")
+            .where(ModelSource.supports_responses.is_(True))
+            .where(ModelSourceModel.model == model)
+            .where(_enablement_filter(only_disabled))
+            .order_by(ModelSource.name, ModelSource.id)
+            .limit(1)
+        )
+        if require_streaming:
+            stmt = stmt.where(ModelSourceModel.supports_streaming.is_(True))
+        if allowed_source_ids is not None:
+            if not allowed_source_ids:
+                return None
+            stmt = stmt.where(ModelSource.id.in_(allowed_source_ids))
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def find_audio_transcriptions_source_for_model(
+        self,
+        model: str,
+        *,
+        allowed_source_ids: set[str] | None = None,
+    ) -> ModelSource | None:
+        stmt = (
+            select(ModelSource)
+            .options(selectinload(ModelSource.models))
+            .join(ModelSourceModel, ModelSourceModel.source_id == ModelSource.id)
+            .where(ModelSource.kind == "openai_compatible")
+            .where(ModelSource.is_enabled.is_(True))
+            .where(ModelSource.supports_audio_transcriptions.is_(True))
+            .where(ModelSourceModel.model == model)
+            .where(ModelSourceModel.is_enabled.is_(True))
+            .order_by(ModelSource.name, ModelSource.id)
+            .limit(1)
+        )
+        if allowed_source_ids is not None:
+            if not allowed_source_ids:
+                return None
+            stmt = stmt.where(ModelSource.id.in_(allowed_source_ids))
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def find_embeddings_source_for_model(
+        self,
+        model: str,
+        *,
+        allowed_source_ids: set[str] | None = None,
+    ) -> ModelSource | None:
+        stmt = (
+            select(ModelSource)
+            .options(selectinload(ModelSource.models))
+            .join(ModelSourceModel, ModelSourceModel.source_id == ModelSource.id)
+            .where(ModelSource.kind == "openai_compatible")
+            .where(ModelSource.is_enabled.is_(True))
+            .where(ModelSource.supports_embeddings.is_(True))
+            .where(ModelSourceModel.model == model)
+            .where(ModelSourceModel.is_enabled.is_(True))
+            .order_by(ModelSource.name, ModelSource.id)
+            .limit(1)
+        )
+        if allowed_source_ids is not None:
+            if not allowed_source_ids:
+                return None
+            stmt = stmt.where(ModelSource.id.in_(allowed_source_ids))
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def create(self, row: ModelSource, *, commit: bool = True) -> ModelSource:
+        self._session.add(row)
+        if commit:
+            await self._session.commit()
+            await self._session.refresh(row, attribute_names=["models"])
+        return row
+
+    async def delete(self, source_id: str) -> bool:
+        result = await self._session.execute(
+            select(ModelSource)
+            .options(
+                selectinload(ModelSource.models),
+                selectinload(ModelSource.api_key_assignments),
+            )
+            .where(ModelSource.id == source_id)
+        )
+        row = result.scalar_one_or_none()
+        if row is None:
+            return False
+        await self._session.delete(row)
+        await self._session.commit()
+        return True
+
+    async def replace_models(
+        self,
+        source: ModelSource,
+        models: list[ModelSourceModel],
+        *,
+        commit: bool = True,
+    ) -> None:
+        await self._session.execute(delete(ModelSourceModel).where(ModelSourceModel.source_id == source.id))
+        for model in models:
+            model.source_id = source.id
+            self._session.add(model)
+        if commit:
+            await self._session.commit()
+            await self._session.refresh(source, attribute_names=["models"])
+
+    async def refresh_models(self, source: ModelSource) -> None:
+        await self._session.refresh(source, attribute_names=["models"])
+
+    async def commit(self) -> None:
+        await self._session.commit()
+
+    async def rollback(self) -> None:
+        await self._session.rollback()

app/modules/model_sources/schemas.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from datetime import datetime
+
+from pydantic import Field
+
+from app.modules.shared.schemas import DashboardModel
+
+
+class ModelSourceModelInput(DashboardModel):
+    model: str = Field(min_length=1, max_length=255)
+    display_name: str | None = Field(default=None, max_length=255)
+    context_window: int | None = Field(default=None, ge=1)
+    max_output_tokens: int | None = Field(default=None, ge=1)
+    supports_streaming: bool = True
+    supports_tools: bool = False
+    supports_vision: bool = False
+    input_per_1m: float | None = Field(default=None, ge=0)
+    cached_input_per_1m: float | None = Field(default=None, ge=0)
+    output_per_1m: float | None = Field(default=None, ge=0)
+    audio_per_minute: float | None = Field(default=None, ge=0)
+    raw_metadata_json: str | None = None
+    is_enabled: bool = True
+
+
+class ModelSourceModelResponse(ModelSourceModelInput):
+    id: int
+    source_id: str
+    created_at: datetime
+    updated_at: datetime
+
+
+class ModelSourceCreateRequest(DashboardModel):
+    name: str = Field(min_length=1, max_length=128)
+    base_url: str = Field(min_length=1, max_length=2048)
+    api_key: str | None = Field(default=None, min_length=1)
+    supports_chat_completions: bool = True
+    supports_responses: bool = False
+    supports_audio_transcriptions: bool = False
+    supports_embeddings: bool = False
+    timeout_seconds: int | None = Field(default=None, ge=1)
+    max_concurrency: int | None = Field(default=None, ge=1)
+    models: list[ModelSourceModelInput] = Field(default_factory=list)
+
+
+class ModelSourceUpdateRequest(DashboardModel):
+    name: str | None = Field(default=None, min_length=1, max_length=128)
+    base_url: str | None = Field(default=None, min_length=1, max_length=2048)
+    api_key: str | None = Field(default=None, min_length=1)
+    is_enabled: bool | None = None
+    supports_chat_completions: bool | None = None
+    supports_responses: bool | None = None
+    supports_audio_transcriptions: bool | None = None
+    supports_embeddings: bool | None = None
+    timeout_seconds: int | None = Field(default=None, ge=1)
+    max_concurrency: int | None = Field(default=None, ge=1)
+    models: list[ModelSourceModelInput] | None = None
+
+
+class ModelSourceResponse(DashboardModel):
+    id: str
+    name: str
+    kind: str
+    base_url: str
+    is_enabled: bool
+    health_status: str
+    supports_chat_completions: bool
+    supports_responses: bool
+    supports_audio_transcriptions: bool
+    supports_embeddings: bool
+    timeout_seconds: int | None
+    max_concurrency: int | None
+    created_at: datetime
+    updated_at: datetime
+    models: list[ModelSourceModelResponse] = Field(default_factory=list)
+
+
+class ModelSourcesResponse(DashboardModel):
+    sources: list[ModelSourceResponse] = Field(default_factory=list)

app/modules/model_sources/selection.py

@@ -0,0 +1,166 @@
+"""Shared model-source selection helpers.
+
+Both the HTTP request handlers and the WebSocket session path need to decide
+whether a requested model is served by an OpenAI-compatible model source.
+Keeping that decision in one module stops the two transports from drifting
+apart: previously only the HTTP handlers consulted model sources, so a
+source-owned model requested over WebSocket fell through to subscription
+account selection and was rejected upstream.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from app.core.openai.model_registry import get_model_registry
+from app.db.models import ModelSource
+from app.db.session import detach_session_objects, get_background_session
+from app.modules.api_keys.service import ApiKeyData
+from app.modules.model_sources.repository import ModelSourcesRepository
+
+logger = logging.getLogger(__name__)
+
+
+def allowed_source_ids_for_api_key(api_key: ApiKeyData | None) -> set[str] | None:
+    """Source ids an API key may use, or ``None`` when scoping is disabled."""
+    if api_key is None or not api_key.source_assignment_scope_enabled:
+        return None
+    return set(api_key.assigned_source_ids)
+
+
+async def select_responses_model_source(
+    model: str,
+    api_key: ApiKeyData | None,
+    *,
+    raw_model: str | None = None,
+    require_streaming: bool = False,
+    only_disabled: bool = False,
+) -> tuple[ModelSource, str] | None:
+    """Resolve ``model`` to a Responses-capable model source, if any.
+
+    ``only_disabled`` inverts the enabled-state filter and leaves every other
+    rule -- candidate order, API key model allowlist, source assignment scope,
+    subscription-registry precedence, route shape, streaming -- untouched, so
+    the answer is exactly "the source this request would have used, had the
+    operator not switched it off". Callers use it after the ordinary lookup
+    misses, to refuse the request instead of handing a source-owned model to a
+    subscription account that cannot serve it. Sharing one function keeps the
+    two lookups from drifting apart the way the transports once did.
+    """
+    assigned_source_ids = allowed_source_ids_for_api_key(api_key)
+    exact_allowed_models = set(api_key.allowed_models) if api_key and api_key.allowed_models else None
+    candidates = [candidate for candidate in (raw_model, model) if candidate]
+    if not candidates:
+        return None
+    deduped_candidates = list(dict.fromkeys(candidates))
+    registry_models = get_model_registry().get_models_with_fallback()
+    async with get_background_session() as session:
+        repository = ModelSourcesRepository(session)
+        for candidate in deduped_candidates:
+            if exact_allowed_models is not None and candidate not in exact_allowed_models:
+                continue
+            subscription_model = registry_models.get(candidate)
+            if assigned_source_ids is None and subscription_model is not None:
+                continue
+            source = await repository.find_responses_source_for_model(
+                candidate,
+                allowed_source_ids=assigned_source_ids,
+                require_streaming=require_streaming,
+                only_disabled=only_disabled,
+            )
+            if source is not None:
+                break
+        else:
+            source = None
+        # ``close_session`` rolls back the read transaction, which would
+        # expire the loaded row; detach it so the forwarding path can read
+        # its attributes after this session boundary.
+        detach_session_objects(session)
+        return (source, candidate) if source is not None else None
+
+
+def effective_model_for_api_key(api_key: ApiKeyData | None, requested_model: str | None) -> str | None:
+    """The model an API key forces, falling back to the requested one."""
+    if api_key is None or api_key.enforced_model is None:
+        return requested_model
+    return api_key.enforced_model
+
+
+async def responses_model_is_source_owned(
+    model: str | None,
+    api_key: ApiKeyData | None,
+    *,
+    raw_model: str | None = None,
+) -> bool:
+    """True when ``model`` is served by a Responses-capable source, even a disabled one.
+
+    Used by the WebSocket path, which cannot forward to a model source and must
+    fail the session so the client falls back to the HTTP transport.
+
+    Disabled-source ownership counts as ownership here: the WebSocket transport
+    cannot serve a source-owned model regardless of the source's enabled state,
+    and dispatching the turn to a subscription account instead produces the
+    opaque upstream 400 (``The '<model>' model is not supported when using
+    Codex with a ChatGPT account.``). Bouncing the turn to HTTP lands it on the
+    disabled-source denial there, which answers 503 ``model_source_disabled``
+    and names the actual condition. The second, ``only_disabled`` probe runs
+    only after the ordinary lookup misses, so enabled sources keep resolving in
+    a single query.
+
+    The API key's ``enforced_model`` is considered alongside the requested
+    model, matching how the HTTP handlers build their candidate list: an
+    enforced model that resolves to a source must not slip through to
+    subscription-account selection.
+
+    ``raw_model`` is the client's requested model captured before request
+    preparation normalized aliases (``gpt-5-high`` -> ``gpt-5``), mirroring the
+    HTTP path's ``raw_source_model``: the caller has already substituted the
+    API key's ``enforced_model`` and applied the fast-mode correction, so it is
+    used verbatim as the leading source candidate. When omitted (request states
+    that predate preparation, e.g. replayed turns), the raw candidate is
+    derived from ``enforced_model``/``model`` as before.
+
+    Resolution failures fail open to ``False``. This helper only gates the
+    WebSocket transport, where the alternative is worse: the lookup runs after
+    the turn's usage reservation is acquired but before it is registered for
+    cleanup, so a propagating database error tears the whole session down and
+    strands the reservation until the stale reaper runs. Failing open degrades
+    to the pre-guard behaviour (the subscription upstream rejects the model),
+    and source forwarding could not have worked anyway — it needs the same
+    database for the source's credentials. The HTTP handlers deliberately do
+    not use this helper: they call ``select_responses_model_source`` directly
+    and must keep surfacing resolution errors rather than silently routing
+    source traffic to a subscription account.
+    """
+    raw = raw_model if raw_model is not None else effective_model_for_api_key(api_key, model)
+    if not model and not raw:
+        return False
+    try:
+        if (
+            await select_responses_model_source(
+                model or raw or "",
+                api_key,
+                raw_model=raw,
+                require_streaming=True,
+            )
+            is not None
+        ):
+            return True
+        return (
+            await select_responses_model_source(
+                model or raw or "",
+                api_key,
+                raw_model=raw,
+                require_streaming=True,
+                only_disabled=True,
+            )
+            is not None
+        )
+    except Exception:
+        logger.warning(
+            "model_source_resolution_failed_open model=%s raw_model=%s",
+            model,
+            raw,
+            exc_info=True,
+        )
+        return False