mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
fix: address bugbot findings on user-fields MCP credential paths
- delete_user_credential: refuse to delete a row that holds a user-fields payload, so the BYOK delete endpoint can't silently destroy a user's saved field values (mirrors the overwrite guards on the write paths). - _build_stdio_env: restore the pre-PR contract of returning None for non-stdio transports so HTTP/SSE clients don't receive stdio env dicts. - _annotate_user_credential_flags: decrypt each credential row once and classify in-process, instead of paying the crypto cost twice per row. Co-authored-by: Yassin Kortam <yassin@berri.ai>
This commit is contained in:
parent
43501183f4
commit
d9e1ce2894
3 changed files with 45 additions and 11 deletions
|
|
@ -5,7 +5,7 @@ import json
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
|
||||
|
||||
from prisma.errors import UniqueViolationError
|
||||
from prisma.errors import RecordNotFoundError, UniqueViolationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -668,7 +668,23 @@ async def delete_user_credential(
|
|||
user_id: str,
|
||||
server_id: str,
|
||||
) -> None:
|
||||
"""Delete the user's stored credential for a BYOK MCP server."""
|
||||
"""Delete the user's stored credential for a BYOK MCP server.
|
||||
|
||||
BYOK, OAuth2, and user-fields payloads share the same
|
||||
``(user_id, server_id)`` row. Refuse to delete a row that holds a
|
||||
user-fields payload so the BYOK delete endpoint does not silently
|
||||
destroy the user's saved field values — mirroring the overwrite
|
||||
guards on the write paths.
|
||||
"""
|
||||
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 None:
|
||||
raise RecordNotFoundError()
|
||||
if _decode_user_fields_payload(existing.credential_b64) is not None:
|
||||
# Treat as "no BYOK credential present" so the endpoint reports
|
||||
# has_credential=False without clobbering the user-fields row.
|
||||
raise RecordNotFoundError()
|
||||
await prisma_client.db.litellm_mcpusercredentials.delete(
|
||||
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
|
||||
)
|
||||
|
|
@ -681,11 +697,13 @@ async def delete_user_credential(
|
|||
# distinguish formats without an extra column.
|
||||
|
||||
|
||||
def _decode_user_fields_payload(stored: str) -> Optional[Dict[str, str]]:
|
||||
"""Return the field-values dict if ``stored`` holds a user-fields payload."""
|
||||
decoded = _decode_user_credential(stored)
|
||||
if decoded is None:
|
||||
return None
|
||||
def _parse_user_fields_plaintext(decoded: str) -> Optional[Dict[str, str]]:
|
||||
"""Return the field-values dict if ``decoded`` is a user-fields JSON payload.
|
||||
|
||||
Takes already-decrypted plaintext so callers that have already paid
|
||||
the decryption cost (e.g. annotating server-list responses) can avoid
|
||||
a redundant decryption round-trip.
|
||||
"""
|
||||
try:
|
||||
parsed = json.loads(decoded)
|
||||
except (ValueError, TypeError):
|
||||
|
|
@ -698,6 +716,14 @@ def _decode_user_fields_payload(stored: str) -> Optional[Dict[str, str]]:
|
|||
return {str(k): str(v) for k, v in values.items()}
|
||||
|
||||
|
||||
def _decode_user_fields_payload(stored: str) -> Optional[Dict[str, str]]:
|
||||
"""Return the field-values dict if ``stored`` holds a user-fields payload."""
|
||||
decoded = _decode_user_credential(stored)
|
||||
if decoded is None:
|
||||
return None
|
||||
return _parse_user_fields_plaintext(decoded)
|
||||
|
||||
|
||||
async def store_user_field_values(
|
||||
prisma_client: PrismaClient,
|
||||
user_id: str,
|
||||
|
|
|
|||
|
|
@ -1309,7 +1309,10 @@ class MCPServerManager:
|
|||
"""
|
||||
|
||||
if server.transport != MCPTransport.stdio:
|
||||
return user_field_env or None
|
||||
# Non-stdio transports (HTTP/SSE) don't take an env dict; user
|
||||
# fields with env_var_name are stdio-only. Match the legacy
|
||||
# contract of always returning None for non-stdio servers.
|
||||
return None
|
||||
if not server.env:
|
||||
return user_field_env or None
|
||||
|
||||
|
|
|
|||
|
|
@ -860,7 +860,7 @@ if MCP_AVAILABLE:
|
|||
batched query."""
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
_decode_user_credential,
|
||||
_decode_user_fields_payload,
|
||||
_parse_user_fields_plaintext,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client as _byok_prisma_client
|
||||
|
||||
|
|
@ -879,10 +879,15 @@ if MCP_AVAILABLE:
|
|||
byok_set: set = set()
|
||||
user_fields_by_server: Dict[str, Dict[str, str]] = {}
|
||||
for row in cred_rows:
|
||||
payload = _decode_user_fields_payload(row.credential_b64)
|
||||
# Decrypt once and classify, instead of paying the crypto
|
||||
# cost twice (once for user-fields detection, once for BYOK).
|
||||
decoded = _decode_user_credential(row.credential_b64)
|
||||
if not decoded:
|
||||
continue
|
||||
payload = _parse_user_fields_plaintext(decoded)
|
||||
if payload is not None:
|
||||
user_fields_by_server[row.server_id] = payload
|
||||
elif _decode_user_credential(row.credential_b64):
|
||||
else:
|
||||
byok_set.add(row.server_id)
|
||||
for server in servers:
|
||||
if getattr(server, "is_byok", False):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue