feat(mcp): implement the api_key arm end-to-end with a per-user CredentialStore

Builds the api_key resolve() arm fully, including the injected credential pull. ApiKeyConfig
now carries a key_source discriminated union (SharedKey | PerUserKey): a shared key lives in
config, while a per-user / BYOK key is seeded per-subject and not static. header_for(value)
applies the v1 scheme prefix (Bearer/ApiKey/Basic/token/raw) to whichever value the arm
resolves.

UpstreamCredentialProvider gains a constructor that injects a CredentialStore (the new port
in credential_store.py, with an InMemoryCredentialStore body for tests and local wiring; the
durable DB-backed body lands later behind the same port). The api_key arm matches on
key_source: shared builds the header from config; per_user pulls the subject's secret via
self._credential_store.get(CredentialKey(tenant, subject, server)) and fails closed
(unauthorized) when absent. The pull lives inside resolve(); only the storage mechanics are
injected.

Tests cover the shared path across every scheme, the per-user hit, the fail-closed miss, and
per-subject isolation (another subject never receives u1's key), all constructed clean-room
with an in-memory store.

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)

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)
 
 

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

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)