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] 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 542b6e45d2b..08e32ae21e8 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 ca6992419ce..8bf5f9367ae 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 @@ 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
+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()