scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

feat(mcp/v2): scaffold resolve() on a per-mode discriminated config union

Browse source 
Phase 1, step 1. ServerSpec now carries one per-mode `config` (the AuthConfig discriminated
union of seven frozen models, one per AuthSpecKind) and derives auth_spec_kind from it, so
the mode is a single source of truth and an illegal combination (e.g. an aws_sigv4 server
holding OAuth fields, or a mode missing required fields) is rejected at construction rather
than at call time.

resolve() moves to upstream_credentials.py and dispatches on the config variant via a
wildcard-free class-pattern match with an assert_never tail; basedpyright gates this shape
exactly like the enum (removing an arm fails reportMatchNotExhaustive + assert_never). Each
arm receives its own fully-typed config, so there are no None-checks.

The self-contained arms are real: none (NoOpAuth), api_key (StaticHeaderAuth carrying the
v1 scheme prefix), and passthrough (forwards the inbound token, fails closed when absent).
The OAuth-flow and aws_sigv4 arms are typed fail-closed stubs awaiting their collaborators
(token store, OAuth providers, RFC 8693 exchanger, SigV4 signer), which are the next step.

Tests construct Subject/ServerSpec directly with zero v1 fixtures: union rejection, the
three real arms across every api_key scheme, the no-inbound-token-replay invariant, and
fail-closed stubs.
This commit is contained in:
Tin Chi Lo 2026-06-17 15:01:22 -07:00
parent 1163711bb4
commit d2c56ed656
4 changed files with 406 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm/proxy/_experimental/mcp_server/v2/oauth/httpx_auth.py Normal file
View file

@ -0,0 +1,38 @@
"""Concrete `httpx.Auth` objects the resolver returns for the self-contained modes.
These are the egress credential as the SDK consumes it: an `httpx.Auth` attached to the
upstream `AsyncClient`. The OAuth-flow modes (`authorization_code`, `client_credentials`,
`token_exchange`) return SDK-provided auth objects instead and land later.
`auth_flow` mutating the outbound request is the `httpx.Auth` contract, not a house-style
violation: the request is httpx's object, and these carry no state of their own.
"""
from __future__ import annotations
from collections.abc import Generator
import httpx
class NoOpAuth(httpx.Auth):
"""Attaches nothing — the `none` mode (and the seam-level default)."""
def auth_flow(
self, request: httpx.Request
) -> Generator[httpx.Request, httpx.Response, None]:
yield request
class StaticHeaderAuth(httpx.Auth):
"""Sets one fixed header on every request — the `api_key` family and `passthrough`."""
def __init__(self, header_value: str, header_name: str = "Authorization") -> None:
self.header_name = header_name
self.header_value = header_value
def auth_flow(
self, request: httpx.Request
) -> Generator[httpx.Request, httpx.Response, None]:
request.headers[self.header_name] = self.header_value
yield request

217
litellm/proxy/_experimental/mcp_server/v2/oauth/types.py
View file

@ -1,18 +1,19 @@
"""The OAuth / upstream-credential vocabulary — the v2 typed seam (Phase 0).
"""The OAuth / upstream-credential vocabulary — the v2 typed seam.
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.
This module is the *contract* the credential build implements and the spec tests assert
against. It ships the data types only; the resolver lives in `upstream_credentials.py`.
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 new 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.
- **Mode is the single source of truth.** A server declares exactly one per-mode `config`
(the `AuthConfig` discriminated union); `auth_spec_kind` is *derived* from it, never a
second field that can drift. The resolver dispatches on the config variant, one arm per
mode. No field-presence inference, no precedence cascade.
- **Illegal states unrepresentable.** Each mode's config is its own frozen model holding
only that mode's fields, all required — an `aws_sigv4` server cannot hold OAuth fields,
and bad config is rejected at construction, not at call time.
- **Fail-closed at the boundary.** A raw mode string can only enter through
`parse_auth_spec_kind()`, which returns a typed `CredError`.
- **Errors as values.** Every seam returns `Result[_, CredError]`; only edge adapters raise.
- **Clean-room.** No imports from v1.
@ -24,18 +25,17 @@ discriminated on a `Literal` `tag`, matched via `self.tag` with an `assert_never
from __future__ import annotations
from enum import Enum
from typing import Literal
from typing import Annotated, Literal
import httpx
from expression import case, tag, tagged_union
from pydantic import BaseModel, ConfigDict
from pydantic import BaseModel, ConfigDict, Field
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.
"""The server's statically-declared upstream-auth mode — derived from its `config`.
Covers v1's full `MCPAuth` surface, not only OAuth grants: the three grant modes, the
collapsed static-header family, client passthrough, no-auth, and AWS request signing.
@ -45,9 +45,9 @@ class AuthSpecKind(str, Enum):
scheme is a parameter the arm carries, not its own mode.
"""
authorization_code = "authorization_code" # per-user 3LO; gateway is OAuth client
authorization_code = "authorization_code" # per-user 3LO; gateway-stored token
client_credentials = "client_credentials" # gateway service account (M2M)
token_exchange = "token_exchange" # RFC 8693 on-behalf-of
token_exchange = "token_exchange" # RFC 8693: token endpoint + subject_token (OBO)
api_key = "api_key" # static header, any scheme (BYOK = per-user-seeded source)
passthrough = "passthrough" # client forwards an upstream-audience token
none = "none" # no upstream credential; resolve yields a no-op auth, never an error
@ -111,6 +111,108 @@ class CredError:
assert_never(self.tag)
ApiKeyScheme = Literal["bearer", "apikey", "basic", "token", "raw"]
class AuthorizationCodeConfig(BaseModel):
"""Per-user 3LO; the gateway is the OAuth client and stores the user's token."""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.authorization_code] = AuthSpecKind.authorization_code
client_id: str
client_secret: str
authorization_url: str
token_url: str
scopes: tuple[str, ...] = ()
class ClientCredentialsConfig(BaseModel):
"""M2M service account; one upstream identity for every user."""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.client_credentials] = AuthSpecKind.client_credentials
client_id: str
client_secret: str
token_url: str
scopes: tuple[str, ...] = ()
class TokenExchangeConfig(BaseModel):
"""RFC 8693 OBO; swap the caller's live subject_token for an upstream-audience token."""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.token_exchange] = AuthSpecKind.token_exchange
token_exchange_endpoint: str
audience: str
subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token"
class ApiKeyConfig(BaseModel):
"""A fixed credential injected as a header. BYOK = the same arm, key seeded per-user."""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.api_key] = AuthSpecKind.api_key
value: str
scheme: ApiKeyScheme = "bearer"
def header_value(self) -> str:
# Reconstructs v1's per-scheme Authorization prefixes (mcp_server_manager.py:877-884).
match self.scheme:
case "bearer":
return f"Bearer {self.value}"
case "apikey":
return f"ApiKey {self.value}"
case "basic":
return f"Basic {self.value}"
case "token":
return f"token {self.value}"
case "raw":
return self.value
assert_never(self.scheme)
class PassthroughConfig(BaseModel):
"""Client-driven upstream OAuth; the gateway forwards the client's upstream token."""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.passthrough] = AuthSpecKind.passthrough
class NoneConfig(BaseModel):
"""No upstream credential; the request is sent unauthenticated."""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.none] = AuthSpecKind.none
class AwsSigV4Config(BaseModel):
"""AWS SigV4 per-request signing. Creds come from static keys, an assumed role, or
the ambient environment; that source is left loose here and tightened when the arm lands.
"""
model_config = ConfigDict(frozen=True)
kind: Literal[AuthSpecKind.aws_sigv4] = AuthSpecKind.aws_sigv4
region: str
service: str = "bedrock-agentcore"
access_key_id: str | None = None
secret_access_key: str | None = None
session_token: str | None = None
role_arn: str | None = None
session_name: str | None = None
AuthConfig = Annotated[
AuthorizationCodeConfig
| ClientCredentialsConfig
| TokenExchangeConfig
| ApiKeyConfig
| PassthroughConfig
| NoneConfig
| AwsSigV4Config,
Field(discriminator="kind"),
]
class Subject(BaseModel):
"""The validated inbound principal. NOT the v1 request object and NOT the LiteLLM key."""
@ -128,90 +230,21 @@ class ServerSpec(BaseModel):
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
config: AuthConfig
@property
def auth_spec_kind(self) -> AuthSpecKind:
return self.config.kind
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
Inside the core the mode is always a valid `AuthSpecKind`, so the resolver 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)
case AuthSpecKind.none:
return self._none(subject, server)
case AuthSpecKind.aws_sigv4:
return self._aws_sigv4(subject, server)
assert_never(server.auth_spec_kind)
# --- 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 _none(
self, subject: Subject, server: ServerSpec
) -> Result[httpx.Auth, CredError]:
return _todo(AuthSpecKind.none)
def _aws_sigv4(
self, subject: Subject, server: ServerSpec
) -> Result[httpx.Auth, CredError]:
return _todo(AuthSpecKind.aws_sigv4)
def _todo(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]:
return Error(
CredError.of_misconfigured(
f"{kind.value}: resolver arm not implemented (Phase 0 skeleton)"
)
)

109
litellm/proxy/_experimental/mcp_server/v2/oauth/upstream_credentials.py Normal file
View file

@ -0,0 +1,109 @@
"""The ONE credential resolver: `(subject, server) -> Result[httpx.Auth, CredError]`.
`resolve()` selects exactly one arm off the server's declared per-mode `config` and either
produces an `httpx.Auth` or fails closed with a typed `CredError`. There is no precedence
cascade and no silent downgrade; a header an attacker strips cannot change the mode.
The match is over the `AuthConfig` variant (not a separate enum), so each arm receives its
own fully-typed config with every field guaranteed present — no `None`-checks. The match is
wildcard-free with an `assert_never` tail, so adding a mode without an arm fails the type
gate, and a bypassed gate fails loudly at runtime instead of returning `None`.
Self-contained modes (`none`, `passthrough`, `api_key`) are implemented. The OAuth-flow and
signing modes are typed stubs that fail closed until their collaborators (token store, OAuth
providers, RFC 8693 exchanger, SigV4 signer) are injected.
"""
from __future__ import annotations
import httpx
from typing_extensions import assert_never
from ..result import Error, Ok, Result
from .httpx_auth import NoOpAuth, StaticHeaderAuth
from .types import (
ApiKeyConfig,
AuthorizationCodeConfig,
AuthSpecKind,
AwsSigV4Config,
ClientCredentialsConfig,
CredError,
NoneConfig,
PassthroughConfig,
ServerSpec,
Subject,
TokenExchangeConfig,
)
class UpstreamCredentialProvider:
"""Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode."""
def resolve(
self, subject: Subject, server: ServerSpec
) -> Result[httpx.Auth, CredError]:
match server.config:
case AuthorizationCodeConfig() as config:
return self._authorization_code(subject, server, config)
case ClientCredentialsConfig() as config:
return self._client_credentials(subject, server, config)
case TokenExchangeConfig() as config:
return self._token_exchange(subject, server, config)
case ApiKeyConfig() as config:
return self._api_key(subject, server, config)
case PassthroughConfig() as config:
return self._passthrough(subject, server, config)
case NoneConfig() as config:
return self._none(subject, server, config)
case AwsSigV4Config() as config:
return self._aws_sigv4(subject, server, config)
assert_never(server.config)
# --- self-contained arms (no injected collaborator) -----------------------------------
def _none(
self, subject: Subject, server: ServerSpec, config: NoneConfig
) -> Result[httpx.Auth, CredError]:
return Ok(NoOpAuth())
def _api_key(
self, subject: Subject, server: ServerSpec, config: ApiKeyConfig
) -> Result[httpx.Auth, CredError]:
return Ok(StaticHeaderAuth(config.header_value()))
def _passthrough(
self, subject: Subject, server: ServerSpec, config: PassthroughConfig
) -> Result[httpx.Auth, CredError]:
# The one arm that forwards a caller-supplied token, and only one the client obtained
# for the upstream's audience — never a gateway-bound credential.
if subject.inbound_token is None:
return Error(
CredError.of_unauthorized("passthrough: no inbound token to forward")
)
return Ok(StaticHeaderAuth(f"Bearer {subject.inbound_token}"))
# --- arms awaiting their collaborators (typed stubs, fail closed) ----------------------
def _authorization_code(
self, subject: Subject, server: ServerSpec, config: AuthorizationCodeConfig
) -> Result[httpx.Auth, CredError]:
return _todo(AuthSpecKind.authorization_code)
def _client_credentials(
self, subject: Subject, server: ServerSpec, config: ClientCredentialsConfig
) -> Result[httpx.Auth, CredError]:
return _todo(AuthSpecKind.client_credentials)
def _token_exchange(
self, subject: Subject, server: ServerSpec, config: TokenExchangeConfig
) -> Result[httpx.Auth, CredError]:
return _todo(AuthSpecKind.token_exchange)
def _aws_sigv4(
self, subject: Subject, server: ServerSpec, config: AwsSigV4Config
) -> Result[httpx.Auth, CredError]:
return _todo(AuthSpecKind.aws_sigv4)
def _todo(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]:
return Error(
CredError.of_misconfigured(f"{kind.value}: resolver arm not implemented yet")
)

134
tests/mcp_tests/v2/test_upstream_credentials.py Normal file
View file

@ -0,0 +1,134 @@
"""Spec tests for the v2 upstream-credential resolver scaffold.
Clean-room litmus: every case constructs `Subject` / `ServerSpec` directly, with zero v1
fixtures. If an arm could not be exercised without a v1 request object, the seam has leaked.
"""
import httpx
import pytest
from pydantic import ValidationError
from litellm.proxy._experimental.mcp_server.v2.oauth.httpx_auth import (
NoOpAuth,
StaticHeaderAuth,
)
from litellm.proxy._experimental.mcp_server.v2.oauth.types import (
ApiKeyConfig,
AuthSpecKind,
NoneConfig,
PassthroughConfig,
ServerSpec,
Subject,
)
from litellm.proxy._experimental.mcp_server.v2.oauth.upstream_credentials import (
UpstreamCredentialProvider,
)
from litellm.proxy._experimental.mcp_server.v2.result import Error, Ok
PROVIDER = UpstreamCredentialProvider()
SUBJECT = Subject(tenant_id="t1", subject_id="u1")
def _spec(config: object) -> ServerSpec:
return ServerSpec(server_id="s1", resource="https://up.example/mcp", config=config) # type: ignore[arg-type]
def _applied_headers(auth: httpx.Auth) -> httpx.Headers:
request = httpx.Request("POST", "https://up.example/mcp")
return next(auth.auth_flow(request)).headers
def test_auth_spec_kind_is_derived_from_config():
spec = _spec(ApiKeyConfig(value="k"))
assert spec.auth_spec_kind is AuthSpecKind.api_key
def test_discriminated_union_rejects_config_missing_required_fields():
# authorization_code requires client_id/secret/urls; an empty body must fail at construction.
with pytest.raises(ValidationError):
_spec({"kind": "authorization_code"})
def test_discriminated_union_picks_the_variant_by_kind():
spec = _spec({"kind": "none"})
assert isinstance(spec.config, NoneConfig)
def test_none_attaches_no_credential():
result = PROVIDER.resolve(SUBJECT, _spec(NoneConfig()))
assert isinstance(result, Ok)
assert isinstance(result.ok, NoOpAuth)
assert "Authorization" not in _applied_headers(result.ok)
@pytest.mark.parametrize(
"scheme,expected",
[
("bearer", "Bearer k"),
("apikey", "ApiKey k"),
("basic", "Basic k"),
("token", "token k"),
("raw", "k"),
],
)
def test_api_key_emits_the_right_scheme(scheme: str, expected: str):
result = PROVIDER.resolve(SUBJECT, _spec(ApiKeyConfig(value="k", scheme=scheme))) # type: ignore[arg-type]
assert isinstance(result, Ok)
assert isinstance(result.ok, StaticHeaderAuth)
assert _applied_headers(result.ok)["Authorization"] == expected
def test_passthrough_forwards_the_inbound_token():
subject = Subject(tenant_id="t1", subject_id="u1", inbound_token="upstream-tok")
result = PROVIDER.resolve(subject, _spec(PassthroughConfig()))
assert isinstance(result, Ok)
assert _applied_headers(result.ok)["Authorization"] == "Bearer upstream-tok"
def test_passthrough_without_a_token_fails_closed():
result = PROVIDER.resolve(SUBJECT, _spec(PassthroughConfig()))
assert isinstance(result, Error)
assert result.error.tag == "unauthorized"
def test_self_contained_arms_never_read_the_inbound_token():
# The #30559 guard: none/api_key must produce the identical credential whether or not a
# caller bearer is present, proving they never forward the gateway-bound token.
with_token = Subject(tenant_id="t1", subject_id="u1", inbound_token="leak-me")
for config in (NoneConfig(), ApiKeyConfig(value="k")):
without = PROVIDER.resolve(SUBJECT, _spec(config))
present = PROVIDER.resolve(with_token, _spec(config))
assert isinstance(without, Ok) and isinstance(present, Ok)
assert _applied_headers(without.ok).get("Authorization") == _applied_headers(
present.ok
).get("Authorization")
@pytest.mark.parametrize(
"config",
[
{
"kind": "authorization_code",
"client_id": "c",
"client_secret": "s",
"authorization_url": "https://idp/auth",
"token_url": "https://idp/token",
},
{
"kind": "client_credentials",
"client_id": "c",
"client_secret": "s",
"token_url": "https://idp/token",
},
{
"kind": "token_exchange",
"token_exchange_endpoint": "https://idp/token",
"audience": "https://up.example",
},
{"kind": "aws_sigv4", "region": "us-east-1"},
],
)
def test_unimplemented_arms_fail_closed(config: dict):
result = PROVIDER.resolve(SUBJECT, _spec(config))
assert isinstance(result, Error)
assert result.error.tag == "misconfigured"