Merge pull request #32815 from BerriAI/litellm_mcp_credential_class_merge

fix(mcp): merge credentials within the client-forwarded class on an auth-type switch

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

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

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: