From f3000bda366c6902ec68b2c27866862199e8cc0d Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 00:27:15 +0000 Subject: [PATCH 1/3] chore(mcp): encrypt user-scoped credentials at rest LiteLLM_MCPUserCredentials.credential_b64 stored both BYOK API keys and OAuth2 access tokens as plain urlsafe-base64 of the raw value. Any DB read could recover the upstream-provider key. Run all writes through encrypt_value_helper (nacl SecretBox, the same helper used for the server-level credentials column) and read back via a small dual-path helper that tries decryption first, then falls back to plain base64 so existing rows keep working until they get rewritten. Folds the three near-identical "decode -> json.loads -> check type == oauth2" sites into _decode_oauth_payload, which simplifies the BYOK guard inside store_user_oauth_credential. --- litellm/proxy/_experimental/mcp_server/db.py | 99 +++--- .../mcp_server/test_db_credentials.py | 293 ++++++++++++++++++ 2 files changed, 350 insertions(+), 42 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 21abaa3f981d..542b6e45d2b8 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,4 +1,5 @@ import base64 +import binascii import json from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast @@ -498,6 +499,47 @@ async def rotate_mcp_server_credentials_master_key( ) +def _decode_user_credential(stored: str) -> Optional[str]: + """Read back a value persisted in ``LiteLLM_MCPUserCredentials.credential_b64``. + + Tries nacl decryption first (current write format). Falls back to a + plain ``urlsafe_b64decode`` for rows persisted by older code that wrote + the credential without encryption. Returns ``None`` when neither path + yields a valid string. + """ + decrypted = decrypt_value_helper( + value=stored, + key="mcp_user_credential", + exception_type="debug", + return_original_value=False, + ) + if decrypted is not None: + return decrypted + try: + return base64.urlsafe_b64decode(stored).decode() + except (binascii.Error, UnicodeDecodeError, ValueError): + return None + + +def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: + """Return the OAuth2 payload dict if ``stored`` holds one, else ``None``. + + A row is considered an OAuth2 credential iff its decoded value parses as + a JSON object with ``"type": "oauth2"``. Plain BYOK credentials (which + share the same column) decode to a non-JSON string and return ``None``. + """ + decoded = _decode_user_credential(stored) + if decoded is None: + return None + try: + parsed = json.loads(decoded) + except (ValueError, TypeError): + return None + if isinstance(parsed, dict) and parsed.get("type") == "oauth2": + return parsed + return None + + async def store_user_credential( prisma_client: PrismaClient, user_id: str, @@ -506,7 +548,7 @@ async def store_user_credential( ) -> None: """Store a user credential for a BYOK MCP server.""" - encoded = base64.urlsafe_b64encode(credential.encode()).decode() + encoded = encrypt_value_helper(credential) await prisma_client.db.litellm_mcpusercredentials.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ @@ -532,16 +574,7 @@ async def get_user_credential( ) if row is None: return None - try: - return base64.urlsafe_b64decode(row.credential_b64).decode() - except Exception: - # Fall back to nacl decryption for credentials stored by older code - return decrypt_value_helper( - value=row.credential_b64, - key="byok_credential", - exception_type="debug", - return_original_value=False, - ) + return _decode_user_credential(row.credential_b64) async def has_user_credential( @@ -582,7 +615,7 @@ async def store_user_oauth_credential( ) -> None: """Persist an OAuth2 access token for a user+server pair. - The payload is JSON-serialised and stored base64-encoded in the same + The payload is JSON-serialised and stored encrypted in the same ``credential_b64`` column used by BYOK. A ``"type": "oauth2"`` key differentiates it from plain BYOK API keys. """ @@ -606,29 +639,22 @@ async def store_user_oauth_credential( payload["scopes"] = scopes # Guard against silently overwriting a BYOK credential with an OAuth token. - # BYOK credentials lack a "type" field (or use a non-"oauth2" type). # Skip the guard when the caller knows the row is already an OAuth2 credential # (e.g. during token refresh), saving an extra DB round-trip. if not skip_byok_guard: existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - if existing is not None: - _byok_error = ValueError( + if ( + existing is not None + and _decode_oauth_payload(existing.credential_b64) is None + ): + raise ValueError( f"A non-OAuth2 credential already exists for user {user_id} " f"and server {server_id}. Refusing to overwrite." ) - try: - raw = json.loads( - base64.urlsafe_b64decode(existing.credential_b64).decode() - ) - except Exception: - # Credential is not base64+JSON — it's a plain-text BYOK key. - raise _byok_error - if raw.get("type") != "oauth2": - raise _byok_error - encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + encoded = encrypt_value_helper(json.dumps(payload)) await prisma_client.db.litellm_mcpusercredentials.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ @@ -672,15 +698,7 @@ async def get_user_oauth_credential( ) if row is None: return None - try: - decoded = base64.urlsafe_b64decode(row.credential_b64).decode() - parsed = json.loads(decoded) - if isinstance(parsed, dict) and parsed.get("type") == "oauth2": - return parsed - # Row exists but is a BYOK (plain string), not an OAuth token - return None - except Exception: - return None + return _decode_oauth_payload(row.credential_b64) async def list_user_oauth_credentials( @@ -694,14 +712,11 @@ async def list_user_oauth_credentials( ) results: List[Dict[str, Any]] = [] for row in rows: - try: - decoded = base64.urlsafe_b64decode(row.credential_b64).decode() - parsed = json.loads(decoded) - if isinstance(parsed, dict) and parsed.get("type") == "oauth2": - parsed["server_id"] = row.server_id - results.append(parsed) - except Exception: - pass # Skip non-OAuth rows (BYOK plain strings) + payload = _decode_oauth_payload(row.credential_b64) + if payload is None: + continue + payload["server_id"] = row.server_id + results.append(payload) return results 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 new file mode 100644 index 000000000000..ca6992419ce5 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -0,0 +1,293 @@ +""" +Tests for the encrypted-at-rest persistence of MCP user credentials. + +The ``LiteLLM_MCPUserCredentials.credential_b64`` column previously stored +both BYOK API keys and OAuth2 access tokens as plain ``urlsafe_b64encode`` +of the raw value, leaving credentials readable from any DB read. The fix +runs every write through ``encrypt_value_helper`` (nacl SecretBox) and +keeps a plain-base64 fallback on read so existing rows continue to work. +""" + +import base64 +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.proxy._experimental.mcp_server.db import ( # noqa: E402 + _decode_user_credential, + get_user_credential, + get_user_oauth_credential, + list_user_oauth_credentials, + store_user_credential, + store_user_oauth_credential, +) + + +SALT_KEY = "test-salt-key-for-byok-credential-tests-1234" + + +@pytest.fixture(autouse=True) +def _set_salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY) + + +def _make_prisma_with_existing(row): + """Build a MagicMock prisma_client whose user-credentials table returns ``row`` + for find_unique and behaves async-correctly for upsert/find_many.""" + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_mcpusercredentials.upsert = AsyncMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) + return prisma + + +def _legacy_row(payload: str): + """A row exactly as the pre-fix code would have written it: plain + ``urlsafe_b64encode`` of the raw payload, no encryption.""" + row = MagicMock() + row.credential_b64 = base64.urlsafe_b64encode(payload.encode()).decode() + row.user_id = "alice" + row.server_id = "srv-1" + return row + + +def _stored_value(prisma) -> str: + """Pull the credential_b64 value passed to the most recent upsert call.""" + call = prisma.db.litellm_mcpusercredentials.upsert.call_args + data = call.kwargs["data"] + create_value = data["create"]["credential_b64"] + update_value = data["update"]["credential_b64"] + assert create_value == update_value, "create/update must agree" + return create_value + + +# ── BYOK round-trip ─────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_store_user_credential_does_not_persist_plaintext(): + # Stored bytes must not just be base64 of the secret — that's the regression. + secret = "sk-proj-very-secret-byok-key" + prisma = _make_prisma_with_existing(row=None) + + await store_user_credential(prisma, "alice", "srv-1", secret) + + stored = _stored_value(prisma) + plain_b64 = base64.urlsafe_b64encode(secret.encode()).decode() + assert stored != plain_b64 + # And the secret must not appear anywhere in a plain-b64 decode of the column. + try: + decoded_bytes = base64.urlsafe_b64decode(stored) + except Exception: + decoded_bytes = b"" + assert secret.encode() not in decoded_bytes + + +@pytest.mark.asyncio +async def test_byok_round_trip_returns_plaintext(): + secret = "sk-proj-very-secret-byok-key" + prisma = _make_prisma_with_existing(row=None) + await store_user_credential(prisma, "alice", "srv-1", secret) + + stored = _stored_value(prisma) + row = MagicMock() + row.credential_b64 = stored + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + + result = await get_user_credential(prisma, "alice", "srv-1") + assert result == secret + + +@pytest.mark.asyncio +async def test_byok_get_returns_plaintext_for_legacy_row(): + # Backward-compat: rows persisted by the pre-fix code (plain base64) must + # still decrypt-or-decode cleanly. + legacy_secret = "legacy-byok-key" + prisma = _make_prisma_with_existing(row=_legacy_row(legacy_secret)) + + result = await get_user_credential(prisma, "alice", "srv-1") + assert result == legacy_secret + + +@pytest.mark.asyncio +async def test_byok_get_returns_none_for_missing_row(): + prisma = _make_prisma_with_existing(row=None) + result = await get_user_credential(prisma, "alice", "srv-1") + assert result is None + + +# ── OAuth2 round-trip ───────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_store_user_oauth_credential_does_not_persist_plaintext(): + access_token = "ya29.a0AfH6SMBverysecretaccesstoken" + prisma = _make_prisma_with_existing(row=None) + + await store_user_oauth_credential( + prisma, "alice", "srv-1", access_token, refresh_token="rfr-xyz" + ) + + stored = _stored_value(prisma) + try: + decoded_bytes = base64.urlsafe_b64decode(stored) + except Exception: + decoded_bytes = b"" + assert access_token.encode() not in decoded_bytes + assert b"rfr-xyz" not in decoded_bytes + + +@pytest.mark.asyncio +async def test_oauth_round_trip_returns_payload(): + access_token = "ya29.a0AfH6SMBverysecretaccesstoken" + prisma = _make_prisma_with_existing(row=None) + await store_user_oauth_credential( + prisma, + "alice", + "srv-1", + access_token, + refresh_token="rfr-xyz", + scopes=["a", "b"], + ) + + stored = _stored_value(prisma) + row = MagicMock() + row.credential_b64 = stored + row.server_id = "srv-1" + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is not None + assert result["type"] == "oauth2" + assert result["access_token"] == access_token + assert result["refresh_token"] == "rfr-xyz" + assert result["scopes"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_oauth_get_returns_payload_for_legacy_row(): + payload = { + "type": "oauth2", + "access_token": "legacy-token", + "connected_at": "2024-01-01T00:00:00Z", + } + legacy_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + row = MagicMock() + row.credential_b64 = legacy_b64 + row.server_id = "srv-1" + prisma = _make_prisma_with_existing(row=row) + + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is not None + assert result["access_token"] == "legacy-token" + + +@pytest.mark.asyncio +async def test_oauth_get_returns_none_for_byok_row(): + # A row that holds a BYOK string must not leak as an OAuth payload. + prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-not-json")) + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is None + + +# ── BYOK guard inside store_user_oauth_credential ───────────────────────────── + + +@pytest.mark.asyncio +async def test_byok_guard_rejects_overwriting_legacy_byok(): + prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-key")) + with pytest.raises(ValueError, match="non-OAuth2 credential"): + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") + + +@pytest.mark.asyncio +async def test_byok_guard_rejects_overwriting_encrypted_byok(): + # Simulate a row written by the new (encrypted) code path: write a BYOK, + # then attempt to overwrite with an OAuth token. + prisma = _make_prisma_with_existing(row=None) + await store_user_credential(prisma, "alice", "srv-1", "sk-secret-byok") + + encrypted_row = MagicMock() + encrypted_row.credential_b64 = _stored_value(prisma) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock( + return_value=encrypted_row + ) + + with pytest.raises(ValueError, match="non-OAuth2 credential"): + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") + + +@pytest.mark.asyncio +async def test_byok_guard_allows_overwriting_existing_oauth(): + # Refresh path: row already holds an OAuth payload, write must succeed. + prisma = _make_prisma_with_existing(row=None) + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok-1") + + oauth_row = MagicMock() + oauth_row.credential_b64 = _stored_value(prisma) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=oauth_row) + + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok-2") + # Final upsert wrote a new payload (different from the first) + assert _stored_value(prisma) != oauth_row.credential_b64 + + +# ── list_user_oauth_credentials ─────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_list_oauth_credentials_filters_byok_and_returns_payloads(): + # Three rows: one encrypted OAuth, one legacy-plaintext OAuth, one BYOK. + # Only the two OAuth rows should come back. + prisma = _make_prisma_with_existing(row=None) + + await store_user_oauth_credential(prisma, "alice", "srv-encrypted", "tok-enc") + encrypted_b64 = _stored_value(prisma) + encrypted_row = MagicMock() + encrypted_row.credential_b64 = encrypted_b64 + encrypted_row.server_id = "srv-encrypted" + + legacy_payload = { + "type": "oauth2", + "access_token": "tok-legacy", + "connected_at": "2024-01-01T00:00:00Z", + } + legacy_row = MagicMock() + legacy_row.credential_b64 = base64.urlsafe_b64encode( + json.dumps(legacy_payload).encode() + ).decode() + legacy_row.server_id = "srv-legacy" + + byok_row = MagicMock() + byok_row.credential_b64 = base64.urlsafe_b64encode(b"plain-byok-key").decode() + byok_row.server_id = "srv-byok" + + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[encrypted_row, legacy_row, byok_row] + ) + + results = await list_user_oauth_credentials(prisma, "alice") + + server_ids = {r["server_id"] for r in results} + assert server_ids == {"srv-encrypted", "srv-legacy"} + tokens = {r["access_token"] for r in results} + assert tokens == {"tok-enc", "tok-legacy"} + + +# ── _decode_user_credential helper ──────────────────────────────────────────── + + +def test_decode_user_credential_handles_garbage(): + # Malformed input must return None, not raise. + assert _decode_user_credential("not-base64-and-not-encrypted!!!") is None + + +def test_decode_user_credential_legacy_path(): + plain = "legacy-secret" + stored = base64.urlsafe_b64encode(plain.encode()).decode() + assert _decode_user_credential(stored) == plain From c76c300392e09c60713dd72fb1318af4db523003 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 00:44:37 +0000 Subject: [PATCH 2/3] fix(mcp): address Greptile P2s on credential encoding helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three minor fixes from Greptile review: 1. _decode_user_credential now also catches TypeError so a null credential_b64 value returns None instead of propagating, matching the documented "returns None when neither path yields a valid string" contract. 2. The OAuth2 BYOK guard error no longer claims the existing row is a BYOK credential — after a salt-key rotation, an OAuth2 row can fail to decrypt and reach the same guard. Reword to "could not be verified as an OAuth2 token", which is accurate for both cases. 3. Drop the no-op sys.path.insert in the new test file (other tests in the directory don't need it; pytest picks up the package via the installed editable wheel). Adds a regression test for the None-input case. --- litellm/proxy/_experimental/mcp_server/db.py | 11 ++++++++--- .../mcp_server/test_db_credentials.py | 15 ++++++++------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 542b6e45d2b8..08e32ae21e88 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -517,7 +517,7 @@ def _decode_user_credential(stored: str) -> Optional[str]: return decrypted try: return base64.urlsafe_b64decode(stored).decode() - except (binascii.Error, UnicodeDecodeError, ValueError): + except (binascii.Error, UnicodeDecodeError, ValueError, TypeError): return None @@ -649,9 +649,14 @@ async def store_user_oauth_credential( existing is not None and _decode_oauth_payload(existing.credential_b64) is None ): + # Existing row is either a BYOK secret or an OAuth2 row that no + # longer decrypts (e.g. after a salt-key rotation). In either + # case, refuse to overwrite — the caller would clobber data + # that may still be recoverable. raise ValueError( - f"A non-OAuth2 credential already exists for user {user_id} " - f"and server {server_id}. Refusing to overwrite." + f"Existing credential for user {user_id} and server " + f"{server_id} could not be verified as an OAuth2 token. " + f"Refusing to overwrite." ) encoded = encrypt_value_helper(json.dumps(payload)) 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 ca6992419ce5..8bf5f9367aef 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 @@ -10,15 +10,11 @@ import base64 import json -import os -import sys from unittest.mock import AsyncMock, MagicMock import pytest -sys.path.insert(0, os.path.abspath("../../../../..")) - -from litellm.proxy._experimental.mcp_server.db import ( # noqa: E402 +from litellm.proxy._experimental.mcp_server.db import ( _decode_user_credential, get_user_credential, get_user_oauth_credential, @@ -201,7 +197,7 @@ async def test_oauth_get_returns_none_for_byok_row(): @pytest.mark.asyncio async def test_byok_guard_rejects_overwriting_legacy_byok(): prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-key")) - with pytest.raises(ValueError, match="non-OAuth2 credential"): + with pytest.raises(ValueError, match="could not be verified as an OAuth2"): await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") @@ -218,7 +214,7 @@ async def test_byok_guard_rejects_overwriting_encrypted_byok(): return_value=encrypted_row ) - with pytest.raises(ValueError, match="non-OAuth2 credential"): + with pytest.raises(ValueError, match="could not be verified as an OAuth2"): await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") @@ -287,6 +283,11 @@ def test_decode_user_credential_handles_garbage(): assert _decode_user_credential("not-base64-and-not-encrypted!!!") is None +def test_decode_user_credential_handles_none(): + # Defensive: a null DB value must return None, not propagate TypeError. + assert _decode_user_credential(None) is None + + def test_decode_user_credential_legacy_path(): plain = "legacy-secret" stored = base64.urlsafe_b64encode(plain.encode()).decode() From e0b32eb1cfa068acb76d7e2bb72f985d7867bb82 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 01:58:26 +0000 Subject: [PATCH 3/3] fix(mcp): re-encrypt user credentials during master-key rotation Greptile P1: this PR encrypts LiteLLM_MCPUserCredentials rows under the salt key, but the /key/regenerate rotation endpoint had no corresponding step for that table. Rotating the master key would leave every BYOK and OAuth2 user credential permanently unreadable. Adds rotate_mcp_user_credentials_master_key, mirroring the existing rotate_mcp_server_credentials_master_key pattern: read each row with the current key (via _decode_user_credential, which also handles unmigrated legacy plaintext rows), re-encrypt under the new master key, write back. One bad row is logged and skipped instead of aborting the whole rotation. Wired into key_management_endpoints.py as step 4b, alongside the existing server-credentials rotation, with the same try/except shape so a transient DB error on this table doesn't kill the whole regenerate-key flow. Tests cover: round-trip through rotation under a new key, automatic re-encryption of legacy plaintext rows (rotation also acts as a migration trigger), and a corrupt row not aborting the rotation. --- litellm/proxy/_experimental/mcp_server/db.py | 35 ++++++ .../key_management_endpoints.py | 12 ++ .../mcp_server/test_db_credentials.py | 108 ++++++++++++++++++ 3 files changed, 155 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 08e32ae21e88..a6f0d145e9b1 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 a424f1558fbd..b67710ef0a0c 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.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 8bf5f9367aef..078adf72d4c3 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 @@ 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"