diff --git a/docker/Dockerfile.custom_ui b/docker/Dockerfile.custom_ui index f836190a49a..cc44893bf92 100644 --- a/docker/Dockerfile.custom_ui +++ b/docker/Dockerfile.custom_ui @@ -71,8 +71,16 @@ WORKDIR /app RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh +# Run as non-root user +RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser \ + && chown -R appuser:appuser /app +USER appuser + # Expose the necessary port EXPOSE 4000/tcp +HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \ + CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:4000/health')"] + # Override the CMD instruction with your desired command and arguments CMD ["--port", "4000", "--config", "config.yaml", "--detailed_debug"] \ No newline at end of file diff --git a/docker/Dockerfile.health_check b/docker/Dockerfile.health_check index fb9cc201d2f..e968bec340e 100644 --- a/docker/Dockerfile.health_check +++ b/docker/Dockerfile.health_check @@ -13,12 +13,12 @@ RUN pip install --no-cache-dir -r requirements.txt RUN chmod +x /app/health_check_client.py # Run as non-root user -RUN adduser --disabled-password --gecos "" --uid 1001 healthcheck -USER healthcheck +RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser +USER appuser # Health check -HEALTHCHECK --interval=30s --timeout=5s --retries=3 \ - CMD python /app/health_check_client.py --help || exit 1 +HEALTHCHECK --interval=30s --timeout=5s --start-period=5s --retries=3 \ + CMD ["python", "/app/health_check_client.py", "--help"] # Set entrypoint ENTRYPOINT ["python", "/app/health_check_client.py"] diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 88cbcac52cc..c64d475fdaa 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -602,6 +602,8 @@ router_settings: | MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE | Maximum number of entries in MCP OAuth2 token cache. Default is 200 | MCP_OAUTH2_TOKEN_CACHE_MIN_TTL | Minimum TTL in seconds for MCP OAuth2 token cache. Default is 10 | MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from token expiry when computing cache TTL. Default is 60 +| MCP_PER_USER_TOKEN_DEFAULT_TTL | Default TTL in seconds for per-user MCP OAuth tokens stored in Redis. Default is 43200 (12 hours) +| MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from per-user MCP OAuth token expiry when computing Redis TTL. Default is 60 | DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20 | DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10 | DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602 diff --git a/litellm/constants.py b/litellm/constants.py index a7d86ddb16b..337cb1243fb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -135,6 +135,15 @@ MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int( MCP_NPM_CACHE_DIR = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache") MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10")) +# Per-user OAuth token Redis cache (for server-side token storage) +MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX = "mcp:per_user_token" +MCP_PER_USER_TOKEN_DEFAULT_TTL = int( + os.getenv("MCP_PER_USER_TOKEN_DEFAULT_TTL", "43200") # 12 hours +) +MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int( + os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60") +) + # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0")) MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c9abaad908..7a8820a8785 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,5 +1,6 @@ import json import ssl +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse from typing import ( TYPE_CHECKING, Any, @@ -5027,6 +5028,16 @@ class BaseLLMHTTPHandler: litellm_params={}, ) ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + # OpenAI's WebSocket responses endpoint requires ?model= in the URL, + # matching the Realtime API convention (wss://.../v1/realtime?model=...). + # Use urllib.parse so existing query params (e.g. api-version) are preserved. + _parsed = urlparse(ws_url) + _qs = parse_qs(_parsed.query) + if "model" not in _qs: + _qs["model"] = [model] + ws_url = urlunparse( + _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) + ) try: ssl_context = get_shared_realtime_ssl_context() diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index fbef33c32ed..e9bd41bb951 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -21,7 +21,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.utils import PrismaClient +from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPCredentials @@ -576,6 +578,7 @@ async def store_user_oauth_credential( refresh_token: Optional[str] = None, expires_in: Optional[int] = None, scopes: Optional[List[str]] = None, + skip_byok_guard: bool = False, ) -> None: """Persist an OAuth2 access token for a user+server pair. @@ -604,21 +607,26 @@ async def store_user_oauth_credential( # Guard against silently overwriting a BYOK credential with an OAuth token. # BYOK credentials lack a "type" field (or use a non-"oauth2" type). - existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} - ) - if existing is not None: - _byok_error = ValueError( - f"A non-OAuth2 credential already exists for user {user_id} " - f"and server {server_id}. Refusing to overwrite." + # Skip the guard when the caller knows the row is already an OAuth2 credential + # (e.g. during token refresh), saving an extra DB round-trip. + if not skip_byok_guard: + existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - try: - raw = json.loads(base64.urlsafe_b64decode(existing.credential_b64).decode()) - except Exception: - # Credential is not base64+JSON — it's a plain-text BYOK key. - raise _byok_error - if raw.get("type") != "oauth2": - raise _byok_error + if existing is not None: + _byok_error = ValueError( + f"A non-OAuth2 credential already exists for user {user_id} " + f"and server {server_id}. Refusing to overwrite." + ) + try: + raw = json.loads( + base64.urlsafe_b64decode(existing.credential_b64).decode() + ) + except Exception: + # Credential is not base64+JSON — it's a plain-text BYOK key. + raise _byok_error + if raw.get("type") != "oauth2": + raise _byok_error encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() await prisma_client.db.litellm_mcpusercredentials.upsert( @@ -697,6 +705,115 @@ async def list_user_oauth_credentials( return results +async def refresh_user_oauth_token( + prisma_client: PrismaClient, + user_id: str, + server: Any, + cred: Dict[str, Any], +) -> Optional[Dict[str, Any]]: + """Attempt to refresh a per-user OAuth2 token using its stored refresh_token. + + POSTs to ``server.token_url`` with ``grant_type=refresh_token``. + + On success: persists the new credential via ``store_user_oauth_credential`` + and returns the updated payload dict. + On failure (network error, invalid_grant, missing refresh_token, …): logs a + warning and returns ``None`` — the caller is responsible for clearing the + stale credential and triggering re-authentication. + """ + refresh_token: Optional[str] = cred.get("refresh_token") + token_url: Optional[str] = getattr(server, "token_url", None) + server_id: str = getattr(server, "server_id", "") + client_id: Optional[str] = getattr(server, "client_id", None) + client_secret: Optional[str] = getattr(server, "client_secret", None) + + if not refresh_token: + verbose_proxy_logger.debug( + "refresh_user_oauth_token: no refresh_token stored for user=%s server=%s", + user_id, + server_id, + ) + return None + if not token_url: + verbose_proxy_logger.debug( + "refresh_user_oauth_token: server=%s has no token_url configured", + server_id, + ) + return None + + token_data: Dict[str, str] = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + } + if client_id: + token_data["client_id"] = client_id + if client_secret: + token_data["client_secret"] = client_secret + + try: + async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.Oauth2Check + ) + response = await async_client.post( + token_url, + headers={"Accept": "application/json"}, + data=token_data, + ) + response.raise_for_status() + body: Dict[str, Any] = response.json() + except Exception as exc: + verbose_proxy_logger.warning( + "refresh_user_oauth_token: refresh request failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + return None + + access_token: Optional[str] = body.get("access_token") + if not access_token: + verbose_proxy_logger.warning( + "refresh_user_oauth_token: token response missing access_token for " + "user=%s server=%s", + user_id, + server_id, + ) + return None + + expires_in: Optional[int] = None + raw_expires = body.get("expires_in") + try: + expires_in = int(raw_expires) if raw_expires is not None else None + except (TypeError, ValueError): + pass + + # Rotate refresh token when the provider returns a new one + new_refresh_token: Optional[str] = body.get("refresh_token") or refresh_token + + raw_scope = body.get("scope") + scopes: Optional[List[str]] = ( + raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + ) or cred.get("scopes") + + await store_user_oauth_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=server_id, + access_token=access_token, + refresh_token=new_refresh_token, + expires_in=expires_in, + scopes=scopes, + skip_byok_guard=True, # Row is already OAuth2; skip the extra find_unique check + ) + + verbose_proxy_logger.info( + "refresh_user_oauth_token: refreshed token for user=%s server=%s", + user_id, + server_id, + ) + return await get_user_oauth_credential(prisma_client, user_id, server_id) + + async def approve_mcp_server( prisma_client: PrismaClient, server_id: str, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 07309eb57f2..d0d61986322 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,10 +1,11 @@ import json -from typing import Optional +from typing import Any, Dict, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -147,6 +148,160 @@ def _resolve_oauth2_server_for_root_endpoints( return None +def _validate_token_response( + token_response: Dict[str, Any], + validation_rules: Dict[str, Any], + server_id: str, +) -> None: + """Raise HTTPException 403 if any validation rule doesn't match the token response. + + Supports dot-notation for nested fields (e.g. ``"team.enterprise_id"`` checks + ``token_response["team"]["enterprise_id"]``). Top-level keys are tried first, + then dot-split traversal. All comparisons are string-coerced so that numeric + values in the response (e.g. ``"org_id": 12345``) match string rules + (``"org_id": "12345"``). + """ + for key, expected in validation_rules.items(): + actual: Any = token_response.get(key) + # Try dot-notation traversal when top-level lookup returns None + if actual is None and "." in key: + obj: Any = token_response + for part in key.split("."): + if isinstance(obj, dict): + obj = obj.get(part) + else: + obj = None + break + actual = obj + # Treat absent fields as a distinct failure from a mismatched value + if actual is None: + raise HTTPException( + status_code=403, + detail={ + "error": "token_validation_failed", + "server_id": server_id, + "field": key, + "message": ( + f"OAuth token rejected: required field '{key}' is absent" + ), + }, + ) + if str(actual) != str(expected): + raise HTTPException( + status_code=403, + detail={ + "error": "token_validation_failed", + "server_id": server_id, + "field": key, + "message": ( + f"OAuth token rejected: '{key}' = '{actual}', " + f"expected '{expected}'" + ), + }, + ) + + +async def _extract_user_id_from_request(request: Request) -> Optional[str]: + """Best-effort extraction of LiteLLM user_id from the request's Authorization header. + + Called at the OAuth token endpoint so that per-user tokens can be stored + server-side. Uses a read-only cache lookup to avoid re-running the full + auth pipeline (which has side effects such as rate-limit increments and + spend logging). Returns ``None`` if no cached credential is found. + """ + auth_header = request.headers.get("Authorization") or request.headers.get( + "authorization" + ) + if not auth_header: + return None + lower = auth_header.lower() + if not lower.startswith("bearer "): + return None + token = auth_header[7:].strip() + try: + from litellm.proxy._types import hash_token # noqa: PLC0415 + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + cached = await user_api_key_cache.async_get_cache(hash_token(token)) + return getattr(cached, "user_id", None) + except Exception: + return None + + +async def _store_per_user_token_server_side( + server: MCPServer, + user_id: str, + token_response: Dict[str, Any], +) -> None: + """Persist the OAuth token server-side and warm the Redis cache. + + Called from the token endpoint after a successful code exchange or refresh. + Errors are logged but NOT re-raised — the token is always returned to the + client even when server-side storage fails. + """ + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415 + _compute_per_user_token_ttl, + mcp_per_user_token_cache, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + access_token: Optional[str] = token_response.get("access_token") + if not access_token: + return + + raw_expires = token_response.get("expires_in") + try: + expires_in: Optional[int] = int(raw_expires) if raw_expires is not None else None + except (TypeError, ValueError): + expires_in = None + + refresh_token: Optional[str] = token_response.get("refresh_token") or None + raw_scope = token_response.get("scope") + scopes: Optional[list] = ( + raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + ) + + try: + prisma_client = get_prisma_client_or_throw( + "Database not connected. Cannot store per-user OAuth token." + ) + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + store_user_oauth_credential, + ) + + await store_user_oauth_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=server.server_id, + access_token=access_token, + refresh_token=refresh_token, + expires_in=expires_in, + scopes=scopes, + ) + verbose_logger.info( + "_store_per_user_token_server_side: stored token for user=%s server=%s", + user_id, + server.server_id, + ) + except Exception as exc: + verbose_logger.warning( + "_store_per_user_token_server_side: DB storage failed for user=%s server=%s: %s", + user_id, + server.server_id, + exc, + ) + return # Don't warm Redis if DB write failed + + # Warm the Redis cache so the first subsequent MCP call is a cache hit + ttl = _compute_per_user_token_ttl(server, expires_in) + await mcp_per_user_token_cache.set( + user_id=user_id, + server_id=server.server_id, + access_token=access_token, + ttl=ttl, + ) + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -266,6 +421,44 @@ async def exchange_token_with_server( token_response = response.json() access_token = token_response["access_token"] + # Validate token response against server-configured rules before any storage. + # This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc. + if mcp_server.token_validation and isinstance(mcp_server.token_validation, dict): + _validate_token_response( + token_response=token_response, + validation_rules=mcp_server.token_validation, + server_id=mcp_server.server_id, + ) + + # Store server-side when the server is configured for per-user OAuth and + # the calling client has provided a valid LiteLLM identity. + # Errors are non-fatal: the token is still returned to the client. + if mcp_server.needs_user_oauth_token: + user_id = await _extract_user_id_from_request(request) + if user_id: + try: + await _store_per_user_token_server_side( + server=mcp_server, + user_id=user_id, + token_response=token_response, + ) + except Exception as exc: + verbose_logger.warning( + "exchange_token_with_server: server-side storage failed " + "for user=%s server=%s: %s", + user_id, + mcp_server.server_id, + exc, + ) + else: + verbose_logger.debug( + "exchange_token_with_server: no LiteLLM user_id found in request; " + "per-user token for server=%s will not be stored server-side. " + "The client should call POST /mcp/server/{id}/oauth-user-credential " + "to store it manually.", + mcp_server.server_id, + ) + result = { "access_token": access_token, "token_type": token_response.get("token_type", "Bearer"), diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 402e12d9356..8d3831e75fb 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2455,6 +2455,37 @@ class MCPServerManager: ) tasks.append(during_hook_task) + # For per-user OAuth servers: if the client didn't supply a token in + # oauth2_headers, look up the stored token from Redis / DB. This is the + # call_tool equivalent of _get_user_oauth_extra_headers_from_db used in + # list_tools. + if ( + mcp_server.needs_user_oauth_token + and not oauth2_headers + and user_api_key_auth is not None + ): + user_id = getattr(user_api_key_auth, "user_id", None) + if user_id: + try: + from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + _get_user_oauth_extra_headers_from_db, + ) + + stored_headers = await _get_user_oauth_extra_headers_from_db( + server=mcp_server, + user_api_key_auth=user_api_key_auth, + ) + if stored_headers: + oauth2_headers = stored_headers + except Exception as _lookup_exc: + verbose_logger.debug( + "call_tool: per-user token lookup failed for " + "user=%s server=%s: %s", + user_id, + mcp_server.server_id, + _lookup_exc, + ) + # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: verbose_logger.debug( diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 84a2e94467b..476e215666e 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -17,8 +17,15 @@ from litellm.constants import ( MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_PER_USER_TOKEN_DEFAULT_TTL, + MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -152,6 +159,107 @@ class MCPOAuth2TokenCache(InMemoryCache): mcp_oauth2_token_cache = MCPOAuth2TokenCache() +def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int: + """Compute Redis TTL for a per-user token. + + Uses server.token_storage_ttl_seconds when configured; otherwise derives + TTL from expires_in minus the expiry buffer; falls back to the default TTL. + """ + if server.token_storage_ttl_seconds is not None: + return max(server.token_storage_ttl_seconds, 1) + if expires_in is not None: + return max( + expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, + 1, + ) + return MCP_PER_USER_TOKEN_DEFAULT_TTL + + +class MCPPerUserTokenCache: + """Redis-backed cache for per-user OAuth2 access tokens. + + Uses LiteLLM's existing ``user_api_key_cache`` (DualCache with optional + Redis backend). Tokens are NaCl-encrypted with ``encrypt_value_helper`` + before storage so they are safe at rest in Redis. + + Redis key format: ``mcp:per_user_token:{user_id}:{server_id}`` + Redis value: ``encrypt_value_helper(access_token)`` — URL-safe base64 + """ + + def _cache_key(self, user_id: str, server_id: str) -> str: + return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + + async def get(self, user_id: str, server_id: str) -> Optional[str]: + """Return the plaintext access_token, or None on miss/error.""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + encrypted = await user_api_key_cache.async_get_cache(key) + if encrypted is None: + return None + plaintext = decrypt_value_helper( + encrypted, + key="mcp_per_user_token", + exception_type="debug", + ) + return plaintext or None + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.get failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + return None + + async def set( + self, + user_id: str, + server_id: str, + access_token: str, + ttl: int, + ) -> None: + """Store NaCl-encrypted access_token in Redis with the given TTL.""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + encrypted = encrypt_value_helper(access_token) + await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl) + verbose_logger.debug( + "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds", + user_id, + server_id, + ttl, + ) + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.set failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + + async def delete(self, user_id: str, server_id: str) -> None: + """Invalidate the cached token (removes from both in-memory and Redis layers).""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + await user_api_key_cache.async_delete_cache(key) + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.delete failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + + +mcp_per_user_token_cache = MCPPerUserTokenCache() + + async def resolve_mcp_auth( server: "MCPServer", mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7fc28b68e9c..99578d006e1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -896,11 +896,17 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, ) -> Optional[Dict[str, str]]: - """Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict. + """Look up stored OAuth2 token for (user, server) and return as extra_headers dict. + + Lookup order: + 1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied + 2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query + 3. Auto-refresh when the stored token is expired and a refresh_token exists Args: prefetched_creds: Optional dict keyed by server_id with credential payloads. - When provided, avoids a per-server DB round-trip. + When provided, the Redis and individual DB lookups are + skipped in favour of the pre-fetched batch result. """ if server.auth_type != MCPAuth.oauth2: return None @@ -914,8 +920,27 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 get_user_oauth_credential, is_oauth_credential_expired, + refresh_user_oauth_token, + ) + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415 + _compute_per_user_token_ttl, + mcp_per_user_token_cache, ) + # ── Fast path: Redis cache ──────────────────────────────────────── + # Only used when prefetched_creds is not supplied (individual lookup). + if prefetched_creds is None: + cached_token = await mcp_per_user_token_cache.get(user_id, server_id) + if cached_token is not None: + verbose_logger.debug( + "_get_user_oauth_extra_headers_from_db: Redis hit for " + "user=%s server=%s", + user_id, + server_id, + ) + return {"Authorization": f"Bearer {cached_token}"} + + # ── Slow path: DB lookup ────────────────────────────────────────── if prefetched_creds is not None: cred = prefetched_creds.get(server_id) else: @@ -929,18 +954,83 @@ if MCP_AVAILABLE: cred = await get_user_oauth_credential( prisma_client, user_id, server_id ) - if cred and cred.get("access_token"): - if is_oauth_credential_expired(cred): - verbose_logger.debug( - f"_get_user_oauth_extra_headers_from_db: token expired for " - f"user={user_id} server={server_id}" - ) + + if not cred or not cred.get("access_token"): + return None + + if is_oauth_credential_expired(cred): + verbose_logger.debug( + "_get_user_oauth_extra_headers_from_db: token expired for " + "user=%s server=%s — attempting refresh", + user_id, + server_id, + ) + # Attempt token refresh; requires a DB client (not available from prefetch) + if cred.get("refresh_token"): + try: + from litellm.proxy.utils import ( # noqa: PLC0415 + get_prisma_client_or_throw, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Cannot refresh OAuth token." + ) + cred = await refresh_user_oauth_token( + prisma_client=prisma_client, + user_id=user_id, + server=server, + cred=cred, + ) + except Exception as refresh_exc: + verbose_logger.warning( + "_get_user_oauth_extra_headers_from_db: refresh failed " + "for user=%s server=%s: %s", + user_id, + server_id, + refresh_exc, + ) + cred = None + + if not cred or not cred.get("access_token"): + # Clear stale Redis/cache entry so we don't serve it again. + # Do this for both the individual and prefetch paths so the + # next request doesn't get a stale cache hit. + await mcp_per_user_token_cache.delete(user_id, server_id) return None - return {"Authorization": f"Bearer {cred['access_token']}"} + + access_token: str = cred["access_token"] + + # Warm (or re-warm) the Redis cache from the DB result. + # Always write regardless of whether expires_at is present — tokens + # without an expiry are still valid and should be cached using the + # server/default TTL so subsequent requests are fast. + if prefetched_creds is None: + raw_expires = None + expires_at = cred.get("expires_at") + if expires_at: + from datetime import datetime, timezone # noqa: PLC0415 + + try: + exp_dt = datetime.fromisoformat(expires_at) + if exp_dt.tzinfo is None: + exp_dt = exp_dt.replace(tzinfo=timezone.utc) + remaining = int( + (exp_dt - datetime.now(timezone.utc)).total_seconds() + ) + raw_expires = max(remaining, 0) if remaining > 0 else None + except (ValueError, TypeError): + pass + ttl = _compute_per_user_token_ttl(server, raw_expires) + await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl) + + return {"Authorization": f"Bearer {access_token}"} except Exception as e: verbose_logger.warning( - f"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " - f"user={user_id} server={server_id}: {e}" + "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " + "user=%s server=%s: %s", + user_id, + server_id, + e, ) return None @@ -2504,6 +2594,14 @@ if MCP_AVAILABLE: server_name, client_ip=_client_ip ) if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: + # For servers that store per-user tokens server-side, skip the + # pre-emptive 401 — the call_tool / list_tools dispatch will look + # up the stored token from Redis / DB and only fail at the MCP + # protocol level if none is found, giving the client a proper + # tool-execution error rather than an HTTP 401. + if server.needs_user_oauth_token: + continue + request = StarletteRequest(scope) base_url = get_request_base_url(request) diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index db7657a0174..a7d0968c0ef 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -71,6 +71,15 @@ class MCPServer(BaseModel): # OAuth2 flow type. Defaults to None (interactive / authorization_code). # Set to "client_credentials" to enable M2M token fetching. oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None + # Per-user OAuth server-side storage config. + # token_validation: key-value pairs that must match fields in the OAuth token + # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). + # Tokens that fail validation are rejected before storage. + token_validation: Optional[Dict[str, Any]] = None + # Optional TTL override (seconds) for the Redis per-user token cache. + # Defaults to the token's expires_in minus the expiry buffer, or + # MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent. + token_storage_ttl_seconds: Optional[int] = None model_config = ConfigDict(arbitrary_types_allowed=True) @property diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index a1684e23769..001b9464006 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -178,7 +178,6 @@ async def test_get_response(): async def test_aavertex_ai_anthropic_async(): # load_vertex_ai_credentials() try: - model = "claude-3-5-sonnet@20240620" vertex_ai_project = "pathrise-convert-1606954137718" @@ -351,7 +350,6 @@ def test_avertex_ai_stream(): @pytest.mark.flaky(retries=3, delay=1) @pytest.mark.asyncio async def test_async_vertexai_response_basic(): - load_vertex_ai_credentials() try: user_message = "Hello, how are you?" @@ -1382,7 +1380,6 @@ async def test_gemini_pro_json_schema_args_sent_httpx( ] ) elif resp is not None: - assert resp.model == model.split("/")[1] @@ -2291,6 +2288,8 @@ def test_prompt_factory_nested(): async def test_completion_fine_tuned_model(): load_vertex_ai_credentials() mock_response = AsyncMock() + mock_response.headers = {} + mock_response.status_code = 200 def return_val(): return { @@ -2326,7 +2325,6 @@ async def test_completion_fine_tuned_model(): } mock_response.json = return_val - mock_response.status_code = 200 expected_payload = { "contents": [ diff --git a/tests/mcp_tests/test_per_user_oauth_cache.py b/tests/mcp_tests/test_per_user_oauth_cache.py new file mode 100644 index 00000000000..36c26a5a505 --- /dev/null +++ b/tests/mcp_tests/test_per_user_oauth_cache.py @@ -0,0 +1,527 @@ +""" +Unit tests for per-user MCP OAuth token storage: +- MCPPerUserTokenCache (NaCl-encrypted Redis cache) +- _validate_token_response (token validation rules) +- _compute_per_user_token_ttl (TTL computation) +- refresh_user_oauth_token (token refresh flow) +""" + +import sys +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# Stub out modules that aren't available in the unit-test environment +# so we can import the targets without a full proxy stack. +for _mod in ("orjson",): + if _mod not in sys.modules: + sys.modules[_mod] = MagicMock() + +from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: E402 + MCPPerUserTokenCache, + _compute_per_user_token_ttl, + mcp_per_user_token_cache, +) +from litellm.types.mcp import MCPAuth, MCPTransport # noqa: E402 +from litellm.types.mcp_server.mcp_server_manager import MCPServer # noqa: E402 + + +def _import_validate(): + """Lazy import to avoid pulling orjson at collection time.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _validate_token_response, + ) + + return _validate_token_response + + +# ── Fixtures ───────────────────────────────────────────────────────────────── + + +def _make_server(**kwargs) -> MCPServer: + defaults: Dict[str, Any] = { + "server_id": "slack-test", + "name": "Slack", + "server_name": "slack", + "url": "https://slack-mcp.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "client_id": "SLACK_CLIENT_ID", + "client_secret": "SLACK_CLIENT_SECRET", + "token_url": "https://slack.com/api/oauth.v2.access", + "authorization_url": "https://slack.com/oauth/v2/authorize", + } + defaults.update(kwargs) + return MCPServer(**defaults) + + +# ── _validate_token_response ────────────────────────────────────────────────── + + +class TestValidateTokenResponse: + def test_passes_when_all_rules_match(self): + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "enterprise_id": "E04XXXXXX", + "team": {"id": "T123", "name": "Acme"}, + } + # Should not raise + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + + def test_raises_on_mismatch(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "xoxb-123", "enterprise_id": "E99999999"} + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + detail = exc_info.value.detail + assert detail["error"] == "token_validation_failed" + assert detail["field"] == "enterprise_id" + + def test_raises_when_field_absent(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "xoxb-123"} + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + # Absent field should produce a distinct "absent" message, not str(None) + assert "absent" in exc_info.value.detail["message"] + + def test_absent_field_does_not_match_string_none(self): + """str(None)='None' must NOT match the string rule value 'None'.""" + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "tok"} # enterprise_id absent + # Even if admin writes validation_rules={"enterprise_id": "None"}, absent + # field should raise, not pass. + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "None"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + assert "absent" in exc_info.value.detail["message"] + + def test_dot_notation_nested_field(self): + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "team": {"enterprise_id": "E04XXXXXX"}, + } + # Should not raise — dot-notation traverses nested dict + _validate_token_response( + token_response=token_response, + validation_rules={"team.enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + + def test_dot_notation_mismatch(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "team": {"enterprise_id": "WRONG"}, + } + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"team.enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["field"] == "team.enterprise_id" + + def test_numeric_value_string_coercion(self): + """Numeric values in token response should match string rules.""" + _validate_token_response = _import_validate() + token_response = {"access_token": "tok", "org_id": 12345} + # Should not raise — str(12345) == "12345" + _validate_token_response( + token_response=token_response, + validation_rules={"org_id": "12345"}, + server_id="test", + ) + + def test_multiple_rules_all_must_match(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = { + "access_token": "tok", + "enterprise_id": "E04XXXXXX", + "cloud_id": "WRONG_CLOUD", + } + with pytest.raises(HTTPException): + _validate_token_response( + token_response=token_response, + validation_rules={ + "enterprise_id": "E04XXXXXX", + "cloud_id": "abc-123", + }, + server_id="atlassian", + ) + + +# ── _compute_per_user_token_ttl ────────────────────────────────────────────── + + +class TestComputePerUserTokenTtl: + def test_uses_server_override_when_set(self): + server = _make_server(token_storage_ttl_seconds=7200) + assert _compute_per_user_token_ttl(server, expires_in=99999) == 7200 + + def test_uses_expires_in_minus_buffer(self): + server = _make_server() + # Default buffer is 60s + ttl = _compute_per_user_token_ttl(server, expires_in=3600) + assert ttl == 3600 - 60 + + def test_minimum_ttl_is_1(self): + server = _make_server() + # expires_in smaller than buffer → clamp to 1 + ttl = _compute_per_user_token_ttl(server, expires_in=30) + assert ttl == 1 + + def test_default_ttl_when_expires_in_none(self): + from litellm.constants import MCP_PER_USER_TOKEN_DEFAULT_TTL + + server = _make_server() + ttl = _compute_per_user_token_ttl(server, expires_in=None) + assert ttl == MCP_PER_USER_TOKEN_DEFAULT_TTL + + +# ── MCPPerUserTokenCache ────────────────────────────────────────────────────── + + +class TestMCPPerUserTokenCache: + """Tests for Redis-backed per-user token cache. + + Patches ``user_api_key_cache`` to avoid needing a real Redis instance. + Patches ``encrypt_value_helper`` / ``decrypt_value_helper`` to verify + encryption is applied before Redis writes and decryption after reads. + """ + + @pytest.fixture + def cache(self): + return MCPPerUserTokenCache() + + @pytest.fixture + def mock_dual_cache(self): + dc = MagicMock() + dc.async_get_cache = AsyncMock(return_value=None) + dc.async_set_cache = AsyncMock() + return dc + + @pytest.mark.asyncio + async def test_get_returns_none_on_miss(self, cache, mock_dual_cache): + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper" + ) as mock_decrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = None + result = await cache.get("alice", "slack-test") + assert result is None + mock_decrypt.assert_not_called() + + @pytest.mark.asyncio + async def test_get_decrypts_cached_value(self, cache, mock_dual_cache): + fake_encrypted = "encrypted_blob_abc123" + fake_plaintext = "xoxb-slack-token" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper", + return_value=fake_plaintext, + ) as mock_decrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = fake_encrypted + result = await cache.get("alice", "slack-test") + + assert result == fake_plaintext + mock_decrypt.assert_called_once_with( + fake_encrypted, + key="mcp_per_user_token", + exception_type="debug", + ) + + @pytest.mark.asyncio + async def test_set_encrypts_before_storing(self, cache, mock_dual_cache): + fake_encrypted = "encrypted_blob_xyz" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value=fake_encrypted, + ) as mock_encrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.set("alice", "slack-test", "xoxb-token", ttl=3540) + + mock_encrypt.assert_called_once_with("xoxb-token") + mock_dual_cache.async_set_cache.assert_called_once() + call_kwargs = mock_dual_cache.async_set_cache.call_args + assert call_kwargs[0][1] == fake_encrypted # encrypted value stored + assert call_kwargs[1]["ttl"] == 3540 + + @pytest.mark.asyncio + async def test_set_uses_correct_cache_key(self, cache, mock_dual_cache): + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value="enc", + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.set("bob", "github-server", "ghp_token", ttl=3600) + + key_used = mock_dual_cache.async_set_cache.call_args[0][0] + assert key_used == "mcp:per_user_token:bob:github-server" + + @pytest.mark.asyncio + async def test_delete_calls_async_delete_cache(self, cache, mock_dual_cache): + mock_dual_cache.async_delete_cache = AsyncMock() + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.delete("alice", "slack-test") + + mock_dual_cache.async_delete_cache.assert_called_once_with( + "mcp:per_user_token:alice:slack-test" + ) + mock_dual_cache.async_set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_get_returns_none_on_decrypt_failure(self, cache, mock_dual_cache): + """Cache misses and decrypt errors should both return None without raising.""" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper", + return_value=None, # decrypt returns None on failure + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = "bad_encrypted_data" + result = await cache.get("alice", "slack-test") + + assert result is None + + @pytest.mark.asyncio + async def test_set_is_noop_on_cache_error(self, cache, mock_dual_cache): + """Errors in the cache layer must not propagate to the caller.""" + mock_dual_cache.async_set_cache.side_effect = RuntimeError("Redis down") + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value="enc", + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + # Should not raise + await cache.set("alice", "slack-test", "token", ttl=3600) + + +# ── refresh_user_oauth_token ────────────────────────────────────────────────── + + +class TestRefreshUserOauthToken: + """Tests for the DB-level token refresh helper.""" + + @pytest.fixture + def server(self): + return _make_server() + + @pytest.fixture + def cred(self): + return { + "type": "oauth2", + "access_token": "OLD_TOKEN", + "refresh_token": "REFRESH_TOKEN_123", + "expires_at": ( + datetime.now(timezone.utc) - timedelta(hours=1) + ).isoformat(), + } + + @pytest.mark.asyncio + async def test_returns_none_when_no_refresh_token(self, server): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + cred = {"type": "oauth2", "access_token": "OLD"} # no refresh_token + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_returns_none_when_no_token_url(self, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + server = _make_server(token_url=None) + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_returns_none_on_http_error(self, server, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + mock_client = AsyncMock() + mock_client.post.side_effect = Exception("Connection refused") + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ): + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_stores_and_returns_new_credential(self, server, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + new_token_response = MagicMock() + new_token_response.json.return_value = { + "access_token": "NEW_TOKEN", + "expires_in": 3600, + "refresh_token": "NEW_REFRESH", + "scope": "channels:read chat:write", + } + new_token_response.raise_for_status = MagicMock() + + mock_client = AsyncMock() + mock_client.post.return_value = new_token_response + + stored_cred = { + "type": "oauth2", + "access_token": "NEW_TOKEN", + "refresh_token": "NEW_REFRESH", + } + mock_prisma = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ), patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new_callable=AsyncMock, + ) as mock_store, patch( + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new_callable=AsyncMock, + return_value=stored_cred, + ): + result = await refresh_user_oauth_token( + prisma_client=mock_prisma, + user_id="alice", + server=server, + cred=cred, + ) + + assert result == stored_cred + mock_store.assert_called_once() + call_kwargs = mock_store.call_args[1] + assert call_kwargs["access_token"] == "NEW_TOKEN" + assert call_kwargs["refresh_token"] == "NEW_REFRESH" + assert call_kwargs["expires_in"] == 3600 + assert call_kwargs["scopes"] == ["channels:read", "chat:write"] + # Refresh path must skip the BYOK guard (row is already OAuth2) + assert call_kwargs.get("skip_byok_guard") is True + + @pytest.mark.asyncio + async def test_falls_back_to_old_refresh_token_when_not_rotated( + self, server, cred + ): + """When provider doesn't return a new refresh_token, keep the old one.""" + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + new_token_response = MagicMock() + new_token_response.json.return_value = { + "access_token": "NEW_TOKEN", + "expires_in": 3600, + # No refresh_token in response + } + new_token_response.raise_for_status = MagicMock() + + mock_client = AsyncMock() + mock_client.post.return_value = new_token_response + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ), patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new_callable=AsyncMock, + ) as mock_store, patch( + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new_callable=AsyncMock, + return_value={"type": "oauth2", "access_token": "NEW_TOKEN"}, + ): + await refresh_user_oauth_token( + prisma_client=AsyncMock(), + user_id="alice", + server=server, + cred=cred, + ) + + call_kwargs = mock_store.call_args[1] + # Old refresh_token preserved when provider doesn't rotate + assert call_kwargs["refresh_token"] == "REFRESH_TOKEN_123" + + +# ── MCPServer new fields ────────────────────────────────────────────────────── + + +class TestMCPServerNewFields: + def test_token_validation_default_none(self): + server = _make_server() + assert server.token_validation is None + + def test_token_validation_set(self): + server = _make_server(token_validation={"enterprise_id": "E04XXXXXX"}) + assert server.token_validation == {"enterprise_id": "E04XXXXXX"} + + def test_token_storage_ttl_default_none(self): + server = _make_server() + assert server.token_storage_ttl_seconds is None + + def test_token_storage_ttl_set(self): + server = _make_server(token_storage_ttl_seconds=7200) + assert server.token_storage_ttl_seconds == 7200 + + def test_needs_user_oauth_token_true_for_oauth2_without_m2m(self): + server = _make_server(auth_type=MCPAuth.oauth2) + assert server.needs_user_oauth_token is True + + def test_needs_user_oauth_token_false_for_m2m(self): + server = _make_server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + ) + assert server.needs_user_oauth_token is False diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 0d83b9f88de..efec841cc4d 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -971,3 +971,106 @@ class TestWebSocketChunkTypes: ) assert len(messages) == 1 assert messages[0]["content"][0]["text"] == "Part 1Part 2" + + +class TestNativeWebSocketUrlConstruction: + """Test that native WebSocket URLs include the model query parameter. + + These tests mock websockets.connect so they exercise the actual URL-building + code inside BaseLLMHTTPHandler.async_responses_websocket rather than + reimplementing the logic themselves. + """ + + @pytest.mark.asyncio + async def test_openai_ws_url_includes_model(self): + """Handler must pass ?model= in the URL to the backend WebSocket.""" + from unittest.mock import AsyncMock, MagicMock, patch + + captured_urls = [] + + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) + + async def __aenter__(self): + raise Exception("stop") + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = "https://api.openai.com/v1/responses" + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o-mini", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o-mini"], f"Expected model in URL, got: {captured_urls[0]}" + + @pytest.mark.asyncio + async def test_ws_url_preserves_existing_params_and_adds_model(self): + """When api_base already has query params, model is added alongside them.""" + from unittest.mock import AsyncMock, MagicMock, patch + + captured_urls = [] + + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) + + async def __aenter__(self): + raise Exception("stop") + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = ( + "https://custom.example.com/v1/responses?api-version=2024-05-01" + ) + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o"], f"model missing from URL: {captured_urls[0]}" + assert qs.get("api-version") == ["2024-05-01"], f"existing param lost: {captured_urls[0]}" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx index e1b3b358300..3c5101fc2dc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx @@ -131,7 +131,7 @@ describe("ModelsAndEndpointsView", () => { , ); expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument(); - }, 15000); + }); it("should show Missing provider banner by default", async () => { localStorageMock.clear(); @@ -149,7 +149,7 @@ describe("ModelsAndEndpointsView", () => { , ); expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument(); - }, 15000); + }); it("should hide Missing provider banner when dismiss button is clicked and persist to localStorage", async () => { localStorageMock.clear(); @@ -180,7 +180,7 @@ describe("ModelsAndEndpointsView", () => { // LocalStorage should be updated expect(localStorageMock.getItem("hideMissingProviderBanner")).toBe("true"); - }, 15000); + }); it("should show compact Request Provider button when banner is dismissed", async () => { // Set localStorage to hide banner @@ -209,7 +209,7 @@ describe("ModelsAndEndpointsView", () => { const requestProviderLinks = document.querySelectorAll('a[href="https://models.litellm.ai/?request=true"]'); // There should be a compact button when banner is hidden expect(requestProviderLinks.length).toBeGreaterThan(0); - }, 15000); + }); it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => { mockHealthCheckComponent.mockClear(); diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 202820a11a2..7ad3e32ef5c 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -46,10 +46,17 @@ function LoginPageContent() { // Cross-origin SSO: worker redirected back with a single-use code. // Exchange it for the JWT via the worker's /v3/login/exchange endpoint. const params = new URLSearchParams(window.location.search); - const ssoCode = params.get("code"); + const rawSsoCode = params.get("code"); + // Validate the SSO code is a plausible OAuth authorization code (alphanumeric + // plus common URL-safe chars) so that arbitrary user input cannot trigger the + // exchange endpoint. + const ssoCode = + rawSsoCode && /^[a-zA-Z0-9._~+/=-]+$/.test(rawSsoCode) ? rawSsoCode : null; if (ssoCode) { - // codeql[js/user-controlled-bypass] - const workerUrl = localStorage.getItem("litellm_worker_url"); + const rawWorkerUrl = localStorage.getItem("litellm_worker_url"); + // Validate the stored worker URL: only allow http(s) URLs. + const workerUrl = + rawWorkerUrl && /^https?:\/\/.+/.test(rawWorkerUrl) ? rawWorkerUrl : null; exchangeLoginCode(ssoCode, workerUrl).then(() => { params.delete("code"); const cleanSearch = params.toString(); diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx index 0c4cad8cb0b..0539d6d8f19 100644 --- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx +++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx @@ -2,6 +2,7 @@ import { Suspense, useEffect, useMemo } from "react"; import { useSearchParams } from "next/navigation"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; // Written to sessionStorage so both the admin hook (useMcpOAuthFlow) and the // user hook (useUserMcpOAuthFlow) can pick up the result. Each hook reads @@ -52,14 +53,24 @@ const McpOAuthCallbackContent = () => { // Write to both namespace keys (admin and user) so whichever hook is // active can consume the result. sessionStorage only — no localStorage. const serialized = JSON.stringify(payload); - window.sessionStorage.setItem(ADMIN_RESULT_KEY, serialized); - window.sessionStorage.setItem(USER_RESULT_KEY, serialized); + setSecureItem(ADMIN_RESULT_KEY, serialized); + setSecureItem(USER_RESULT_KEY, serialized); } catch (err) { // Silently ignore storage errors } - const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY); - const destination = returnUrl || resolveDefaultRedirect(); + const returnUrl = getSecureItem(RETURN_URL_STORAGE_KEY); + let destination = resolveDefaultRedirect(); + if (returnUrl) { + try { + const parsed = new URL(returnUrl, window.location.origin); + if (parsed.origin === window.location.origin) { + destination = parsed.href; + } + } catch { + // invalid URL — fall through to default + } + } window.location.replace(destination); }, [payload]); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 44df1b5bd41..d3fab5cf5bb 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -277,13 +277,18 @@ function CreateKeyPageContent() { // Check for a stored return URL const returnUrl = consumeReturnUrl(); if (returnUrl && isValidReturnUrl(returnUrl)) { + // Inline origin check: only redirect to same-origin URLs to prevent open redirect. + const safeUrl = new URL(returnUrl, window.location.origin); + if (safeUrl.origin !== window.location.origin) { + return; + } const currentUrl = window.location.href; const normalizedReturnUrl = normalizeUrlForCompare(returnUrl); const normalizedCurrentUrl = normalizeUrlForCompare(currentUrl); // Only redirect if the return URL is different from the current URL // This prevents infinite redirect loops if (normalizedReturnUrl !== normalizedCurrentUrl) { - window.location.replace(returnUrl); + window.location.replace(safeUrl.href); } } }, [authLoading, token]); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 9a4659da9d3..03d982ca7c2 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -51,7 +51,7 @@ function renderWithProviders(ui: React.ReactElement) { return render({ui}); } -describe("CreateUserButton", { timeout: 20000 }, () => { +describe("CreateUserButton", () => { beforeEach(() => { vi.clearAllMocks(); mockGetProxyUISettings.mockResolvedValue({ @@ -62,288 +62,296 @@ describe("CreateUserButton", { timeout: 20000 }, () => { }); }); - it("should render the create user form when embedded", () => { - renderWithProviders( - , - ); - expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument(); - }); + describe("rendering and visibility", () => { + it("should render the create user form when embedded", () => { + renderWithProviders( + , + ); + expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument(); + }); - it("should render the invite user button when not embedded", async () => { - renderWithProviders(); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + it("should render the invite user button when not embedded", async () => { + renderWithProviders(); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + }); + + it("should open the invite modal when invite user button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + expect(dialog).toBeInTheDocument(); + expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument(); + }); + + it("should display email invitations info message in embedded mode", () => { + renderWithProviders(); + expect(screen.getByText("Email invitations")).toBeInTheDocument(); + }); + + it("should display user role options when possibleUIRoles is provided", async () => { + const possibleUIRoles = { + proxy_admin: { ui_label: "Admin", description: "Full access" }, + proxy_user: { ui_label: "User", description: "Limited access" }, + }; + renderWithProviders( + , + ); + await userEvent.click(screen.getByRole("combobox", { name: /user role/i })); + expect(screen.getByText("Admin")).toBeInTheDocument(); + expect(screen.getByText("User")).toBeInTheDocument(); + }); + + it("should close modal when cancel is clicked in standalone mode", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument(); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.click(within(dialog).getByRole("button", { name: /close/i })); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); }); }); - it("should open the invite modal when invite user button is clicked", async () => { - const user = userEvent.setup(); - renderWithProviders(); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + describe("embedded mode submission", () => { + it("should call userCreateCall when form is submitted in embedded mode", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-1", + user_id: "new-user-123", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "test@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ + user_email: "test@example.com", + user_role: "proxy_user", + })); + }); }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - expect(dialog).toBeInTheDocument(); - expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument(); - }); - it("should display email invitations info message in embedded mode", () => { - renderWithProviders(); - expect(screen.getByText("Email invitations")).toBeInTheDocument(); - }); + it("should call onUserCreated callback when user is created in embedded mode", async () => { + const user = userEvent.setup(); + const onUserCreated = vi.fn(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } }); - it("should display user role options when possibleUIRoles is provided", async () => { - const possibleUIRoles = { - proxy_admin: { ui_label: "Admin", description: "Full access" }, - proxy_user: { ui_label: "User", description: "Limited access" }, - }; - renderWithProviders( - , - ); - await userEvent.click(screen.getByRole("combobox", { name: /user role/i })); - expect(screen.getByText("Admin")).toBeInTheDocument(); - expect(screen.getByText("User")).toBeInTheDocument(); - }); + renderWithProviders( + , + ); - it("should call userCreateCall when form is submitted in embedded mode", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-1", - user_id: "new-user-123", - has_user_setup_sso: false, - } as any); + await user.type(screen.getByLabelText(/user email/i), "embedded@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); - renderWithProviders( - , - ); + await waitFor(() => { + expect(onUserCreated).toHaveBeenCalledWith("new-user-456"); + }); + }); - await user.type(screen.getByLabelText(/user email/i), "test@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); + it("should show error notification when user creation fails", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } }); - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ - user_email: "test@example.com", - user_role: "proxy_user", - })); + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists"); + }); + }); + + it("should show info notification when making API call", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-3", + user_id: "new-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "info@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call"); + }); }); }); - it("should call onUserCreated callback when user is created in embedded mode", async () => { - const user = userEvent.setup(); - const onUserCreated = vi.fn(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } }); + describe("standalone mode submission", () => { + it("should show success notification when user is created successfully in standalone mode", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-2", + user_id: "new-user-789", + has_user_setup_sso: false, + } as any); - renderWithProviders( - , - ); + renderWithProviders( + , + ); - await user.type(screen.getByLabelText(/user email/i), "embedded@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - await waitFor(() => { - expect(onUserCreated).toHaveBeenCalledWith("new-user-456"); + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + }); + }); + + it("should show onboarding modal when user is created and SSO is disabled", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-sso", + user_id: "sso-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user"); + }); + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + }); }); }); - it("should show success notification when user is created successfully in standalone mode", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-2", - user_id: "new-user-789", - has_user_setup_sso: false, - } as any); + describe("organizations", () => { + it("should send organizations list in POST body when organizations are selected", async () => { + const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); + vi.mocked(useOrganizations).mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "My Org" }], + isLoading: false, + } as any); - renderWithProviders( - , - ); + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-org", + user_id: "org-user", + has_user_setup_sso: false, + } as any); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + + // Select org from the dropdown + const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i }); + await user.click(orgSelect); + await user.click(screen.getByText("My Org (org-1)")); + + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ + organizations: ["org-1"], + })); + }); }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + it("should not call organizationMemberAddCall after user creation", async () => { + const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); + vi.mocked(useOrganizations).mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "My Org" }], + isLoading: false, + } as any); - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-nma", + user_id: "no-member-add-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalled(); + }); + expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled(); }); }); - - it("should show error notification when user creation fails", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } }); - - renderWithProviders( - , - ); - - await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); - - await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists"); - }); - }); - - it("should show info notification when making API call", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-3", - user_id: "new-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await user.type(screen.getByLabelText(/user email/i), "info@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); - - await waitFor(() => { - expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call"); - }); - }); - - it("should close modal when cancel is clicked in standalone mode", async () => { - const user = userEvent.setup(); - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument(); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.click(within(dialog).getByRole("button", { name: /close/i })); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - }); - - it("should show onboarding modal when user is created and SSO is disabled", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-sso", - user_id: "sso-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user"); - }); - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); - }); - }); - - it("should send organizations list in POST body when organizations are selected", async () => { - const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); - vi.mocked(useOrganizations).mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "My Org" }], - isLoading: false, - } as any); - - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-org", - user_id: "org-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - - // Select org from the dropdown - const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i }); - await user.click(orgSelect); - await user.click(screen.getByText("My Org (org-1)")); - - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ - organizations: ["org-1"], - })); - }); - }); - - it("should not call organizationMemberAddCall after user creation", async () => { - const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); - vi.mocked(useOrganizations).mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "My Org" }], - isLoading: false, - } as any); - - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-nma", - user_id: "no-member-add-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalled(); - }); - expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled(); - }); }); diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index 651d1495c61..4b89820bad6 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -843,7 +843,7 @@ describe("OldTeams - access_group_ids in team create", () => { }), ); }); - }, { timeout: 30000 }); + }); }); describe("OldTeams - models dropdown options", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 197bcd6569f..0a8d2d9124e 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -175,7 +175,7 @@ describe("Add Model Tab", () => { ); expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument(); - }, 10000); // This test is flaky, adding a timeout until we find a better solution + }); it("should display both Add Model and Add Auto Router tabs", async () => { const props = createTestProps(); @@ -269,7 +269,7 @@ describe("Add Model Tab", () => { }, { timeout: 10000 }, ); - }, 15000); // 15 second timeout to allow waitFor to complete + }); it("should show team selection when team-only switch is enabled", async () => { const props = createTestProps(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx new file mode 100644 index 00000000000..888f3066252 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx @@ -0,0 +1,208 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen, waitFor, act, fireEvent } from "@testing-library/react"; +import { Form } from "antd"; +import OAuthFormFields from "./OAuthFormFields"; + +// ── helpers ────────────────────────────────────────────────────────────────── + +/** Minimal Ant Form wrapper so Form.Item registers correctly. */ +const WithForm: React.FC<{ children: React.ReactNode; onFinish?: (values: any) => void }> = ({ + children, + onFinish, +}) => { + const [form] = Form.useForm(); + return ( +
+ {children} + +
+ ); +}; + +// ── tests ───────────────────────────────────────────────────────────────────── + +describe("OAuthFormFields", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + // ── visibility by flow type ───────────────────────────────────────────────── + + describe("interactive mode (isM2M=false)", () => { + it("renders Token Validation Rules field", () => { + render( + + + , + ); + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + it("renders Token Storage TTL field", () => { + render( + + + , + ); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + + it("renders standard interactive fields alongside the new fields", () => { + render( + + + , + ); + expect(screen.getByText("Authorization URL (optional)")).toBeInTheDocument(); + expect(screen.getByText("Registration URL (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + }); + + describe("M2M mode (isM2M=true)", () => { + it("does NOT render Token Validation Rules field", () => { + render( + + + , + ); + expect(screen.queryByText("Token Validation Rules (optional)")).not.toBeInTheDocument(); + }); + + it("does NOT render Token Storage TTL field", () => { + render( + + + , + ); + expect(screen.queryByText("Token Storage TTL (seconds, optional)")).not.toBeInTheDocument(); + }); + + it("still renders M2M-specific fields", () => { + render( + + + , + ); + expect(screen.getByText("Client ID")).toBeInTheDocument(); + expect(screen.getByText("Token URL")).toBeInTheDocument(); + }); + }); + + // ── token_validation_json inline JSON validator ────────────────────────────── + + describe("token_validation_json validation", () => { + it("accepts empty value without error", async () => { + const onFinish = vi.fn(); + render( + + + , + ); + + // Leave the textarea empty and submit + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + + it("accepts a valid JSON object without error", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: '{"organization": "my-org"}' } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + + it("shows 'Must be valid JSON' error for malformed JSON", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "not-valid-json{" } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + }); + + it("shows error for a plain string value (not a JSON object)", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + // A bare string is valid JSON but we still want to accept it; only truly + // unparseable text should fail. Bare "hello" is actually invalid JSON + // (no quotes), so it should fail. + fireEvent.change(textarea, { target: { value: "hello" } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + }); + + it("whitespace-only value is treated as empty and passes validation", async () => { + const onFinish = vi.fn(); + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: " " } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx index 85487a8a479..4a808ca489d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Form, Select, Tooltip } from "antd"; +import { Form, Input, InputNumber, Select, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { OAUTH_FLOW } from "./types"; @@ -151,6 +151,50 @@ const OAuthFormFields: React.FC = ({ > + + } + name="token_validation_json" + rules={[ + { + validator: (_: any, value: string) => { + if (!value || value.trim() === "") return Promise.resolve(); + try { + JSON.parse(value); + return Promise.resolve(); + } catch { + return Promise.reject(new Error("Must be valid JSON")); + } + }, + }, + ]} + > + + + + } + name="token_storage_ttl_seconds" + > + + {oauthFlow && (

diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index d49c49446bb..b4251267137 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -150,151 +150,139 @@ describe("CreateMCPServer", () => { }); }); - it( - "should not require auth value when creating a server with API Key auth type", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should not require auth value when creating a server with API Key auth type", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - // Fill in server name (use id to avoid duplicate placeholder) - const nameInput = getServerNameInput(); - await user.type(nameInput, "Test_Server"); + // Fill in server name (use id to avoid duplicate placeholder) + const nameInput = getServerNameInput(); + await user.type(nameInput, "Test_Server"); - // Fill in URL - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + // Fill in URL + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - // Select API Key auth type - await selectAntOption("Authentication", "API Key"); + // Select API Key auth type + await selectAntOption("Authentication", "API Key"); - await waitFor(() => { - expect(screen.getByText("Authentication Value")).toBeInTheDocument(); - }); + await waitFor(() => { + expect(screen.getByText("Authentication Value")).toBeInTheDocument(); + }); - // Leave auth value empty and submit - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "Test_Server", - alias: "Test_Server", - url: "https://example.com/mcp", - transport: "http", - auth_type: "api_key", - created_at: "2024-01-01T00:00:00Z", - created_by: "user-1", - updated_at: "2024-01-01T00:00:00Z", - updated_by: "user-1", - }); + // Leave auth value empty and submit + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "Test_Server", + alias: "Test_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "api_key", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - // The form should submit without validation error on auth_value - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); - }, - ); + // The form should submit without validation error on auth_value + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + }); - it( - "should not require auth value when creating a server with Bearer Token auth type", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should not require auth value when creating a server with Bearer Token auth type", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - const nameInput = getServerNameInput(); - await user.type(nameInput, "Test_Server"); + const nameInput = getServerNameInput(); + await user.type(nameInput, "Test_Server"); - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - await selectAntOption("Authentication", "Bearer Token"); + await selectAntOption("Authentication", "Bearer Token"); - await waitFor(() => { - expect(screen.getByText("Authentication Value")).toBeInTheDocument(); - }); + await waitFor(() => { + expect(screen.getByText("Authentication Value")).toBeInTheDocument(); + }); - // Leave auth value empty and submit - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "Test_Server", - alias: "Test_Server", - url: "https://example.com/mcp", - transport: "http", - auth_type: "bearer_token", - created_at: "2024-01-01T00:00:00Z", - created_by: "user-1", - updated_at: "2024-01-01T00:00:00Z", - updated_by: "user-1", - }); + // Leave auth value empty and submit + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "Test_Server", + alias: "Test_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "bearer_token", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); - }, - ); + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + }); - it( - "should successfully create a server when auth value is provided", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should successfully create a server when auth value is provided", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - const nameInput = getServerNameInput(); - await user.type(nameInput, "My_Server"); + const nameInput = getServerNameInput(); + await user.type(nameInput, "My_Server"); - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - await selectAntOption("Authentication", "API Key"); + await selectAntOption("Authentication", "API Key"); - await waitFor(() => { - expect(screen.getByText("Authentication Value")).toBeInTheDocument(); - }); + await waitFor(() => { + expect(screen.getByText("Authentication Value")).toBeInTheDocument(); + }); - // Fill in auth value - const authInput = screen.getByPlaceholderText("Enter token or secret"); - await user.type(authInput, "my-secret-key"); + // Fill in auth value + const authInput = screen.getByPlaceholderText("Enter token or secret"); + await user.type(authInput, "my-secret-key"); - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "My_Server", - alias: "My_Server", - url: "https://example.com/mcp", - transport: "http", - auth_type: "api_key", - created_at: "2024-01-01T00:00:00Z", - created_by: "user-1", - updated_at: "2024-01-01T00:00:00Z", - updated_by: "user-1", - }); + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "My_Server", + alias: "My_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "api_key", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); - const [token, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; - expect(token).toBe("test-token"); - expect(payload.credentials).toEqual({ auth_value: "my-secret-key" }); - }, - ); + const [token, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(token).toBe("test-token"); + expect(payload.credentials).toEqual({ auth_value: "my-secret-key" }); + }); it("should not show auth value field when None auth type is selected", async () => { await selectHttpTransport(); @@ -307,50 +295,187 @@ describe("CreateMCPServer", () => { }); }); - it( - "should successfully create a server with no auth", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should successfully create a server with no auth", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - const nameInput = getServerNameInput(); - await user.type(nameInput, "No_Auth_Server"); + const nameInput = getServerNameInput(); + await user.type(nameInput, "No_Auth_Server"); - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - await selectAntOption("Authentication", "None"); + await selectAntOption("Authentication", "None"); - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "No_Auth_Server", - alias: "No_Auth_Server", - url: "https://example.com/mcp", - transport: "http", - auth_type: "none", - created_at: "2024-01-01T00:00:00Z", - created_by: "user-1", - updated_at: "2024-01-01T00:00:00Z", - updated_by: "user-1", - }); + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "No_Auth_Server", + alias: "No_Auth_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "none", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); - const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; - expect(payload.auth_type).toBe("none"); - // No credentials should be sent for "none" auth - expect(payload.credentials).toBeUndefined(); - }, - ); + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.auth_type).toBe("none"); + // No credentials should be sent for "none" auth + expect(payload.credentials).toBeUndefined(); + }); + }); + + describe("when OAuth interactive auth is selected", () => { + /** Select HTTP transport + OAuth auth, then wait for the OAuth form to appear. */ + async function setupOAuthInteractive() { + render(); + await selectAntOption("Transport Type", "Streamable HTTP"); + + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + + await selectAntOption("Authentication", "OAuth"); + + // Wait for OAuthFormFields to render (OAuth Flow Type selector is the sentinel) + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + // OAuthFormFields defaults to INTERACTIVE, so the new fields should appear + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + } + + it("shows Token Validation Rules and Token Storage TTL fields", async () => { + await setupOAuthInteractive(); + // Asserted in setupOAuthInteractive + }); + + it("includes token_validation in payload when token_validation_json is filled with valid JSON", async () => { + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-oauth", + server_name: "OAuth_Server", + alias: "OAuth_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "oauth2", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); + + await setupOAuthInteractive(); + + // Fill required form fields + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "OAuth_Server" } }); + }); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } }); + }); + + // Fill in the token_validation_json textarea + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: '{"organization": "my-org", "team.id": "42"}' } }); + }); + + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.token_validation).toEqual({ organization: "my-org", "team.id": "42" }); + }); + + it("omits token_validation from payload when token_validation_json is empty", async () => { + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-oauth", + server_name: "OAuth_Server", + alias: "OAuth_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "oauth2", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); + + await setupOAuthInteractive(); + + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "OAuth_Server" } }); + }); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } }); + }); + + // Leave token_validation_json empty + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.token_validation).toBeUndefined(); + }); + + it("does not submit and shows validation error for invalid JSON in token_validation_json", async () => { + await setupOAuthInteractive(); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "not-valid-json{" } }); + }); + + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "OAuth_Server" } }); + }); + + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Either the inline form validation message or the notification fires — + // both indicate the submit was blocked. + await waitFor(() => { + const inlineError = screen.queryByText("Must be valid JSON"); + const notCalled = !vi.mocked(networking.createMCPServer).mock.calls.length; + expect(inlineError !== null || notCalled).toBe(true); + }); + }); }); describe("when modal is cancelled", () => { 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 4c824fcee0b..17bcd59c43e 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 @@ -17,6 +17,7 @@ import { validateMCPServerUrl, validateMCPServerName } from "./utils"; import NotificationsManager from "../molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { useTestMCPConnection } from "@/hooks/useTestMCPConnection"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; const asset_logos_folder = "../ui/assets/logos/"; export const mcpLogoImg = `${asset_logos_folder}mcp_logo.png`; @@ -94,8 +95,7 @@ const CreateMCPServer: React.FC = ({ } try { const values = form.getFieldsValue(true); - // codeql[js/clear-text-storage-of-sensitive-data] - window.sessionStorage.setItem( + setSecureItem( CREATE_OAUTH_UI_STATE_KEY, JSON.stringify({ modalVisible: isModalVisible, @@ -178,7 +178,7 @@ const CreateMCPServer: React.FC = ({ if (typeof window === "undefined") { return; } - const storedState = window.sessionStorage.getItem(CREATE_OAUTH_UI_STATE_KEY); + const storedState = getSecureItem(CREATE_OAUTH_UI_STATE_KEY); if (!storedState) { return; } @@ -284,6 +284,7 @@ const CreateMCPServer: React.FC = ({ credentials: credentialValues, allow_all_keys: allowAllKeysRaw, available_on_public_internet: availableOnPublicInternetRaw, + token_validation_json: rawTokenValidationJson, ...restValues } = values; @@ -356,6 +357,18 @@ const CreateMCPServer: React.FC = ({ restValues.transport = "http"; } + // Parse token_validation JSON if provided + let tokenValidation: Record | null = null; + if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") { + try { + tokenValidation = JSON.parse(rawTokenValidationJson); + } catch { + NotificationsManager.fromBackend("Invalid JSON in Token Validation Rules"); + setIsLoading(false); + return; + } + } + // Prepare the payload with cost configuration and allowed tools const payload: Record = { ...restValues, @@ -376,6 +389,7 @@ const CreateMCPServer: React.FC = ({ allow_all_keys: Boolean(allowAllKeysRaw), available_on_public_internet: Boolean(availableOnPublicInternetRaw), static_headers: staticHeaders, + ...(tokenValidation !== null && { token_validation: tokenValidation }), }; payload.static_headers = staticHeaders; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index e33e2fff491..aba2a3d9222 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -3,6 +3,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { render, screen, waitFor, fireEvent, act } from "@testing-library/react"; import MCPServerEdit from "./mcp_server_edit"; import * as networking from "../networking"; +import NotificationsManager from "../molecules/notifications_manager"; vi.mock("../networking", () => ({ updateMCPServer: vi.fn(), @@ -37,6 +38,29 @@ vi.mock("./mcp_tool_configuration", () => ({ default: () =>

, })); +// ── fixtures ────────────────────────────────────────────────────────────────── + +const interactiveOAuthServer = { + server_id: "oauth_server_1", + server_name: "OAuthServer", + alias: "oauth_server", // underscores: hyphens fail validateMCPServerName + description: "Interactive OAuth MCP server", + transport: "http", + url: "https://example.com/mcp", + auth_type: "oauth2", + // No token_url → edit form defaults to INTERACTIVE flow + token_url: null, + authorization_url: null, + registration_url: null, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + mcp_access_groups: [], +}; + +// ── test suites ─────────────────────────────────────────────────────────────── + describe("MCPServerEdit (stdio)", () => { beforeEach(() => { vi.clearAllMocks(); @@ -152,3 +176,228 @@ describe("MCPServerEdit (stdio)", () => { expect(payload.env).toEqual({ CIRCLECI_TOKEN: "new-token", CIRCLECI_BASE_URL: "https://circleci.com" }); }); }); + +describe("MCPServerEdit (interactive OAuth)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("renders Token Validation Rules and Token Storage TTL fields for interactive OAuth server", async () => { + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + }); + + // Note: The M2M flow hiding logic is tested via OAuthFormFields.test.tsx (isM2M prop directly), + // since Form.useWatch doesn't synchronously reflect initialValues in jsdom. + + it("pre-populates token_validation_json from existing server token_validation", async () => { + const tokenValidation = { organization: "my-org", "team.id": "123" }; + + render( + , + ); + + await waitFor(() => { + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + expect(textarea).not.toBeNull(); + const parsed = JSON.parse(textarea.value); + expect(parsed).toEqual(tokenValidation); + }); + }); + + it("includes token_validation in update payload when token_validation_json is filled", async () => { + const onSuccess = vi.fn(); + vi.mocked(networking.updateMCPServer).mockResolvedValue({ + ...interactiveOAuthServer, + token_validation: { organization: "my-org" }, + }); + + render( + , + ); + + // Wait for the form to mount and the token_validation_json field to appear + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: '{"organization": "my-org"}' } }); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + expect(payload.token_validation).toEqual({ organization: "my-org" }); + }); + + it("does not include token_validation in payload when field is empty and server had none", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue(interactiveOAuthServer); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + // Leave token_validation_json empty + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + expect(payload.token_validation).toBeUndefined(); + }); + + it("sends token_validation: null to clear an existing value when textarea is cleared", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue({ + ...interactiveOAuthServer, + token_validation: null, + }); + + render( + , + ); + + await waitFor(() => { + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + expect(textarea?.value).toContain("old-org"); + }); + + // Clear the textarea + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "" } }); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + // null signals the backend to clear the existing validation rules + expect(payload.token_validation).toBeNull(); + }); + + it("shows inline validation error and does not submit on invalid JSON in token_validation_json", async () => { + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "{ bad json" } }); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + // The Form.Item inline validator intercepts invalid JSON before handleSave runs, + // so the inline error message appears and updateMCPServer is never called. + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + }); + + it("includes token_storage_ttl_seconds in payload when set", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue({ + ...interactiveOAuthServer, + token_storage_ttl_seconds: 7200, + }); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + expect(payload.token_storage_ttl_seconds).toBe(7200); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index e81d2f3960e..574e7871759 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Form, Select, Button as AntdButton, Tooltip, Input } from "antd"; +import { Form, Select, Button as AntdButton, Tooltip, Input, InputNumber } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; import { AUTH_TYPE, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types"; @@ -12,6 +12,7 @@ import MCPLogoSelector from "./MCPLogoSelector"; import { validateMCPServerUrl, validateMCPServerName } from "./utils"; import NotificationsManager from "../molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; interface MCPServerEditProps { mcpServer: MCPServer; @@ -73,8 +74,7 @@ const MCPServerEdit: React.FC = ({ } try { const values = form.getFieldsValue(true); - // codeql[js/clear-text-storage-of-sensitive-data] - window.sessionStorage.setItem( + setSecureItem( EDIT_OAUTH_UI_STATE_KEY, JSON.stringify({ serverId: mcpServer.server_id, @@ -190,6 +190,9 @@ const MCPServerEdit: React.FC = ({ transport: effectiveTransport, static_headers: initialStaticHeaders, oauth_flow_type: mcpServer.token_url ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE, + token_validation_json: mcpServer.token_validation + ? JSON.stringify(mcpServer.token_validation, null, 2) + : undefined, }), [mcpServer, effectiveTransport, initialStaticHeaders, initialEnvJson], ); @@ -214,7 +217,7 @@ const MCPServerEdit: React.FC = ({ if (typeof window === "undefined") { return; } - const storedState = window.sessionStorage.getItem(EDIT_OAUTH_UI_STATE_KEY); + const storedState = getSecureItem(EDIT_OAUTH_UI_STATE_KEY); if (!storedState) { return; } @@ -400,6 +403,7 @@ const MCPServerEdit: React.FC = ({ args: rawArgs, allow_all_keys: allowAllKeysRaw, available_on_public_internet: availableOnPublicInternetRaw, + token_validation_json: rawTokenValidationJson, ...restValues } = values; @@ -522,6 +526,17 @@ const MCPServerEdit: React.FC = ({ restValues.transport = "http"; } + // Parse token_validation JSON if provided + let tokenValidation: Record | null = null; + if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") { + try { + tokenValidation = JSON.parse(rawTokenValidationJson); + } catch { + NotificationsManager.fromBackend("Invalid JSON in Token Validation Rules"); + return; + } + } + // Prepare the payload with cost configuration and permission fields const mcpInfoServerName = restValues.server_name || @@ -556,6 +571,10 @@ const MCPServerEdit: React.FC = ({ static_headers: staticHeaders, allow_all_keys: Boolean(allowAllKeysRaw ?? mcpServer.allow_all_keys), available_on_public_internet: Boolean(availableOnPublicInternetRaw ?? mcpServer.available_on_public_internet), + // Include token_validation when it is set (non-null) or when clearing an existing value + ...(tokenValidation !== null || mcpServer.token_validation + ? { token_validation: tokenValidation } + : {}), }; const includeCredentials = restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); @@ -863,6 +882,58 @@ const MCPServerEdit: React.FC = ({ className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500" /> + {!isM2MFlow && ( + <> + + Token Validation Rules (optional) + + + + + } + name="token_validation_json" + rules={[ + { + validator: (_: any, value: string) => { + if (!value || value.trim() === "") return Promise.resolve(); + try { + JSON.parse(value); + return Promise.resolve(); + } catch { + return Promise.reject(new Error("Must be valid JSON")); + } + }, + }, + ]} + > + + + + Token Storage TTL (seconds, optional) + + + + + } + name="token_storage_ttl_seconds" + > + + + + )}

Use OAuth to fetch a fresh access token and temporarily save it in the session as the authentication value.