From 6bae09bdeb20770016940135592ef5e226efc093 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 17 Jun 2026 16:38:54 -0700 Subject: [PATCH] feat(mcp): distinguish BYOK vs per-user env-var api_key by missing-status The per-user api_key source was one PerUserKey variant that returned a single CredError on a missing credential, collapsing v1's two distinct behaviours. Split it into Byok and PerUserEnvVar (the ApiKeySource union is now SharedKey | PerUserEnvVar | Byok). Both still pull the subject's value from the injected CredentialStore; they differ only when it is absent: BYOK returns unauthorized (401 + WWW-Authenticate, the user must provide it) while the env-var case returns the new precondition_required (412, a setup precondition). The edge mapping in http_status renders those as 401 vs 412, and a comment marks that SharedKey is read straight from ServerSpec.config rather than the per-user store. Adds tests for both missing paths and the distinct status mapping. --- .../gateway/mcp/_spike_exhaustiveness.py | 2 + .../mcp/outbound_credentials/resolver.py | 33 ++++++++++++----- .../gateway/mcp/outbound_credentials/types.py | 37 ++++++++++++++++--- tests/mcp_tests/gateway/test_resolver.py | 31 ++++++++++++---- 4 files changed, 81 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/gateway/mcp/_spike_exhaustiveness.py b/litellm/proxy/gateway/mcp/_spike_exhaustiveness.py index 1779633ba2d..623d3e5946f 100644 --- a/litellm/proxy/gateway/mcp/_spike_exhaustiveness.py +++ b/litellm/proxy/gateway/mcp/_spike_exhaustiveness.py @@ -61,6 +61,8 @@ def http_status(err: CredError) -> int: return 503 case "unsupported_mode": return 500 + case "precondition_required": + return 412 assert_never(err.tag) diff --git a/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py b/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py index d9314f213b8..658f55a13d1 100644 --- a/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py +++ b/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py @@ -28,11 +28,12 @@ from .types import ( AuthorizationCodeConfig, AuthSpecKind, AwsSigV4Config, + Byok, ClientCredentialsConfig, CredError, NoneConfig, PassthroughConfig, - PerUserKey, + PerUserEnvVar, ServerSpec, SharedKey, Subject, @@ -77,24 +78,38 @@ class UpstreamCredentialProvider: ) -> Result[httpx.Auth, CredError]: match config.key_source: case SharedKey() as source: + # The shared key is read straight from ServerSpec.config, not the per-user + # store; it is the same credential for every caller. 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, + case PerUserEnvVar(): + value = self._per_user_value(subject, server) + if value is None: + return Error( + CredError.of_precondition_required( + "api_key: per-user env var not set for this subject" + ) ) - ) + return Ok(StaticHeaderAuth(config.header_for(value))) + case Byok(): + value = self._per_user_value(subject, server) if value is None: return Error( CredError.of_unauthorized( - "api_key: no per-user credential for this subject" + "api_key: no BYOK credential for this subject" ) ) return Ok(StaticHeaderAuth(config.header_for(value))) assert_never(config.key_source) + def _per_user_value(self, subject: Subject, server: ServerSpec) -> str | None: + return self._credential_store.get( + CredentialKey( + tenant_id=subject.tenant_id, + subject_id=subject.subject_id, + server_id=server.server_id, + ) + ) + def _passthrough( self, subject: Subject, server: ServerSpec, config: PassthroughConfig ) -> Result[httpx.Auth, CredError]: diff --git a/litellm/proxy/gateway/mcp/outbound_credentials/types.py b/litellm/proxy/gateway/mcp/outbound_credentials/types.py index dd9d23213b7..f48609c7ea9 100644 --- a/litellm/proxy/gateway/mcp/outbound_credentials/types.py +++ b/litellm/proxy/gateway/mcp/outbound_credentials/types.py @@ -63,7 +63,11 @@ class CredError: """ tag: Literal[ - "unauthorized", "misconfigured", "upstream_unavailable", "unsupported_mode" + "unauthorized", + "misconfigured", + "upstream_unavailable", + "unsupported_mode", + "precondition_required", ] = tag() unauthorized: str = ( @@ -78,6 +82,9 @@ class CredError: unsupported_mode: str = ( case() ) # a raw mode string did not parse into AuthSpecKind (boundary) + precondition_required: str = ( + case() + ) # a required per-user value (e.g. an env var) has not been provided -> 412 @staticmethod def of_unauthorized(detail: str) -> CredError: @@ -95,6 +102,10 @@ class CredError: def of_unsupported_mode(detail: str) -> CredError: return CredError(unsupported_mode=detail) + @staticmethod + def of_precondition_required(detail: str) -> CredError: + return CredError(precondition_required=detail) + @property def summary(self) -> str: # Exhaustiveness: every Literal tag has an arm; the trailing assert_never typechecks @@ -108,6 +119,8 @@ class CredError: return f"upstream unavailable: {self.upstream_unavailable}" case "unsupported_mode": return self.unsupported_mode + case "precondition_required": + return f"precondition required: {self.precondition_required}" assert_never(self.tag) @@ -155,15 +168,27 @@ class SharedKey(BaseModel): 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.""" +class PerUserEnvVar(BaseModel): + """A per-user value the admin templated as an env var; the user fills it in. Pulled from + the credential store at resolve time. Missing means the user has not completed setup, a + precondition (412) rather than an auth failure.""" model_config = ConfigDict(frozen=True) - source: Literal["per_user"] = "per_user" + source: Literal["per_user_env_var"] = "per_user_env_var" -ApiKeySource = Annotated[SharedKey | PerUserKey, Field(discriminator="source")] +class Byok(BaseModel): + """A key the user brings via the entry flow, stored per-user. Pulled from the credential + store at resolve time. Missing means the user must provide it, a 401 + WWW-Authenticate + challenge.""" + + model_config = ConfigDict(frozen=True) + source: Literal["byok"] = "byok" + + +ApiKeySource = Annotated[ + SharedKey | PerUserEnvVar | Byok, Field(discriminator="source") +] class ApiKeyConfig(BaseModel): diff --git a/tests/mcp_tests/gateway/test_resolver.py b/tests/mcp_tests/gateway/test_resolver.py index aaa41276abb..508e153e900 100644 --- a/tests/mcp_tests/gateway/test_resolver.py +++ b/tests/mcp_tests/gateway/test_resolver.py @@ -16,12 +16,15 @@ from litellm.proxy.gateway.mcp.outbound_credentials.httpx_auth import ( NoOpAuth, StaticHeaderAuth, ) +from litellm.proxy.gateway.mcp._spike_exhaustiveness import http_status from litellm.proxy.gateway.mcp.outbound_credentials.types import ( ApiKeyConfig, AuthSpecKind, + Byok, + CredError, NoneConfig, PassthroughConfig, - PerUserKey, + PerUserEnvVar, ServerSpec, SharedKey, Subject, @@ -85,23 +88,31 @@ def test_api_key_emits_the_right_scheme(scheme: str, expected: str): assert _applied_headers(result.ok)["Authorization"] == expected -def test_api_key_per_user_pulls_the_subject_credential(): +@pytest.mark.parametrize("source", [Byok(), PerUserEnvVar()]) +def test_api_key_per_user_pulls_the_subject_credential(source: object): 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()))) + result = provider.resolve(SUBJECT, _spec(ApiKeyConfig(key_source=source))) # type: ignore[arg-type] 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()))) +def test_api_key_byok_missing_returns_unauthorized(): + # Missing BYOK credential -> 401 + WWW-Authenticate (the user must provide it). + result = PROVIDER.resolve(SUBJECT, _spec(ApiKeyConfig(key_source=Byok()))) assert isinstance(result, Error) assert result.error.tag == "unauthorized" +def test_api_key_env_var_missing_returns_precondition_required(): + # Missing per-user env var -> 412 (a setup precondition), distinct from BYOK's 401. + result = PROVIDER.resolve(SUBJECT, _spec(ApiKeyConfig(key_source=PerUserEnvVar()))) + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + + 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( @@ -109,11 +120,17 @@ def test_api_key_per_user_isolated_by_subject(): ) provider = UpstreamCredentialProvider(store) other = Subject(tenant_id="t1", subject_id="u2") - result = provider.resolve(other, _spec(ApiKeyConfig(key_source=PerUserKey()))) + result = provider.resolve(other, _spec(ApiKeyConfig(key_source=Byok()))) assert isinstance(result, Error) assert result.error.tag == "unauthorized" +def test_missing_status_maps_byok_401_distinct_from_env_var_412(): + # The two per-user sources surface different HTTP statuses at the edge. + assert http_status(CredError.of_unauthorized("byok missing")) == 401 + assert http_status(CredError.of_precondition_required("env var missing")) == 412 + + 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()))