Merge pull request #32735 from BerriAI/litellm_mcp_no_dcr_persist_for_passthrough

fix(mcp): stop persisting the DCR client onto true_passthrough and oauth_delegate server rows
This commit is contained in:
tin-berri 2026-07-10 11:03:38 -07:00 committed by GitHub
commit 11aeeea1fb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 237 additions and 2 deletions

View file

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

View file

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

View file

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

View file

@ -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."""