diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 4fd47a97066..f746a42080a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -471,7 +471,8 @@ def _raise_if_not_oauth2(mcp_server: MCPServer) -> None: through: the caller owns the upstream token, and this relayed flow is how a browser obtains one against the upstream IdP (the admin UI's browser-only Authorize uses it). The minted token is upstream-audienced and held by the caller; the gateway persists nothing for these - modes (DCR persistence is opt-in and never enabled on this path). + modes (``_persist_dcr_client_registration`` skips them unconditionally, so even the admin + Authorize path with ``persist_credentials`` enabled writes nothing to the server row). """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, @@ -807,7 +808,7 @@ async def _reuse_persisted_dcr_client_if_available(mcp_server: MCPServer) -> boo return bool(mcp_server.client_id) -DcrRegistrationPersistenceResult = Literal["persisted", "reused", "failed"] +DcrRegistrationPersistenceResult = Literal["persisted", "reused", "skipped", "failed"] async def _persist_dcr_client_registration( @@ -821,7 +822,16 @@ async def _persist_dcr_client_registration( full re-authorization instead of a silent refresh. Mirrors the ``encrypt_credentials`` write that ``client_credentials`` and token exchange already use. Failures are logged, never raised: registration still returns to the caller even when persistence fails. + + The client-forwarded token modes (``true_passthrough`` / ``oauth_delegate``) are skipped + unconditionally: the caller holds the upstream token and the gateway must hold no OAuth + client identity for these servers. Persisting here would stamp ``oauth2_flow`` and a + ``client_id`` onto a server whose mode promises the gateway stores nothing, making a + fresh pass-through server read as gateway-authorized. """ + if mcp_server.is_true_passthrough or mcp_server.is_oauth_delegate: + return "skipped" + try: registration = _DcrClientRegistration.model_validate(registration_response) except ValidationError as exc: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 64226eea821..c88027abcd4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -149,6 +149,29 @@ async def test_authorization_code_isolates_by_subject(): assert isinstance(bob, Error) and bob.error.tag == "unauthorized" +@pytest.mark.asyncio +async def test_authorization_code_isolates_by_server_id_even_when_servers_share_a_url(): + """A token stored for one server must be invisible to a different server_id pointing at the + same upstream URL: credentials bind to the server entry they were authorized for, so a + recreated or duplicated server starts unauthorized instead of inheriting the old grant. Guards + against any future token lookup keyed on the resource URL instead of (user_id, server_id) -- + both the egress resolve and the has_user_token discovery check must agree.""" + shared_url = "https://upstream.example.com" + store = _FakeTokenStore({("alice", "server-a"): OAuthToken(access_token="at-alice")}) + provider = UpstreamCredentialProvider(oauth_token_store=store) + subject = Subject(tenant_id="", subject_id="alice") + spec_a = ServerSpec(server_id="server-a", resource=shared_url, config=AuthorizationCodeConfig()) + spec_b = ServerSpec(server_id="server-b", resource=shared_url, config=AuthorizationCodeConfig()) + + granted = await provider.resolve_credentials(subject, spec_a) + fresh = await provider.resolve_credentials(subject, spec_b) + + assert isinstance(granted, Ok) and _emitted(granted.ok)["Authorization"] == "Bearer at-alice" + assert isinstance(fresh, Error) and fresh.error.tag == "unauthorized" + assert await provider.has_user_token(subject, spec_a) is True + assert await provider.has_user_token(subject, spec_b) is False + + @pytest.mark.asyncio async def test_has_user_token_reflects_the_stored_token(): present = UpstreamCredentialProvider( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 9de342eabd7..c8e871b3bc1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -665,6 +665,178 @@ async def test_register_client_persists_dcr_client_identity(): mock_update_server.assert_called_once() +async def _register_persistence_attempted_for_auth_type(auth_type: MCPAuth) -> bool: + """Run register_client_with_server with persist_credentials=True for a server of ``auth_type`` + and report whether the DCR result was persisted onto the server row. The client-forwarded token + modes must skip the persist even on the admin path: writing it stamps oauth2_flow and a + client_id onto a server whose contract is that the gateway stores nothing, which makes a fresh + pass-through server read as gateway-authorized. The upstream registration must still be relayed + to the browser either way, since the caller needs the minted client to run its own flow.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="pt_server", + name="pt_server", + server_name="pt_server", + alias="pt_server", + transport=MCPTransport.http, + auth_type=auth_type, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + mock_response = MagicMock() + mock_response.json.return_value = { + "client_id": "generated-client", + "client_secret": "generated-secret", + "token_endpoint_auth_method": "none", + } + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + mock_update = AsyncMock(return_value=MagicMock()) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + ): + response = await register_client_with_server( + request=mock_request, + mcp_server=server, + client_name="Litellm Proxy", + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", + persist_credentials=True, + ) + + assert json.loads(response.body.decode("utf-8")) == mock_response.json.return_value + return mock_update.await_count > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type", [MCPAuth.true_passthrough, MCPAuth.oauth_delegate]) +async def test_register_client_does_not_persist_for_client_forwarded_modes(auth_type): + """The admin Authorize path (persist_credentials=True) must not write the DCR client onto a + true_passthrough / oauth_delegate server row: the browser still receives the registration, but + the gateway keeps no OAuth client identity for these modes.""" + assert await _register_persistence_attempted_for_auth_type(auth_type) is False + + +@pytest.mark.asyncio +async def test_register_client_persist_discriminator_oauth2_persists(): + """Guard the no-persist assertion above against vacuity: the same helper run against a genuine + oauth2 server DOES persist, so a regression that silently disables persistence everywhere (or a + helper that never reaches the persist) fails here instead of passing both.""" + assert await _register_persistence_attempted_for_auth_type(MCPAuth.oauth2) is True + + +@pytest.mark.asyncio +async def test_register_client_persists_only_to_its_own_row_when_another_server_shares_the_url(): + """A fresh server must mint and persist its OWN DCR client even when another server row with + the same upstream URL already holds one: both the reuse lookup and the persist are keyed by + server_id, never by URL, so OAuth client identity is not transferable between server entries. + If either side ever falls back to a URL match, this fails: the fresh server would skip the + upstream registration (adopting the sibling's client) or persist onto the wrong row.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client_with_server, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + shared_url = "https://provider.example/mcp" + fresh_server = MCPServer( + server_id="server-b", + name="server-b", + server_name="server-b", + alias="server-b", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url=shared_url, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + sibling_row_with_client = MagicMock(server_id="server-a", url=shared_url) + sibling_row_with_client.credentials = {"client_id": "client-a-do-not-adopt"} + own_row_without_client = MagicMock(server_id="server-b", url=shared_url) + own_row_without_client.credentials = {} + rows_by_server_id = {"server-a": sibling_row_with_client, "server-b": own_row_without_client} + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + mock_response = MagicMock() + mock_response.json.return_value = {"client_id": "fresh-client-b", "token_endpoint_auth_method": "none"} + mock_response.raise_for_status = MagicMock() + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + mock_update = AsyncMock(return_value=MagicMock()) + + async def _get_row(prisma_client, server_id): + return rows_by_server_id.get(server_id) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=mock_async_client, + ), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(side_effect=_get_row)), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + ): + response = await register_client_with_server( + request=mock_request, + mcp_server=fresh_server, + client_name="Litellm Proxy", + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", + persist_credentials=True, + ) + + mock_async_client.post.assert_called_once() + body = json.loads(response.body.decode("utf-8")) + assert body["client_id"] == "fresh-client-b" + + mock_update.assert_called_once() + update_data = mock_update.call_args.kwargs["data"] + assert update_data.server_id == "server-b" + assert update_data.credentials["client_id"] == "fresh-client-b" + + @pytest.mark.asyncio async def test_register_client_does_not_clobber_token_url_when_absent(): """When the in-memory server has no token_url, the DCR persist must omit it from the diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py index a60dab9148d..7d2cf4442a5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py @@ -92,6 +92,36 @@ async def test_token_cached_across_calls(): assert mock_client.post.call_count == 1 +@pytest.mark.asyncio +async def test_m2m_token_not_shared_across_server_ids_with_identical_config(): + """Two servers with byte-identical client_credentials config but different server_ids must not + share a cached M2M token: the cache is keyed by server_id, so a new server entry (even one + recreated with the same URL and credentials) mints its own token instead of inheriting the + sibling's. Guards against the cache key ever collapsing to the URL or the client config.""" + cache = MCPOAuth2TokenCache() + server_a = _server(server_id="srv-a") + server_b = _server(server_id="srv-b") + mock_client = AsyncMock() + mock_client.post.side_effect = [_token_response("tok-for-a"), _token_response("tok-for-b")] + + with ( + patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", + return_value=mock_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.mcp_oauth2_token_cache", + cache, + ), + ): + token_a = await resolve_mcp_auth(server_a) + token_b = await resolve_mcp_auth(server_b) + + assert token_a == "tok-for-a" + assert token_b == "tok-for-b" + assert mock_client.post.call_count == 2 + + @pytest.mark.asyncio async def test_per_request_header_beats_oauth2(): """An explicit mcp_auth_header takes priority over the OAuth2 token."""