From 7d64f9d26b71fc9e432eb1c578c0bee31082e862 Mon Sep 17 00:00:00 2001 From: Tin Date: Fri, 10 Jul 2026 12:58:44 -0700 Subject: [PATCH] fix(mcp): merge credentials within the client-forwarded class on an auth-type switch --- litellm/proxy/_experimental/mcp_server/db.py | 56 +++++++-- .../mcp_server/test_mcp_partial_update.py | 108 +++++++++++++++++- .../mcp_server/test_mcp_sigv4_auth.py | 97 +++++----------- 3 files changed, 183 insertions(+), 78 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index c7b32614c61..97cefb3f2cb 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -76,6 +76,34 @@ _TOKEN_EXCHANGE_COLUMN_FIELDS: frozenset = frozenset( } ) +# The client-forwarded token modes share one stored-credential shape: the admin-declared upstream +# OAuth app (client_id/client_secret) plus the same authorize relay, and neither mints anything the +# gateway keeps. So a switch WITHIN this class must preserve the stored app, unlike a cross-class +# switch (e.g. an oauth2 row whose client may be DCR-minted and is not reusable elsewhere). +_CLIENT_FORWARDED_AUTH_TYPES: frozenset = frozenset({"true_passthrough", "oauth_delegate"}) + +# Minted token material that must never survive a client rotation on a persisted row. +_MINTED_TOKEN_CREDENTIAL_FIELDS: frozenset = frozenset({"access_token", "refresh_token", "expires_in"}) + + +def _credential_auth_class(auth_type: Optional[str]) -> Optional[str]: + """Collapse the client-forwarded modes to one credential class; every other auth_type is its own + class. Used so credential handling keys off whether the stored-credential shape actually changed, + not off a raw auth_type inequality that treats true_passthrough<->oauth_delegate as a full reset.""" + if auth_type in _CLIENT_FORWARDED_AUTH_TYPES: + return "client_forwarded" + return auth_type + + +def _drop_stale_minted_on_client_rotation(merged: Dict[str, Any], new_creds: Dict[str, Any]) -> Dict[str, Any]: + """When the update rotates the client, drop stale minted token keys it did not itself set, so an old + app's access/refresh token never rides forward under the new client. A no-op when no client key changed.""" + if "client_id" not in new_creds and "client_secret" not in new_creds: + return merged + return { + key: value for key, value in merged.items() if key not in _MINTED_TOKEN_CREDENTIAL_FIELDS or key in new_creds + } + def _is_global_env_var_scope(scope: Any) -> bool: """``scope="user"`` entries are placeholders the user fills in; everything @@ -679,7 +707,9 @@ async def update_mcp_server( existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id}) auth_type_changed = bool( - data.auth_type and existing and existing.auth_type is not None and existing.auth_type != data.auth_type + data.auth_type + and existing + and _credential_auth_class(existing.auth_type) != _credential_auth_class(data.auth_type) ) # Clear stale credentials when auth_type changes but no new credentials provided @@ -712,11 +742,12 @@ async def update_mcp_server( # would wipe encrypted secrets that the UI cannot display back. if "credentials" in data_dict and data_dict["credentials"] is not None: if existing and existing.credentials: - # Only merge when auth_type is unchanged. Switching auth types - # (e.g. oauth2 → api_key) should replace credentials entirely - # to avoid stale secrets from the previous auth type lingering. - auth_type_unchanged = data.auth_type is None or data.auth_type == existing.auth_type - if auth_type_unchanged: + # Only merge when the credential CLASS is unchanged. A cross-class switch + # (e.g. oauth2 → api_key, or oauth2 → true_passthrough) replaces credentials + # entirely to avoid stale secrets from the previous class lingering; a switch + # within the client-forwarded class (true_passthrough ↔ oauth_delegate) keeps + # the same declared app and so must merge, not replace. + if not auth_type_changed: existing_creds = ( json.loads(existing.credentials) if isinstance(existing.credentials, str) @@ -727,8 +758,9 @@ async def update_mcp_server( if isinstance(data_dict["credentials"], str) else dict(data_dict["credentials"]) ) - # New values override existing; existing keys not in update are preserved - merged = {**existing_creds, **new_creds} + # New values override existing; existing keys not in update are preserved. A client + # rotation additionally drops the previous app's stale minted token keys. + merged = _drop_stale_minted_on_client_rotation({**existing_creds, **new_creds}, new_creds) # Migrate-on-write for legacy rows: token-exchange settings the # old blob shape carried move to their dedicated columns (unless # the caller set the column this update, or the row already has @@ -747,6 +779,14 @@ async def update_mcp_server( # Add audit fields data_dict["updated_by"] = touched_by + # prisma-python rejects a raw ``None`` for a ``Json?`` field ("value is required but not set"); the + # clear paths above use ``None`` as the merge-skip sentinel, so translate it here to ``Json(None)``, + # which writes SQL null and reads back as ``None``. Done at the edge so the merge guards stay simple. + if "credentials" in data_dict and data_dict["credentials"] is None: + from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools + + data_dict["credentials"] = Json(None) + updated_mcp_server = await MCPServerRepository(prisma_client).table.update( where={"server_id": data.server_id}, data=data_dict, # type: ignore diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index 00d96f975bb..41d61c8f508 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -11,6 +11,7 @@ import json from unittest.mock import AsyncMock, MagicMock import pytest +from prisma import Json from litellm.proxy._experimental.mcp_server.db import ( create_mcp_server, @@ -19,6 +20,11 @@ from litellm.proxy._experimental.mcp_server.db import ( from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest +def _credentials_cleared(value) -> bool: + """The clear sentinel after the edge translation: prisma Json(None) (SQL null) or a bare None.""" + return value is None or (isinstance(value, Json) and getattr(value, "data", "x") is None) + + def _mock_prisma(): mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable = AsyncMock() @@ -208,7 +214,7 @@ async def test_auth_type_switch_clears_stale_flow_scoped_fields(): "token_exchange_profile", ): assert data_dict[stale_field] is None, f"{stale_field} must be cleared on auth_type switch" - assert data_dict["credentials"] is None + assert _credentials_cleared(data_dict["credentials"]) @pytest.mark.asyncio @@ -558,3 +564,103 @@ async def test_te_update_without_blob_te_keys_leaves_credentials_untouched(): assert data_dict["token_exchange_endpoint"] == "https://new.example.com/token" assert "credentials" not in data_dict + + +# ── client-forwarded credential class: true_passthrough <-> oauth_delegate share one +# stored-app shape, so a switch between them must MERGE (keep the declared app), not REPLACE ── + + +@pytest.mark.asyncio +async def test_cf_pair_switch_without_credentials_keeps_stored_app_and_endpoints(): + """true_passthrough -> oauth_delegate with no credentials in the update must not clear the + stored client or null the endpoint columns: both modes use the same declared app and relay.""" + mock_prisma = _mock_prisma() + existing = _existing_row("true_passthrough", credentials={"client_id": "enc-A", "client_secret": "enc-B"}) + existing.authorization_url = "https://provider.example/authorize" + existing.token_url = "https://provider.example/token" + existing.registration_url = "https://provider.example/register" + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="cf-server", auth_type="oauth_delegate") + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + assert "credentials" not in data_dict + for scoped_field in ("authorization_url", "token_url", "registration_url", "oauth2_flow"): + assert scoped_field not in data_dict, f"{scoped_field} must not be nulled within the CF class" + + +@pytest.mark.asyncio +async def test_cf_pair_switch_with_partial_credentials_merges_not_replaces(): + """oauth_delegate update carrying only client_id onto a true_passthrough row must MERGE, so the + stored client_secret survives instead of being dropped by a REPLACE.""" + mock_prisma = _mock_prisma() + existing = _existing_row("true_passthrough", credentials={"client_id": "enc-A", "client_secret": "enc-B"}) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="cf-server", auth_type="oauth_delegate", credentials={"client_id": "B"}) + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + merged = json.loads(data_dict["credentials"]) + assert merged["client_secret"] == "enc-B" + assert merged["client_id"] != "enc-A" + + +@pytest.mark.asyncio +async def test_null_existing_auth_type_to_cf_counts_as_changed_and_clears_blob(): + """A legacy row with NULL auth_type switched to a client-forwarded mode is a cross-class change, + so the stale blob must be cleared (the two change predicates must agree on this).""" + mock_prisma = _mock_prisma() + existing = _existing_row(None, credentials={"client_id": "enc-old"}) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest(server_id="cf-server", auth_type="true_passthrough") + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + # The clear must reach prisma as Json(None) (SQL null), never a raw None, which prisma rejects. + assert isinstance(data_dict["credentials"], Json) + assert getattr(data_dict["credentials"], "data", "x") is None + + +@pytest.mark.asyncio +async def test_client_rotation_strips_legacy_minted_token_keys(): + """Rotating the client on a same-class row must drop stale minted token material the update did + not set, so an old access_token/refresh_token never rides forward under the new client.""" + mock_prisma = _mock_prisma() + existing = _existing_row( + "oauth2", credentials={"client_id": "A", "access_token": "T", "refresh_token": "R", "expires_in": 3600} + ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing) + + data = UpdateMCPServerRequest( + server_id="oauth2-server", auth_type="oauth2", credentials={"client_id": "B", "client_secret": "S"} + ) + await update_mcp_server(mock_prisma, data, "test-user") + data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + + merged = json.loads(data_dict["credentials"]) + assert "access_token" not in merged + assert "refresh_token" not in merged + assert "expires_in" not in merged + assert "client_secret" in merged + + +@pytest.mark.asyncio +async def test_cf_to_non_cf_switch_clears_dcr_bridge(): + """A cross-class switch OUT of a client-forwarded mode (true_passthrough -> api_key) must clear + dcr_bridge: the switch is cross-class so the flow-scoped sweep runs and nulls it, leaving no stale + dcr_bridge=True on a row that no longer supports it.""" + data = UpdateMCPServerRequest(server_id="s", auth_type="api_key") + data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") + assert data_dict["dcr_bridge"] is None + + +@pytest.mark.asyncio +async def test_cf_pair_switch_does_not_clear_dcr_bridge(): + """A within-class switch (true_passthrough <-> oauth_delegate) is not a credential-class change, so + the flow-scoped sweep does not run and dcr_bridge is left intact (both modes use the DCR bridge).""" + data = UpdateMCPServerRequest(server_id="s", auth_type="oauth_delegate") + data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") + assert "dcr_bridge" not in data_dict diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index f5a48b5ec34..5992fd1814f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -399,9 +399,7 @@ class TestMCPServerManagerSigV4: server = next(iter(manager.config_mcp_servers.values())) assert server.auth_type == MCPAuth.aws_sigv4 assert server.aws_access_key_id == "AKIAIOSFODNN7EXAMPLE" - assert ( - server.aws_secret_access_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" - ) + assert server.aws_secret_access_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" assert server.aws_region_name == "us-east-1" assert server.aws_service_name == "bedrock-agentcore" @@ -531,9 +529,7 @@ class TestMCPServerManagerSigV4: "aws_session_name": "my-session", } - result = manager._extract_aws_credentials( - creds, credentials_are_encrypted=False - ) + result = manager._extract_aws_credentials(creds, credentials_are_encrypted=False) assert result["aws_role_name"] == "arn:aws:iam::123456789012:role/TestRole" assert result["aws_session_name"] == "my-session" @@ -561,10 +557,7 @@ class TestSigV4CredentialEncryption: # Secrets should be encrypted assert result["aws_access_key_id"] == "enc:AKIAIOSFODNN7EXAMPLE" - assert ( - result["aws_secret_access_key"] - == "enc:wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" - ) + assert result["aws_secret_access_key"] == "enc:wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" assert result["aws_session_token"] == "enc:FwoGZX..." # Non-secrets should be unchanged assert result["aws_region_name"] == "us-east-1" @@ -606,12 +599,8 @@ class TestCredentialMergeOnUpdate: ) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -650,9 +639,7 @@ class TestCredentialMergeOnUpdate: from litellm.proxy._types import UpdateMCPServerRequest mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -679,12 +666,8 @@ class TestCredentialMergeOnUpdate: existing_record.credentials = None mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -725,12 +708,8 @@ class TestCredentialMergeOnUpdate: ) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -772,12 +751,8 @@ class TestCredentialMergeOnUpdate: ) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -819,9 +794,7 @@ class TestSigV4BuildFromTable: table_record.server_name = "sigv4_server" table_record.alias = None table_record.description = None - table_record.url = ( - "https://bedrock-agentcore.us-east-1.amazonaws.com/invocations" - ) + table_record.url = "https://bedrock-agentcore.us-east-1.amazonaws.com/invocations" table_record.spec_path = None table_record.transport = "http" table_record.auth_type = "aws_sigv4" @@ -867,9 +840,7 @@ class TestSigV4BuildFromTable: with patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.decrypt_value_helper", - side_effect=lambda value, key, exception_type, return_original_value: value.replace( - "enc:", "" - ), + side_effect=lambda value, key, exception_type, return_original_value: value.replace("enc:", ""), ): server = await manager.build_mcp_server_from_table(table_record) @@ -930,9 +901,7 @@ class TestSigV4BuildFromTable: with patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.decrypt_value_helper", - side_effect=lambda value, key, exception_type, return_original_value: value.replace( - "enc:", "" - ), + side_effect=lambda value, key, exception_type, return_original_value: value.replace("enc:", ""), ): server = await manager.build_mcp_server_from_table(table_record) @@ -1018,9 +987,7 @@ class TestRotateCredentials: server.env_vars = None mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock( - return_value=[server] - ) + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() with ( @@ -1039,9 +1006,7 @@ class TestRotateCredentials: side_effect=lambda value, new_encryption_key: f"enc_new:{value}", ), ): - await rotate_mcp_server_credentials_master_key( - mock_prisma, "admin", "new-key" - ) + await rotate_mcp_server_credentials_master_key(mock_prisma, "admin", "new-key") update_call = mock_prisma.db.litellm_mcpservertable.update assert update_call.called @@ -1069,9 +1034,7 @@ class TestRotateCredentials: ] mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock( - return_value=[server] - ) + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() with ( @@ -1090,9 +1053,7 @@ class TestRotateCredentials: side_effect=lambda value, new_encryption_key: f"enc_new:{value}", ), ): - await rotate_mcp_server_credentials_master_key( - mock_prisma, "admin", "new-key" - ) + await rotate_mcp_server_credentials_master_key(mock_prisma, "admin", "new-key") update_call = mock_prisma.db.litellm_mcpservertable.update assert update_call.called @@ -1116,17 +1077,11 @@ class TestAuthTypeSwitchClearsCredentials: existing_record = MagicMock() existing_record.auth_type = "oauth2" - existing_record.credentials = json.dumps( - {"client_id": "enc:cid", "client_secret": "enc:csec"} - ) + existing_record.credentials = json.dumps({"client_id": "enc:cid", "client_secret": "enc:csec"}) mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( - return_value=existing_record - ) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock( - return_value=MagicMock() - ) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) data = UpdateMCPServerRequest( server_id="test-server", @@ -1141,8 +1096,12 @@ class TestAuthTypeSwitchClearsCredentials: await update_mcp_server(mock_prisma, data, "test-user") data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] - # Credentials should be cleared (set to None) - assert data_dict.get("credentials") is None + # Credentials should be cleared. The clear reaches prisma as Json(None) (SQL null), which + # prisma-python requires for a Json? field; a bare None is also accepted for older callers. + from prisma import Json + + cleared = data_dict.get("credentials") + assert cleared is None or (isinstance(cleared, Json) and getattr(cleared, "data", "x") is None) class TestInheritCredentials: