diff --git a/litellm/proxy/gateway/mcp/oauth/credential_store.py b/litellm/proxy/gateway/mcp/oauth/credential_store.py new file mode 100644 index 00000000000..02129833b75 --- /dev/null +++ b/litellm/proxy/gateway/mcp/oauth/credential_store.py @@ -0,0 +1,38 @@ +"""The per-subject credential store port for the `api_key` per-user / BYOK source. + +`resolve()` owns the *pull* (which key, fail-closed on a miss); this is only the storage +mechanics behind a `Protocol`, so the resolver stays testable without a database. The +durable body (over `LiteLLM_MCPUserEnvVars` and the BYOK credential table) lands later; +`InMemoryCredentialStore` is a working body for tests and local wiring. +""" + +from __future__ import annotations + +from typing import Protocol + +from pydantic import BaseModel, ConfigDict + + +class CredentialKey(BaseModel): + """Identifies a per-user secret. Per-tenant / per-user isolation is the key shape.""" + + model_config = ConfigDict(frozen=True) + tenant_id: str + subject_id: str + server_id: str + + +class CredentialStore(Protocol): + """Fetches the per-subject secret for an `api_key` per-user / BYOK server.""" + + def get(self, key: CredentialKey) -> str | None: ... + + +class InMemoryCredentialStore: + """A working in-memory `CredentialStore` for tests and local wiring.""" + + def __init__(self, seeded: dict[CredentialKey, str] | None = None) -> None: + self._values: dict[CredentialKey, str] = dict(seeded or {}) + + def get(self, key: CredentialKey) -> str | None: + return self._values.get(key) diff --git a/litellm/proxy/gateway/mcp/oauth/types.py b/litellm/proxy/gateway/mcp/oauth/types.py index 8104c49eccb..b563d702693 100644 --- a/litellm/proxy/gateway/mcp/oauth/types.py +++ b/litellm/proxy/gateway/mcp/oauth/types.py @@ -147,27 +147,47 @@ class TokenExchangeConfig(BaseModel): subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token" +class SharedKey(BaseModel): + """A fixed key configured on the server, identical for every caller.""" + + model_config = ConfigDict(frozen=True) + source: Literal["shared"] = "shared" + value: str + + +class PerUserKey(BaseModel): + """A key seeded per-user (per-user env var or BYOK). The value is not static; it is + pulled from the credential store at resolve time, keyed by the subject.""" + + model_config = ConfigDict(frozen=True) + source: Literal["per_user"] = "per_user" + + +ApiKeySource = Annotated[SharedKey | PerUserKey, Field(discriminator="source")] + + class ApiKeyConfig(BaseModel): - """A fixed credential injected as a header. BYOK = the same arm, key seeded per-user.""" + """A fixed credential injected as a header. The value is shared (in config) or seeded + per-user (pulled from the store); `scheme` is how it is written into the header.""" model_config = ConfigDict(frozen=True) kind: Literal[AuthSpecKind.api_key] = AuthSpecKind.api_key - value: str scheme: ApiKeyScheme = "bearer" + key_source: ApiKeySource - def header_value(self) -> str: + def header_for(self, value: str) -> str: # Reconstructs v1's per-scheme Authorization prefixes (mcp_server_manager.py:877-884). match self.scheme: case "bearer": - return f"Bearer {self.value}" + return f"Bearer {value}" case "apikey": - return f"ApiKey {self.value}" + return f"ApiKey {value}" case "basic": - return f"Basic {self.value}" + return f"Basic {value}" case "token": - return f"token {self.value}" + return f"token {value}" case "raw": - return self.value + return value assert_never(self.scheme) diff --git a/litellm/proxy/gateway/mcp/oauth/upstream_credentials.py b/litellm/proxy/gateway/mcp/oauth/upstream_credentials.py index 90ad1e4cccd..d9314f213b8 100644 --- a/litellm/proxy/gateway/mcp/oauth/upstream_credentials.py +++ b/litellm/proxy/gateway/mcp/oauth/upstream_credentials.py @@ -9,9 +9,10 @@ own fully-typed config with every field guaranteed present — no `None`-checks. 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. +Implemented: `none`, `passthrough`, and `api_key` (shared from config, or per-user / BYOK +pulled from the injected `CredentialStore`). 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 @@ -20,6 +21,7 @@ import httpx from typing_extensions import assert_never from ..result import Error, Ok, Result +from .credential_store import CredentialKey, CredentialStore from .httpx_auth import NoOpAuth, StaticHeaderAuth from .types import ( ApiKeyConfig, @@ -30,7 +32,9 @@ from .types import ( CredError, NoneConfig, PassthroughConfig, + PerUserKey, ServerSpec, + SharedKey, Subject, TokenExchangeConfig, ) @@ -39,6 +43,9 @@ from .types import ( class UpstreamCredentialProvider: """Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode.""" + def __init__(self, credential_store: CredentialStore) -> None: + self._credential_store = credential_store + def resolve( self, subject: Subject, server: ServerSpec ) -> Result[httpx.Auth, CredError]: @@ -59,7 +66,7 @@ class UpstreamCredentialProvider: return self._aws_sigv4(subject, server, config) assert_never(server.config) - # --- self-contained arms (no injected collaborator) ----------------------------------- + # --- implemented arms ----------------------------------------------------------------- def _none( self, subject: Subject, server: ServerSpec, config: NoneConfig ) -> Result[httpx.Auth, CredError]: @@ -68,7 +75,25 @@ class UpstreamCredentialProvider: def _api_key( self, subject: Subject, server: ServerSpec, config: ApiKeyConfig ) -> Result[httpx.Auth, CredError]: - return Ok(StaticHeaderAuth(config.header_value())) + match config.key_source: + case SharedKey() as source: + return Ok(StaticHeaderAuth(config.header_for(source.value))) + case PerUserKey(): + value = self._credential_store.get( + CredentialKey( + tenant_id=subject.tenant_id, + subject_id=subject.subject_id, + server_id=server.server_id, + ) + ) + if value is None: + return Error( + CredError.of_unauthorized( + "api_key: no per-user credential for this subject" + ) + ) + return Ok(StaticHeaderAuth(config.header_for(value))) + assert_never(config.key_source) def _passthrough( self, subject: Subject, server: ServerSpec, config: PassthroughConfig diff --git a/tests/mcp_tests/gateway/test_upstream_credentials.py b/tests/mcp_tests/gateway/test_upstream_credentials.py index 5170c9b971e..e1fdbed476e 100644 --- a/tests/mcp_tests/gateway/test_upstream_credentials.py +++ b/tests/mcp_tests/gateway/test_upstream_credentials.py @@ -8,6 +8,10 @@ import httpx import pytest from pydantic import ValidationError +from litellm.proxy.gateway.mcp.oauth.credential_store import ( + CredentialKey, + InMemoryCredentialStore, +) from litellm.proxy.gateway.mcp.oauth.httpx_auth import ( NoOpAuth, StaticHeaderAuth, @@ -17,7 +21,9 @@ from litellm.proxy.gateway.mcp.oauth.types import ( AuthSpecKind, NoneConfig, PassthroughConfig, + PerUserKey, ServerSpec, + SharedKey, Subject, ) from litellm.proxy.gateway.mcp.oauth.upstream_credentials import ( @@ -25,7 +31,7 @@ from litellm.proxy.gateway.mcp.oauth.upstream_credentials import ( ) from litellm.proxy.gateway.mcp.result import Error, Ok -PROVIDER = UpstreamCredentialProvider() +PROVIDER = UpstreamCredentialProvider(InMemoryCredentialStore()) SUBJECT = Subject(tenant_id="t1", subject_id="u1") @@ -39,7 +45,7 @@ def _applied_headers(auth: httpx.Auth) -> httpx.Headers: def test_auth_spec_kind_is_derived_from_config(): - spec = _spec(ApiKeyConfig(value="k")) + spec = _spec(ApiKeyConfig(key_source=SharedKey(value="k"))) assert spec.auth_spec_kind is AuthSpecKind.api_key @@ -72,12 +78,42 @@ def test_none_attaches_no_credential(): ], ) 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] + config = ApiKeyConfig(scheme=scheme, key_source=SharedKey(value="k")) # type: ignore[arg-type] + result = PROVIDER.resolve(SUBJECT, _spec(config)) assert isinstance(result, Ok) assert isinstance(result.ok, StaticHeaderAuth) assert _applied_headers(result.ok)["Authorization"] == expected +def test_api_key_per_user_pulls_the_subject_credential(): + store = InMemoryCredentialStore( + {CredentialKey(tenant_id="t1", subject_id="u1", server_id="s1"): "user-secret"} + ) + provider = UpstreamCredentialProvider(store) + result = provider.resolve(SUBJECT, _spec(ApiKeyConfig(key_source=PerUserKey()))) + assert isinstance(result, Ok) + assert _applied_headers(result.ok)["Authorization"] == "Bearer user-secret" + + +def test_api_key_per_user_missing_credential_fails_closed(): + # Empty store -> the per-user arm fails closed rather than sending no/garbage auth. + result = PROVIDER.resolve(SUBJECT, _spec(ApiKeyConfig(key_source=PerUserKey()))) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + +def test_api_key_per_user_isolated_by_subject(): + # The stored key belongs to (t1,u1,s1); a different subject must not receive it. + store = InMemoryCredentialStore( + {CredentialKey(tenant_id="t1", subject_id="u1", server_id="s1"): "u1-secret"} + ) + provider = UpstreamCredentialProvider(store) + other = Subject(tenant_id="t1", subject_id="u2") + result = provider.resolve(other, _spec(ApiKeyConfig(key_source=PerUserKey()))) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + 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())) @@ -95,7 +131,7 @@ 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")): + for config in (NoneConfig(), ApiKeyConfig(key_source=SharedKey(value="k"))): without = PROVIDER.resolve(SUBJECT, _spec(config)) present = PROVIDER.resolve(with_token, _spec(config)) assert isinstance(without, Ok) and isinstance(present, Ok)