From 3d2ef7472f0bcdd202cf2922bfae7ec8769d8eb2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 17 Jun 2026 13:45:31 -0700 Subject: [PATCH] feat(mcp/v2): stand up Mini Phase 0 typed OAuth credential seam Pulls forward the minimum from S0 and S8 needed for the clean-room v2 mini-chassis, the typed OAuth credential seam that Phase 1 gives a body. Nothing is wired into v1 yet; this phase has no runtime callers, so it is kept separate from the Phase 1 OAuth/credential build that follows. Adds expression to the proxy extra and basedpyright to the dev group so the FP spine is a declared dependency rather than relying on an ambient install. Ships the scoped types.py slice: the CredError tagged union, frozen Subject/ServerSpec, the boundary parser, and the resolve((subject, server)) -> Result[httpx.Auth, CredError] signature dispatched exhaustively on the declared auth_spec_kind with stub arms (no credential logic). Vendors Ok|Error as Result so wrong-side access is a type error. Includes the match-exhaustiveness spike that de-risks the whole approach: it proves Expression's @tagged_union plus match satisfies basedpyright strict reportMatchNotExhaustive, and README.md records the verdict plus the remove-an-arm reproduction showing the gate bites --- .../_experimental/mcp_server/v2/README.md | 65 ++++++ .../_experimental/mcp_server/v2/__init__.py | 16 ++ .../mcp_server/v2/_spike_exhaustiveness.py | 63 ++++++ .../mcp_server/v2/oauth/__init__.py | 1 + .../mcp_server/v2/oauth/types.py | 199 ++++++++++++++++++ .../mcp_server/v2/pyrightconfig.json | 10 + .../_experimental/mcp_server/v2/result.py | 73 +++++++ pyproject.toml | 2 + 8 files changed, 429 insertions(+) create mode 100644 litellm/proxy/_experimental/mcp_server/v2/README.md create mode 100644 litellm/proxy/_experimental/mcp_server/v2/__init__.py create mode 100644 litellm/proxy/_experimental/mcp_server/v2/_spike_exhaustiveness.py create mode 100644 litellm/proxy/_experimental/mcp_server/v2/oauth/__init__.py create mode 100644 litellm/proxy/_experimental/mcp_server/v2/oauth/types.py create mode 100644 litellm/proxy/_experimental/mcp_server/v2/pyrightconfig.json create mode 100644 litellm/proxy/_experimental/mcp_server/v2/result.py diff --git a/litellm/proxy/_experimental/mcp_server/v2/README.md b/litellm/proxy/_experimental/mcp_server/v2/README.md new file mode 100644 index 00000000000..1353f594a75 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/README.md @@ -0,0 +1,65 @@ +# MCP Gateway v2 — Mini Phase 0 (the typed OAuth seam) + +This directory is the clean-room v2 mini-chassis. It is the pre-Phase-1 graft that stands up +just enough of the typed/FP spine for the OAuth credential subdomain, so the highest-risk +security code is v2-shaped from line one instead of rewritten when the full S0 chassis lands in +Phase 2. It imports nothing from v1; v1 will only ever reach it through a thin adapter built in +Phase 1. + +Scope is deliberately small: the typed credential seam (`oauth/types.py`), the vendored `Result` +(`result.py`), and the basedpyright match-exhaustiveness spike (`_spike_exhaustiveness.py`). No +transport, registry, CI gate, semgrep rules, import-linter layers, LOC caps, or composition root; +those land with the full S0 in Phase 2. + +## The spike verdict: the bet HOLDS + +The one genuine technical risk of the whole FP approach is whether Expression's `@tagged_union` +plus `match` actually satisfies basedpyright strict `reportMatchNotExhaustive` at compile time. +Mini Phase 0 answers that before any OAuth logic commits to the style. The answer is yes, with two +non-obvious rules taken from the sibling v2 effort. `_spike_exhaustiveness.py` pins the only two +`match` shapes we rely on, and both pass strict: + +1. `match` over a closed `Enum` (what `resolve()` dispatches on against `AuthSpecKind`). +2. `match self.tag` over an Expression `@tagged_union` whose `tag` is a `Literal` (what `CredError` + uses for its `summary`). + +Both end in `assert_never(...)`. The exhaustiveness guarantee is real precisely because deleting an +arm breaks that tail: the scrutinee stops narrowing to `Never`, so `assert_never` becomes a type +error. That is the property we want, so the gate is load-bearing rather than decorative. + +### Rejected patterns (do not reintroduce) + +- `tag: str` plus `match err:` over the union object is NOT exhaustiveness-checked. One class with a + `str` tag never narrows to `Never`, so basedpyright demands a `case _:` and silently ignores + missing tags. We discriminate on a `Literal` tag instead. +- `expression.Result` is a single class carrying both `.ok` and `.error`, so an unguarded `.ok` + access is invisible to the checker. We vendor an `Ok | Error` union (`result.py`) so reaching for + the wrong side before a `match`/`isinstance` is a type error. + +## Reproducing the proof + +Run the gate (passes clean, 0 errors): + +``` +basedpyright --project litellm/proxy/_experimental/mcp_server/v2/pyrightconfig.json +``` + +Now prove the gate actually bites. Delete any one arm from `label_enum` in +`_spike_exhaustiveness.py` (for example the `AuthSpecKind.api_key` case) and re-run. basedpyright +reports two errors, which is the proof the exhaustiveness check is doing its job: + +``` +_spike_exhaustiveness.py:33:11 - error: Cases within match statement do not exhaustively handle all values + Unhandled type: "Literal[AuthSpecKind.api_key]" (reportMatchNotExhaustive) +_spike_exhaustiveness.py:44:18 - error: Argument of type "Literal[AuthSpecKind.api_key]" cannot be + assigned to parameter "arg" of type "Never" in function "assert_never" (reportArgumentType) +``` + +Restore the arm and the errors disappear. This is why adding a sixth `AuthSpecKind` member without a +`resolve()` arm fails the type gate rather than failing at runtime. + +## Toolchain notes + +`expression` is declared in the litellm `proxy` extra and `basedpyright` in the dev group +(`pyproject.toml`); the base SDK install is untouched. The spike was validated against +`expression 5.6.0` and `basedpyright 1.39.8` (pyright 1.1.410) on Python 3.10. diff --git a/litellm/proxy/_experimental/mcp_server/v2/__init__.py b/litellm/proxy/_experimental/mcp_server/v2/__init__.py new file mode 100644 index 00000000000..b09e8c5e3d8 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/__init__.py @@ -0,0 +1,16 @@ +"""MCP Gateway v2 — clean-room subdomain (Phase 0 mini-chassis). + +Written v2-native: nothing here imports from v1 (the rest of +`litellm.proxy._experimental.mcp_server.*`). v1 only ever reaches v2 through a thin +adapter built in Phase 1, never the other way round. + +Phase 0 scope = the typed OAuth-credential seam only (`oauth/types.py`), the vendored +`Result` (`result.py`), and the basedpyright match-exhaustiveness spike +(`_spike_exhaustiveness.py`). No transport, registry, or CI/semgrep/composition-root +infrastructure — those land in Phase 2 (S0). + +House style is adopted from the sibling v2 effort in `litellm/translation/` (see its +`CLAUDE.md`): product types are frozen dataclasses / frozen pydantic models, sum types are +Expression `@tagged_union` discriminated on a `Literal` `tag` (match on `self.tag`, +`assert_never` the tail), failures are values (`Result`, nothing raises in-tree). +""" diff --git a/litellm/proxy/_experimental/mcp_server/v2/_spike_exhaustiveness.py b/litellm/proxy/_experimental/mcp_server/v2/_spike_exhaustiveness.py new file mode 100644 index 00000000000..a2cc96d2a4d --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/_spike_exhaustiveness.py @@ -0,0 +1,63 @@ +"""Phase 0 spike — does `match` give us *compile-time* exhaustiveness under basedpyright strict? + +This is the one genuine technical risk of the FP approach (Full Migration Plan, Mini Phase 0, +goal 1). We answer it before any OAuth logic commits to the style. Verdict (see README.md): +the bet HOLDS, with two non-obvious rules taken from the sibling v2 package `litellm/translation`. + +Two `match` shapes are exhaustiveness-checked and are the only ones we use: + + (A) `match` over a closed `Enum` -> what `resolve()` dispatches on (`oauth/types.py`). + (B) `match self.tag` over an Expression `@tagged_union` whose `tag` is a `Literal` + -> what `CredError` uses for its `summary`. + +Both end in `assert_never(...)`: if any arm is deleted, the scrutinee no longer narrows to +`Never`, so `assert_never` becomes a type error — that "remove-an-arm" failure is the proof +the gate bites (reproduced in README.md). + +Rejected patterns and why (do NOT use; recorded so the finding is not re-litigated): + - `tag: str` + `match err:` over the union object -> NOT exhaustiveness-checked: one class, + `str` tag never narrows to `Never`. basedpyright demands `case _:` and ignores missing tags. + - `expression.Result` -> single class carrying both `.ok` and `.error`; unguarded `.ok` + access is invisible to the checker. Replaced by the vendored `Ok | Error` union (result.py). +""" + +from __future__ import annotations + +from typing_extensions import assert_never + +from .oauth.types import AuthSpecKind, CredError + + +# (A) Enum dispatch — the load-bearing case (`resolve()` uses exactly this shape). +def label_enum(kind: AuthSpecKind) -> str: + match kind: + case AuthSpecKind.authorization_code: + return "per-user 3LO" + case AuthSpecKind.client_credentials: + return "service account" + case AuthSpecKind.token_exchange: + return "on-behalf-of" + case AuthSpecKind.api_key: + return "static header" + case AuthSpecKind.passthrough: + return "client-forwarded" + # Reached only if an enum member has no arm above. basedpyright then narrows `kind` to that + # uncovered member (not `Never`), so this `assert_never` is a type error => the gate bit. + assert_never(kind) + + +# (B) @tagged_union discriminated on a Literal `tag` — matched via the tag, not the object. +def http_status(err: CredError) -> int: + match err.tag: + case "unauthorized": + return 401 + case "misconfigured": + return 500 + case "upstream_unavailable": + return 503 + case "unsupported_mode": + return 500 + assert_never(err.tag) + + +__all__ = ["label_enum", "http_status"] diff --git a/litellm/proxy/_experimental/mcp_server/v2/oauth/__init__.py b/litellm/proxy/_experimental/mcp_server/v2/oauth/__init__.py new file mode 100644 index 00000000000..bdbdb978e8a --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/oauth/__init__.py @@ -0,0 +1 @@ +"""OAuth / upstream-credential subdomain (Phase 0 typed seam).""" diff --git a/litellm/proxy/_experimental/mcp_server/v2/oauth/types.py b/litellm/proxy/_experimental/mcp_server/v2/oauth/types.py new file mode 100644 index 00000000000..3781ddef35c --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/oauth/types.py @@ -0,0 +1,199 @@ +"""The OAuth / upstream-credential vocabulary — the v2 typed seam (Phase 0). + +This module is the *contract* the Phase 1 build implements and the spec tests assert +against. It ships types and one stub `resolve()`; no credential logic lands in Phase 0. + +Design invariants encoded here: + +- **Mode is the single source of truth.** `resolve()` dispatches on the server's declared + `AuthSpecKind`, one arm per mode. No field-presence inference, no precedence cascade. +- **Exhaustive dispatch, no wildcard.** `auth_spec_kind` is a closed enum, so the `match` in + `resolve()` has no `_` arm — basedpyright (`reportMatchNotExhaustive`) guarantees every mode + is handled. Adding a sixth mode without an arm fails the type gate, not at runtime. +- **Fail-closed at the boundary.** An unknown mode string can only enter through + `parse_auth_spec_kind()`, which returns a typed `CredError`. Inside the core the mode is + always valid, so illegal states are unrepresentable. +- **Errors as values.** Every seam returns `Result[_, CredError]`; only edge adapters raise. +- **Clean-room.** No imports from v1. + +House style follows `litellm/translation/` (sum types = Expression `@tagged_union` +discriminated on a `Literal` `tag`, matched via `self.tag` with an `assert_never` tail; +`Result` is the vendored `Ok | Error` union, not `expression.Result`). +""" + +from __future__ import annotations + +from enum import Enum +from typing import Literal + +import httpx +from expression import case, tag, tagged_union +from pydantic import BaseModel, ConfigDict +from typing_extensions import assert_never + +from ..result import Error, Ok, Result + + +class AuthSpecKind(str, Enum): + """The server's statically-declared upstream-auth mode — the single source of truth. + + Canonical names follow the OAuth-grant vocabulary (see the Notion "resolve() Per-Mode + Behavior" page). BYOK is *not* a member: it is the `api_key` mode seeded per-user, a + source selector inside that arm. + """ + + authorization_code = ( + "authorization_code" # per-user 3LO; gateway is the OAuth client + ) + client_credentials = "client_credentials" # gateway service account (M2M) + token_exchange = "token_exchange" # RFC 8693 on-behalf-of + api_key = "api_key" # fixed/static header (BYOK = per-user-seeded source) + passthrough = "passthrough" # client forwards an upstream-audience token + + +@tagged_union(frozen=True) +class CredError: + """Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx.Auth`. + + Discriminated on the `Literal` `tag`; consumers `match self.tag` (see `summary`) so the + type checker can prove exhaustiveness. Construct via the `of_*` factories. + """ + + tag: Literal[ + "unauthorized", "misconfigured", "upstream_unavailable", "unsupported_mode" + ] = tag() + + unauthorized: str = ( + case() + ) # no usable credential for this (subject, server) -> 401 challenge + misconfigured: str = ( + case() + ) # the declared mode is missing required config -> 5xx (operator) + upstream_unavailable: str = ( + case() + ) # the IdP / token endpoint could not be reached -> 503 + unsupported_mode: str = ( + case() + ) # a raw mode string did not parse into AuthSpecKind (boundary) + + @staticmethod + def of_unauthorized(detail: str) -> CredError: + return CredError(unauthorized=detail) + + @staticmethod + def of_misconfigured(detail: str) -> CredError: + return CredError(misconfigured=detail) + + @staticmethod + def of_upstream_unavailable(detail: str) -> CredError: + return CredError(upstream_unavailable=detail) + + @staticmethod + def of_unsupported_mode(detail: str) -> CredError: + return CredError(unsupported_mode=detail) + + @property + def summary(self) -> str: + # Exhaustiveness: every Literal tag has an arm; the trailing assert_never typechecks + # only while that stays true (a `case _` would defeat reportMatchNotExhaustive). + match self.tag: + case "unauthorized": + return f"unauthorized: {self.unauthorized}" + case "misconfigured": + return f"misconfigured: {self.misconfigured}" + case "upstream_unavailable": + return f"upstream unavailable: {self.upstream_unavailable}" + case "unsupported_mode": + return self.unsupported_mode + assert_never(self.tag) + + +class Subject(BaseModel): + """The validated inbound principal. NOT the v1 request object and NOT the LiteLLM key.""" + + model_config = ConfigDict(frozen=True) + + tenant_id: str + subject_id: str + # Opaque, already-validated inbound identity. Only `token_exchange` / `passthrough` read it. + inbound_token: str | None = None + + +class ServerSpec(BaseModel): + """The declared upstream. A v2-native type; the v1->v2 adapter maps onto this.""" + + model_config = ConfigDict(frozen=True) + + server_id: str + auth_spec_kind: AuthSpecKind + resource: str # RFC 8707 audience URI this upstream's tokens are bound to + + +def parse_auth_spec_kind(raw: str) -> Result[AuthSpecKind, CredError]: + """Boundary parser — the *only* place an unknown mode is handled, and it fails closed. + + Inside the core the mode is always a valid `AuthSpecKind`, so `resolve()` never needs a + wildcard arm and basedpyright can prove its `match` exhaustive. + """ + try: + return Ok(AuthSpecKind(raw)) + except ValueError: + return Error(CredError.of_unsupported_mode(f"unknown auth_spec_kind: {raw!r}")) + + +class UpstreamCredentialProvider: + """The ONE credential resolver. Phase 0 ships the seam; arms are stubs filled in Phase 1.""" + + def resolve( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + """Select exactly one provider off the declared mode, or fail closed. + + The `match` is intentionally wildcard-free: every `AuthSpecKind` member must have an + arm or the type gate fails. This is the property the Phase 0 spike verifies. + """ + match server.auth_spec_kind: + case AuthSpecKind.authorization_code: + return self._authorization_code(subject, server) + case AuthSpecKind.client_credentials: + return self._client_credentials(subject, server) + case AuthSpecKind.token_exchange: + return self._token_exchange(subject, server) + case AuthSpecKind.api_key: + return self._api_key(subject, server) + case AuthSpecKind.passthrough: + return self._passthrough(subject, server) + + # --- arms: Phase 0 stubs (errors-as-values, no raise). Filled in Phase 1. ------------- + def _authorization_code( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + return _todo(AuthSpecKind.authorization_code) + + def _client_credentials( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + return _todo(AuthSpecKind.client_credentials) + + def _token_exchange( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + return _todo(AuthSpecKind.token_exchange) + + def _api_key( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + return _todo(AuthSpecKind.api_key) + + def _passthrough( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + return _todo(AuthSpecKind.passthrough) + + +def _todo(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: + return Error( + CredError.of_misconfigured( + f"{kind.value}: resolver arm not implemented (Phase 0 skeleton)" + ) + ) diff --git a/litellm/proxy/_experimental/mcp_server/v2/pyrightconfig.json b/litellm/proxy/_experimental/mcp_server/v2/pyrightconfig.json new file mode 100644 index 00000000000..6ad94a4ab1b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/pyrightconfig.json @@ -0,0 +1,10 @@ +{ + "include": ["."], + "typeCheckingMode": "strict", + "pythonVersion": "3.10", + "venvPath": "../../../../..", + "venv": ".venv", + "reportMissingTypeStubs": false, + "reportUnknownMemberType": true, + "reportMatchNotExhaustive": true +} diff --git a/litellm/proxy/_experimental/mcp_server/v2/result.py b/litellm/proxy/_experimental/mcp_server/v2/result.py new file mode 100644 index 00000000000..0cf2fa52621 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/result.py @@ -0,0 +1,73 @@ +"""A tagged union the type checker can actually see. + +``Ok`` and ``Error`` are separate classes joined by a ``Union`` alias, so reaching for +``result.ok`` before eliminating the ``Error`` arm (via ``isinstance`` or a ``match`` +pattern) is a pyright error, not a runtime ``AttributeError``. The single-class +``expression.Result`` this replaces declared both payload fields on one class, which made +unguarded access invisible to tooling. + +Both variants are covariant and frozen; the absent side defaults to ``Never`` so a bare +``Ok(value)`` or ``Error(err)`` infers fully and is assignable to any ``Result`` whose +matching side fits. + +``is_ok``/``is_error`` are runtime predicates only; they do not narrow the union for the +type checker. Inside the strictly-checked v2 tree, discriminate with ``match`` or +``isinstance``. + +Adopted verbatim from the sibling v2 package ``litellm/translation/result.py`` so the two +clean-room efforts share one ``Result`` shape and one rationale. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Generic, Literal, TypeAlias + +from typing_extensions import Never, TypeVar + +_TOk_co = TypeVar("_TOk_co", covariant=True, default=Never) +_TError_co = TypeVar("_TError_co", covariant=True, default=Never) +_TMapped = TypeVar("_TMapped") +_TBindError = TypeVar("_TBindError") + + +@dataclass(frozen=True) +class Ok(Generic[_TOk_co, _TError_co]): + ok: _TOk_co + + def is_ok(self) -> Literal[True]: + return True + + def is_error(self) -> Literal[False]: + return False + + def map(self, mapper: Callable[[_TOk_co], _TMapped]) -> Ok[_TMapped, _TError_co]: + return Ok(mapper(self.ok)) + + def bind( + self, mapper: Callable[[_TOk_co], Result[_TMapped, _TBindError]] + ) -> Result[_TMapped, _TBindError]: + return mapper(self.ok) + + +@dataclass(frozen=True) +class Error(Generic[_TOk_co, _TError_co]): + error: _TError_co + + def is_ok(self) -> Literal[False]: + return False + + def is_error(self) -> Literal[True]: + return True + + def map(self, mapper: Callable[[_TOk_co], _TMapped]) -> Error[_TMapped, _TError_co]: + return Error(self.error) + + def bind( + self, mapper: Callable[[_TOk_co], Result[_TMapped, _TBindError]] + ) -> Error[_TMapped, _TError_co]: + return Error(self.error) + + +Result: TypeAlias = Ok[_TOk_co, _TError_co] | Error[_TOk_co, _TError_co] diff --git a/pyproject.toml b/pyproject.toml index 8ee2840b573..77c462b79ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,7 @@ proxy = [ "soundfile>=0.12.1,<1.0", "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", "pydantic-settings>=2.14.1,<3.0", + "expression>=5.6.0,<6.0", ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy # imports (fastapi, cryptography, ...) are all guarded, so it runs on the base @@ -163,6 +164,7 @@ dev = [ "responses==0.26.0", "respx==0.22.0", "ruff==0.15.3", + "basedpyright==1.39.8", "types-requests==2.32.4.20260107", "types-setuptools==75.8.0.20250225", "types-redis==4.6.0.20241004",