diff --git a/litellm/proxy/_experimental/mcp_server/v2/oauth/httpx_auth.py b/litellm/proxy/_experimental/mcp_server/v2/oauth/httpx_auth.py new file mode 100644 index 00000000000..d355dd246d6 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/oauth/httpx_auth.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/v2/oauth/types.py b/litellm/proxy/_experimental/mcp_server/v2/oauth/types.py index fec69d41334..8104c49eccb 100644 --- a/litellm/proxy/_experimental/mcp_server/v2/oauth/types.py +++ b/litellm/proxy/_experimental/mcp_server/v2/oauth/types.py @@ -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)" - ) - ) diff --git a/litellm/proxy/_experimental/mcp_server/v2/oauth/upstream_credentials.py b/litellm/proxy/_experimental/mcp_server/v2/oauth/upstream_credentials.py new file mode 100644 index 00000000000..90ad1e4cccd --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/v2/oauth/upstream_credentials.py @@ -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") + ) diff --git a/tests/mcp_tests/v2/test_upstream_credentials.py b/tests/mcp_tests/v2/test_upstream_credentials.py new file mode 100644 index 00000000000..0824fc1961b --- /dev/null +++ b/tests/mcp_tests/v2/test_upstream_credentials.py @@ -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"