From 49e12c872d55b06742a89e16198c2ffcb65bf4e5 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 18 Jun 2026 17:49:07 -0700 Subject: [PATCH] feat(mcp): graft per-user api_key (BYOK) to the v2 resolver BYOK is the api_key mode with the key seeded per-user; it grafts on the same header seam as the shared key, riding the Subject edge. The v2 _api_key arm already resolves the Byok key_source via the CredentialStore port, so this is graft wiring only: - V1ByokCredentialStore (v2_port_bodies) bridges the CredentialStore port to v1's existing read (db.get_user_credential -> find_unique + credential_b64 decrypt), keyed by (subject_id == user_id, server_id). Missing row -> Ok(None) (the arm returns 401); DB outage -> upstream_unavailable. An empty subject_id short-circuits to Ok(None) so an identity-less caller never shares a credential slot (fail closed). The reader is injected so it stays unit-testable without a DB. - _to_server_spec maps an is_byok api_key server to ApiKeyConfig(key_source=Byok); the arm pulls the per-user value from the store. _provider injects V1ByokCredentialStore. Tests cover the store body (present / missing / empty-subject-skips-store / DB-error) and the is_byok mapping. 92 tests pass; bridge typechecks clean, no new errors in the port bodies. Live e2e is deferred: the store reads LiteLLM_MCPUserCredentials, so it needs a DB-backed proxy with a seeded per-user credential (the local config-only harness has no DB). --- .../mcp_server/v2_port_bodies.py | 44 +++++++++++++++++- .../mcp_server/v2_resolver_bridge.py | 18 ++++++-- .../mcp_server/test_v2_port_bodies.py | 45 +++++++++++++++++++ .../mcp_server/test_v2_resolver_bridge.py | 21 +++++++++ 4 files changed, 123 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py b/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py index f1347978ca1..8216f209ac7 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py +++ b/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py @@ -9,7 +9,7 @@ from __future__ import annotations import asyncio from datetime import timedelta -from typing import Optional +from typing import Awaitable, Callable, Optional import httpx from pydantic import BaseModel, ConfigDict, SecretStr, ValidationError @@ -17,6 +17,9 @@ from pydantic import BaseModel, ConfigDict, SecretStr, ValidationError from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.gateway.mcp.outbound_credentials.clock import Clock, SystemClock +from litellm.proxy.gateway.mcp.outbound_credentials.credential_store import ( + CredentialKey, +) from litellm.proxy.gateway.mcp.outbound_credentials.token_store import StoredToken from litellm.proxy.gateway.mcp.outbound_credentials.types import ( AssumeRole, @@ -156,6 +159,45 @@ def _build_sigv4_auth(config: AwsSigV4Config) -> httpx.Auth: return MCPSigV4Auth(aws_region_name=config.region, aws_service_name=config.service) +async def _read_v1_user_credential(subject_id: str, server_id: str) -> Optional[str]: + from litellm.proxy._experimental.mcp_server.db import get_user_credential + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise RuntimeError("no DB client available for BYOK credential lookup") + return await get_user_credential(prisma_client, subject_id, server_id) + + +class V1ByokCredentialStore: + """Per-user BYOK credential store backed by v1's LiteLLM_MCPUserCredentials. + + Bridges the v2 CredentialStore port to v1's existing read (`db.get_user_credential`, which + does the find_unique + credential_b64 decrypt), keyed by (subject_id == user_id, server_id). + A missing row is `Ok(None)` (the arm turns that into a 401); a DB outage is + `upstream_unavailable`. The reader is injected so the arm stays unit-testable without a DB. + """ + + def __init__( + self, + reader: Callable[ + [str, str], Awaitable[Optional[str]] + ] = _read_v1_user_credential, + ) -> None: + self._reader = reader + + async def get(self, key: CredentialKey) -> Result[Optional[str], CredError]: + if not key.subject_id: + # No authenticated identity -> no per-user credential; never share one slot. + return Ok(None) + try: + value = await self._reader(key.subject_id, key.server_id) + except Exception as e: + return Error( + CredError.of_upstream_unavailable(f"BYOK credential lookup failed: {e}") + ) + return Ok(value) + + def _classify_sigv4_error(error: Exception) -> CredError: # botocore is an optional dependency without type stubs; match its connection-error classes # by name rather than importing it just to isinstance-check. diff --git a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py index bc1c6c8eefe..3641911d9f2 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py +++ b/litellm/proxy/_experimental/mcp_server/v2_resolver_bridge.py @@ -23,11 +23,9 @@ from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.v2_port_bodies import ( HttpxClientCredentialsFetcher, HttpxSigV4Signer, + V1ByokCredentialStore, ) from litellm.proxy.gateway.mcp.outbound_credentials.clock import SystemClock -from litellm.proxy.gateway.mcp.outbound_credentials.credential_store import ( - InMemoryCredentialStore, -) from litellm.proxy.gateway.mcp.outbound_credentials.resolver import ( UpstreamCredentialProvider, ) @@ -44,6 +42,7 @@ from litellm.proxy.gateway.mcp.outbound_credentials.types import ( AssumeRole, AuthorizationCodeConfig, AwsSigV4Config, + Byok, ClientCredentialsConfig, CredError, NoneConfig, @@ -97,7 +96,7 @@ def _provider() -> UpstreamCredentialProvider: # real bodies get wired in as their modes are grafted. unwired = _Unwired() return UpstreamCredentialProvider( - credential_store=InMemoryCredentialStore(), + credential_store=V1ByokCredentialStore(), token_store=InMemoryTokenStore(), token_refresher=unwired, clock=SystemClock(), @@ -115,6 +114,17 @@ def _to_server_spec(server: MCPServer) -> Optional[ServerSpec]: server_id=server.server_id, resource=resource, config=NoneConfig() ) if server.auth_type == MCPAuth.api_key: + if server.is_byok: + # Per-user key; the arm pulls it from the CredentialStore keyed by the Subject. + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=ApiKeyConfig( + header_name="X-API-Key", + value_prefix="", + key_source=Byok(), + ), + ) token = server.authentication_token if not token: return None # api_key with no key: let v1 handle it (parity-safe) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py index f3d1dd11412..4c4ad7af6bf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py @@ -11,8 +11,12 @@ from litellm.proxy._experimental.mcp_server import v2_port_bodies from litellm.proxy._experimental.mcp_server.v2_port_bodies import ( HttpxClientCredentialsFetcher, HttpxSigV4Signer, + V1ByokCredentialStore, _classify_sigv4_error, ) +from litellm.proxy.gateway.mcp.outbound_credentials.credential_store import ( + CredentialKey, +) from litellm.proxy.gateway.mcp.outbound_credentials.types import ( AwsSigV4Config, ClientCredentialsConfig, @@ -162,3 +166,44 @@ async def test_classify_sigv4_connection_error_is_upstream_unavailable(): async def test_classify_sigv4_other_error_is_misconfigured(): assert _classify_sigv4_error(ValueError("no creds")).tag == "misconfigured" + + +def _cred_key(subject_id="u1", server_id="s1"): + return CredentialKey(tenant_id="org1", subject_id=subject_id, server_id=server_id) + + +async def test_byok_store_returns_user_credential(): + async def reader(subject_id, server_id): + assert (subject_id, server_id) == ("u1", "s1") + return "user-byok-key" + + result = await V1ByokCredentialStore(reader=reader).get(_cred_key()) + assert isinstance(result, Ok) + assert result.ok == "user-byok-key" + + +async def test_byok_store_missing_credential_is_ok_none(): + async def reader(subject_id, server_id): + return None + + result = await V1ByokCredentialStore(reader=reader).get(_cred_key()) + assert isinstance(result, Ok) + assert result.ok is None + + +async def test_byok_store_empty_subject_skips_the_store(): + async def reader(subject_id, server_id): + raise AssertionError("store must not be queried for an empty subject") + + result = await V1ByokCredentialStore(reader=reader).get(_cred_key(subject_id="")) + assert isinstance(result, Ok) + assert result.ok is None + + +async def test_byok_store_db_error_is_upstream_unavailable(): + async def reader(subject_id, server_id): + raise RuntimeError("db down") + + result = await V1ByokCredentialStore(reader=reader).get(_cred_key()) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py index 31c20c556fb..0f3d1b542be 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_resolver_bridge.py @@ -274,3 +274,24 @@ async def test_resolve_v2_auth_value_threads_identity_without_breaking_static(v2 server, user_api_key_auth=auth, subject_token="jwt" ) assert result == {"X-API-Key": "up-secret"} + + +async def test_byok_server_maps_to_byok_key_source(): + from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( + _to_server_spec, + ) + from litellm.proxy.gateway.mcp.outbound_credentials.types import ApiKeyConfig, Byok + + server = MCPServer( + server_id="byok1", + name="byok1", + transport=MCPTransport.http, + url="https://up.example/mcp", + auth_type=MCPAuth.api_key, + is_byok=True, + ) + spec = _to_server_spec(server) + assert spec is not None + assert isinstance(spec.config, ApiKeyConfig) + assert isinstance(spec.config.key_source, Byok) + assert spec.config.header_name == "X-API-Key"