diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 08e32ae21e..a6f0d145e9 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -540,6 +540,41 @@ def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: return None +async def rotate_mcp_user_credentials_master_key( + prisma_client: PrismaClient, new_master_key: str +): + """Re-encrypt every ``LiteLLM_MCPUserCredentials`` row with ``new_master_key``. + + Reads each ``credential_b64`` with the current salt key (falling back to + legacy plain base64 for unmigrated rows) and writes it back encrypted + under the new master key. Rows that are unreadable under both paths + are logged and skipped so one corrupt row does not abort the rotation. + """ + rows = await prisma_client.db.litellm_mcpusercredentials.find_many() + for row in rows: + plaintext = _decode_user_credential(row.credential_b64) + if plaintext is None: + verbose_proxy_logger.warning( + "rotate_mcp_user_credentials_master_key: could not decode " + "credential for user_id=%s server_id=%s, skipping", + row.user_id, + row.server_id, + ) + continue + re_encrypted = encrypt_value_helper( + plaintext, new_encryption_key=new_master_key + ) + await prisma_client.db.litellm_mcpusercredentials.update( + where={ + "user_id_server_id": { + "user_id": row.user_id, + "server_id": row.server_id, + } + }, + data={"credential_b64": re_encrypted}, + ) + + async def store_user_credential( prisma_client: PrismaClient, user_id: str, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a424f1558f..b67710ef0a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._experimental.mcp_server.db import ( rotate_mcp_server_credentials_master_key, + rotate_mcp_user_credentials_master_key, ) from litellm.proxy._types import * from litellm.proxy._types import LiteLLM_VerificationToken @@ -3709,6 +3710,17 @@ async def _rotate_master_key( # noqa: PLR0915 "Failed to rotate MCP server credentials: %s", str(e) ) + # 4b. process MCP user-scoped credentials table (BYOK + OAuth2 tokens) + try: + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma_client, + new_master_key=new_master_key, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to rotate MCP user credentials: %s", str(e) + ) + # 5. process credentials table try: credentials = await prisma_client.db.litellm_credentialstable.find_many() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 8bf5f9367a..078adf72d4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -19,9 +19,11 @@ from litellm.proxy._experimental.mcp_server.db import ( get_user_credential, get_user_oauth_credential, list_user_oauth_credentials, + rotate_mcp_user_credentials_master_key, store_user_credential, store_user_oauth_credential, ) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper SALT_KEY = "test-salt-key-for-byok-credential-tests-1234" @@ -292,3 +294,109 @@ def test_decode_user_credential_legacy_path(): plain = "legacy-secret" stored = base64.urlsafe_b64encode(plain.encode()).decode() assert _decode_user_credential(stored) == plain + + +# ── master-key rotation ─────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_rotate_re_encrypts_byok_with_new_key(monkeypatch): + # Encrypt a row under the current salt, then rotate to a new key, then + # confirm the stored ciphertext decrypts under the NEW key — and not under + # the old one. + prisma = _make_prisma_with_existing(row=None) + secret = "sk-original-byok-key" + await store_user_credential(prisma, "alice", "srv-1", secret) + encrypted_old = _stored_value(prisma) + + row = MagicMock() + row.user_id = "alice" + row.server_id = "srv-1" + row.credential_b64 = encrypted_old + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_mcpusercredentials.update = AsyncMock() + + new_master_key = "rotated-salt-key-9999-9999-9999-9999" + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma, new_master_key=new_master_key + ) + + update_call = prisma.db.litellm_mcpusercredentials.update.call_args + new_stored = update_call.kwargs["data"]["credential_b64"] + assert new_stored != encrypted_old, "rotation must produce different ciphertext" + + # Decrypt the rotated value under the NEW salt key — round-trips to plaintext. + monkeypatch.setenv("LITELLM_SALT_KEY", new_master_key) + assert ( + decrypt_value_helper( + value=new_stored, + key="mcp_user_credential", + exception_type="debug", + return_original_value=False, + ) + == secret + ) + + +@pytest.mark.asyncio +async def test_rotate_migrates_legacy_plaintext_rows(monkeypatch): + # A legacy plain-base64 row must also get re-encrypted under the new key. + prisma = _make_prisma_with_existing(row=None) + + legacy_row = MagicMock() + legacy_row.user_id = "alice" + legacy_row.server_id = "srv-legacy" + legacy_row.credential_b64 = base64.urlsafe_b64encode(b"legacy-plain").decode() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[legacy_row] + ) + prisma.db.litellm_mcpusercredentials.update = AsyncMock() + + new_key = "another-rotation-key-aaaa-bbbb-cccc-dddd" + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma, new_master_key=new_key + ) + + new_stored = prisma.db.litellm_mcpusercredentials.update.call_args.kwargs["data"][ + "credential_b64" + ] + monkeypatch.setenv("LITELLM_SALT_KEY", new_key) + assert ( + decrypt_value_helper( + value=new_stored, + key="mcp_user_credential", + exception_type="debug", + return_original_value=False, + ) + == "legacy-plain" + ) + + +@pytest.mark.asyncio +async def test_rotate_skips_undecodable_rows(): + # One bad row must not abort the rotation for the rest. + prisma = _make_prisma_with_existing(row=None) + + bad_row = MagicMock() + bad_row.user_id = "alice" + bad_row.server_id = "srv-corrupt" + bad_row.credential_b64 = "!!! not base64 and not encrypted !!!" + + good_row = MagicMock() + good_row.user_id = "bob" + good_row.server_id = "srv-ok" + good_row.credential_b64 = base64.urlsafe_b64encode(b"good-byok").decode() + + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[bad_row, good_row] + ) + prisma.db.litellm_mcpusercredentials.update = AsyncMock() + + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma, new_master_key="new-key-xxxx" + ) + + # Only one update call — the good row. + assert prisma.db.litellm_mcpusercredentials.update.call_count == 1 + where = prisma.db.litellm_mcpusercredentials.update.call_args.kwargs["where"] + assert where["user_id_server_id"]["server_id"] == "srv-ok"