From e30e4639271dbe461352dfe60cb70991fd199ce2 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 19 May 2026 02:48:57 +0000 Subject: [PATCH 01/23] feat(mcp): add admin-declared per-user fields for MCP servers Generalizes the existing BYOK pattern (one user-provided credential per server) to N admin-declared fields. Each field can target an HTTP header (http/sse transports) or env var (stdio), with an optional value template for prefixes like "Bearer {value}". Field values are encrypted at rest, stored in the existing LiteLLM_MCPUserCredentials table with a "type": "user_fields" discriminator so they don't collide with BYOK strings or OAuth2 blobs. Demo flow: 1. Admin adds a server with one or more user fields. 2. The user dashboard shows a red "N missing fields" badge until each required field has a value. 3. Calling the server via Claude Code (or any MCP client) before saving returns HTTP 401 with error="user_fields_missing", the list of missing field descriptors, and a config_url pointing at the dashboard. 4. After the user saves their values, the badge clears and tool calls dispatch with the user's values injected as the configured headers or env vars. Backend - New JSONB column LiteLLM_MCPServerTable.user_fields with a migration. - MCPUserField / MCPUserFieldValuesRequest / MCPUserFieldsStatus types on the existing create/update/read models. - DB helpers store/get/delete user-field values via the same encryption path as BYOK; a "type" discriminator keeps the three formats apart. - New endpoints GET/POST/DELETE /v1/mcp/server/{id}/user-field-values and GET /v1/mcp/user-field-values (aggregated for dashboard badges). - GET /v1/mcp/server is annotated per-caller with missing_user_field_keys. - execute_mcp_tool enforces required fields with a friendly 401 carrying the dashboard config_url; the managed MCP dispatch path injects the resolved headers and stdio env vars. UI - Admin "Add MCP Server" form gains a dynamic User Fields section. - Dashboard servers list shows a red badge with the missing-field count for each affected server and opens a new UserFieldsModal where the end-user fills in their values. --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 5 + litellm/proxy/_experimental/mcp_server/db.py | 105 +++ .../mcp_server/mcp_server_manager.py | 147 ++++- .../proxy/_experimental/mcp_server/server.py | 109 ++++ .../_experimental/mcp_server/user_fields.py | 194 ++++++ litellm/proxy/_types.py | 65 ++ .../mcp_management_endpoints.py | 350 +++++++++- litellm/proxy/schema.prisma | 5 + .../types/mcp_server/mcp_server_manager.py | 4 + schema.prisma | 5 + .../mcp_server/test_mcp_user_fields.py | 608 ++++++++++++++++++ .../components/mcp_tools/UserFieldsModal.tsx | 243 +++++++ .../mcp_tools/create_mcp_server.tsx | 219 ++++++- .../mcp_tools/mcp_server_columns.tsx | 43 ++ .../src/components/mcp_tools/mcp_servers.tsx | 17 + .../src/components/mcp_tools/types.tsx | 35 + 17 files changed, 2144 insertions(+), 12 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260519120000_add_mcp_user_fields/migration.sql create mode 100644 litellm/proxy/_experimental/mcp_server/user_fields.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/UserFieldsModal.tsx diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260519120000_add_mcp_user_fields/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260519120000_add_mcp_user_fields/migration.sql new file mode 100644 index 000000000000..2c3ad5663047 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260519120000_add_mcp_user_fields/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "user_fields" JSONB NOT NULL DEFAULT '[]'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 78143fe0411c..b16714b0f040 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -328,6 +328,11 @@ model LiteLLM_MCPServerTable { is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? + // Admin-defined per-user fields (e.g. bearer tokens, workspace IDs) that + // each end-user must supply via the dashboard before they can use the + // server. Each entry is a JSON object describing how the field is rendered + // in the dashboard and injected at request time. + user_fields Json @default("[]") source_url String? // BYOM submission lifecycle approval_status String? @default("active") diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index a6f0d145e9b1..94314357d2a3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -85,6 +85,13 @@ def _prepare_mcp_server_data( # but be explicit to ensure a False value is always written to the DB). data_dict["is_byok"] = getattr(data, "is_byok", False) + # user_fields is a list of MCPUserField models. exclude_none=True will + # already have dict-ified them, but the JSONB column expects a JSON + # string when written through Prisma. + user_fields = data_dict.get("user_fields") + if user_fields is not None: + data_dict["user_fields"] = safe_dumps(user_fields) + return data_dict @@ -635,6 +642,104 @@ async def delete_user_credential( ) +# ── User-fields helpers ─────────────────────────────────────────────────────── +# +# User-fields share the credential_b64 column with BYOK and OAuth2. We tag +# the JSON payload with ``"type": "user_fields"`` so the read path can +# 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 + try: + parsed = json.loads(decoded) + except (ValueError, TypeError): + return None + if not isinstance(parsed, dict) or parsed.get("type") != "user_fields": + return None + values = parsed.get("values") + if not isinstance(values, dict): + return {} + return {str(k): str(v) for k, v in values.items()} + + +async def store_user_field_values( + prisma_client: PrismaClient, + user_id: str, + server_id: str, + values: Dict[str, str], +) -> None: + """Persist the calling user's values for an MCP server's user fields. + + The full set of values is encoded into a single encrypted JSON blob in + ``LiteLLM_MCPUserCredentials.credential_b64``. A ``"type"`` discriminator + lets ``get_user_field_values`` tell user-fields rows apart from BYOK + strings and OAuth2 payloads sharing the same column. + """ + + payload = json.dumps({"type": "user_fields", "values": values}) + encoded = encrypt_value_helper(payload) + await prisma_client.db.litellm_mcpusercredentials.upsert( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, + data={ + "create": { + "user_id": user_id, + "server_id": server_id, + "credential_b64": encoded, + }, + "update": {"credential_b64": encoded}, + }, + ) + + +async def get_user_field_values( + prisma_client: PrismaClient, + user_id: str, + server_id: str, +) -> Optional[Dict[str, str]]: + """Return the user's stored field values, or ``None`` if not stored. + + Returns ``None`` both when no row exists and when the row holds a + different credential type (BYOK / OAuth2) — callers should treat both + as "no user-fields configured for this user yet". + """ + + row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} + ) + if row is None: + return None + return _decode_user_fields_payload(row.credential_b64) + + +async def delete_user_field_values( + prisma_client: PrismaClient, + user_id: str, + server_id: str, +) -> bool: + """Delete the user's stored field values. + + Only removes the row when it actually holds a user-fields payload, so a + co-located BYOK / OAuth2 credential for the same (user, server) pair is + not accidentally wiped. Returns True if a row was deleted. + """ + + row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} + ) + if row is None: + return False + if _decode_user_fields_payload(row.credential_b64) is None: + return False + await prisma_client.db.litellm_mcpusercredentials.delete( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} + ) + return True + + # ── OAuth2 user-credential helpers ──────────────────────────────────────────── diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d1b49039e8ec..0f0bf1285018 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -190,6 +190,25 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]: return data +def _deserialize_user_fields(data: Any) -> List[Dict[str, Any]]: + """Decode the JSON-encoded ``user_fields`` blob from the MCP server row. + + Always returns a list — falsy or malformed values become ``[]`` so callers + can iterate without a None check. + """ + if not data: + return [] + if isinstance(data, list): + return data + if isinstance(data, str): + try: + decoded = json.loads(data) + except (json.JSONDecodeError, TypeError): + return [] + return decoded if isinstance(decoded, list) else [] + return [] + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -812,6 +831,9 @@ async def build_mcp_server_from_table( is_byok=bool(getattr(mcp_server, "is_byok", False)), byok_description=getattr(mcp_server, "byok_description", None) or [], byok_api_key_help_url=getattr(mcp_server, "byok_api_key_help_url", None), + user_fields=_deserialize_user_fields( + getattr(mcp_server, "user_fields", None) + ), # AWS SigV4 fields aws_access_key_id=aws_creds.get("aws_access_key_id"), aws_secret_access_key=aws_creds.get("aws_secret_access_key"), @@ -1276,11 +1298,20 @@ def _build_stdio_env( self, server: MCPServer, raw_headers: Optional[Dict[str, str]] = None, + user_field_env: Optional[Dict[str, str]] = None, ) -> Optional[Dict[str, str]]: - """Resolve stdio env values, supporting header-driven placeholders.""" + """Resolve stdio env values, supporting header-driven placeholders. - if server.transport != MCPTransport.stdio or not server.env: - return None + ``user_field_env`` carries values resolved from admin-declared + user_fields with an ``env_var_name`` set. They take precedence + over the static server.env entries so a user's stored value + always overrides any placeholder default. + """ + + if server.transport != MCPTransport.stdio: + return user_field_env or None + if not server.env: + return user_field_env or None resolved_env: Dict[str, str] = {} normalized_headers = {k.lower(): v for k, v in (raw_headers or {}).items()} @@ -1297,8 +1328,99 @@ def _build_stdio_env( else: resolved_env[env_key] = env_value + if user_field_env: + resolved_env.update(user_field_env) + # Preserve the legacy contract: when an env dict is present on the + # server we always return a dict (even if empty after template + # resolution). Only return None when no env is configured at all. return resolved_env + async def _resolve_user_field_values( + self, + mcp_server: MCPServer, + user_api_key_auth: Optional["UserAPIKeyAuth"], + ) -> Dict[str, str]: + """Read the calling user's stored user-field values for ``mcp_server``. + + Returns ``{}`` when the user has no row, no user_id, or there's no + DB. The caller decides what to do with missing required fields — + ``server.execute_mcp_tool`` raises a 401 with a config_url before + we even get here, so by the time injection happens we already + know the values are present. + """ + from litellm.proxy._experimental.mcp_server.user_fields import ( + coerce_user_fields, + ) + + if not coerce_user_fields(mcp_server): + return {} + if user_api_key_auth is None or not getattr(user_api_key_auth, "user_id", None): + return {} + + # Reuse the cache and lookup function from server.py so we don't + # pay a second DB round-trip after the enforcement check populated + # the same cache key. + try: + from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + _USER_FIELDS_CACHE_TTL, + _user_fields_cache, + _write_user_fields_cache, + ) + except ImportError: + _user_fields_cache = {} # type: ignore[assignment] + _USER_FIELDS_CACHE_TTL = 60 + _write_user_fields_cache = None # type: ignore[assignment] + + user_id = user_api_key_auth.user_id or "" + cache_key = (user_id, mcp_server.server_id) + cached = _user_fields_cache.get(cache_key) if _user_fields_cache else None + if cached is not None: + values, ts = cached + if time.monotonic() - ts < _USER_FIELDS_CACHE_TTL: + return values or {} + + from litellm.proxy._experimental.mcp_server.db import get_user_field_values + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return {} + values = await get_user_field_values( + prisma_client=prisma_client, + user_id=user_id, + server_id=mcp_server.server_id, + ) + if _write_user_fields_cache is not None: + _write_user_fields_cache(user_id, mcp_server.server_id, values) + return values or {} + + async def _resolve_user_field_headers( + self, + mcp_server: MCPServer, + user_api_key_auth: Optional["UserAPIKeyAuth"], + ) -> Dict[str, str]: + from litellm.proxy._experimental.mcp_server.user_fields import ( + resolve_user_field_headers, + ) + + stored = await self._resolve_user_field_values(mcp_server, user_api_key_auth) + if not stored: + return {} + return resolve_user_field_headers(mcp_server, stored) + + async def _resolve_user_field_env( + self, + mcp_server: MCPServer, + user_api_key_auth: Optional["UserAPIKeyAuth"], + ) -> Dict[str, str]: + from litellm.proxy._experimental.mcp_server.user_fields import ( + resolve_user_field_env, + ) + + stored = await self._resolve_user_field_values(mcp_server, user_api_key_auth) + if not stored: + return {} + return resolve_user_field_env(mcp_server, stored) + async def _create_mcp_client( self, server: MCPServer, @@ -2633,6 +2755,7 @@ async def _call_regular_mcp_tool( # noqa: PLR0915 proxy_logging_obj: Optional[ProxyLogging], host_progress_callback: Optional[Callable] = None, hook_extra_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> CallToolResult: """ Call a regular MCP tool using the MCP client. @@ -2717,6 +2840,16 @@ async def _call_regular_mcp_tool( # noqa: PLR0915 extra_headers = {} extra_headers.update(mcp_server.static_headers) + # User-fields: inject each declared user field's stored value either + # as a header (http/sse) or env var (stdio path; handled below). + user_field_headers = await self._resolve_user_field_headers( + mcp_server, user_api_key_auth + ) + if user_field_headers: + if extra_headers is None: + extra_headers = {} + extra_headers.update(user_field_headers) + if hook_extra_headers: if extra_headers is None: extra_headers = {} @@ -2746,7 +2879,12 @@ async def _call_regular_mcp_tool( # noqa: PLR0915 if extra_headers is not None and len(extra_headers) == 0: extra_headers = None - stdio_env = self._build_stdio_env(mcp_server, raw_headers) + user_field_env = await self._resolve_user_field_env( + mcp_server, user_api_key_auth + ) + stdio_env = self._build_stdio_env( + mcp_server, raw_headers, user_field_env=user_field_env or None + ) client = await self._create_mcp_client( server=mcp_server, @@ -2923,6 +3061,7 @@ async def call_tool( proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, hook_extra_headers=hook_result.get("extra_headers"), + user_api_key_auth=user_api_key_auth, ) # For OpenAPI tools, await outside the client context diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0a74a92f9ce8..e92c68be954b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -44,6 +44,11 @@ _mcp_gateway_initialize_instructions, ) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug +from litellm.proxy._experimental.mcp_server.user_fields import ( + build_user_fields_missing_error, + compute_missing_user_fields, + server_has_user_fields, +) from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, @@ -75,6 +80,14 @@ _BYOK_CRED_CACHE_TTL = 60 # seconds _BYOK_CRED_CACHE_MAX_SIZE = 4096 # cap to prevent unbounded growth +# Short-lived in-memory cache for user-field values, mirroring the BYOK cache. +# Value is (values_dict_or_None, monotonic_timestamp). None means "we checked +# and the user has no row yet" — still cached to avoid hammering the DB on +# tool calls that will fail the required-fields check. +_user_fields_cache: Dict[Tuple[str, str], Tuple[Optional[Dict[str, str]], float]] = {} +_USER_FIELDS_CACHE_TTL = 60 # seconds, matches BYOK +_USER_FIELDS_CACHE_MAX_SIZE = 4096 + def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: """Remove a (user_id, server_id) entry from the BYOK credential cache. @@ -94,6 +107,20 @@ def _write_byok_cred_cache( _byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic()) +def _invalidate_user_fields_cache(user_id: str, server_id: str) -> None: + """Drop a cached user-fields entry after a user updates their values.""" + _user_fields_cache.pop((user_id, server_id), None) + + +def _write_user_fields_cache( + user_id: str, server_id: str, values: Optional[Dict[str, str]] +) -> None: + """Cache stored user-field values, capping total entries.""" + if len(_user_fields_cache) >= _USER_FIELDS_CACHE_MAX_SIZE: + _user_fields_cache.clear() + _user_fields_cache[(user_id, server_id)] = (values, time.monotonic()) + + # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 # We're making this conditional import to avoid breaking users who use python 3.8. @@ -2039,6 +2066,81 @@ async def _check_byok_credential( }, ) + async def _get_user_field_values_cached( + mcp_server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> Tuple[Optional[str], Optional[Dict[str, str]]]: + """Return (user_id, stored_values) for the calling user. + + ``stored_values`` is None when either no user_id is available or + the user has not yet saved any values. Reads through a 60s + in-memory cache so back-to-back tool calls don't hammer the DB. + """ + user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or "" + if not user_id: + return None, None + + cache_key = (user_id, mcp_server.server_id) + cached = _user_fields_cache.get(cache_key) + if cached is not None: + values, ts = cached + if time.monotonic() - ts < _USER_FIELDS_CACHE_TTL: + return user_id, values + + from litellm.proxy._experimental.mcp_server.db import get_user_field_values + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return user_id, None + + values = await get_user_field_values( + prisma_client=prisma_client, + user_id=user_id, + server_id=mcp_server.server_id, + ) + _write_user_fields_cache(user_id, mcp_server.server_id, values) + return user_id, values + + async def _enforce_user_fields( + mcp_server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> None: + """Raise a friendly 401 when required user-fields are missing. + + Looks at the server's declared ``user_fields`` and the calling + user's saved values. If any ``required`` field is unset, raises + with ``error=user_fields_missing`` and a ``config_url`` pointing + at the dashboard. The error body is what Claude Code / other MCP + clients surface to the end-user — we put the URL in both the + ``message`` and as a structured field so simple clients still + show a clickable link. + """ + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + _resolve_proxy_base_url_env, + ) + + user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or "" + if not user_id: + # No identity means we can't look up stored values; treat as + # missing and let the user log in via the dashboard. + detail = build_user_fields_missing_error( + mcp_server, + compute_missing_user_fields(mcp_server, None), + _resolve_proxy_base_url_env(), + ) + raise HTTPException(status_code=401, detail=detail) + + _, stored_values = await _get_user_field_values_cached( + mcp_server, user_api_key_auth + ) + missing = compute_missing_user_fields(mcp_server, stored_values) + if not missing: + return + detail = build_user_fields_missing_error( + mcp_server, missing, _resolve_proxy_base_url_env() + ) + raise HTTPException(status_code=401, detail=detail) + async def execute_mcp_tool( # noqa: PLR0915 name: str, arguments: Dict[str, Any], @@ -2155,6 +2257,13 @@ async def execute_mcp_tool( # noqa: PLR0915 # External auth header supplied; still enforce user-identity check. await _check_byok_credential(mcp_server, user_api_key_auth) + # User fields: required admin-declared per-user values must be + # present before we dispatch. The helper raises a friendly 401 + # with a config_url pointing at the dashboard when any required + # field is missing. + if server_has_user_fields(mcp_server): + await _enforce_user_fields(mcp_server, user_api_key_auth) + # Check if tool exists in local registry first (for OpenAPI-based tools) # These tools are registered with their prefixed names ######################################################### diff --git a/litellm/proxy/_experimental/mcp_server/user_fields.py b/litellm/proxy/_experimental/mcp_server/user_fields.py new file mode 100644 index 000000000000..5684f004e199 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/user_fields.py @@ -0,0 +1,194 @@ +"""Helpers for resolving admin-declared MCP user fields at request time. + +User fields are per-user values (e.g. bearer tokens, workspace IDs) the +admin declares when adding an MCP server. End-users fill them in via the +dashboard; this module retrieves the stored values and injects them into +outbound MCP requests as HTTP headers (http/sse) or env vars (stdio). + +The retrieval result is cached in process so a tool call does not pay a +DB round-trip for every step. See ``_user_fields_cache`` in +``server.py`` for the cache itself; this module only contains the pure +resolution / injection logic. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, List, Optional, Tuple + +from litellm._logging import verbose_logger +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def coerce_user_fields(server: MCPServer) -> List[Dict[str, Any]]: + """Return the server's declared user fields as a list of plain dicts. + + The column is stored as JSONB but Prisma sometimes hands it back as a + string; this normalises both shapes and silently drops malformed + entries (the admin form filters these out, but DB writes from + external tooling might not). + """ + raw = getattr(server, "user_fields", None) + if not raw: + return [] + if isinstance(raw, list): + return [e for e in raw if isinstance(e, dict)] + if isinstance(raw, str): + import json + + try: + parsed = json.loads(raw) + except (ValueError, TypeError): + return [] + if isinstance(parsed, list): + return [e for e in parsed if isinstance(e, dict)] + return [] + + +def server_has_user_fields(server: MCPServer) -> bool: + """True iff the server declares any user fields at all.""" + return bool(coerce_user_fields(server)) + + +def compute_missing_user_fields( + server: MCPServer, stored_values: Optional[Dict[str, str]] +) -> List[Dict[str, Any]]: + """Return the declared field definitions the user has yet to fill in. + + Only ``required`` fields count as missing — optional fields without a + stored value are still considered satisfied so the user can save + partial configurations. + """ + stored = stored_values or {} + missing: List[Dict[str, Any]] = [] + for entry in coerce_user_fields(server): + field_key = entry.get("field_key") + if not isinstance(field_key, str) or not field_key: + continue + if not entry.get("required", True): + continue + if not stored.get(field_key): + missing.append(entry) + return missing + + +def build_user_fields_missing_error( + server: MCPServer, + missing: List[Dict[str, Any]], + base_dashboard_url: Optional[str], +) -> Dict[str, Any]: + """Construct the FastAPI ``detail`` payload for a missing-fields 401. + + Includes a ``config_url`` pointing the end-user (or their agent) at + the dashboard page where they can fill in the missing fields. The + URL is built defensively — if no base URL is known we emit a relative + path so curl-style clients still surface something useful. + """ + server_id = server.server_id + display_name = server.alias or server.server_name or server.name or server_id + + config_path = f"/ui?page=mcp-servers&server_id={server_id}" + if base_dashboard_url: + # Strip trailing slash to avoid double-slash in the joined URL. + base = base_dashboard_url.rstrip("/") + config_url = f"{base}{config_path}" + else: + config_url = config_path + + missing_summary = [ + { + "field_key": f.get("field_key"), + "display_name": f.get("display_name") or f.get("field_key"), + "description": f.get("description"), + } + for f in missing + ] + field_names = ", ".join( + str(m["display_name"]) for m in missing_summary if m.get("display_name") + ) + plural = "fields" if len(missing) != 1 else "field" + message = ( + f"This MCP server ({display_name}) needs your {plural} before it can run: " + f"{field_names}. Open {config_url} to fill them in, then retry the tool call." + ) + return { + "error": "user_fields_missing", + "server_id": server_id, + "server_name": server.server_name or server.name, + "missing_fields": missing_summary, + "config_url": config_url, + "message": message, + } + + +def resolve_user_field_headers( + server: MCPServer, stored_values: Dict[str, str] +) -> Dict[str, str]: + """Build the HTTP header dict to inject for an http/sse MCP server. + + Each declared field with a ``header_name`` contributes one header. + ``header_value_template`` (default ``"{value}"``) lets admins inject + well-known prefixes like ``"Bearer {value}"`` without making the user + re-type them. + """ + headers: Dict[str, str] = {} + for entry in coerce_user_fields(server): + field_key = entry.get("field_key") + header_name = entry.get("header_name") + if not field_key or not header_name: + continue + value = stored_values.get(field_key) + if not value: + continue + template = entry.get("header_value_template") or "{value}" + try: + headers[header_name] = template.format(value=value) + except (KeyError, IndexError, ValueError): + # Malformed template — fall back to raw value rather than crashing. + verbose_logger.warning( + "MCP user_fields: invalid header_value_template %r for field %r " + "on server %s; falling back to raw value.", + template, + field_key, + server.server_id, + ) + headers[header_name] = value + return headers + + +def resolve_user_field_env( + server: MCPServer, stored_values: Dict[str, str] +) -> Dict[str, str]: + """Build the env-var dict to inject for a stdio MCP server.""" + env: Dict[str, str] = {} + for entry in coerce_user_fields(server): + field_key = entry.get("field_key") + env_var_name = entry.get("env_var_name") + if not field_key or not env_var_name: + continue + value = stored_values.get(field_key) + if not value: + continue + env[env_var_name] = value + return env + + +def lookup_cached_user_fields( + cache: Dict[Tuple[str, str], Tuple[Optional[Dict[str, str]], float]], + user_id: str, + server_id: str, + ttl_seconds: int, +) -> Tuple[bool, Optional[Dict[str, str]]]: + """Return (cache_hit, values) for the (user, server) cache pair. + + Pulled out so server.py can pass its own cache dict in; keeping the + cache as module-level state in server.py preserves the existing + invalidation hooks called from the management endpoints. + """ + cached = cache.get((user_id, server_id)) + if cached is None: + return False, None + values, ts = cached + if time.monotonic() - ts >= ttl_seconds: + return False, None + return True, values diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c989f5dff130..c89e6ef0f350 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1245,6 +1245,44 @@ class MCPApprovalStatus(str, enum.Enum): rejected = "rejected" +class MCPUserField(LiteLLMPydanticObjectBase): + """Describes a single admin-declared per-user field on an MCP server. + + The admin declares a list of these when creating/editing the server. + Each end-user then supplies their own value(s) via the dashboard, which + are injected at request time as either HTTP headers (for http/sse + transports) or environment variables (for stdio transport). + """ + + field_key: str # storage key, must be unique within the server + display_name: Optional[str] = None # human-readable label + description: Optional[str] = None # help text for the dashboard + required: bool = True + # HTTP injection: when non-empty, the value is added to outbound headers. + # header_value_template defaults to "{value}"; use e.g. "Bearer {value}" + # to prefix the credential automatically. + header_name: Optional[str] = None + header_value_template: Optional[str] = None + # Stdio injection: when non-empty, the value is forwarded as an env var + # to the spawned MCP process. + env_var_name: Optional[str] = None + + +class MCPUserFieldValuesRequest(LiteLLMPydanticObjectBase): + """Body for storing the calling user's values for an MCP server's user fields.""" + + values: Dict[str, str] + + +class MCPUserFieldsStatus(LiteLLMPydanticObjectBase): + """Per-server view of the calling user's user-field completeness.""" + + server_id: str + user_fields: List[MCPUserField] = Field(default_factory=list) + stored_field_keys: List[str] = Field(default_factory=list) + missing_field_keys: List[str] = Field(default_factory=list) + + # MCP Proxy Request Types class NewMCPServerRequest(LiteLLMPydanticObjectBase): server_id: Optional[str] = None @@ -1278,6 +1316,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): is_byok: bool = False byok_description: List[str] = Field(default_factory=list) byok_api_key_help_url: Optional[str] = None + user_fields: List[MCPUserField] = Field(default_factory=list) source_url: Optional[str] = None # BYOM submission fields — set by the endpoint, not by the caller. # Any caller-provided values are silently overridden before persistence. @@ -1361,6 +1400,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): is_byok: bool = False byok_description: List[str] = Field(default_factory=list) byok_api_key_help_url: Optional[str] = None + user_fields: List[MCPUserField] = Field(default_factory=list) source_url: Optional[str] = None @model_validator(mode="before") @@ -1391,6 +1431,27 @@ def validate_transport_fields(cls, values): class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): """Represents a LiteLLM_MCPServerTable record""" + @model_validator(mode="before") + @classmethod + def _coerce_user_fields(cls, values): + """Accept ``user_fields`` as either a list or a JSON string. + + Prisma can hand JSONB columns back as a serialized string when the + row was written via ``safe_dumps`` (see ``_prepare_mcp_server_data``). + Coercing here keeps the rest of the pipeline able to assume a list. + """ + if isinstance(values, dict): + raw = values.get("user_fields") + if isinstance(raw, str): + try: + parsed = json.loads(raw) + except (json.JSONDecodeError, TypeError): + parsed = [] + if not isinstance(parsed, list): + parsed = [] + values["user_fields"] = parsed + return values + server_id: str server_name: Optional[str] = None alias: Optional[str] = None @@ -1434,6 +1495,10 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): byok_description: List[str] = Field(default_factory=list) byok_api_key_help_url: Optional[str] = None has_user_credential: Optional[bool] = None + user_fields: List[MCPUserField] = Field(default_factory=list) + # Computed per-request for the calling user. Populated only on the + # list/get endpoints that resolve credentials for the caller. + missing_user_field_keys: Optional[List[str]] = None source_url: Optional[str] = None # BYOM submission fields approval_status: Optional[str] = Field( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e9d9c243e7cd..0c392791d88e 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -112,13 +112,17 @@ def validate_tool_name(name: str) -> _ToolNameValidationResult: # type: ignore[ delete_mcp_server, delete_user_credential, get_all_mcp_servers_for_user, + delete_user_field_values, + get_all_mcp_servers, get_mcp_server, get_mcp_servers, get_mcp_submissions, + get_user_field_values, get_user_oauth_credential, list_user_oauth_credentials, reject_mcp_server, store_user_credential, + store_user_field_values, store_user_oauth_credential, update_mcp_server, ) @@ -146,6 +150,9 @@ def validate_tool_name(name: str) -> _ToolNameValidationResult: # type: ignore[ MCPUserCredentialListItem, MCPUserCredentialRequest, MCPUserCredentialResponse, + MCPUserField, + MCPUserFieldsStatus, + MCPUserFieldValuesRequest, NewMCPServerRequest, RejectMCPServerRequest, SpecialMCPServerName, @@ -473,6 +480,47 @@ def _redact_mcp_credentials_list( ) -> List[LiteLLM_MCPServerTable]: return [_redact_mcp_credentials(server) for server in mcp_servers] + def _coerce_user_fields_list(raw: Any) -> List[Dict[str, Any]]: + """Normalize a server.user_fields value (JSON string or list) to a list of dicts. + + Used by the BYOK-list-annotation block and any other site that needs + to inspect declared fields without instantiating MCPUserField models. + """ + if not raw: + return [] + if isinstance(raw, list): + return [e for e in raw if isinstance(e, dict)] + if isinstance(raw, str): + try: + parsed = json.loads(raw) + except (ValueError, TypeError): + return [] + if isinstance(parsed, list): + return [e for e in parsed if isinstance(e, dict)] + return [] + + def _has_required_user_fields(raw: Any) -> bool: + """True iff the server declares at least one required user field.""" + for entry in _coerce_user_fields_list(raw): + if entry.get("required", True): + return True + return False + + def _compute_missing_user_field_keys( + raw: Any, stored_values: Dict[str, str] + ) -> List[str]: + """Field keys the user must still supply (required + currently empty).""" + missing: List[str] = [] + for entry in _coerce_user_fields_list(raw): + field_key = entry.get("field_key") + if not isinstance(field_key, str) or not field_key: + continue + if not entry.get("required", True): + continue + if not stored_values.get(field_key): + missing.append(field_key) + return missing + def _is_restricted_virtual_key_request(user_api_key_dict: UserAPIKeyAuth) -> bool: """Best-effort detection for route-restricted virtual keys. @@ -944,26 +992,57 @@ async def fetch_all_mcp_servers( server.mcp_info = {} server.mcp_info["is_public"] = True - # Annotate has_user_credential for BYOK servers (single batched query) + # Annotate has_user_credential (BYOK) and missing_user_field_keys + # (user_fields) for the calling user. Both flags share the same + # storage table, so we batch a single query covering both feature + # sets. + from litellm.proxy._experimental.mcp_server.db import ( + _decode_user_credential, + _decode_user_fields_payload, + ) from litellm.proxy.proxy_server import prisma_client as _byok_prisma_client user_id = user_api_key_dict.user_id or "" if user_id and _byok_prisma_client is not None: - byok_server_ids = [ + relevant_server_ids = [ s.server_id for s in redacted_mcp_servers if getattr(s, "is_byok", False) + or _has_required_user_fields(getattr(s, "user_fields", None)) ] - if byok_server_ids: + if relevant_server_ids: cred_rows = ( await _byok_prisma_client.db.litellm_mcpusercredentials.find_many( - where={"user_id": user_id, "server_id": {"in": byok_server_ids}} + where={ + "user_id": user_id, + "server_id": {"in": relevant_server_ids}, + } ) ) - cred_set = {r.server_id for r in cred_rows} + # Build two indexes: one for BYOK presence, one mapping + # server_id → stored user-field values. + 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) + if payload is not None: + user_fields_by_server[row.server_id] = payload + else: + # Anything that isn't a user-fields payload counts + # as a BYOK credential for has_user_credential. + decoded = _decode_user_credential(row.credential_b64) + if decoded: + byok_set.add(row.server_id) for server in redacted_mcp_servers: if getattr(server, "is_byok", False): - server.has_user_credential = server.server_id in cred_set + server.has_user_credential = server.server_id in byok_set + if _has_required_user_fields(getattr(server, "user_fields", None)): + stored = user_fields_by_server.get(server.server_id, {}) + server.missing_user_field_keys = ( + _compute_missing_user_field_keys( + getattr(server, "user_fields", None), stored + ) + ) # Virtual keys only get a sanitized discovery view. if is_restricted_virtual_key: @@ -2055,6 +2134,265 @@ async def get_mcp_oauth_user_credential_status( connected_at=cred.get("connected_at"), ) + # ── User-fields endpoints ───────────────────────────────────────────────── + # + # Admin-declared user fields (e.g. per-user bearer tokens) are stored on + # the MCP server row as a list of definitions. Each end-user then supplies + # their own values via these endpoints; values are injected at request + # time as either HTTP headers (http/sse) or env vars (stdio). + + def _build_user_fields_status( + server: "LiteLLM_MCPServerTable", + stored_values: Optional[Dict[str, str]], + ) -> MCPUserFieldsStatus: + """Compose the dashboard-facing per-server status object. + + ``stored_values`` is the dict of values the calling user has saved + (None when no row exists yet). ``missing_field_keys`` enumerates only + the ``required`` fields the user has yet to fill in. + """ + raw_fields = getattr(server, "user_fields", None) or [] + if isinstance(raw_fields, str): + try: + raw_fields = json.loads(raw_fields) + except (ValueError, TypeError): + raw_fields = [] + user_fields: List[MCPUserField] = [] + for entry in raw_fields: + if not isinstance(entry, dict): + continue + try: + user_fields.append(MCPUserField(**entry)) + except Exception: # noqa: BLE001 — drop malformed entries silently + continue + stored: Dict[str, str] = stored_values or {} + stored_keys = [k for k, v in stored.items() if v] + missing_keys = [ + f.field_key + for f in user_fields + if f.required and not stored.get(f.field_key) + ] + return MCPUserFieldsStatus( + server_id=server.server_id, + user_fields=user_fields, + stored_field_keys=stored_keys, + missing_field_keys=missing_keys, + ) + + @router.get( + "/server/{server_id}/user-field-values", + description=( + "Return the calling user's saved values (presence only — values are " + "never echoed back) for an MCP server's admin-declared user fields." + ), + dependencies=[Depends(user_api_key_auth)], + response_model=MCPUserFieldsStatus, + ) + @management_endpoint_wrapper + async def get_mcp_user_field_values( + server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + server = await get_mcp_server(prisma_client, server_id) + if server is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP Server {server_id} not found"}, + ) + stored = await get_user_field_values(prisma_client, user_id, server_id) + return _build_user_fields_status(server, stored) + + @router.post( + "/server/{server_id}/user-field-values", + description=( + "Store the calling user's values for an MCP server's admin-declared " + "user fields. Values are encrypted at rest." + ), + dependencies=[Depends(user_api_key_auth)], + response_model=MCPUserFieldsStatus, + ) + @management_endpoint_wrapper + async def store_mcp_user_field_values( + server_id: str, + payload: MCPUserFieldValuesRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + server = await get_mcp_server(prisma_client, server_id) + if server is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP Server {server_id} not found"}, + ) + + # Pre-compute allowed keys from the server's declared fields, then + # filter the incoming payload to only those keys. This prevents + # callers from polluting the storage blob with arbitrary keys. + raw_fields = getattr(server, "user_fields", None) or [] + if isinstance(raw_fields, str): + try: + raw_fields = json.loads(raw_fields) + except (ValueError, TypeError): + raw_fields = [] + declared_keys = { + entry.get("field_key") + for entry in raw_fields + if isinstance(entry, dict) and entry.get("field_key") + } + if not declared_keys: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "This MCP server has no user fields declared", + "server_id": server_id, + }, + ) + + # Merge with anything already stored — partial saves should preserve + # previously-supplied values rather than wipe them. + existing = await get_user_field_values(prisma_client, user_id, server_id) or {} + merged: Dict[str, str] = {**existing} + for key, value in payload.values.items(): + if key not in declared_keys: + continue + if value == "": + merged.pop(key, None) + else: + merged[key] = value + + await store_user_field_values(prisma_client, user_id, server_id, merged) + + # Invalidate the BYOK credential cache for this (user, server) pair + # so the next tool call re-reads the row. We piggyback on the + # existing cache because user-fields, BYOK, and OAuth2 all share + # the same DB row. + from litellm.proxy._experimental.mcp_server.server import ( + _invalidate_byok_cred_cache, + _invalidate_user_fields_cache, + ) + + _invalidate_byok_cred_cache(user_id, server_id) + _invalidate_user_fields_cache(user_id, server_id) + return _build_user_fields_status(server, merged) + + @router.delete( + "/server/{server_id}/user-field-values", + description="Clear all of the calling user's saved user-field values for an MCP server.", + dependencies=[Depends(user_api_key_auth)], + response_model=MCPUserFieldsStatus, + ) + @management_endpoint_wrapper + async def delete_mcp_user_field_values( + server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + server = await get_mcp_server(prisma_client, server_id) + if server is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP Server {server_id} not found"}, + ) + try: + await delete_user_field_values(prisma_client, user_id, server_id) + except RecordNotFoundError: + pass + from litellm.proxy._experimental.mcp_server.server import ( + _invalidate_byok_cred_cache, + _invalidate_user_fields_cache, + ) + + _invalidate_byok_cred_cache(user_id, server_id) + _invalidate_user_fields_cache(user_id, server_id) + return _build_user_fields_status(server, None) + + @router.get( + "/user-field-values", + description=( + "Aggregate user-fields status across every MCP server the calling user " + "can see. Used by the dashboard to render 'needs setup' badges." + ), + dependencies=[Depends(user_api_key_auth)], + response_model=List[MCPUserFieldsStatus], + ) + @management_endpoint_wrapper + async def list_mcp_user_field_values( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + all_servers = await get_all_mcp_servers(prisma_client) + # Pre-filter to servers that actually declare user_fields, then + # batch-fetch every credential row for the calling user in one + # query. The N+1 alternative (per-server get_user_field_values) + # would scale linearly with server count. + relevant_servers: List["LiteLLM_MCPServerTable"] = [] + relevant_ids: List[str] = [] + for server in all_servers: + raw_fields = getattr(server, "user_fields", None) or [] + if isinstance(raw_fields, str): + try: + raw_fields = json.loads(raw_fields) + except (ValueError, TypeError): + raw_fields = [] + if not raw_fields: + continue + relevant_servers.append(server) + relevant_ids.append(server.server_id) + if not relevant_servers: + return [] + + # Single batched read; we only need rows for the calling user. + from litellm.proxy._experimental.mcp_server.db import ( + _decode_user_fields_payload, + ) + + cred_rows = await prisma_client.db.litellm_mcpusercredentials.find_many( + where={"user_id": user_id, "server_id": {"in": relevant_ids}} + ) + stored_by_server: Dict[str, Dict[str, str]] = {} + for row in cred_rows: + decoded = _decode_user_fields_payload(row.credential_b64) + if decoded is not None: + stored_by_server[row.server_id] = decoded + + return [ + _build_user_fields_status(server, stored_by_server.get(server.server_id)) + for server in relevant_servers + ] + @router.get( "/user-credentials", description="List all OAuth2 MCP credentials stored for the calling user", diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 78143fe0411c..b16714b0f040 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -328,6 +328,11 @@ model LiteLLM_MCPServerTable { is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? + // Admin-defined per-user fields (e.g. bearer tokens, workspace IDs) that + // each end-user must supply via the dashboard before they can use the + // server. Each entry is a JSON object describing how the field is rendered + // in the dashboard and injected at request time. + user_fields Json @default("[]") source_url String? // BYOM submission lifecycle approval_status String? @default("active") diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 776c7fa67a6e..4d79a3c6075d 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -77,6 +77,10 @@ class MCPServer(BaseModel): is_byok: bool = False byok_description: List[str] = [] byok_api_key_help_url: Optional[str] = None + # Raw user_fields blob (list of dicts) as stored on the server. Decoded + # at request time to inject per-user headers/env-vars. Kept as plain + # dicts here to avoid a dependency on proxy/_types from this module. + user_fields: List[Dict[str, Any]] = [] created_at: Optional[datetime] = None updated_at: Optional[datetime] = None # OAuth2 flow type. Defaults to None (interactive / authorization_code). diff --git a/schema.prisma b/schema.prisma index 78143fe0411c..b16714b0f040 100644 --- a/schema.prisma +++ b/schema.prisma @@ -328,6 +328,11 @@ model LiteLLM_MCPServerTable { is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? + // Admin-defined per-user fields (e.g. bearer tokens, workspace IDs) that + // each end-user must supply via the dashboard before they can use the + // server. Each entry is a JSON object describing how the field is rendered + // in the dashboard and injected at request time. + user_fields Json @default("[]") source_url String? // BYOM submission lifecycle approval_status String? @default("active") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py new file mode 100644 index 000000000000..f947849509ae --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -0,0 +1,608 @@ +"""Unit tests for admin-declared MCP user-fields. + +Covers the per-user fields feature end-to-end: + - Pydantic shapes for MCPUserField on the create/update/read models + - Encrypted JSON round-trip in LiteLLM_MCPUserCredentials.credential_b64 + - Storage helpers honour the type discriminator (don't collide with BYOK / OAuth2) + - Pure helpers in user_fields.py compute the right missing-field list and + inject the right headers / env vars + - execute_mcp_tool raises the friendly 401 when required fields are absent, + and proceeds when they're present +""" + +import json +import os +from typing import Any, Dict, List, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# Tests run with a fixed salt so encrypt_value_helper has a stable key. +os.environ.setdefault("LITELLM_SALT_KEY", "test-salt-for-user-fields") + +from litellm.proxy._experimental.mcp_server.db import ( + _decode_user_fields_payload, +) +from litellm.proxy._experimental.mcp_server.user_fields import ( + build_user_fields_missing_error, + coerce_user_fields, + compute_missing_user_fields, + resolve_user_field_env, + resolve_user_field_headers, + server_has_user_fields, +) +from litellm.proxy._types import ( + LiteLLM_MCPServerTable, + MCPTransport, + MCPUserField, + NewMCPServerRequest, + UpdateMCPServerRequest, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +# --------------------------------------------------------------------------- +# Pydantic shapes +# --------------------------------------------------------------------------- + + +def test_new_mcp_server_request_accepts_user_fields(): + """user_fields should serialize as a list of MCPUserField models.""" + req = NewMCPServerRequest( + server_name="gmail", + url="https://gmail.example.com", + transport=MCPTransport.http, + user_fields=[ + MCPUserField( + field_key="TOKEN", + display_name="Gmail OAuth Token", + header_name="Authorization", + header_value_template="Bearer {value}", + required=True, + ), + ], + ) + assert len(req.user_fields) == 1 + field = req.user_fields[0] + assert field.field_key == "TOKEN" + assert field.required is True + assert field.header_value_template == "Bearer {value}" + + +def test_update_mcp_server_request_user_fields_default_empty(): + req = UpdateMCPServerRequest( + server_id="s1", + server_name="gmail", + url="https://gmail.example.com", + transport=MCPTransport.http, + ) + assert req.user_fields == [] + + +def test_litellm_mcp_server_table_user_fields_default_empty(): + row = LiteLLM_MCPServerTable( + server_id="s1", + server_name="gmail", + transport=MCPTransport.http, + ) + assert row.user_fields == [] + # missing_user_field_keys is opt-in (set per-request by the list endpoint) + assert row.missing_user_field_keys is None + + +# --------------------------------------------------------------------------- +# Storage round-trip +# --------------------------------------------------------------------------- + + +def _encrypt_user_fields_blob(values: Dict[str, str]) -> str: + """Mirror what store_user_field_values writes into credential_b64.""" + payload = json.dumps({"type": "user_fields", "values": values}) + return encrypt_value_helper(payload) + + +def test_decode_user_fields_payload_round_trip(): + encoded = _encrypt_user_fields_blob({"BEARER": "tok", "WS": "ws1"}) + decoded = _decode_user_fields_payload(encoded) + assert decoded == {"BEARER": "tok", "WS": "ws1"} + + +def test_decode_user_fields_payload_rejects_oauth2_blob(): + """OAuth2 rows live in the same column — must not be misread as user-fields.""" + oauth_blob = encrypt_value_helper( + json.dumps({"type": "oauth2", "access_token": "abc"}) + ) + assert _decode_user_fields_payload(oauth_blob) is None + + +def test_decode_user_fields_payload_rejects_byok_string(): + """Plain BYOK strings must not parse as a user-fields payload.""" + byok_blob = encrypt_value_helper("plain-api-key") + assert _decode_user_fields_payload(byok_blob) is None + + +def test_decode_user_fields_payload_handles_garbage(): + assert _decode_user_fields_payload("not-base64-at-all") is None + + +# --------------------------------------------------------------------------- +# user_fields helpers +# --------------------------------------------------------------------------- + + +def _gmail_server(user_fields: Optional[List[Dict[str, Any]]] = None) -> MCPServer: + """Build a minimal MCPServer fixture for the helper tests.""" + return MCPServer( + server_id="s1", + name="Gmail", + alias="gmail-prod", + transport=MCPTransport.http, + user_fields=( + user_fields + if user_fields is not None + else [ + { + "field_key": "GMAIL_TOKEN", + "display_name": "Gmail Token", + "description": "Your Gmail OAuth bearer token", + "header_name": "Authorization", + "header_value_template": "Bearer {value}", + "required": True, + }, + { + "field_key": "WORKSPACE", + "display_name": "Workspace ID", + "header_name": "X-Workspace", + "required": False, + }, + ] + ), + ) + + +def test_coerce_user_fields_accepts_list(): + srv = _gmail_server() + assert len(coerce_user_fields(srv)) == 2 + + +def test_coerce_user_fields_accepts_json_string(): + srv = _gmail_server(user_fields=None) + # Prisma occasionally hands JSONB columns back as a string. + srv.user_fields = json.loads(json.dumps([{"field_key": "X", "required": True}])) + assert coerce_user_fields(srv) == [{"field_key": "X", "required": True}] + + +def test_coerce_user_fields_empty_when_missing(): + srv = MCPServer(server_id="s2", name="empty", transport=MCPTransport.http) + assert coerce_user_fields(srv) == [] + assert server_has_user_fields(srv) is False + + +def test_compute_missing_required_only(): + srv = _gmail_server() + missing = compute_missing_user_fields(srv, None) + # Only the required field is reported as missing. + assert [f["field_key"] for f in missing] == ["GMAIL_TOKEN"] + + +def test_compute_missing_empty_when_all_filled(): + srv = _gmail_server() + missing = compute_missing_user_fields( + srv, {"GMAIL_TOKEN": "tok", "WORKSPACE": "ws1"} + ) + assert missing == [] + + +def test_compute_missing_optional_field_unaffected(): + """Optional fields without a value should not appear in missing.""" + srv = _gmail_server() + missing = compute_missing_user_fields(srv, {"GMAIL_TOKEN": "tok"}) + assert missing == [] + + +def test_resolve_user_field_headers_applies_template(): + srv = _gmail_server() + headers = resolve_user_field_headers( + srv, {"GMAIL_TOKEN": "tok", "WORKSPACE": "ws1"} + ) + assert headers == {"Authorization": "Bearer tok", "X-Workspace": "ws1"} + + +def test_resolve_user_field_headers_skips_unset_values(): + srv = _gmail_server() + headers = resolve_user_field_headers(srv, {"WORKSPACE": "ws1"}) + # GMAIL_TOKEN has no stored value → no Authorization header. + assert headers == {"X-Workspace": "ws1"} + + +def test_resolve_user_field_headers_falls_back_on_bad_template(): + """A broken template must not crash the request path.""" + srv = _gmail_server( + user_fields=[ + { + "field_key": "TOKEN", + "header_name": "Authorization", + "header_value_template": "Bearer {unknown_placeholder}", + "required": True, + } + ] + ) + headers = resolve_user_field_headers(srv, {"TOKEN": "raw"}) + # Falls back to the raw value rather than raising. + assert headers == {"Authorization": "raw"} + + +def test_resolve_user_field_env_for_stdio(): + srv = MCPServer( + server_id="s3", + name="local", + transport=MCPTransport.stdio, + user_fields=[ + {"field_key": "GH_TOKEN", "env_var_name": "GITHUB_TOKEN", "required": True} + ], + ) + env = resolve_user_field_env(srv, {"GH_TOKEN": "ghs_xxx"}) + assert env == {"GITHUB_TOKEN": "ghs_xxx"} + + +def test_build_user_fields_missing_error_uses_proxy_base_url(): + srv = _gmail_server() + missing = compute_missing_user_fields(srv, None) + err = build_user_fields_missing_error(srv, missing, "https://proxy.example.com/") + assert err["error"] == "user_fields_missing" + assert err["server_id"] == "s1" + assert ( + err["config_url"] + == "https://proxy.example.com/ui?page=mcp-servers&server_id=s1" + ) + # Friendly message contains the URL and the field display name. + assert "Gmail Token" in err["message"] + assert err["config_url"] in err["message"] + assert err["missing_fields"][0]["field_key"] == "GMAIL_TOKEN" + + +def test_build_user_fields_missing_error_falls_back_to_relative_path(): + srv = _gmail_server() + err = build_user_fields_missing_error( + srv, compute_missing_user_fields(srv, None), None + ) + # When no base URL is known, the config_url is a relative path so it + # still surfaces something useful in error logs. + assert err["config_url"] == "/ui?page=mcp-servers&server_id=s1" + + +# --------------------------------------------------------------------------- +# execute_mcp_tool enforcement +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_enforce_user_fields_raises_when_missing(): + """When a required field is missing, the helper raises HTTP 401 with the + structured error payload Claude Code can render to the end-user.""" + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.server import ( + _enforce_user_fields, + _user_fields_cache, + ) + from litellm.proxy._types import UserAPIKeyAuth + + _user_fields_cache.clear() + srv = _gmail_server() + user = UserAPIKeyAuth(api_key="hashed", user_id="user-1") + + # Simulate "no stored values" by patching the cache directly so we don't + # need a live prisma_client. + _user_fields_cache[("user-1", "s1")] = (None, 1e18) # fresh entry, value=None + + with pytest.raises(HTTPException) as exc_info: + await _enforce_user_fields(srv, user) + assert exc_info.value.status_code == 401 + detail = exc_info.value.detail + assert detail["error"] == "user_fields_missing" + assert detail["server_id"] == "s1" + assert "GMAIL_TOKEN" in [m["field_key"] for m in detail["missing_fields"]] + assert "config_url" in detail + + +@pytest.mark.asyncio +async def test_enforce_user_fields_passes_when_all_required_present(): + from litellm.proxy._experimental.mcp_server.server import ( + _enforce_user_fields, + _user_fields_cache, + ) + from litellm.proxy._types import UserAPIKeyAuth + + _user_fields_cache.clear() + srv = _gmail_server() + user = UserAPIKeyAuth(api_key="hashed", user_id="user-2") + _user_fields_cache[("user-2", "s1")] = ({"GMAIL_TOKEN": "tok"}, 1e18) + + # Should not raise. + await _enforce_user_fields(srv, user) + + +@pytest.mark.asyncio +async def test_enforce_user_fields_no_user_id_raises(): + """A request without a user identity cannot be satisfied — surface the + same friendly error so the client knows where to send the user.""" + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.server import _enforce_user_fields + from litellm.proxy._types import UserAPIKeyAuth + + srv = _gmail_server() + user = UserAPIKeyAuth(api_key="hashed") # no user_id + + with pytest.raises(HTTPException) as exc_info: + await _enforce_user_fields(srv, user) + assert exc_info.value.status_code == 401 + assert exc_info.value.detail["error"] == "user_fields_missing" + + +# --------------------------------------------------------------------------- +# Manager wiring +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_build_stdio_env_merges_user_field_env_over_static(): + """Stored user_field env vars must take precedence over static server.env.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + srv = MCPServer( + server_id="s4", + name="local", + transport=MCPTransport.stdio, + env={"GITHUB_TOKEN": "static-default", "OTHER": "keep-me"}, + ) + mgr = MCPServerManager() + merged = mgr._build_stdio_env( + srv, raw_headers=None, user_field_env={"GITHUB_TOKEN": "user-value"} + ) + assert merged == {"GITHUB_TOKEN": "user-value", "OTHER": "keep-me"} + + +@pytest.mark.asyncio +async def test_build_stdio_env_returns_user_env_when_no_static_env(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + srv = MCPServer(server_id="s5", name="local", transport=MCPTransport.stdio) + mgr = MCPServerManager() + merged = mgr._build_stdio_env(srv, raw_headers=None, user_field_env={"X": "1"}) + assert merged == {"X": "1"} + + +# --------------------------------------------------------------------------- +# HTTP endpoints — POST / GET / DELETE /v1/mcp/server/{id}/user-field-values +# --------------------------------------------------------------------------- + + +def _server_row_with_user_fields() -> Any: + """Mock Prisma row matching what get_mcp_server returns. + + The endpoint sees `user_fields` as a list, mirroring how Prisma decodes + JSONB in the happy path. + """ + now = "2026-05-19T00:00:00" + return MagicMock( + server_id="srv-1", + server_name="Gmail", + alias="gmail", + transport=MCPTransport.http, + url="https://gmail.example.com", + created_at=now, + updated_at=now, + is_byok=False, + byok_description=[], + byok_api_key_help_url=None, + user_fields=[ + { + "field_key": "GMAIL_TOKEN", + "display_name": "Gmail Token", + "header_name": "Authorization", + "header_value_template": "Bearer {value}", + "required": True, + }, + { + "field_key": "WORKSPACE", + "display_name": "Workspace", + "header_name": "X-Workspace", + "required": False, + }, + ], + credentials=None, + mcp_info={}, + mcp_access_groups=[], + allowed_tools=[], + tool_name_to_display_name={}, + tool_name_to_description={}, + extra_headers=[], + static_headers={}, + status="unknown", + last_health_check=None, + health_check_error=None, + command=None, + args=[], + env={}, + authorization_url=None, + token_url=None, + registration_url=None, + allow_all_keys=False, + available_on_public_internet=True, + delegate_auth_to_upstream=False, + source_url=None, + approval_status="active", + submitted_by=None, + submitted_at=None, + reviewed_at=None, + review_notes=None, + spec_path=None, + instructions=None, + created_by="admin", + updated_by="admin", + auth_type=None, + ) + + +@pytest.mark.asyncio +async def test_get_user_field_values_endpoint_reports_missing(): + """Initial GET (no stored row) should show all required fields as missing.""" + from litellm.proxy._types import UserAPIKeyAuth + + server_row = _server_row_with_user_fields() + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=server_row + ) + prisma_client.db.litellm_mcpusercredentials.find_unique = AsyncMock( + return_value=None + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_user_field_values", + AsyncMock(return_value=None), + ), + ): + # Call the endpoint function directly to skip FastAPI auth wiring. + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + # The endpoint is closed over `router` inside `setup_mcp_management_routes`; + # exercise it via a TestClient. + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-1" + ) + client = TestClient(app) + res = client.get("/v1/mcp/server/srv-1/user-field-values") + assert res.status_code == 200, res.text + body = res.json() + assert body["server_id"] == "srv-1" + assert "GMAIL_TOKEN" in body["missing_field_keys"] + # Optional field is never reported as missing. + assert "WORKSPACE" not in body["missing_field_keys"] + assert body["stored_field_keys"] == [] + # Declared field descriptors are echoed back so the UI can render + # input fields without a second round-trip. + keys = {f["field_key"] for f in body["user_fields"]} + assert keys == {"GMAIL_TOKEN", "WORKSPACE"} + + +@pytest.mark.asyncio +async def test_post_user_field_values_rejects_undeclared_keys(): + """Attempting to save a key the server didn't declare must be dropped silently.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + server_row = _server_row_with_user_fields() + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=server_row + ) + + captured = {} + + async def fake_store(prisma, user_id, server_id, values): + captured["user_id"] = user_id + captured["server_id"] = server_id + captured["values"] = values + + async def fake_get(prisma, user_id, server_id): + return None # no existing values + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_field_values", + new=fake_store, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_field_values", + new=fake_get, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._invalidate_byok_cred_cache" + ), + patch( + "litellm.proxy._experimental.mcp_server.server._invalidate_user_fields_cache" + ), + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-2" + ) + client = TestClient(app) + res = client.post( + "/v1/mcp/server/srv-1/user-field-values", + json={ + "values": { + "GMAIL_TOKEN": "tok", + "WORKSPACE": "ws1", + "EVIL_KEY": "should-be-dropped", + } + }, + ) + assert res.status_code == 200, res.text + assert captured["values"] == {"GMAIL_TOKEN": "tok", "WORKSPACE": "ws1"} + body = res.json() + assert body["missing_field_keys"] == [] + + +@pytest.mark.asyncio +async def test_post_user_field_values_rejects_server_with_no_declared_fields(): + """Storing values on a server that doesn't use user_fields should 400.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + server_row = _server_row_with_user_fields() + server_row.user_fields = [] # no declared fields + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=server_row + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-3" + ) + client = TestClient(app) + res = client.post( + "/v1/mcp/server/srv-1/user-field-values", json={"values": {"X": "y"}} + ) + assert res.status_code == 400 diff --git a/ui/litellm-dashboard/src/components/mcp_tools/UserFieldsModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/UserFieldsModal.tsx new file mode 100644 index 000000000000..2d77eaa7496d --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/UserFieldsModal.tsx @@ -0,0 +1,243 @@ +"use client"; + +import React, { useEffect, useState } from "react"; +import { Input, Modal, Tag, Typography } from "antd"; +import { + CheckOutlined, + CloseOutlined, + LockOutlined, +} from "@ant-design/icons"; +import MessageManager from "@/components/molecules/message_manager"; +import { MCPServer, MCPUserField, MCPUserFieldsStatus } from "./types"; + +interface UserFieldsModalProps { + server: MCPServer; + open: boolean; + onClose: () => void; + onSuccess: (status: MCPUserFieldsStatus) => void; + accessToken: string; +} + +/** + * Dashboard modal where an end-user fills in the per-user fields declared by + * the admin for an MCP server. On submit, posts the values to + * /v1/mcp/server/{server_id}/user-field-values and returns the new status so + * the caller can clear the red badge. + */ +export const UserFieldsModal: React.FC = ({ + server, + open, + onClose, + onSuccess, + accessToken, +}) => { + const [values, setValues] = useState>({}); + const [storedFieldKeys, setStoredFieldKeys] = useState>(new Set()); + const [loading, setLoading] = useState(false); + const [fetching, setFetching] = useState(false); + + const declaredFields: MCPUserField[] = server.user_fields ?? []; + const displayName = server.alias || server.server_name || "Service"; + + // When the modal opens, fetch the current status so the user can see which + // fields they've already saved (we don't echo back the values themselves). + useEffect(() => { + if (!open || !accessToken) return; + let cancelled = false; + const load = async () => { + setFetching(true); + try { + const res = await fetch( + `/v1/mcp/server/${server.server_id}/user-field-values`, + { + method: "GET", + headers: { Authorization: `Bearer ${accessToken}` }, + }, + ); + if (!res.ok) return; + const status = (await res.json()) as MCPUserFieldsStatus; + if (!cancelled) { + setStoredFieldKeys(new Set(status.stored_field_keys ?? [])); + } + } catch (err) { + // Non-fatal; just leave storedFieldKeys empty. + } finally { + if (!cancelled) setFetching(false); + } + }; + void load(); + return () => { + cancelled = true; + }; + }, [open, server.server_id, accessToken]); + + const handleClose = () => { + setValues({}); + setLoading(false); + onClose(); + }; + + const handleSave = async () => { + // Trim and drop empty strings so optional fields aren't overwritten with "". + const payloadValues: Record = {}; + for (const field of declaredFields) { + const v = values[field.field_key]; + if (typeof v === "string" && v.trim().length > 0) { + payloadValues[field.field_key] = v.trim(); + } + } + // Validate required fields the user hasn't already saved. + const missingRequired = declaredFields.filter( + (f) => + (f.required ?? true) && + !storedFieldKeys.has(f.field_key) && + !payloadValues[f.field_key], + ); + if (missingRequired.length > 0) { + MessageManager.error( + `Please fill in: ${missingRequired + .map((f) => f.display_name || f.field_key) + .join(", ")}`, + ); + return; + } + setLoading(true); + try { + const res = await fetch( + `/v1/mcp/server/${server.server_id}/user-field-values`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${accessToken}`, + }, + body: JSON.stringify({ values: payloadValues }), + }, + ); + if (!res.ok) { + const err = await res.json().catch(() => ({})); + throw new Error(err?.detail?.error || err?.detail || "Failed to save fields"); + } + const status = (await res.json()) as MCPUserFieldsStatus; + MessageManager.success(`Saved your fields for ${displayName}`); + onSuccess(status); + handleClose(); + } catch (e: any) { + MessageManager.error(e?.message || "Failed to save fields"); + } finally { + setLoading(false); + } + }; + + return ( + +
+
+ + Connect {displayName} + + +
+ + + This server needs a few personal values from you before you can use it. + Your values are encrypted at rest and only sent to {displayName} on + your behalf. + + + {declaredFields.length === 0 ? ( + + This server has no user fields to configure. + + ) : ( +
+ {declaredFields.map((field) => { + const isSaved = storedFieldKeys.has(field.field_key); + const required = field.required ?? true; + return ( +
+
+ + {isSaved && ( + + Saved + + )} +
+ {field.description && ( + + {field.description} + + )} + + setValues((prev) => ({ + ...prev, + [field.field_key]: e.target.value, + })) + } + autoComplete="off" + /> +
+ ); + })} +
+ )} + +
+ + + Your values are encrypted with the proxy's salt key and never + logged. Other users on this proxy cannot see them. + +
+ +
+ + +
+
+ + ); +}; + +export default UserFieldsModal; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index f8b0141b25d1..4307789fac77 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -1,6 +1,21 @@ import React, { useState } from "react"; -import { Modal, Tooltip, Form, Select, Input, Switch, Collapse } from "antd"; -import { InfoCircleOutlined } from "@ant-design/icons"; +import { + Card, + Collapse, + Form, + Input, + Modal, + Select, + Space, + Switch, + Tooltip, + Typography, +} from "antd"; +import { + DeleteOutlined, + InfoCircleOutlined, + PlusOutlined, +} from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { createMCPServer, registerMCPServer } from "../networking"; import { AUTH_TYPE, DiscoverableMCPServer, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types"; @@ -370,6 +385,31 @@ const CreateMCPServer: React.FC = ({ } } + // Strip incomplete user_fields entries (e.g. an empty row added by an + // admin who forgot to remove it). The backend validator would reject + // these with a 422; filtering client-side gives a clearer experience. + const rawUserFields = (restValues as any).user_fields; + let cleanedUserFields: any[] | undefined; + if (Array.isArray(rawUserFields)) { + cleanedUserFields = rawUserFields + .filter( + (entry: any) => + entry && typeof entry.field_key === "string" && entry.field_key.trim().length > 0, + ) + .map((entry: any) => { + const result: Record = { field_key: entry.field_key.trim() }; + if (entry.display_name) result.display_name = entry.display_name; + if (entry.description) result.description = entry.description; + if (entry.header_name) result.header_name = entry.header_name; + if (entry.header_value_template) + result.header_value_template = entry.header_value_template; + if (entry.env_var_name) result.env_var_name = entry.env_var_name; + result.required = entry.required !== false; + return result; + }); + } + (restValues as any).user_fields = cleanedUserFields ?? []; + // Prepare the payload with cost configuration and allowed tools const payload: Record = { ...restValues, @@ -679,6 +719,181 @@ const CreateMCPServer: React.FC = ({ /> )} + {/* User fields - admin-declared per-user values (bearer tokens, etc.) */} + {transportType !== "" && ( + + Per-User Fields + + + + + ), + children: ( + + {(fields, { add, remove }) => ( + <> + {fields.length === 0 && ( + + Declare fields like BEARER_TOKEN or{" "} + WORKSPACE_ID that each user must fill in + before this server works for them. The dashboard will + highlight servers with missing fields, and tool calls + will fail with a friendly error pointing the user to + the dashboard. + + )} + {fields.map(({ key, name, ...restField }) => ( + + Field #{name + 1} + + } + extra={ + + } + > + + Field key + + + + + } + rules={[ + { required: true, message: "field_key is required" }, + { + pattern: /^[A-Za-z][A-Za-z0-9_]*$/, + message: "Use letters, digits, underscores; start with a letter", + }, + ]} + > + + + + Display name + + } + > + + + + Description (help text shown on dashboard) + + } + > + + + + HTTP header name (http/sse) + + + + + } + > + + + + Header value template + + + + + } + > + + + + Env var name (stdio) + + + + + } + > + + + + Required + + } + valuePropName="checked" + initialValue={true} + > + + + + ))} + + + )} + + ), + }, + ]} + /> + )} + {/* BYOK toggle - only for OpenAPI */} {transportType === TRANSPORT.OPENAPI && ( <> diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx index ea5ccf1c847e..01ab4375a7e8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx @@ -92,6 +92,7 @@ export const mcpServerColumns = ( onByokConnect?: (server: MCPServer) => void, onRecheckHealth?: (serverId: string) => void, recheckingServerIds?: Set, + onUserFieldsConnect?: (server: MCPServer) => void, ): ColumnDef[] => [ { accessorKey: "server_id", @@ -266,6 +267,48 @@ export const mcpServerColumns = ( header: "Credential", cell: ({ row }) => { const server = row.original; + // User-fields take priority over BYOK display because the most + // critical feedback for the user is "you still need to fill this in". + const declaredUserFields = server.user_fields ?? []; + const hasUserFields = declaredUserFields.length > 0; + if (hasUserFields) { + // missing_user_field_keys is populated by the proxy for the calling + // user on the list endpoint. Falling back to an empty array assumes + // "nothing missing" so old API responses don't show a false alarm. + const missing = server.missing_user_field_keys ?? []; + if (missing.length > 0) { + return onUserFieldsConnect ? ( + + ) : ( + + + {missing.length} missing field{missing.length === 1 ? "" : "s"} + + ); + } + return ( +
+ + Ready + + {onUserFieldsConnect && ( + + )} +
+ ); + } if (!server.is_byok) { return —; } diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index 72d5e4b5aa8e..412360a40ad8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -20,6 +20,7 @@ import MCPSemanticFilterSettings from "../Settings/AdminSettings/MCPSemanticFilt import MCPNetworkSettings from "./MCPNetworkSettings"; import MCPDiscovery from "./mcp_discovery"; import { ByokCredentialModal } from "./ByokCredentialModal"; +import { UserFieldsModal } from "./UserFieldsModal"; import { getSecureItem } from "@/utils/secureStorage"; const { Text: AntdText, Title: AntdTitle } = Typography; @@ -64,6 +65,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const [prefillData, setPrefillData] = useState(null); const [isDeletingServer, setIsDeletingServer] = useState(false); const [byokModalServer, setByokModalServer] = useState(null); + const [userFieldsModalServer, setUserFieldsModalServer] = useState(null); const isInternalUser = userRole === "Internal User"; useEffect(() => { @@ -173,6 +175,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) (server: MCPServer) => setByokModalServer(server), recheckServerHealth, recheckingServerIds, + (server: MCPServer) => setUserFieldsModalServer(server), ), [userRole, isLoadingHealth, recheckServerHealth, recheckingServerIds], ); @@ -461,6 +464,20 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) accessToken={accessToken || ""} /> )} + {userFieldsModalServer && ( + setUserFieldsModalServer(null)} + onSuccess={() => { + // Refetch the server list so missing_user_field_keys updates + // and the red badge clears (or stays red if still incomplete). + refetch(); + setUserFieldsModalServer(null); + }} + accessToken={accessToken || ""} + /> + )} ); }; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 7cfe08d9ee57..a54eb002583a 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -169,6 +169,33 @@ export interface MCPToolsViewerProps { extraHeaders?: string[] | null; } +/** + * One admin-declared per-user field on an MCP server. The admin lists these + * when creating the server; each end-user supplies their own values via the + * dashboard. Values get injected as HTTP headers (http/sse) or env vars (stdio) + * at request time. + */ +export interface MCPUserField { + field_key: string; + display_name?: string | null; + description?: string | null; + required?: boolean; + /** HTTP injection target (http/sse transports). */ + header_name?: string | null; + /** Defaults to "{value}". Use e.g. "Bearer {value}" to prefix. */ + header_value_template?: string | null; + /** Env var injection target (stdio transport). */ + env_var_name?: string | null; +} + +/** Response shape from /v1/mcp/server/{server_id}/user-field-values. */ +export interface MCPUserFieldsStatus { + server_id: string; + user_fields: MCPUserField[]; + stored_field_keys: string[]; + missing_field_keys: string[]; +} + export interface MCPServer { server_id: string; server_name?: string | null; @@ -215,6 +242,14 @@ export interface MCPServer { byok_api_key_help_url?: string | null; has_user_credential?: boolean | null; + /** Admin-declared per-user fields (e.g. bearer tokens) */ + user_fields?: MCPUserField[] | null; + /** + * Populated by the proxy on list/get for the calling user — the field_keys + * the user has not yet filled in. Drives the red badge on the dashboard. + */ + missing_user_field_keys?: string[] | null; + /** GitHub / source repository URL */ source_url?: string | null; From fa671c6b327fb566228a968fea53ac217ceaed93 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 05:50:21 +0000 Subject: [PATCH 02/23] fix(mcp): address security and correctness bugs in user-fields PR - Add missing 'import time' in mcp_server_manager.py; _resolve_user_field_values() was calling time.monotonic() without the module imported, raising NameError on every cached tool call against a server with user_fields. - list_mcp_user_field_values now filters via get_all_mcp_servers_for_user instead of get_all_mcp_servers so non-admin callers no longer see server IDs and user-field metadata (header_name/env_var_name) for servers they cannot access. - store_user_field_values now refuses to overwrite a non-user-fields credential (BYOK / OAuth2) sharing the same (user_id, server_id) row, mirroring the existing skip_byok_guard pattern in store_user_oauth_credential. The POST endpoint surfaces this as a 409 conflict. - Deduplicate user-fields parsing helpers: replace _coerce_user_fields_list / _has_required_user_fields / _compute_missing_user_field_keys with the existing coerce_user_fields / server_has_user_fields / compute_missing_user_fields from user_fields.py. - Remove dead lookup_cached_user_fields() (and its now-unused 'time' import) from user_fields.py. Co-authored-by: Yassin Kortam --- litellm/proxy/_experimental/mcp_server/db.py | 19 ++++ .../mcp_server/mcp_server_manager.py | 1 + .../_experimental/mcp_server/user_fields.py | 24 +---- .../mcp_management_endpoints.py | 102 ++++++------------ 4 files changed, 56 insertions(+), 90 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 94314357d2a3..54deb8c202e5 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -678,8 +678,27 @@ async def store_user_field_values( ``LiteLLM_MCPUserCredentials.credential_b64``. A ``"type"`` discriminator lets ``get_user_field_values`` tell user-fields rows apart from BYOK strings and OAuth2 payloads sharing the same column. + + BYOK and OAuth2 credentials share the same ``(user_id, server_id)`` row. + Refuse to overwrite a non-user-fields credential so saving user-field + values does not silently destroy a stored BYOK API key or OAuth2 token. """ + # Guard against silently overwriting a BYOK or OAuth2 credential that + # shares the same (user_id, server_id) row. + 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 + and _decode_user_fields_payload(existing.credential_b64) is None + ): + raise ValueError( + f"Existing credential for user {user_id} and server " + f"{server_id} is not a user-fields payload (likely BYOK or " + f"OAuth2). Refusing to overwrite." + ) + payload = json.dumps({"type": "user_fields", "values": values}) encoded = encrypt_value_helper(payload) await prisma_client.db.litellm_mcpusercredentials.upsert( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0f0bf1285018..efc5228335a6 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -12,6 +12,7 @@ import json import os import re +import time from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast from urllib.parse import urlparse diff --git a/litellm/proxy/_experimental/mcp_server/user_fields.py b/litellm/proxy/_experimental/mcp_server/user_fields.py index 5684f004e199..a3c910fbe73c 100644 --- a/litellm/proxy/_experimental/mcp_server/user_fields.py +++ b/litellm/proxy/_experimental/mcp_server/user_fields.py @@ -13,8 +13,7 @@ from __future__ import annotations -import time -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional from litellm._logging import verbose_logger from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -171,24 +170,3 @@ def resolve_user_field_env( continue env[env_var_name] = value return env - - -def lookup_cached_user_fields( - cache: Dict[Tuple[str, str], Tuple[Optional[Dict[str, str]], float]], - user_id: str, - server_id: str, - ttl_seconds: int, -) -> Tuple[bool, Optional[Dict[str, str]]]: - """Return (cache_hit, values) for the (user, server) cache pair. - - Pulled out so server.py can pass its own cache dict in; keeping the - cache as module-level state in server.py preserves the existing - invalidation hooks called from the management endpoints. - """ - cached = cache.get((user_id, server_id)) - if cached is None: - return False, None - values, ts = cached - if time.monotonic() - ts >= ttl_seconds: - return False, None - return True, values diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 0c392791d88e..320549717e0e 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -113,7 +113,6 @@ def validate_tool_name(name: str) -> _ToolNameValidationResult: # type: ignore[ delete_user_credential, get_all_mcp_servers_for_user, delete_user_field_values, - get_all_mcp_servers, get_mcp_server, get_mcp_servers, get_mcp_submissions, @@ -126,6 +125,11 @@ def validate_tool_name(name: str) -> _ToolNameValidationResult: # type: ignore[ store_user_oauth_credential, update_mcp_server, ) + from litellm.proxy._experimental.mcp_server.user_fields import ( + coerce_user_fields, + compute_missing_user_fields, + server_has_user_fields, + ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize_with_server, exchange_token_with_server, @@ -480,47 +484,6 @@ def _redact_mcp_credentials_list( ) -> List[LiteLLM_MCPServerTable]: return [_redact_mcp_credentials(server) for server in mcp_servers] - def _coerce_user_fields_list(raw: Any) -> List[Dict[str, Any]]: - """Normalize a server.user_fields value (JSON string or list) to a list of dicts. - - Used by the BYOK-list-annotation block and any other site that needs - to inspect declared fields without instantiating MCPUserField models. - """ - if not raw: - return [] - if isinstance(raw, list): - return [e for e in raw if isinstance(e, dict)] - if isinstance(raw, str): - try: - parsed = json.loads(raw) - except (ValueError, TypeError): - return [] - if isinstance(parsed, list): - return [e for e in parsed if isinstance(e, dict)] - return [] - - def _has_required_user_fields(raw: Any) -> bool: - """True iff the server declares at least one required user field.""" - for entry in _coerce_user_fields_list(raw): - if entry.get("required", True): - return True - return False - - def _compute_missing_user_field_keys( - raw: Any, stored_values: Dict[str, str] - ) -> List[str]: - """Field keys the user must still supply (required + currently empty).""" - missing: List[str] = [] - for entry in _coerce_user_fields_list(raw): - field_key = entry.get("field_key") - if not isinstance(field_key, str) or not field_key: - continue - if not entry.get("required", True): - continue - if not stored_values.get(field_key): - missing.append(field_key) - return missing - def _is_restricted_virtual_key_request(user_api_key_dict: UserAPIKeyAuth) -> bool: """Best-effort detection for route-restricted virtual keys. @@ -1007,8 +970,7 @@ async def fetch_all_mcp_servers( relevant_server_ids = [ s.server_id for s in redacted_mcp_servers - if getattr(s, "is_byok", False) - or _has_required_user_fields(getattr(s, "user_fields", None)) + if getattr(s, "is_byok", False) or server_has_user_fields(s) ] if relevant_server_ids: cred_rows = ( @@ -1036,13 +998,13 @@ async def fetch_all_mcp_servers( for server in redacted_mcp_servers: if getattr(server, "is_byok", False): server.has_user_credential = server.server_id in byok_set - if _has_required_user_fields(getattr(server, "user_fields", None)): + if server_has_user_fields(server): stored = user_fields_by_server.get(server.server_id, {}) - server.missing_user_field_keys = ( - _compute_missing_user_field_keys( - getattr(server, "user_fields", None), stored - ) - ) + server.missing_user_field_keys = [ + f["field_key"] + for f in compute_missing_user_fields(server, stored) + if isinstance(f.get("field_key"), str) + ] # Virtual keys only get a sanitized discovery view. if is_restricted_virtual_key: @@ -2245,16 +2207,10 @@ async def store_mcp_user_field_values( # Pre-compute allowed keys from the server's declared fields, then # filter the incoming payload to only those keys. This prevents # callers from polluting the storage blob with arbitrary keys. - raw_fields = getattr(server, "user_fields", None) or [] - if isinstance(raw_fields, str): - try: - raw_fields = json.loads(raw_fields) - except (ValueError, TypeError): - raw_fields = [] declared_keys = { entry.get("field_key") - for entry in raw_fields - if isinstance(entry, dict) and entry.get("field_key") + for entry in coerce_user_fields(server) + if entry.get("field_key") } if not declared_keys: raise HTTPException( @@ -2277,7 +2233,19 @@ async def store_mcp_user_field_values( else: merged[key] = value - await store_user_field_values(prisma_client, user_id, server_id, merged) + try: + await store_user_field_values(prisma_client, user_id, server_id, merged) + except ValueError as e: + # The (user, server) row already holds a BYOK or OAuth2 credential. + # Refuse rather than silently destroying the existing credential. + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": "credential_conflict", + "message": str(e), + "server_id": server_id, + }, + ) # Invalidate the BYOK credential cache for this (user, server) pair # so the next tool call re-reads the row. We piggyback on the @@ -2353,7 +2321,13 @@ async def list_mcp_user_field_values( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) - all_servers = await get_all_mcp_servers(prisma_client) + # Only consider servers the calling user is permitted to see. This + # mirrors the per-user filtering used by the main MCP list endpoint + # and prevents leaking server IDs / user-field metadata (header + # names, env var names) to users without access. + all_servers = await get_all_mcp_servers_for_user( + prisma_client, user_api_key_dict + ) # Pre-filter to servers that actually declare user_fields, then # batch-fetch every credential row for the calling user in one # query. The N+1 alternative (per-server get_user_field_values) @@ -2361,13 +2335,7 @@ async def list_mcp_user_field_values( relevant_servers: List["LiteLLM_MCPServerTable"] = [] relevant_ids: List[str] = [] for server in all_servers: - raw_fields = getattr(server, "user_fields", None) or [] - if isinstance(raw_fields, str): - try: - raw_fields = json.loads(raw_fields) - except (ValueError, TypeError): - raw_fields = [] - if not raw_fields: + if not server_has_user_fields(server): continue relevant_servers.append(server) relevant_ids.append(server.server_id) From a432772748f53515dfef7f9e4c28e3ca542b99b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 06:31:21 +0000 Subject: [PATCH 03/23] fix(lint): extract user-credential annotation helper to satisfy PLR0915 fetch_all_mcp_servers was 55 statements (limit 50) after adding the user-fields annotation. Move the BYOK + missing-user-fields batch annotation into _annotate_user_credential_flags. --- .../mcp_management_endpoints.py | 95 +++++++++---------- 1 file changed, 46 insertions(+), 49 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 320549717e0e..1a3df4ef3016 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -851,6 +851,50 @@ async def _get_team_scoped_mcp_server_list( return _redact_mcp_credentials_list(servers) + async def _annotate_user_credential_flags( + servers: List[LiteLLM_MCPServerTable], + user_id: str, + ) -> None: + """Populate ``has_user_credential`` and ``missing_user_field_keys`` + for the calling user across BYOK / user-fields servers in a single + batched query.""" + from litellm.proxy._experimental.mcp_server.db import ( + _decode_user_credential, + _decode_user_fields_payload, + ) + from litellm.proxy.proxy_server import prisma_client as _byok_prisma_client + + if not user_id or _byok_prisma_client is None: + return + relevant_server_ids = [ + s.server_id + for s in servers + if getattr(s, "is_byok", False) or server_has_user_fields(s) + ] + if not relevant_server_ids: + return + cred_rows = await _byok_prisma_client.db.litellm_mcpusercredentials.find_many( + where={"user_id": user_id, "server_id": {"in": relevant_server_ids}} + ) + 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) + if payload is not None: + user_fields_by_server[row.server_id] = payload + elif _decode_user_credential(row.credential_b64): + byok_set.add(row.server_id) + for server in servers: + if getattr(server, "is_byok", False): + server.has_user_credential = server.server_id in byok_set + if server_has_user_fields(server): + stored = user_fields_by_server.get(server.server_id, {}) + server.missing_user_field_keys = [ + f["field_key"] + for f in compute_missing_user_fields(server, stored) + if isinstance(f.get("field_key"), str) + ] + @router.get( "/server", description="Returns the mcp server list with associated teams", @@ -955,56 +999,9 @@ async def fetch_all_mcp_servers( server.mcp_info = {} server.mcp_info["is_public"] = True - # Annotate has_user_credential (BYOK) and missing_user_field_keys - # (user_fields) for the calling user. Both flags share the same - # storage table, so we batch a single query covering both feature - # sets. - from litellm.proxy._experimental.mcp_server.db import ( - _decode_user_credential, - _decode_user_fields_payload, + await _annotate_user_credential_flags( + redacted_mcp_servers, user_api_key_dict.user_id or "" ) - from litellm.proxy.proxy_server import prisma_client as _byok_prisma_client - - user_id = user_api_key_dict.user_id or "" - if user_id and _byok_prisma_client is not None: - relevant_server_ids = [ - s.server_id - for s in redacted_mcp_servers - if getattr(s, "is_byok", False) or server_has_user_fields(s) - ] - if relevant_server_ids: - cred_rows = ( - await _byok_prisma_client.db.litellm_mcpusercredentials.find_many( - where={ - "user_id": user_id, - "server_id": {"in": relevant_server_ids}, - } - ) - ) - # Build two indexes: one for BYOK presence, one mapping - # server_id → stored user-field values. - 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) - if payload is not None: - user_fields_by_server[row.server_id] = payload - else: - # Anything that isn't a user-fields payload counts - # as a BYOK credential for has_user_credential. - decoded = _decode_user_credential(row.credential_b64) - if decoded: - byok_set.add(row.server_id) - for server in redacted_mcp_servers: - if getattr(server, "is_byok", False): - server.has_user_credential = server.server_id in byok_set - if server_has_user_fields(server): - stored = user_fields_by_server.get(server.server_id, {}) - server.missing_user_field_keys = [ - f["field_key"] - for f in compute_missing_user_fields(server, stored) - if isinstance(f.get("field_key"), str) - ] # Virtual keys only get a sanitized discovery view. if is_restricted_virtual_key: From 8255b000620957f6d973939b322cd97e323f5baf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 06:35:33 +0000 Subject: [PATCH 04/23] fix(mcp): scope user-field-values endpoints to caller's allowed servers The per-server GET/POST/DELETE /v1/mcp/server/{server_id}/user-field-values handlers previously called get_mcp_server directly with no access check. A non-admin user who knew (or guessed) another team's server_id could fetch the admin-declared user_fields metadata (display names, header names, env var names) or write/delete user-field values for that server. Gate the three endpoints behind get_all_mcp_servers_for_user (same scoping already used by the aggregated /v1/mcp/user-field-values endpoint and by fetch_mcp_server). Admins bypass the check; everyone else gets a 404 (not 403) so the status code does not leak server existence. --- .../mcp_management_endpoints.py | 31 ++++++++ .../mcp_server/test_mcp_user_fields.py | 75 ++++++++++++++++++- 2 files changed, 103 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 1a3df4ef3016..aa2ba932805d 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2100,6 +2100,28 @@ async def get_mcp_oauth_user_credential_status( # their own values via these endpoints; values are injected at request # time as either HTTP headers (http/sse) or env vars (stdio). + async def _assert_user_can_access_mcp_server( + prisma_client: Any, + user_api_key_dict: UserAPIKeyAuth, + server_id: str, + ) -> None: + """Refuse access to MCP servers the caller is not permitted to see. + + Returns ``404`` (rather than ``403``) on failure so callers can't + probe for the existence of servers outside their allowed set by + comparing status codes. Admin viewers bypass the check. + """ + if _user_has_admin_view(user_api_key_dict): + return + allowed_servers = await get_all_mcp_servers_for_user( + prisma_client, user_api_key_dict + ) + if not does_mcp_server_exist(allowed_servers, server_id): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP Server {server_id} not found"}, + ) + def _build_user_fields_status( server: "LiteLLM_MCPServerTable", stored_values: Optional[Dict[str, str]], @@ -2161,6 +2183,9 @@ async def get_mcp_user_field_values( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) + await _assert_user_can_access_mcp_server( + prisma_client, user_api_key_dict, server_id + ) server = await get_mcp_server(prisma_client, server_id) if server is None: raise HTTPException( @@ -2194,6 +2219,9 @@ async def store_mcp_user_field_values( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) + await _assert_user_can_access_mcp_server( + prisma_client, user_api_key_dict, server_id + ) server = await get_mcp_server(prisma_client, server_id) if server is None: raise HTTPException( @@ -2277,6 +2305,9 @@ async def delete_mcp_user_field_values( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) + await _assert_user_can_access_mcp_server( + prisma_client, user_api_key_dict, server_id + ) server = await get_mcp_server(prisma_client, server_id) if server is None: raise HTTPException( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index f947849509ae..4f0a64fc640e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -474,6 +474,10 @@ async def test_get_user_field_values_endpoint_reports_missing(): "litellm.proxy._experimental.mcp_server.db.get_user_field_values", AsyncMock(return_value=None), ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[server_row]), + ), ): # Call the endpoint function directly to skip FastAPI auth wiring. from litellm.proxy.management_endpoints import mcp_management_endpoints as mod @@ -546,6 +550,10 @@ async def fake_get(prisma, user_id, server_id): patch( "litellm.proxy._experimental.mcp_server.server._invalidate_user_fields_cache" ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[server_row]), + ), ): from fastapi import FastAPI from fastapi.testclient import TestClient @@ -587,9 +595,15 @@ async def test_post_user_field_values_rejects_server_with_no_declared_fields(): return_value=server_row ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=prisma_client, + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[server_row]), + ), ): from fastapi import FastAPI from fastapi.testclient import TestClient @@ -606,3 +620,58 @@ async def test_post_user_field_values_rejects_server_with_no_declared_fields(): "/v1/mcp/server/srv-1/user-field-values", json={"values": {"X": "y"}} ) assert res.status_code == 400 + + +@pytest.mark.asyncio +async def test_user_field_value_endpoints_404_when_server_not_in_allowed_set(): + """Non-admin callers must not be able to probe servers outside their allowed set. + + Without this gate, an authenticated user who knows another team's + ``server_id`` could call GET/POST/DELETE on the user-field-values + endpoints and leak the admin-declared field descriptors (header names, + env var names) for that server. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + server_row = _server_row_with_user_fields() + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=server_row + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[]), + ), + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="outsider" + ) + client = TestClient(app) + + get_res = client.get("/v1/mcp/server/srv-1/user-field-values") + assert get_res.status_code == 404 + post_res = client.post( + "/v1/mcp/server/srv-1/user-field-values", + json={"values": {"GMAIL_TOKEN": "tok"}}, + ) + assert post_res.status_code == 404 + del_res = client.delete("/v1/mcp/server/srv-1/user-field-values") + assert del_res.status_code == 404 + + # find_unique must never be reached — the access gate fails first, + # before any server metadata is read or returned. + prisma_client.db.litellm_mcpservertable.find_unique.assert_not_called() From adbc4c277f1a99f4e03d24b8d3d2f648d1482e33 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 06:43:22 +0000 Subject: [PATCH 05/23] fix(lint): widen user_fields helpers to accept LiteLLM_MCPServerTable The helpers use getattr(server, 'user_fields', None) and work with any server-shaped object. Widen the type so the management endpoint callers that pass LiteLLM_MCPServerTable type-check cleanly. --- .../proxy/_experimental/mcp_server/user_fields.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/user_fields.py b/litellm/proxy/_experimental/mcp_server/user_fields.py index a3c910fbe73c..9f47126e54f5 100644 --- a/litellm/proxy/_experimental/mcp_server/user_fields.py +++ b/litellm/proxy/_experimental/mcp_server/user_fields.py @@ -13,13 +13,20 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_logger from litellm.types.mcp_server.mcp_server_manager import MCPServer +if TYPE_CHECKING: + from litellm.proxy._types import LiteLLM_MCPServerTable -def coerce_user_fields(server: MCPServer) -> List[Dict[str, Any]]: + UserFieldServer = Union[MCPServer, LiteLLM_MCPServerTable] +else: + UserFieldServer = MCPServer + + +def coerce_user_fields(server: UserFieldServer) -> List[Dict[str, Any]]: """Return the server's declared user fields as a list of plain dicts. The column is stored as JSONB but Prisma sometimes hands it back as a @@ -44,13 +51,13 @@ def coerce_user_fields(server: MCPServer) -> List[Dict[str, Any]]: return [] -def server_has_user_fields(server: MCPServer) -> bool: +def server_has_user_fields(server: UserFieldServer) -> bool: """True iff the server declares any user fields at all.""" return bool(coerce_user_fields(server)) def compute_missing_user_fields( - server: MCPServer, stored_values: Optional[Dict[str, str]] + server: UserFieldServer, stored_values: Optional[Dict[str, str]] ) -> List[Dict[str, Any]]: """Return the declared field definitions the user has yet to fill in. From e45eda44b56cbb6473caf3e24653d0c44e8682f0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 06:59:19 +0000 Subject: [PATCH 06/23] fix(mcp): tighten user-fields/BYOK collisions and dedupe value resolution - db.get_user_credential / store_user_credential: refuse to read or overwrite a user-fields payload, so BYOK never injects the raw JSON blob as an Authorization header. - server._enforce_user_fields: when the caller has no user_id but the server only declares optional user_fields, return without raising the spurious 401. - Promote _get_user_field_values_cached to module level so the dispatch path in mcp_server_manager._resolve_user_field_values shares the same cache + DB lookup implementation. - _call_regular_mcp_tool: resolve user-field values once and pass them to resolve_user_field_headers / resolve_user_field_env directly, eliminating the duplicate cache lookup and the two wrapper helpers. Co-authored-by: Yassin Kortam --- litellm/proxy/_experimental/mcp_server/db.py | 33 ++++++- .../mcp_server/mcp_server_manager.py | 94 ++++++------------- .../proxy/_experimental/mcp_server/server.py | 86 +++++++++-------- 3 files changed, 106 insertions(+), 107 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 54deb8c202e5..8122a14881ae 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -588,7 +588,27 @@ async def store_user_credential( server_id: str, credential: str, ) -> None: - """Store a user credential for a BYOK MCP server.""" + """Store a user credential for a BYOK MCP server. + + BYOK, OAuth2, and user-fields payloads share the same ``credential_b64`` + column. Refuse to overwrite a stored user-fields payload so saving a + BYOK credential does not silently destroy the user's saved field values. + """ + + # Guard against silently overwriting a user-fields payload that shares + # the same (user_id, server_id) row. + 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 + and _decode_user_fields_payload(existing.credential_b64) is not None + ): + raise ValueError( + f"Existing credential for user {user_id} and server " + f"{server_id} holds user-fields values. Refusing to overwrite " + f"with a BYOK credential." + ) encoded = encrypt_value_helper(credential) await prisma_client.db.litellm_mcpusercredentials.upsert( @@ -609,13 +629,22 @@ async def get_user_credential( user_id: str, server_id: str, ) -> Optional[str]: - """Return credential for a user+server pair, or None.""" + """Return credential for a user+server pair, or None. + + The ``credential_b64`` column is multiplexed between BYOK strings, + OAuth2 blobs and user-fields blobs. A row holding a user-fields + payload is not a BYOK credential — return ``None`` so the caller + triggers the normal "no credential stored" flow instead of injecting + the raw JSON blob as an Authorization header. + """ row = await prisma_client.db.litellm_mcpusercredentials.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: return None + if _decode_user_fields_payload(row.credential_b64) is not None: + return None return _decode_user_credential(row.credential_b64) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index efc5228335a6..9e21cd42b935 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -12,7 +12,6 @@ import json import os import re -import time from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast from urllib.parse import urlparse @@ -1348,6 +1347,10 @@ async def _resolve_user_field_values( ``server.execute_mcp_tool`` raises a 401 with a config_url before we even get here, so by the time injection happens we already know the values are present. + + Delegates to ``server._get_user_field_values_cached`` so the + enforcement check and dispatch share a single cache + DB-lookup + implementation. """ from litellm.proxy._experimental.mcp_server.user_fields import ( coerce_user_fields, @@ -1358,69 +1361,12 @@ async def _resolve_user_field_values( if user_api_key_auth is None or not getattr(user_api_key_auth, "user_id", None): return {} - # Reuse the cache and lookup function from server.py so we don't - # pay a second DB round-trip after the enforcement check populated - # the same cache key. - try: - from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 - _USER_FIELDS_CACHE_TTL, - _user_fields_cache, - _write_user_fields_cache, - ) - except ImportError: - _user_fields_cache = {} # type: ignore[assignment] - _USER_FIELDS_CACHE_TTL = 60 - _write_user_fields_cache = None # type: ignore[assignment] - - user_id = user_api_key_auth.user_id or "" - cache_key = (user_id, mcp_server.server_id) - cached = _user_fields_cache.get(cache_key) if _user_fields_cache else None - if cached is not None: - values, ts = cached - if time.monotonic() - ts < _USER_FIELDS_CACHE_TTL: - return values or {} - - from litellm.proxy._experimental.mcp_server.db import get_user_field_values - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - return {} - values = await get_user_field_values( - prisma_client=prisma_client, - user_id=user_id, - server_id=mcp_server.server_id, - ) - if _write_user_fields_cache is not None: - _write_user_fields_cache(user_id, mcp_server.server_id, values) - return values or {} - - async def _resolve_user_field_headers( - self, - mcp_server: MCPServer, - user_api_key_auth: Optional["UserAPIKeyAuth"], - ) -> Dict[str, str]: - from litellm.proxy._experimental.mcp_server.user_fields import ( - resolve_user_field_headers, - ) - - stored = await self._resolve_user_field_values(mcp_server, user_api_key_auth) - if not stored: - return {} - return resolve_user_field_headers(mcp_server, stored) - - async def _resolve_user_field_env( - self, - mcp_server: MCPServer, - user_api_key_auth: Optional["UserAPIKeyAuth"], - ) -> Dict[str, str]: - from litellm.proxy._experimental.mcp_server.user_fields import ( - resolve_user_field_env, + from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + _get_user_field_values_cached, ) - stored = await self._resolve_user_field_values(mcp_server, user_api_key_auth) - if not stored: - return {} - return resolve_user_field_env(mcp_server, stored) + _, values = await _get_user_field_values_cached(mcp_server, user_api_key_auth) + return values or {} async def _create_mcp_client( self, @@ -2841,11 +2787,23 @@ async def _call_regular_mcp_tool( # noqa: PLR0915 extra_headers = {} extra_headers.update(mcp_server.static_headers) - # User-fields: inject each declared user field's stored value either - # as a header (http/sse) or env var (stdio path; handled below). - user_field_headers = await self._resolve_user_field_headers( + # User-fields: resolve the calling user's stored values once and + # reuse them for both the header (http/sse) and env (stdio) paths + # below — calling the resolver twice would repeat cache lookups + # and coercion work for the same data. + from litellm.proxy._experimental.mcp_server.user_fields import ( + resolve_user_field_env, + resolve_user_field_headers, + ) + + stored_user_field_values = await self._resolve_user_field_values( mcp_server, user_api_key_auth ) + user_field_headers = ( + resolve_user_field_headers(mcp_server, stored_user_field_values) + if stored_user_field_values + else {} + ) if user_field_headers: if extra_headers is None: extra_headers = {} @@ -2880,8 +2838,10 @@ async def _call_regular_mcp_tool( # noqa: PLR0915 if extra_headers is not None and len(extra_headers) == 0: extra_headers = None - user_field_env = await self._resolve_user_field_env( - mcp_server, user_api_key_auth + user_field_env = ( + resolve_user_field_env(mcp_server, stored_user_field_values) + if stored_user_field_values + else {} ) stdio_env = self._build_stdio_env( mcp_server, raw_headers, user_field_env=user_field_env or None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e92c68be954b..588d02d3b2f2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -121,6 +121,46 @@ def _write_user_fields_cache( _user_fields_cache[(user_id, server_id)] = (values, time.monotonic()) +async def _get_user_field_values_cached( + mcp_server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth], +) -> Tuple[Optional[str], Optional[Dict[str, str]]]: + """Return (user_id, stored_values) for the calling user. + + ``stored_values`` is None when either no user_id is available or + the user has not yet saved any values. Reads through a 60s + in-memory cache so back-to-back tool calls don't hammer the DB. + + Module-level (not nested) so ``mcp_server_manager`` can reuse the + same cache + DB-lookup path during dispatch — keeping two copies + in sync was a maintenance hazard. + """ + user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or "" + if not user_id: + return None, None + + cache_key = (user_id, mcp_server.server_id) + cached = _user_fields_cache.get(cache_key) + if cached is not None: + values, ts = cached + if time.monotonic() - ts < _USER_FIELDS_CACHE_TTL: + return user_id, values + + from litellm.proxy._experimental.mcp_server.db import get_user_field_values + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return user_id, None + + values = await get_user_field_values( + prisma_client=prisma_client, + user_id=user_id, + server_id=mcp_server.server_id, + ) + _write_user_fields_cache(user_id, mcp_server.server_id, values) + return user_id, values + + # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 # We're making this conditional import to avoid breaking users who use python 3.8. @@ -2066,41 +2106,6 @@ async def _check_byok_credential( }, ) - async def _get_user_field_values_cached( - mcp_server: MCPServer, - user_api_key_auth: Optional[UserAPIKeyAuth], - ) -> Tuple[Optional[str], Optional[Dict[str, str]]]: - """Return (user_id, stored_values) for the calling user. - - ``stored_values`` is None when either no user_id is available or - the user has not yet saved any values. Reads through a 60s - in-memory cache so back-to-back tool calls don't hammer the DB. - """ - user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or "" - if not user_id: - return None, None - - cache_key = (user_id, mcp_server.server_id) - cached = _user_fields_cache.get(cache_key) - if cached is not None: - values, ts = cached - if time.monotonic() - ts < _USER_FIELDS_CACHE_TTL: - return user_id, values - - from litellm.proxy._experimental.mcp_server.db import get_user_field_values - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - return user_id, None - - values = await get_user_field_values( - prisma_client=prisma_client, - user_id=user_id, - server_id=mcp_server.server_id, - ) - _write_user_fields_cache(user_id, mcp_server.server_id, values) - return user_id, values - async def _enforce_user_fields( mcp_server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth], @@ -2121,11 +2126,16 @@ async def _enforce_user_fields( user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or "" if not user_id: - # No identity means we can't look up stored values; treat as - # missing and let the user log in via the dashboard. + # No identity means we can't look up stored values. Compute the + # missing-required set against an empty value map and only raise + # when at least one required field is actually missing — servers + # whose user_fields are all optional must not be blocked here. + missing = compute_missing_user_fields(mcp_server, None) + if not missing: + return detail = build_user_fields_missing_error( mcp_server, - compute_missing_user_fields(mcp_server, None), + missing, _resolve_proxy_base_url_env(), ) raise HTTPException(status_code=401, detail=detail) From 6b0af7f89a5b7fe7cf1ec4f038ae3a1b0a965c77 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 07:15:31 +0000 Subject: [PATCH 07/23] Fix MCP user_fields TOCTOU + handle Pydantic-typed user_fields entries - coerce_user_fields / _build_user_fields_status: accept MCPUserField Pydantic instances in addition to raw dicts, so a fully-typed LiteLLM_MCPServerTable passed through these helpers no longer silently drops every entry. - store_user_field_values: replace read-then-write with optimistic concurrency (create + UniqueViolation retry, then compare-and-swap on credential_b64 via update_many) so a concurrent BYOK / OAuth2 write to the same (user_id, server_id) row cannot be silently overwritten. Co-authored-by: Yassin Kortam --- litellm/proxy/_experimental/mcp_server/db.py | 76 +++++++++++++------ .../_experimental/mcp_server/user_fields.py | 31 ++++++-- .../mcp_management_endpoints.py | 3 + 3 files changed, 83 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 8122a14881ae..f0f93846c347 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,9 +1,12 @@ +import asyncio import base64 import binascii 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 litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import ( @@ -711,35 +714,64 @@ async def store_user_field_values( BYOK and OAuth2 credentials share the same ``(user_id, server_id)`` row. Refuse to overwrite a non-user-fields credential so saving user-field values does not silently destroy a stored BYOK API key or OAuth2 token. - """ - # Guard against silently overwriting a BYOK or OAuth2 credential that - # shares the same (user_id, server_id) row. - 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 - and _decode_user_fields_payload(existing.credential_b64) is None - ): - raise ValueError( - f"Existing credential for user {user_id} and server " - f"{server_id} is not a user-fields payload (likely BYOK or " - f"OAuth2). Refusing to overwrite." - ) + Uses optimistic concurrency to close the read-then-write race against + concurrent ``store_user_credential`` / ``store_user_oauth_credential`` + calls: the write is gated on the previously-observed ``credential_b64`` + being unchanged, and we retry from the re-read on contention. + """ payload = json.dumps({"type": "user_fields", "values": values}) encoded = encrypt_value_helper(payload) - await prisma_client.db.litellm_mcpusercredentials.upsert( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, - data={ - "create": { + + # Bound the retry loop; in practice contention is resolved in a single + # extra round-trip but we leave headroom for pathological interleavings. + for attempt in range(5): + 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: + try: + await prisma_client.db.litellm_mcpusercredentials.create( + data={ + "user_id": user_id, + "server_id": server_id, + "credential_b64": encoded, + } + ) + return + except UniqueViolationError: + # A concurrent writer inserted the row first; restart so we + # can inspect what they wrote before clobbering it. + await asyncio.sleep(0) + continue + + if _decode_user_fields_payload(existing.credential_b64) is None: + raise ValueError( + f"Existing credential for user {user_id} and server " + f"{server_id} is not a user-fields payload (likely BYOK or " + f"OAuth2). Refusing to overwrite." + ) + + # Compare-and-swap on credential_b64: only succeed if the row still + # matches what we just inspected, so a concurrent BYOK/OAuth2 write + # between the read and the write is not silently overwritten. + updated = await prisma_client.db.litellm_mcpusercredentials.update_many( + where={ "user_id": user_id, "server_id": server_id, - "credential_b64": encoded, + "credential_b64": existing.credential_b64, }, - "update": {"credential_b64": encoded}, - }, + data={"credential_b64": encoded}, + ) + if updated: + return + await asyncio.sleep(0) + + raise RuntimeError( + f"store_user_field_values: gave up after repeated concurrent " + f"modifications for user {user_id} and server {server_id}" ) diff --git a/litellm/proxy/_experimental/mcp_server/user_fields.py b/litellm/proxy/_experimental/mcp_server/user_fields.py index 9f47126e54f5..af72740fff2b 100644 --- a/litellm/proxy/_experimental/mcp_server/user_fields.py +++ b/litellm/proxy/_experimental/mcp_server/user_fields.py @@ -26,19 +26,40 @@ UserFieldServer = MCPServer +def _entry_to_dict(entry: Any) -> Optional[Dict[str, Any]]: + """Coerce a single user-field entry to a plain dict, or None if malformed. + + Accepts both raw dicts (as Prisma hands JSONB columns back) and Pydantic + ``MCPUserField`` model instances (as a fully-typed + ``LiteLLM_MCPServerTable`` would carry). + """ + if isinstance(entry, dict): + return entry + model_dump = getattr(entry, "model_dump", None) + if callable(model_dump): + try: + dumped = model_dump() + except Exception: # noqa: BLE001 — drop malformed entries silently + return None + if isinstance(dumped, dict): + return dumped + return None + + def coerce_user_fields(server: UserFieldServer) -> List[Dict[str, Any]]: """Return the server's declared user fields as a list of plain dicts. The column is stored as JSONB but Prisma sometimes hands it back as a - string; this normalises both shapes and silently drops malformed - entries (the admin form filters these out, but DB writes from - external tooling might not). + string; fully-typed ``LiteLLM_MCPServerTable`` instances carry it as a + list of ``MCPUserField`` Pydantic models. This normalises all shapes + and silently drops malformed entries (the admin form filters these + out, but DB writes from external tooling might not). """ raw = getattr(server, "user_fields", None) if not raw: return [] if isinstance(raw, list): - return [e for e in raw if isinstance(e, dict)] + return [d for d in (_entry_to_dict(e) for e in raw) if d is not None] if isinstance(raw, str): import json @@ -47,7 +68,7 @@ def coerce_user_fields(server: UserFieldServer) -> List[Dict[str, Any]]: except (ValueError, TypeError): return [] if isinstance(parsed, list): - return [e for e in parsed if isinstance(e, dict)] + return [d for d in (_entry_to_dict(e) for e in parsed) if d is not None] return [] diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index aa2ba932805d..62b35b56fa23 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2140,6 +2140,9 @@ def _build_user_fields_status( raw_fields = [] user_fields: List[MCPUserField] = [] for entry in raw_fields: + if isinstance(entry, MCPUserField): + user_fields.append(entry) + continue if not isinstance(entry, dict): continue try: From 5215fbeeebb7b3060da868b13bbbc95ccb1326a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 07:19:18 +0000 Subject: [PATCH 08/23] fix(mcp): return 409 on BYOK/user-fields conflict; coerce Pydantic user_fields - store_mcp_user_credential now catches the ValueError raised by store_user_credential when the (user, server) row already holds a user-fields payload, returning HTTP 409 instead of leaking a 500. - coerce_user_fields now also accepts MCPUserField Pydantic instances (the shape used by LiteLLM_MCPServerTable.user_fields), so the management-layer annotation and enforcement paths no longer return silently empty when callers pass a parsed server record. --- .../mcp_management_endpoints.py | 16 +++++++++++--- .../mcp_server/test_mcp_user_fields.py | 22 +++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 62b35b56fa23..0b952f95a87b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1918,9 +1918,19 @@ async def store_mcp_user_credential( detail={"error": "User ID not found in token"}, ) if payload.save: - await store_user_credential( - prisma_client, user_id, server_id, payload.credential - ) + try: + await store_user_credential( + prisma_client, user_id, server_id, payload.credential + ) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": "credential_conflict", + "message": str(e), + "server_id": server_id, + }, + ) from litellm.proxy._experimental.mcp_server.server import ( _invalidate_byok_cred_cache, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 4f0a64fc640e..05cb4c4b6dd0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -179,6 +179,28 @@ def test_coerce_user_fields_empty_when_missing(): assert server_has_user_fields(srv) is False +def test_coerce_user_fields_accepts_litellm_mcp_server_table(): + """LiteLLM_MCPServerTable.user_fields is List[MCPUserField] (Pydantic + instances), not List[dict]. The helper must normalise both shapes so + the management-layer annotation / enforcement paths don't silently + return empty results. + """ + table = LiteLLM_MCPServerTable( + server_id="s3", + transport=MCPTransport.http, + user_fields=[ + {"field_key": "TOKEN", "header_name": "Authorization", "required": True}, + {"field_key": "WS", "header_name": "X-Workspace", "required": False}, + ], + ) + assert isinstance(table.user_fields[0], MCPUserField) + coerced = coerce_user_fields(table) + assert [f["field_key"] for f in coerced] == ["TOKEN", "WS"] + assert server_has_user_fields(table) is True + missing = compute_missing_user_fields(table, None) + assert [f["field_key"] for f in missing] == ["TOKEN"] + + def test_compute_missing_required_only(): srv = _gmail_server() missing = compute_missing_user_fields(srv, None) From 43501183f402b78b46c706548e5012be2ffad641 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 07:29:02 +0000 Subject: [PATCH 09/23] test(mcp): cover DELETE and aggregate list user-field-values endpoints --- .../mcp_server/test_mcp_user_fields.py | 215 ++++++++++++++++++ 1 file changed, 215 insertions(+) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 05cb4c4b6dd0..9c50de0c8dff 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -697,3 +697,218 @@ async def test_user_field_value_endpoints_404_when_server_not_in_allowed_set(): # find_unique must never be reached — the access gate fails first, # before any server metadata is read or returned. prisma_client.db.litellm_mcpservertable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_delete_user_field_values_clears_stored_and_returns_status(): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + server_row = _server_row_with_user_fields() + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=server_row + ) + + delete_calls: List[tuple] = [] + + async def fake_delete(prisma, user_id, server_id): + delete_calls.append((user_id, server_id)) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_field_values", + new=fake_delete, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[server_row]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._invalidate_byok_cred_cache" + ), + patch( + "litellm.proxy._experimental.mcp_server.server._invalidate_user_fields_cache" + ), + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-del" + ) + client = TestClient(app) + res = client.delete("/v1/mcp/server/srv-1/user-field-values") + assert res.status_code == 200, res.text + body = res.json() + assert body["server_id"] == "srv-1" + assert body["stored_field_keys"] == [] + assert "GMAIL_TOKEN" in body["missing_field_keys"] + assert delete_calls == [("user-del", "srv-1")] + + +@pytest.mark.asyncio +async def test_delete_user_field_values_swallows_record_not_found(): + """DELETE on a row that was already removed should still succeed.""" + from prisma.errors import RecordNotFoundError + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + server_row = _server_row_with_user_fields() + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=server_row + ) + + async def fake_delete(prisma, user_id, server_id): + raise RecordNotFoundError(data={"error": {"message": "no rows", "meta": {}}}) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_field_values", + new=fake_delete, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[server_row]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._invalidate_byok_cred_cache" + ), + patch( + "litellm.proxy._experimental.mcp_server.server._invalidate_user_fields_cache" + ), + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-del2" + ) + client = TestClient(app) + res = client.delete("/v1/mcp/server/srv-1/user-field-values") + assert res.status_code == 200, res.text + + +@pytest.mark.asyncio +async def test_list_user_field_values_aggregates_per_server_status(): + """GET /v1/mcp/user-field-values returns one entry per server with user_fields.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + gmail_row = _server_row_with_user_fields() + no_fields_row = _server_row_with_user_fields() + no_fields_row.server_id = "srv-2" + no_fields_row.user_fields = [] + + other_row = _server_row_with_user_fields() + other_row.server_id = "srv-3" + other_row.user_fields = [ + { + "field_key": "JIRA_TOKEN", + "header_name": "Authorization", + "header_value_template": "Bearer {value}", + "required": True, + } + ] + + encoded = encrypt_value_helper( + json.dumps({"type": "user_fields", "values": {"JIRA_TOKEN": "jt"}}) + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[MagicMock(server_id="srv-3", credential_b64=encoded)] + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[gmail_row, no_fields_row, other_row]), + ), + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-list" + ) + client = TestClient(app) + res = client.get("/v1/mcp/user-field-values") + assert res.status_code == 200, res.text + body = res.json() + # Only servers with declared user_fields are returned. + ids = {entry["server_id"] for entry in body} + assert ids == {"srv-1", "srv-3"} + + by_id = {entry["server_id"]: entry for entry in body} + # srv-1 has no stored values → required field missing. + assert "GMAIL_TOKEN" in by_id["srv-1"]["missing_field_keys"] + # srv-3 has the stored value → nothing missing. + assert by_id["srv-3"]["missing_field_keys"] == [] + assert by_id["srv-3"]["stored_field_keys"] == ["JIRA_TOKEN"] + + +@pytest.mark.asyncio +async def test_list_user_field_values_returns_empty_when_no_servers_declare_fields(): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.management_endpoints import mcp_management_endpoints as mod + + plain_row = _server_row_with_user_fields() + plain_row.user_fields = [] + + prisma_client = MagicMock() + prisma_client.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", + AsyncMock(return_value=[plain_row]), + ), + ): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = FastAPI() + app.include_router(mod.router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-empty" + ) + client = TestClient(app) + res = client.get("/v1/mcp/user-field-values") + assert res.status_code == 200, res.text + assert res.json() == [] + # Should short-circuit before hitting the credentials table. + prisma_client.db.litellm_mcpusercredentials.find_many.assert_not_called() From d9e1ce289447c04c16d5a219521f3cbd8a39e23b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 07:29:48 +0000 Subject: [PATCH 10/23] 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 --- litellm/proxy/_experimental/mcp_server/db.py | 40 +++++++++++++++---- .../mcp_server/mcp_server_manager.py | 5 ++- .../mcp_management_endpoints.py | 11 +++-- 3 files changed, 45 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index f0f93846c347..175a9e4597a3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -5,7 +5,7 @@ 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, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 9e21cd42b935..deb6b39555b3 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1309,7 +1309,10 @@ def _build_stdio_env( """ 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 diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 0b952f95a87b..9c1410bbae99 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -860,7 +860,7 @@ async def _annotate_user_credential_flags( 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 @@ async def _annotate_user_credential_flags( 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): From 501afd2361f710621647b8f49fec6e07e231aa1d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 07:41:06 +0000 Subject: [PATCH 11/23] fix(mcp): provide data arg to RecordNotFoundError; preserve BYOK has_user_credential on decryption failure - delete_user_credential now constructs RecordNotFoundError with the required 'data' kwarg so the proxy delete endpoints continue to receive a typed exception instead of a TypeError-induced 500. - _annotate_user_credential_flags falls back to row-existence when a credential row cannot be decrypted (e.g. after salt-key rotation), so BYOK servers do not regress to has_user_credential=False when the row is in fact stored. Co-authored-by: Yassin Kortam --- litellm/proxy/_experimental/mcp_server/db.py | 13 +++++++++++-- .../mcp_management_endpoints.py | 8 +++++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 175a9e4597a3..ff4b72002be3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -680,11 +680,20 @@ async def delete_user_credential( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if existing is None: - raise RecordNotFoundError() + raise RecordNotFoundError( + data={"error": {"message": "no BYOK credential row", "meta": {}}} + ) 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() + raise RecordNotFoundError( + data={ + "error": { + "message": "row holds user-fields payload, not a BYOK credential", + "meta": {}, + } + } + ) await prisma_client.db.litellm_mcpusercredentials.delete( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 9c1410bbae99..f3c0fecb1743 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -882,12 +882,14 @@ async def _annotate_user_credential_flags( # 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) + payload = _parse_user_fields_plaintext(decoded) if decoded else None if payload is not None: user_fields_by_server[row.server_id] = payload else: + # Either a BYOK credential or an undecryptable row (e.g. + # after a salt-key rotation). Either way, a credential row + # exists for this (user, server), so surface it as present + # instead of silently telling the user to reconnect. byok_set.add(row.server_id) for server in servers: if getattr(server, "is_byok", False): From 4618e28046f40a3a91fe914cec59626fa45693e4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 07:44:32 +0000 Subject: [PATCH 12/23] fix(mcp): return 409 on user-fields conflict for OAuth credential POST store_user_oauth_credential raises ValueError when the (user, server) row already holds a non-OAuth2 payload, which now includes user-fields blobs. Catch it in the explicit credential-POST endpoint and surface 409 instead of letting management_endpoint_wrapper re-raise as 500, mirroring the existing handling on the BYOK and user-field-values POST endpoints. --- .../mcp_management_endpoints.py | 28 +++++++++++++------ 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f3c0fecb1743..3fd5c1a7f005 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2009,15 +2009,25 @@ async def store_mcp_oauth_user_credential( status_code=status.HTTP_400_BAD_REQUEST, detail={"error": "User ID not found in token"}, ) - await store_user_oauth_credential( - prisma_client, - user_id, - server_id, - payload.access_token, - refresh_token=payload.refresh_token, - expires_in=payload.expires_in, - scopes=payload.scopes, - ) + try: + await store_user_oauth_credential( + prisma_client, + user_id, + server_id, + payload.access_token, + refresh_token=payload.refresh_token, + expires_in=payload.expires_in, + scopes=payload.scopes, + ) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": "credential_conflict", + "message": str(e), + "server_id": server_id, + }, + ) # Read back the persisted record so the response reflects the stored # expires_at rather than recomputing it here (which could diverge by # milliseconds or if the storage logic ever adds a grace period). From cbb13bf53c947cb78ffd45b6f38bdd8c7ff18402 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 08:03:00 +0000 Subject: [PATCH 13/23] Fix concurrent write race + warning message in MCP user-fields MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - store_user_credential / store_user_oauth_credential now use the same compare-and-swap (create + update_many) pattern as store_user_field_values so a concurrent user-fields write between the find_unique check and the upsert cannot silently overwrite the user-fields payload. - store_user_field_values accepts an optional merge_fn callable that is re-evaluated on every CAS retry. The user-fields POST endpoint now passes a merge_fn so a concurrent partial save by the same user from another tab is folded into the result instead of being clobbered with a stale pre-loop merge. - The hook_extra_headers Authorization-overwrite warning no longer claims the existing header came specifically from static_headers — it can now come from static_headers, forwarded raw headers, or user-fields. Co-authored-by: Yassin Kortam --- litellm/proxy/_experimental/mcp_server/db.py | 190 +++++++++++++----- .../mcp_server/mcp_server_manager.py | 4 +- .../mcp_management_endpoints.py | 31 ++- .../mcp_server/test_db_credentials.py | 21 +- .../mcp_server/test_mcp_user_fields.py | 14 +- 5 files changed, 196 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index ff4b72002be3..b0c120ff2ce8 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -3,7 +3,7 @@ import binascii import json from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast +from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Union, cast from prisma.errors import RecordNotFoundError, UniqueViolationError @@ -596,34 +596,61 @@ async def store_user_credential( BYOK, OAuth2, and user-fields payloads share the same ``credential_b64`` column. Refuse to overwrite a stored user-fields payload so saving a BYOK credential does not silently destroy the user's saved field values. + + Uses optimistic concurrency (compare-and-swap on ``credential_b64``) so + a concurrent ``store_user_field_values`` write between the read and the + write cannot be silently overwritten — mirroring the protection that + user-fields writes already have against BYOK writes. """ - # Guard against silently overwriting a user-fields payload that shares - # the same (user_id, server_id) row. - 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 - and _decode_user_fields_payload(existing.credential_b64) is not None - ): - raise ValueError( - f"Existing credential for user {user_id} and server " - f"{server_id} holds user-fields values. Refusing to overwrite " - f"with a BYOK credential." + encoded = encrypt_value_helper(credential) + + for _attempt in range(5): + existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - 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={ - "create": { + if existing is None: + try: + await prisma_client.db.litellm_mcpusercredentials.create( + data={ + "user_id": user_id, + "server_id": server_id, + "credential_b64": encoded, + } + ) + return + except UniqueViolationError: + # A concurrent writer inserted the row first; restart so we + # can inspect what they wrote before clobbering it. + await asyncio.sleep(0) + continue + + if _decode_user_fields_payload(existing.credential_b64) is not None: + raise ValueError( + f"Existing credential for user {user_id} and server " + f"{server_id} holds user-fields values. Refusing to overwrite " + f"with a BYOK credential." + ) + + # Compare-and-swap on credential_b64: only succeed if the row still + # matches what we just inspected, so a concurrent user-fields write + # between the read and the write is not silently overwritten. + updated = await prisma_client.db.litellm_mcpusercredentials.update_many( + where={ "user_id": user_id, "server_id": server_id, - "credential_b64": encoded, + "credential_b64": existing.credential_b64, }, - "update": {"credential_b64": encoded}, - }, + data={"credential_b64": encoded}, + ) + if updated: + return + await asyncio.sleep(0) + + raise RuntimeError( + f"store_user_credential: gave up after repeated concurrent " + f"modifications for user {user_id} and server {server_id}" ) @@ -737,8 +764,10 @@ async def store_user_field_values( prisma_client: PrismaClient, user_id: str, server_id: str, - values: Dict[str, str], -) -> None: + values: Optional[Dict[str, str]] = None, + *, + merge_fn: Optional[Callable[[Dict[str, str]], Dict[str, str]]] = None, +) -> Dict[str, str]: """Persist the calling user's values for an MCP server's user fields. The full set of values is encoded into a single encrypted JSON blob in @@ -750,23 +779,56 @@ async def store_user_field_values( Refuse to overwrite a non-user-fields credential so saving user-field values does not silently destroy a stored BYOK API key or OAuth2 token. + Exactly one of ``values`` or ``merge_fn`` must be supplied: + + * ``values`` writes the dict as-is. + * ``merge_fn`` is invoked with the currently-stored user-fields values + (``{}`` if no row exists) and must return the dict to persist. It is + re-invoked on every CAS retry, so a concurrent partial save from the + same user in another tab is merged into the result instead of being + silently overwritten. + Uses optimistic concurrency to close the read-then-write race against concurrent ``store_user_credential`` / ``store_user_oauth_credential`` calls: the write is gated on the previously-observed ``credential_b64`` being unchanged, and we retry from the re-read on contention. + + Returns the dict that was ultimately persisted. """ - payload = json.dumps({"type": "user_fields", "values": values}) - encoded = encrypt_value_helper(payload) + if (values is None) == (merge_fn is None): + raise ValueError( + "store_user_field_values: exactly one of `values` or `merge_fn` " + "must be provided" + ) + + def _encode(field_values: Dict[str, str]) -> str: + return encrypt_value_helper( + json.dumps({"type": "user_fields", "values": field_values}) + ) + + # Pre-compute encoding for the ``values`` path so we don't pay encryption + # cost on every retry iteration when there is no merge function. + static_encoded: Optional[str] = None + if merge_fn is None: + assert values is not None + static_encoded = _encode(values) # Bound the retry loop; in practice contention is resolved in a single # extra round-trip but we leave headroom for pathological interleavings. - for attempt in range(5): + for _attempt in range(5): 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: + if merge_fn is not None: + current_values = merge_fn({}) + encoded = _encode(current_values) + else: + assert values is not None and static_encoded is not None + current_values = values + encoded = static_encoded try: await prisma_client.db.litellm_mcpusercredentials.create( data={ @@ -775,20 +837,32 @@ async def store_user_field_values( "credential_b64": encoded, } ) - return + return current_values except UniqueViolationError: # A concurrent writer inserted the row first; restart so we # can inspect what they wrote before clobbering it. await asyncio.sleep(0) continue - if _decode_user_fields_payload(existing.credential_b64) is None: + existing_field_values = _decode_user_fields_payload(existing.credential_b64) + if existing_field_values is None: raise ValueError( f"Existing credential for user {user_id} and server " f"{server_id} is not a user-fields payload (likely BYOK or " f"OAuth2). Refusing to overwrite." ) + # When a merge function is supplied, recompute on every retry against + # the freshly-read existing values so a concurrent user-fields write + # from the same user in another tab is not silently dropped. + if merge_fn is not None: + current_values = merge_fn(existing_field_values) + encoded = _encode(current_values) + else: + assert values is not None and static_encoded is not None + current_values = values + encoded = static_encoded + # Compare-and-swap on credential_b64: only succeed if the row still # matches what we just inspected, so a concurrent BYOK/OAuth2 write # between the read and the write is not silently overwritten. @@ -801,7 +875,7 @@ async def store_user_field_values( data={"credential_b64": encoded}, ) if updated: - return + return current_values await asyncio.sleep(0) raise RuntimeError( @@ -893,38 +967,60 @@ async def store_user_oauth_credential( if scopes: payload["scopes"] = scopes - # Guard against silently overwriting a BYOK credential with an OAuth token. - # 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: + encoded = encrypt_value_helper(json.dumps(payload)) + + # Optimistic concurrency: compare-and-swap on ``credential_b64`` so a + # concurrent BYOK or user-fields write between the read and the write is + # not silently overwritten. ``skip_byok_guard`` (token refresh) bypasses + # the type check but still uses CAS to avoid clobbering data. + for _attempt in range(5): 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: + try: + await prisma_client.db.litellm_mcpusercredentials.create( + data={ + "user_id": user_id, + "server_id": server_id, + "credential_b64": encoded, + } + ) + return + except UniqueViolationError: + await asyncio.sleep(0) + continue + if ( - existing is not None + not skip_byok_guard 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. + # Existing row is either a BYOK secret, a user-fields blob, or an + # OAuth2 row that no longer decrypts (e.g. after a salt-key + # rotation). In any case, refuse to overwrite — the caller would + # clobber data that may still be recoverable. raise ValueError( 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)) - await prisma_client.db.litellm_mcpusercredentials.upsert( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, - data={ - "create": { + updated = await prisma_client.db.litellm_mcpusercredentials.update_many( + where={ "user_id": user_id, "server_id": server_id, - "credential_b64": encoded, + "credential_b64": existing.credential_b64, }, - "update": {"credential_b64": encoded}, - }, + data={"credential_b64": encoded}, + ) + if updated: + return + await asyncio.sleep(0) + + raise RuntimeError( + f"store_user_oauth_credential: gave up after repeated concurrent " + f"modifications for user {user_id} and server {server_id}" ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index deb6b39555b3..a2ab9d9f905a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2819,8 +2819,8 @@ async def _call_regular_mcp_tool( # noqa: PLR0915 if "Authorization" in extra_headers: verbose_logger.warning( "MCPServerManager: hook_extra_headers 'Authorization' will overwrite " - "the existing Authorization header from static_headers. " - "The hook JWT will take precedence." + "the existing Authorization header (from static_headers, forwarded raw " + "headers, or user-fields). The hook JWT will take precedence." ) elif server_auth_header is not None: # server_auth_header is passed separately to _create_mcp_client as diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 3fd5c1a7f005..0209535b2c94 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2277,19 +2277,28 @@ async def store_mcp_user_field_values( ) # Merge with anything already stored — partial saves should preserve - # previously-supplied values rather than wipe them. - existing = await get_user_field_values(prisma_client, user_id, server_id) or {} - merged: Dict[str, str] = {**existing} - for key, value in payload.values.items(): - if key not in declared_keys: - continue - if value == "": - merged.pop(key, None) - else: - merged[key] = value + # previously-supplied values rather than wipe them. The merge is run + # inside ``store_user_field_values`` on every CAS retry so a + # concurrent partial save by the same user from another tab is + # folded in rather than silently overwritten. + def _merge_user_fields(existing: Dict[str, str]) -> Dict[str, str]: + merged_local: Dict[str, str] = {**existing} + for key, value in payload.values.items(): + if key not in declared_keys: + continue + if value == "": + merged_local.pop(key, None) + else: + merged_local[key] = value + return merged_local try: - await store_user_field_values(prisma_client, user_id, server_id, merged) + merged = await store_user_field_values( + prisma_client, + user_id, + server_id, + merge_fn=_merge_user_fields, + ) except ValueError as e: # The (user, server) row already holds a BYOK or OAuth2 credential. # Refuse rather than silently destroying the existing credential. 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 078adf72d4c3..a533e4d79c19 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 @@ -36,9 +36,13 @@ def _set_salt_key(monkeypatch): 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.""" + for find_unique and behaves async-correctly for the CAS write paths.""" prisma = MagicMock() prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_mcpusercredentials.create = AsyncMock() + # ``update_many`` returns the number of affected rows for CAS — 1 means + # the swap succeeded so the production loop exits. + prisma.db.litellm_mcpusercredentials.update_many = AsyncMock(return_value=1) prisma.db.litellm_mcpusercredentials.upsert = AsyncMock() prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) return prisma @@ -55,7 +59,20 @@ def _legacy_row(payload: str): def _stored_value(prisma) -> str: - """Pull the credential_b64 value passed to the most recent upsert call.""" + """Pull the credential_b64 value passed to the most recent write. + + The production code uses ``create`` when no row exists and + ``update_many`` (CAS) when a row already exists. Walk both call histories + and return the most recent ``credential_b64`` written. + """ + create_mock = prisma.db.litellm_mcpusercredentials.create + update_many_mock = prisma.db.litellm_mcpusercredentials.update_many + + if update_many_mock.call_args is not None: + return update_many_mock.call_args.kwargs["data"]["credential_b64"] + if create_mock.call_args is not None: + return create_mock.call_args.kwargs["data"]["credential_b64"] + # Fall back to the legacy upsert path for any tests that still rely on it. call = prisma.db.litellm_mcpusercredentials.upsert.call_args data = call.kwargs["data"] create_value = data["create"]["credential_b64"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 9c50de0c8dff..325a0bb9fe76 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -545,10 +545,20 @@ async def test_post_user_field_values_rejects_undeclared_keys(): captured = {} - async def fake_store(prisma, user_id, server_id, values): + async def fake_store( + prisma, user_id, server_id, values=None, *, merge_fn=None, **kwargs + ): captured["user_id"] = user_id captured["server_id"] = server_id - captured["values"] = values + # The endpoint passes a ``merge_fn`` that closes over the request + # payload + declared_keys; resolve it against an empty existing dict + # (mirroring "no row yet") so the test can assert on the final values. + if merge_fn is not None: + result = merge_fn({}) + else: + result = values or {} + captured["values"] = result + return result async def fake_get(prisma, user_id, server_id): return None # no existing values From 3851d8ae8828eedb52f520a98d6172f960620617 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 08:18:33 +0000 Subject: [PATCH 14/23] fix(mcp): forward user_fields on temp records and dedupe normalization Pass through payload.user_fields when building the temporary MCP server record so session-cached servers retain their admin-declared fields. Replace the bespoke list/string/instance handling in _build_user_fields_status with the shared coerce_user_fields helper. --- .../mcp_management_endpoints.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 0209535b2c94..c6f7e65ebd6f 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -670,6 +670,7 @@ def _build_temporary_mcp_server_record( registration_url=payload.registration_url, allow_all_keys=payload.allow_all_keys, available_on_public_internet=payload.available_on_public_internet, + user_fields=payload.user_fields, ) def get_prisma_client_or_throw(message: str): @@ -2159,19 +2160,8 @@ def _build_user_fields_status( (None when no row exists yet). ``missing_field_keys`` enumerates only the ``required`` fields the user has yet to fill in. """ - raw_fields = getattr(server, "user_fields", None) or [] - if isinstance(raw_fields, str): - try: - raw_fields = json.loads(raw_fields) - except (ValueError, TypeError): - raw_fields = [] user_fields: List[MCPUserField] = [] - for entry in raw_fields: - if isinstance(entry, MCPUserField): - user_fields.append(entry) - continue - if not isinstance(entry, dict): - continue + for entry in coerce_user_fields(server): try: user_fields.append(MCPUserField(**entry)) except Exception: # noqa: BLE001 — drop malformed entries silently From 2cce5383052b80bc5710f04edc91d63f5ea21ace Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 08:35:58 +0000 Subject: [PATCH 15/23] test(mcp): cover user_fields helper edge cases for codecov Add unit tests for the previously-uncovered branches in user_fields.py so the patch hits the project's coverage target: - coerce_user_fields with a JSON-encoded string (Prisma JSONB fallback), invalid JSON, non-list JSON, unsupported scalar types, and entries that are neither dicts nor model_dump-able (including a model_dump that raises). - compute_missing_user_fields skipping malformed entries (no key, empty key, non-string key). - resolve_user_field_headers skipping entries without header_name. - resolve_user_field_env skipping entries without env_var_name and empty stored values. Also fixes test_coerce_user_fields_accepts_json_string, which was inadvertently round-tripping through json.loads(json.dumps(...)) and therefore never actually exercising the string-input branch it was documenting. --- .../mcp_server/test_mcp_user_fields.py | 94 ++++++++++++++++++- 1 file changed, 93 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 325a0bb9fe76..981b9d7eb63b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -169,10 +169,50 @@ def test_coerce_user_fields_accepts_list(): def test_coerce_user_fields_accepts_json_string(): srv = _gmail_server(user_fields=None) # Prisma occasionally hands JSONB columns back as a string. - srv.user_fields = json.loads(json.dumps([{"field_key": "X", "required": True}])) + srv.user_fields = json.dumps([{"field_key": "X", "required": True}]) assert coerce_user_fields(srv) == [{"field_key": "X", "required": True}] +def test_coerce_user_fields_returns_empty_for_invalid_json_string(): + srv = _gmail_server(user_fields=None) + srv.user_fields = "not-valid-json{" + assert coerce_user_fields(srv) == [] + + +def test_coerce_user_fields_returns_empty_for_non_list_json_string(): + srv = _gmail_server(user_fields=None) + # Well-formed JSON but not a list — must not crash and must yield []. + srv.user_fields = json.dumps({"field_key": "X"}) + assert coerce_user_fields(srv) == [] + + +def test_coerce_user_fields_returns_empty_for_unsupported_type(): + """Anything that isn't a list or a string (e.g. a stray int from a bad + migration) must drop to the safe empty fallback rather than raise.""" + srv = _gmail_server(user_fields=None) + srv.user_fields = 42 # type: ignore[assignment] + assert coerce_user_fields(srv) == [] + + +def test_coerce_user_fields_drops_entries_without_model_dump(): + """Entries that are neither dicts nor model_dump-able are silently dropped.""" + srv = _gmail_server(user_fields=None) + srv.user_fields = ["a-bare-string", 7] # type: ignore[list-item] + assert coerce_user_fields(srv) == [] + + +def test_coerce_user_fields_drops_entries_when_model_dump_raises(): + """A model_dump that blows up should not bring down request dispatch.""" + + class _Boom: + def model_dump(self): + raise RuntimeError("kaboom") + + srv = _gmail_server(user_fields=None) + srv.user_fields = [_Boom()] # type: ignore[list-item] + assert coerce_user_fields(srv) == [] + + def test_coerce_user_fields_empty_when_missing(): srv = MCPServer(server_id="s2", name="empty", transport=MCPTransport.http) assert coerce_user_fields(srv) == [] @@ -223,6 +263,21 @@ def test_compute_missing_optional_field_unaffected(): assert missing == [] +def test_compute_missing_skips_entries_without_valid_field_key(): + """Malformed entries (no field_key, non-string field_key, empty string) are + dropped instead of being surfaced as missing.""" + srv = _gmail_server( + user_fields=[ + {"field_key": "", "required": True}, + {"required": True}, + {"field_key": 5, "required": True}, + {"field_key": "REAL", "required": True}, + ] + ) + missing = compute_missing_user_fields(srv, None) + assert [f["field_key"] for f in missing] == ["REAL"] + + def test_resolve_user_field_headers_applies_template(): srv = _gmail_server() headers = resolve_user_field_headers( @@ -255,6 +310,21 @@ def test_resolve_user_field_headers_falls_back_on_bad_template(): assert headers == {"Authorization": "raw"} +def test_resolve_user_field_headers_skips_entries_missing_header_name(): + """Stdio-only fields (no header_name) must not produce empty-name headers.""" + srv = MCPServer( + server_id="s4", + name="local", + transport=MCPTransport.http, + user_fields=[ + {"field_key": "ONLY_ENV", "env_var_name": "SECRET", "required": True}, + {"field_key": "WS", "header_name": "X-Workspace", "required": False}, + ], + ) + headers = resolve_user_field_headers(srv, {"ONLY_ENV": "x", "WS": "ws1"}) + assert headers == {"X-Workspace": "ws1"} + + def test_resolve_user_field_env_for_stdio(): srv = MCPServer( server_id="s3", @@ -268,6 +338,28 @@ def test_resolve_user_field_env_for_stdio(): assert env == {"GITHUB_TOKEN": "ghs_xxx"} +def test_resolve_user_field_env_skips_entries_without_env_var_name(): + """HTTP-only fields (no env_var_name) and empty stored values are skipped.""" + srv = MCPServer( + server_id="s5", + name="local", + transport=MCPTransport.stdio, + user_fields=[ + { + "field_key": "HTTP_ONLY", + "header_name": "Authorization", + "required": True, + }, + {"field_key": "GH_TOKEN", "env_var_name": "GITHUB_TOKEN", "required": True}, + {"field_key": "EMPTY", "env_var_name": "EMPTY_VAR", "required": False}, + ], + ) + env = resolve_user_field_env( + srv, {"HTTP_ONLY": "x", "GH_TOKEN": "ghs", "EMPTY": ""} + ) + assert env == {"GITHUB_TOKEN": "ghs"} + + def test_build_user_fields_missing_error_uses_proxy_base_url(): srv = _gmail_server() missing = compute_missing_user_fields(srv, None) From 5437285451c6639c8d35f4977e59d4de04227d6e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 08:45:41 +0000 Subject: [PATCH 16/23] fix(mcp): drop orphan user-field keys + reject duplicate field_keys --- litellm/proxy/_types.py | 28 +++++++++++++++++++ .../mcp_management_endpoints.py | 3 +- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c89e6ef0f350..a2b41095af02 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1268,6 +1268,24 @@ class MCPUserField(LiteLLMPydanticObjectBase): env_var_name: Optional[str] = None +def _validate_unique_user_field_keys( + user_fields: List[MCPUserField], +) -> List[MCPUserField]: + seen: set = set() + duplicates: set = set() + for entry in user_fields: + key = entry.field_key + if key in seen: + duplicates.add(key) + else: + seen.add(key) + if duplicates: + raise ValueError( + f"user_fields contains duplicate field_key values: {sorted(duplicates)}" + ) + return user_fields + + class MCPUserFieldValuesRequest(LiteLLMPydanticObjectBase): """Body for storing the calling user's values for an MCP server's user fields.""" @@ -1368,6 +1386,11 @@ def validate_credentials_requirements(cls, values): """ return values + @field_validator("user_fields") + @classmethod + def _user_fields_unique(cls, v: List[MCPUserField]) -> List[MCPUserField]: + return _validate_unique_user_field_keys(v) + class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): server_id: str @@ -1427,6 +1450,11 @@ def validate_transport_fields(cls, values): ) return values + @field_validator("user_fields") + @classmethod + def _user_fields_unique(cls, v: List[MCPUserField]) -> List[MCPUserField]: + return _validate_unique_user_field_keys(v) + class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): """Represents a LiteLLM_MCPServerTable record""" diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c6f7e65ebd6f..010816db0424 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2166,8 +2166,9 @@ def _build_user_fields_status( user_fields.append(MCPUserField(**entry)) except Exception: # noqa: BLE001 — drop malformed entries silently continue + declared_keys = {f.field_key for f in user_fields} stored: Dict[str, str] = stored_values or {} - stored_keys = [k for k, v in stored.items() if v] + stored_keys = [k for k, v in stored.items() if v and k in declared_keys] missing_keys = [ f.field_key for f in user_fields From 0023194ccd0e1201c94767d14179566cbaefed5f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 08:46:20 +0000 Subject: [PATCH 17/23] fix(mcp): inject user-field headers for OpenAPI/local tool dispatch The OpenAPI/local-tool dispatch path validated required user-field values via _enforce_user_fields but never injected them into the upstream request. The managed MCP path already did this in _call_regular_mcp_tool, so users configuring a user-fields-enabled server backed by an OpenAPI spec would see the dashboard turn green but still hit upstream auth failures. Add a _request_user_field_headers ContextVar mirroring the existing _request_auth_header/_request_extra_headers pattern, populate it in the local-tool branch of execute_mcp_tool, and merge it after static_headers in _merge_openapi_tool_request_headers so admin-declared per-user values win over operator-configured static headers (matching managed-MCP precedence). BYOK Authorization override still wins over the user-field Authorization to mirror the managed path's separate mcp_auth_header parameter. --- .../mcp_server/openapi_to_mcp_generator.py | 35 +++++-- .../proxy/_experimental/mcp_server/server.py | 23 +++++ .../mcp_server/test_mcp_user_fields.py | 63 ++++++++++++ .../mcp_server/test_openapi_tool_auth.py | 97 +++++++++++++++++++ 4 files changed, 210 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 271517bb1e62..c04d70f4272c 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -62,6 +62,14 @@ def sanitize_openapi_tool_name(raw_name: str) -> str: contextvars.ContextVar("_request_extra_headers", default=None) ) +# Per-request user-field headers resolved from MCPServer.user_fields and the +# calling user's stored values. Set this ContextVar before calling a local +# tool handler so admin-declared per-user values reach the upstream API even +# when the tool dispatch goes through the OpenAPI/local registry path. +_request_user_field_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = ( + contextvars.ContextVar("_request_user_field_headers", default=None) +) + def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: """Ensure path params cannot introduce directory traversal.""" @@ -311,23 +319,26 @@ def _merge_openapi_tool_request_headers( Precedence (highest to lowest): 1. ``_request_auth_header`` — BYOK override of ``Authorization`` - 2. ``static_headers`` — operator-configured headers baked into the + 2. ``_request_user_field_headers`` — admin-declared per-user values + resolved from ``MCPServer.user_fields`` and the calling user's + stored values + 3. ``static_headers`` — operator-configured headers baked into the tool closure at registration time - 3. ``_request_extra_headers`` — per-request headers forwarded from + 4. ``_request_extra_headers`` — per-request headers forwarded from the MCP caller (allowlisted by ``MCPServer.extra_headers``) - This matches the existing MCP invariant in - :func:`litellm.proxy._experimental.mcp_server.utils.merge_mcp_headers` - and the managed MCP path, where ``static_headers`` always wins over - caller-forwarded headers. Keeping the same precedence here prevents an - authenticated caller from overriding an operator-configured value - (e.g. a tenant id or upstream API key) by sending the same header name. + User-field headers win over ``static_headers`` because the admin + explicitly declared a field requiring a per-user value for that header + name. This matches the managed MCP dispatch path + (:meth:`MCPServerManager._call_regular_mcp_tool`), where user-field + headers are merged after ``static_headers``. Header names are compared case-insensitively so different casing cannot bypass the precedence rules. """ request_extra = _request_extra_headers.get() or {} static = static_headers or {} + user_field = _request_user_field_headers.get() or {} static_lower_names = {k.lower() for k in static} effective_headers: Dict[str, str] = { @@ -335,6 +346,14 @@ def _merge_openapi_tool_request_headers( } effective_headers.update(static) + if user_field: + existing_lower_names = {k.lower(): k for k in effective_headers} + for header_name, value in user_field.items(): + collision = existing_lower_names.get(header_name.lower()) + if collision is not None and collision != header_name: + del effective_headers[collision] + effective_headers[header_name] = value + override_auth = _request_auth_header.get() if override_auth: for existing in [k for k in effective_headers if k.lower() == "authorization"]: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 588d02d3b2f2..0e9a51014ce6 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -231,6 +231,7 @@ async def _get_user_field_values_cached( from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, + _request_user_field_headers, ) from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport from litellm.proxy._experimental.mcp_server.tool_registry import ( @@ -2363,13 +2364,35 @@ async def execute_mcp_tool( # noqa: PLR0915 forwarded_headers = {} forwarded_headers[header_name] = value + # User-fields headers: enforce_user_fields above guarantees the + # required values are present; resolve them once and inject so + # OpenAPI/local tools see the same user-provided values that the + # managed MCP dispatch path injects. + user_field_headers: Optional[Dict[str, str]] = None + if server_has_user_fields(mcp_server): + from litellm.proxy._experimental.mcp_server.user_fields import ( + resolve_user_field_headers, + ) + + _, stored_field_values = await _get_user_field_values_cached( + mcp_server, user_api_key_auth + ) + if stored_field_values: + resolved = resolve_user_field_headers( + mcp_server, stored_field_values + ) + if resolved: + user_field_headers = resolved + _auth_token = _request_auth_header.set(auth_header_value) _extra_token = _request_extra_headers.set(forwarded_headers) + _user_field_token = _request_user_field_headers.set(user_field_headers) try: local_content = await _handle_local_mcp_tool(name, arguments) finally: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) + _request_user_field_headers.reset(_user_field_token) response = CallToolResult(content=cast(Any, local_content), isError=False) # Try managed MCP server tool (pass the full prefixed name) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 981b9d7eb63b..9c4cfbcb833d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -493,6 +493,69 @@ async def test_build_stdio_env_returns_user_env_when_no_static_env(): assert merged == {"X": "1"} +def test_openapi_tool_merge_injects_user_field_headers_over_static(): + """User-field headers override static_headers in the OpenAPI/local-tool + dispatch path, matching managed-MCP precedence. + + Regression: stored user-field values were validated at enforcement time + but silently dropped for OpenAPI tools, so upstream calls failed even + after users provided required values. + """ + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _merge_openapi_tool_request_headers, + _request_auth_header, + _request_extra_headers, + _request_user_field_headers, + ) + + extra_token = _request_extra_headers.set({"X-Trace": "abc"}) + user_token = _request_user_field_headers.set( + {"Authorization": "Bearer user-tok", "X-Workspace": "ws1"} + ) + auth_token = _request_auth_header.set(None) + try: + merged = _merge_openapi_tool_request_headers( + {"Authorization": "Bearer static-default", "X-Static": "keep"} + ) + finally: + _request_auth_header.reset(auth_token) + _request_user_field_headers.reset(user_token) + _request_extra_headers.reset(extra_token) + + # User-field Authorization wins over static. + assert merged["Authorization"] == "Bearer user-tok" + # Other user-field headers are added. + assert merged["X-Workspace"] == "ws1" + # Static headers without a user-field override survive. + assert merged["X-Static"] == "keep" + # Forwarded extras still come through when not colliding with static. + assert merged["X-Trace"] == "abc" + + +def test_openapi_tool_merge_byok_auth_overrides_user_field_authorization(): + """BYOK `_request_auth_header` is the final word on Authorization, + even when a user_field also targets Authorization. Matches the managed + path's `mcp_auth_header` parameter precedence.""" + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _merge_openapi_tool_request_headers, + _request_auth_header, + _request_extra_headers, + _request_user_field_headers, + ) + + extra_token = _request_extra_headers.set(None) + user_token = _request_user_field_headers.set({"Authorization": "Bearer user-tok"}) + auth_token = _request_auth_header.set("Bearer byok-tok") + try: + merged = _merge_openapi_tool_request_headers({}) + finally: + _request_auth_header.reset(auth_token) + _request_user_field_headers.reset(user_token) + _request_extra_headers.reset(extra_token) + + assert merged["Authorization"] == "Bearer byok-tok" + + # --------------------------------------------------------------------------- # HTTP endpoints — POST / GET / DELETE /v1/mcp/server/{id}/user-field-values # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 3ad01e9c3ece..71112f75d610 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -5,6 +5,7 @@ """ from datetime import datetime, timezone +from typing import Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -159,6 +160,102 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): handle_local.assert_not_awaited() +@pytest.mark.asyncio +async def test_openapi_local_tool_injects_user_field_headers(): + """OpenAPI/local-tool dispatch must inject the calling user's stored + user-field values via `_request_user_field_headers` so upstream calls + see them. Pre-fix this path validated the values at enforcement time + but never set the ContextVar, so the upstream request went out without + them and the user saw "missing credential" errors after configuring + the dashboard. + """ + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_user_field_headers, + ) + + user = UserAPIKeyAuth( + api_key="sk-user", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + fake_server = MagicMock() + fake_server.name = "openapi-petstore" + fake_server.is_byok = False + fake_server.auth_type = None + fake_server.mcp_info = None + fake_server.server_id = "srv-1" + fake_server.server_name = "openapi-petstore" + fake_server.user_fields = [ + { + "field_key": "API_KEY", + "display_name": "API Key", + "header_name": "X-Api-Key", + "required": True, + } + ] + fake_server.extra_headers = None + fake_server.has_client_credentials = False + + fake_tool = MagicMock() + fake_tool.name = "list_pets" + + captured: Dict[str, Optional[Dict[str, str]]] = {"headers": None} + + async def fake_handle_local(name, arguments): + captured["headers"] = _request_user_field_headers.get() + return [] + + pre_call = AsyncMock(return_value={}) + + # Simulate "user has saved their values" by seeding the in-memory cache + # the enforcement / dispatch path consult. + mcp_module._user_fields_cache.clear() + mcp_module._user_fields_cache[("alice", "srv-1")] = ( + {"API_KEY": "secret-token"}, + 1e18, + ) + + try: + with ( + patch.object( + mcp_module.global_mcp_server_manager, + "_get_mcp_server_from_tool_name", + return_value=fake_server, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "pre_call_tool_check", + new=pre_call, + ), + patch.object( + mcp_module.global_mcp_tool_registry, + "get_tool", + return_value=fake_tool, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + new=fake_handle_local, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ), + ): + await mcp_module.execute_mcp_tool( + name="list_pets", + arguments={}, + allowed_mcp_servers=[fake_server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + finally: + mcp_module._user_fields_cache.clear() + + assert captured["headers"] == {"X-Api-Key": "secret-token"} + + @pytest.mark.asyncio async def test_openapi_local_tool_denied_when_server_not_resolvable(): """If the local-registry tool is found but no MCP server resolves From 19b68721c0a417d4af95f138d841088048372b82 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 08:53:09 +0000 Subject: [PATCH 18/23] fix(mcp): security/correctness fixes for user-fields header injection and BYOK annotation - resolve_user_field_headers: switch from str.format to str.replace so admin-supplied templates cannot use Python format-string attribute or item access (e.g. {value.__class__}) to leak object internals into outbound HTTP headers. Update fallback test to assert literal token preservation, and add a regression test for the attribute-access case. - _annotate_user_credential_flags: skip OAuth2 rows when populating has_user_credential for is_byok=True servers. Otherwise a server reconfigured from OAuth2 to BYOK shows a misleading 'Connected' badge while the actual tool call would fail (get_user_credential filters OAuth2 rows out at execution time). Re-uses the already-decrypted plaintext, so no extra crypto cost. Co-authored-by: Yassin Kortam --- .../_experimental/mcp_server/user_fields.py | 12 ++++++-- .../mcp_management_endpoints.py | 26 +++++++++++++---- .../mcp_server/test_mcp_user_fields.py | 28 ++++++++++++++++--- 3 files changed, 53 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/user_fields.py b/litellm/proxy/_experimental/mcp_server/user_fields.py index af72740fff2b..2b6017523bec 100644 --- a/litellm/proxy/_experimental/mcp_server/user_fields.py +++ b/litellm/proxy/_experimental/mcp_server/user_fields.py @@ -168,10 +168,16 @@ def resolve_user_field_headers( if not value: continue template = entry.get("header_value_template") or "{value}" + # Use plain substitution rather than ``str.format``: the latter + # exposes attribute / item access (``{value.__class__}``, ``{value[0]}``) + # which an admin-supplied template could use — intentionally or by + # accident — to leak Python object internals into outbound HTTP + # headers. ``str.replace`` only matches the literal ``{value}`` token. try: - headers[header_name] = template.format(value=value) - except (KeyError, IndexError, ValueError): - # Malformed template — fall back to raw value rather than crashing. + headers[header_name] = template.replace("{value}", value) + except (TypeError, AttributeError): + # Non-string template (e.g. corrupt JSONB row) — fall back to + # raw value rather than crashing the request. verbose_logger.warning( "MCP user_fields: invalid header_value_template %r for field %r " "on server %s; falling back to raw value.", diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 010816db0424..fc4281d734e2 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -859,6 +859,8 @@ async def _annotate_user_credential_flags( """Populate ``has_user_credential`` and ``missing_user_field_keys`` for the calling user across BYOK / user-fields servers in a single batched query.""" + import json as _json + from litellm.proxy._experimental.mcp_server.db import ( _decode_user_credential, _parse_user_fields_plaintext, @@ -886,12 +888,24 @@ async def _annotate_user_credential_flags( payload = _parse_user_fields_plaintext(decoded) if decoded else None if payload is not None: user_fields_by_server[row.server_id] = payload - else: - # Either a BYOK credential or an undecryptable row (e.g. - # after a salt-key rotation). Either way, a credential row - # exists for this (user, server), so surface it as present - # instead of silently telling the user to reconnect. - byok_set.add(row.server_id) + continue + # Skip OAuth2 rows so a server that's been reconfigured from + # OAuth2 to BYOK doesn't show "Connected" while the actual + # tool call would fail (``get_user_credential`` filters OAuth2 + # rows out at execution time). Re-uses the already-decrypted + # plaintext, so no extra crypto cost. + if decoded is not None: + try: + parsed = _json.loads(decoded) + except (ValueError, TypeError): + parsed = None + if isinstance(parsed, dict) and parsed.get("type") == "oauth2": + continue + # Either a BYOK credential or an undecryptable row (e.g. + # after a salt-key rotation). Either way, a credential row + # exists for this (user, server), so surface it as present + # instead of silently telling the user to reconnect. + byok_set.add(row.server_id) for server in servers: if getattr(server, "is_byok", False): server.has_user_credential = server.server_id in byok_set diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 9c4cfbcb833d..57ac58f649d3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -293,8 +293,11 @@ def test_resolve_user_field_headers_skips_unset_values(): assert headers == {"X-Workspace": "ws1"} -def test_resolve_user_field_headers_falls_back_on_bad_template(): - """A broken template must not crash the request path.""" +def test_resolve_user_field_headers_passes_unknown_placeholders_through(): + """Templates with non-``{value}`` placeholders are preserved literally + instead of triggering format-string evaluation. ``str.replace`` only + matches the literal ``{value}`` token, so unrelated braces flow through + untouched and never crash the request path.""" srv = _gmail_server( user_fields=[ { @@ -306,8 +309,25 @@ def test_resolve_user_field_headers_falls_back_on_bad_template(): ] ) headers = resolve_user_field_headers(srv, {"TOKEN": "raw"}) - # Falls back to the raw value rather than raising. - assert headers == {"Authorization": "raw"} + assert headers == {"Authorization": "Bearer {unknown_placeholder}"} + + +def test_resolve_user_field_headers_does_not_evaluate_attribute_access(): + """A template containing ``{value.__class__}`` must NOT evaluate the + attribute traversal — that would leak Python internals into outbound + HTTP headers. With ``str.replace`` it is preserved as a literal token.""" + srv = _gmail_server( + user_fields=[ + { + "field_key": "TOKEN", + "header_name": "Authorization", + "header_value_template": "Bearer {value.__class__}", + "required": True, + } + ] + ) + headers = resolve_user_field_headers(srv, {"TOKEN": "raw"}) + assert headers == {"Authorization": "Bearer {value.__class__}"} def test_resolve_user_field_headers_skips_entries_missing_header_name(): From bd9a6e4872a907c2369f5c2e74ed9863d220ce63 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 09:12:03 +0000 Subject: [PATCH 19/23] fix(mcp): inject user_fields headers on managed OpenAPI dispatch and clear false UI Ready badge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - mcp_server_manager.py: resolve admin-declared per-user field values and forward them as upstream headers through the managed _call_openapi_tool_handler path via the openapi generator's _request_user_field_headers ContextVar. The local-registry dispatch in server.execute_mcp_tool already injects these; the managed fallback was silently dropping them, so a server with spec_path + user_fields would pass _enforce_user_fields but still fail upstream with no credentials. - mcp_server_columns.tsx: distinguish null/undefined (annotation never ran — e.g. empty user_id, no DB connection) from [] (verified complete) when rendering the user-fields credential cell. Previously a server whose missing_user_field_keys was never populated rendered a green 'Ready' badge while enforcement at tool-call time would still return 401, lying to the user. Now show a neutral 'Not verified' badge when the annotation was skipped. Co-authored-by: Yassin Kortam --- .../mcp_server/mcp_server_manager.py | 44 ++++++++++++++++++- .../mcp_tools/mcp_server_columns.tsx | 31 +++++++++++-- 2 files changed, 70 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a2ab9d9f905a..1cd0a7d4ba60 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2482,6 +2482,7 @@ async def _call_openapi_tool_handler( server: MCPServer, tool_name: str, arguments: Dict[str, Any], + user_field_headers: Optional[Dict[str, str]] = None, ) -> CallToolResult: """ Call an OpenAPI tool handler directly. @@ -2493,12 +2494,20 @@ async def _call_openapi_tool_handler( Args: tool_name: The full tool name (with prefix) to call arguments: Tool arguments to pass to the handler + user_field_headers: Optional admin-declared per-user header values + resolved from ``MCPServer.user_fields`` for the calling user. + Forwarded via the openapi generator's request ContextVar so + the closure-baked handler picks them up — mirrors the + local-registry dispatch path in ``server.execute_mcp_tool``. Returns: CallToolResult with the response from the API """ from mcp.types import TextContent + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_user_field_headers, + ) from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -2514,6 +2523,11 @@ async def _call_openapi_tool_handler( isError=True, ) + _user_field_token = ( + _request_user_field_headers.set(user_field_headers) + if user_field_headers + else None + ) try: # Call the tool handler with the arguments # The handler is an async function that makes the HTTP request @@ -2534,6 +2548,9 @@ async def _call_openapi_tool_handler( content=[TextContent(type="text", text=error_msg)], isError=True, ) + finally: + if _user_field_token is not None: + _request_user_field_headers.reset(_user_field_token) async def pre_call_tool_check( self, @@ -3006,9 +3023,34 @@ async def call_tool( "transport to enable hook header injection.", server_name, ) + # User-fields: resolve the calling user's stored values and forward + # them as upstream headers so admin-declared required fields reach + # the OpenAPI handler. Mirrors the local-registry dispatch path in + # ``server.execute_mcp_tool`` — without this, _enforce_user_fields + # would gate the call on the values being present but the values + # would be silently dropped before dispatch. + user_field_headers: Optional[Dict[str, str]] = None + from litellm.proxy._experimental.mcp_server.user_fields import ( + resolve_user_field_headers, + ) + + stored_user_field_values = await self._resolve_user_field_values( + mcp_server, user_api_key_auth + ) + if stored_user_field_values: + resolved = resolve_user_field_headers( + mcp_server, stored_user_field_values + ) + if resolved: + user_field_headers = resolved tasks.append( asyncio.create_task( - self._call_openapi_tool_handler(mcp_server, name, arguments) + self._call_openapi_tool_handler( + mcp_server, + name, + arguments, + user_field_headers=user_field_headers, + ) ) ) else: diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx index 01ab4375a7e8..a17e72c6afd4 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx @@ -273,10 +273,13 @@ export const mcpServerColumns = ( const hasUserFields = declaredUserFields.length > 0; if (hasUserFields) { // missing_user_field_keys is populated by the proxy for the calling - // user on the list endpoint. Falling back to an empty array assumes - // "nothing missing" so old API responses don't show a false alarm. - const missing = server.missing_user_field_keys ?? []; - if (missing.length > 0) { + // user on the list endpoint. Distinguish `null`/`undefined` + // (annotation never ran — e.g. no user_id, no DB) from `[]` + // (annotation ran and all required fields are satisfied): rendering + // a green "Ready" badge for the unannotated case would lie to the + // user, since the actual tool call would still 401 on missing fields. + const missing = server.missing_user_field_keys; + if (missing && missing.length > 0) { return onUserFieldsConnect ? ( + ) : ( + + + Not verified + + ); + } return (
From a6ae106cf132a5eded06a451c223e52325185a9d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 09:20:57 +0000 Subject: [PATCH 20/23] fix(mcp): address bugbot user_fields + BYOK conflict and header dedup - Reject combining is_byok=True with non-empty user_fields at server create/update time: both share the (user_id, server_id) credential row and the store paths refuse cross-type overwrites, so the combination trapped users in an unresolvable 401 loop. - Make the managed MCP dispatch path apply user-field headers with the same case-insensitive precedence as the OpenAPI/local path, so a static_header named 'Authorization' is no longer left alongside a user-field header named 'authorization' on the outbound request. --- .../mcp_server/mcp_server_manager.py | 11 +++++++- .../mcp_management_endpoints.py | 27 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1cd0a7d4ba60..bac327d5d329 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2827,7 +2827,16 @@ async def _call_regular_mcp_tool( # noqa: PLR0915 if user_field_headers: if extra_headers is None: extra_headers = {} - extra_headers.update(user_field_headers) + # Match the OpenAPI/local path's case-insensitive precedence: a + # user-field header replaces any existing header (e.g. from + # static_headers) whose name matches case-insensitively, instead + # of leaving two differently-cased duplicates on the wire. + existing_lower_names = {k.lower(): k for k in extra_headers} + for header_name, value in user_field_headers.items(): + collision = existing_lower_names.get(header_name.lower()) + if collision is not None and collision != header_name: + del extra_headers[collision] + extra_headers[header_name] = value if hook_extra_headers: if extra_headers is None: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index fc4281d734e2..7e66b558d28c 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -216,6 +216,33 @@ def _validate_mcp_server_name_fields(payload: Any) -> None: def validate_and_normalize_mcp_server_payload(payload: Any) -> None: _base_validate_and_normalize_mcp_server_payload(payload) _validate_mcp_server_name_fields(payload) + _validate_mcp_user_fields_byok_exclusive(payload) + + def _validate_mcp_user_fields_byok_exclusive(payload: Any) -> None: + """Reject server configurations that combine is_byok with user_fields. + + Both credential types share the same (user_id, server_id) row in + LiteLLM_MCPUserCredentials, and the store paths refuse to overwrite + the other type — so a user can save BYOK or user-fields for a given + server, never both. Allowing this combination at admin time would + trap end-users in an unresolvable state: every tool call would 401 + on whichever check the user has not (and cannot) satisfy. + """ + if not getattr(payload, "is_byok", False): + return + user_fields = getattr(payload, "user_fields", None) or [] + if user_fields: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": ( + "MCP servers cannot enable both is_byok and user_fields. " + "Use user_fields if the server needs multiple per-user " + "values; use is_byok only for a single legacy BYOK " + "credential." + ) + }, + ) _VALID_MCP_REQUIRED_FIELDS: frozenset = frozenset(NewMCPServerRequest.model_fields) From 56790e43ebd64a44a7f0beaa98fb80aafd46f4ed Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 09:25:01 +0000 Subject: [PATCH 21/23] test(mcp): cover user_fields header template fallback for codecov Adds a regression test for the AttributeError fallback in resolve_user_field_headers when a corrupt JSONB row carries a non-string header_value_template. Bumps patch coverage past the codecov auto target. --- .../mcp_server/test_mcp_user_fields.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 57ac58f649d3..61da42c0089c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -330,6 +330,25 @@ def test_resolve_user_field_headers_does_not_evaluate_attribute_access(): assert headers == {"Authorization": "Bearer {value.__class__}"} +def test_resolve_user_field_headers_falls_back_on_non_string_template(): + """A corrupt JSONB row may carry a non-string ``header_value_template`` + (e.g. an int from an external write that bypassed the API). The resolver + must catch the resulting ``AttributeError`` from ``int.replace`` and fall + back to the raw user value rather than crashing the request path.""" + srv = _gmail_server( + user_fields=[ + { + "field_key": "TOKEN", + "header_name": "Authorization", + "header_value_template": 42, + "required": True, + } + ] + ) + headers = resolve_user_field_headers(srv, {"TOKEN": "raw-tok"}) + assert headers == {"Authorization": "raw-tok"} + + def test_resolve_user_field_headers_skips_entries_missing_header_name(): """Stdio-only fields (no header_name) must not produce empty-name headers.""" srv = MCPServer( From b7fd2ec55063383871b51970afbcffc8a6eb1332 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 09:41:57 +0000 Subject: [PATCH 22/23] fix(mcp): reject oauth2+user_fields admin config to prevent user deadlock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The exclusivity guard already rejected is_byok+user_fields because both serialize into the single credential_b64 row of LiteLLM_MCPUserCredentials and the write paths refuse to overwrite a foreign payload type. The same mutual exclusion applies to auth_type='oauth2' (interactive) — whichever credential the user saves first permanently locks the other, and the user has no way to escape without admin intervention. Extend the check to also reject auth_type=oauth2 combined with non-empty user_fields, and rename the helper to reflect that it covers more than the BYOK case. oauth2_token_exchange remains allowed because its tokens are cached in-memory (token_exchange.py) and never touch credential_b64. --- .../mcp_management_endpoints.py | 44 +++++++++++++------ .../test_mcp_management_endpoints.py | 40 +++++++++++++++++ 2 files changed, 70 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 7e66b558d28c..dc9ad5ed6821 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -216,22 +216,24 @@ def _validate_mcp_server_name_fields(payload: Any) -> None: def validate_and_normalize_mcp_server_payload(payload: Any) -> None: _base_validate_and_normalize_mcp_server_payload(payload) _validate_mcp_server_name_fields(payload) - _validate_mcp_user_fields_byok_exclusive(payload) - - def _validate_mcp_user_fields_byok_exclusive(payload: Any) -> None: - """Reject server configurations that combine is_byok with user_fields. - - Both credential types share the same (user_id, server_id) row in - LiteLLM_MCPUserCredentials, and the store paths refuse to overwrite - the other type — so a user can save BYOK or user-fields for a given - server, never both. Allowing this combination at admin time would - trap end-users in an unresolvable state: every tool call would 401 - on whichever check the user has not (and cannot) satisfy. + _validate_mcp_user_fields_exclusive(payload) + + def _validate_mcp_user_fields_exclusive(payload: Any) -> None: + """Reject server configurations that combine user_fields with another + credential type that shares the same per-user storage row. + + BYOK strings, OAuth2 access tokens, and user_fields blobs all encode + into the single ``credential_b64`` column of LiteLLM_MCPUserCredentials, + and the write paths refuse to overwrite a different type. Mixing + ``is_byok=True`` or ``auth_type=oauth2`` with non-empty ``user_fields`` + would let whichever credential the user saves first permanently block + the other — every tool call would 401 on the unsatisfied check, and + the user has no path to resolve it without admin intervention. """ - if not getattr(payload, "is_byok", False): - return user_fields = getattr(payload, "user_fields", None) or [] - if user_fields: + if not user_fields: + return + if getattr(payload, "is_byok", False): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail={ @@ -243,6 +245,20 @@ def _validate_mcp_user_fields_byok_exclusive(payload: Any) -> None: ) }, ) + if getattr(payload, "auth_type", None) == MCPAuth.oauth2: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": ( + "MCP servers cannot combine auth_type='oauth2' with " + "user_fields. The interactive OAuth2 token and " + "user_fields share the same per-user credential row, " + "so saving one would block saving the other. Use " + "user_fields for per-user secrets or oauth2 for the " + "primary upstream auth, not both." + ) + }, + ) _VALID_MCP_REQUIRED_FIELDS: frozenset = frozenset(NewMCPServerRequest.model_fields) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 5d66c184495e..5ffa0359846f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2468,6 +2468,46 @@ def test_accepts_valid_names(self): assert payload.alias == "valid_server" + def test_rejects_is_byok_with_user_fields(self): + payload = SimpleNamespace( + server_name="valid_server", + alias=None, + is_byok=True, + user_fields=[{"field_key": "TOKEN", "header_name": "Authorization"}], + ) + + with pytest.raises(HTTPException) as exc_info: + mgmt_endpoints.validate_and_normalize_mcp_server_payload(payload) + + assert exc_info.value.status_code == 400 + assert "is_byok and user_fields" in exc_info.value.detail["error"] + + def test_rejects_oauth2_with_user_fields(self): + payload = SimpleNamespace( + server_name="valid_server", + alias=None, + is_byok=False, + auth_type=MCPAuth.oauth2, + user_fields=[{"field_key": "WORKSPACE", "header_name": "X-Workspace"}], + ) + + with pytest.raises(HTTPException) as exc_info: + mgmt_endpoints.validate_and_normalize_mcp_server_payload(payload) + + assert exc_info.value.status_code == 400 + assert "oauth2" in exc_info.value.detail["error"] + + def test_accepts_oauth2_token_exchange_with_user_fields(self): + payload = SimpleNamespace( + server_name="valid_server", + alias=None, + is_byok=False, + auth_type=MCPAuth.oauth2_token_exchange, + user_fields=[{"field_key": "WORKSPACE", "header_name": "X-Workspace"}], + ) + + mgmt_endpoints.validate_and_normalize_mcp_server_payload(payload) + @pytest.mark.asyncio async def test_health_check_view_all_mode(self): """view_all mode should return health info for all MCP servers.""" From 72184582f24fe9bfeff0f7c6e056a9ac0d4624f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 19 May 2026 10:12:22 +0000 Subject: [PATCH 23/23] fix(mcp): enforce required user_fields in MCPServerManager.call_tool The streamable-HTTP entrypoint (server.execute_mcp_tool) gated tool calls on a server's admin-declared user_fields, but the Responses API path (LiteLLM_Proxy_MCP_Handler) bypassed that gate by calling MCPServerManager.call_tool directly. With required fields unset, the manager just resolved an empty value map and dispatched anyway, so the upstream call went out without the configured headers/env. Move the enforcement into MCPServerManager.call_tool itself so every dispatch path raises the same friendly 401 (user_fields_missing + config_url) when required values are absent. --- .../mcp_server/mcp_server_manager.py | 40 +++++++++++ .../mcp_server/test_mcp_user_fields.py | 68 +++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bac327d5d329..fd535023bd9e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1371,6 +1371,44 @@ async def _resolve_user_field_values( _, values = await _get_user_field_values_cached(mcp_server, user_api_key_auth) return values or {} + async def _enforce_required_user_fields( + self, + mcp_server: MCPServer, + user_api_key_auth: Optional["UserAPIKeyAuth"], + ) -> None: + """Raise the user_fields_missing 401 when required values are absent. + + ``server.execute_mcp_tool`` runs the same gate for the streamable-HTTP + entrypoint, but ``call_tool`` is also called directly by the + Responses API path (``LiteLLM_Proxy_MCP_Handler``). Running the check + here covers every dispatch path through the manager so a caller + cannot skip enforcement by going around ``execute_mcp_tool``. + """ + from litellm.proxy._experimental.mcp_server.user_fields import ( + build_user_fields_missing_error, + compute_missing_user_fields, + server_has_user_fields, + ) + + if not server_has_user_fields(mcp_server): + return + + stored_values = await self._resolve_user_field_values( + mcp_server, user_api_key_auth + ) + missing = compute_missing_user_fields(mcp_server, stored_values or None) + if not missing: + return + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + _resolve_proxy_base_url_env, + ) + + detail = build_user_fields_missing_error( + mcp_server, missing, _resolve_proxy_base_url_env() + ) + raise HTTPException(status_code=401, detail=detail) + async def _create_mcp_client( self, server: MCPServer, @@ -2956,6 +2994,8 @@ async def call_tool( if mcp_server is None: raise ValueError(f"Tool {name} not found") + await self._enforce_required_user_fields(mcp_server, user_api_key_auth) + ######################################################### # Pre MCP Tool Call Hook # Allow validation and modification of tool calls before execution diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py index 61da42c0089c..a9531b9858ae 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_user_fields.py @@ -500,6 +500,74 @@ async def test_enforce_user_fields_no_user_id_raises(): # --------------------------------------------------------------------------- +@pytest.mark.asyncio +async def test_manager_call_tool_enforces_required_user_fields(): + """``MCPServerManager.call_tool`` is invoked directly by the Responses + API path, so it must raise the same friendly 401 as ``execute_mcp_tool`` + when a required user-field has no stored value — otherwise a caller + can bypass enforcement and dispatch with the value silently dropped. + """ + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.proxy._experimental.mcp_server.server import _user_fields_cache + from litellm.proxy._types import UserAPIKeyAuth + + _user_fields_cache.clear() + srv = _gmail_server() + mgr = MCPServerManager() + mgr.registry["s1"] = srv + mgr.tool_name_to_mcp_server_name_mapping["gmail-prod-send"] = srv.name + mgr.tool_name_to_mcp_server_name_mapping["send"] = srv.name + user = UserAPIKeyAuth(api_key="hashed", user_id="responses-user") + _user_fields_cache[("responses-user", "s1")] = (None, 1e18) + + with pytest.raises(HTTPException) as exc_info: + await mgr.call_tool( + server_name="gmail-prod", + name="send", + arguments={}, + user_api_key_auth=user, + ) + assert exc_info.value.status_code == 401 + assert exc_info.value.detail["error"] == "user_fields_missing" + assert "GMAIL_TOKEN" in [ + m["field_key"] for m in exc_info.value.detail["missing_fields"] + ] + + +@pytest.mark.asyncio +async def test_manager_call_tool_enforces_when_no_user_id(): + """Anonymous callers cannot satisfy required user-fields, so the manager + must still surface ``user_fields_missing`` rather than silently dispatch + without the configured headers.""" + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.proxy._types import UserAPIKeyAuth + + srv = _gmail_server() + mgr = MCPServerManager() + mgr.registry["s1"] = srv + mgr.tool_name_to_mcp_server_name_mapping["gmail-prod-send"] = srv.name + mgr.tool_name_to_mcp_server_name_mapping["send"] = srv.name + user = UserAPIKeyAuth(api_key="hashed") # no user_id + + with pytest.raises(HTTPException) as exc_info: + await mgr.call_tool( + server_name="gmail-prod", + name="send", + arguments={}, + user_api_key_auth=user, + ) + assert exc_info.value.status_code == 401 + assert exc_info.value.detail["error"] == "user_fields_missing" + + @pytest.mark.asyncio async def test_build_stdio_env_merges_user_field_env_over_static(): """Stored user_field env vars must take precedence over static server.env."""