scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(mcp): distinguish BYOK vs per-user env-var api_key by missing-status

Browse source 
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.
This commit is contained in:
Tin Chi Lo 2026-06-17 16:38:54 -07:00
parent 74baaaf197
commit 6bae09bdeb
4 changed files with 81 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/gateway/mcp/_spike_exhaustiveness.py
View file

@ -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)

33
litellm/proxy/gateway/mcp/outbound_credentials/resolver.py
View file

@ -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]:

37
litellm/proxy/gateway/mcp/outbound_credentials/types.py
View file

@ -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):

31
tests/mcp_tests/gateway/test_resolver.py
View file

@ -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()))