fix(mcp): address Greptile P2s on credential encoding helpers

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.
This commit is contained in:
user 2026-04-30 00:44:37 +00:00
parent f3000bda36
commit c76c300392
2 changed files with 16 additions and 10 deletions

View file

@ -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))

View file

@ -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()