scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge branch 'main' into fix/redundant-decrption

Browse source
This commit is contained in:
yangdx 2026-04-10 14:08:45 +08:00
commit 3cf35a664c
35 changed files with 3112 additions and 1116 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
docker/Dockerfile.custom_ui
View file

@ -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"]

8
docker/Dockerfile.health_check
View file

@ -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"]

2
docs/my-website/docs/proxy/config_settings.md
View file

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

9
litellm/constants.py
View file

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

11
litellm/llms/custom_httpx/llm_http_handler.py
View file

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

145
litellm/proxy/_experimental/mcp_server/db.py
View file

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

195
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

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

31
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

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

108
litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py
View file

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

120
litellm/proxy/_experimental/mcp_server/server.py
View file

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

9
litellm/types/mcp_server/mcp_server_manager.py
View file

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

6
tests/local_testing/test_amazing_vertex_completion.py
View file

@ -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": [

527
tests/mcp_tests/test_per_user_oauth_cache.py Normal file
View file

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

103
tests/test_litellm/responses/test_responses_websocket_all_providers.py
View file

@ -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]}"

8
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx
View file

@ -131,7 +131,7 @@ describe("ModelsAndEndpointsView", () => {
</QueryClientProvider>,
);
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", () => {
</QueryClientProvider>,
);
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();

13
ui/litellm-dashboard/src/app/login/LoginPage.tsx
View file

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

19
ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx
View file

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

7
ui/litellm-dashboard/src/app/page.tsx
View file

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

524
ui/litellm-dashboard/src/components/CreateUserButton.test.tsx
View file

@ -51,7 +51,7 @@ function renderWithProviders(ui: React.ReactElement) {
return render(<QueryClientProvider client={qc}>{ui}</QueryClientProvider>);
}
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(
<CreateUserButton {...defaultProps} isEmbedded />,
);
expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument();
});
describe("rendering and visibility", () => {
it("should render the create user form when embedded", () => {
renderWithProviders(
<CreateUserButton {...defaultProps} isEmbedded />,
);
expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument();
});
it("should render the invite user button when not embedded", async () => {
renderWithProviders(<CreateUserButton {...defaultProps} />);
await waitFor(() => {
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
it("should render the invite user button when not embedded", async () => {
renderWithProviders(<CreateUserButton {...defaultProps} />);
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(<CreateUserButton {...defaultProps} />);
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(<CreateUserButton {...defaultProps} isEmbedded />);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={possibleUIRoles} isEmbedded />,
);
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(<CreateUserButton {...defaultProps} />);
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(<CreateUserButton {...defaultProps} />);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
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(<CreateUserButton {...defaultProps} isEmbedded />);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={possibleUIRoles} isEmbedded />,
);
await userEvent.click(screen.getByRole("combobox", { name: /user role/i }));
expect(screen.getByText("Admin")).toBeInTheDocument();
expect(screen.getByText("User")).toBeInTheDocument();
});
renderWithProviders(
<CreateUserButton {...defaultProps} onUserCreated={onUserCreated} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
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(
<CreateUserButton {...defaultProps} onUserCreated={onUserCreated} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
renderWithProviders(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
);
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(<CreateUserButton {...defaultProps} />);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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(
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
);
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();
});
});

2
ui/litellm-dashboard/src/components/OldTeams.test.tsx
View file

@ -843,7 +843,7 @@ describe("OldTeams - access_group_ids in team create", () => {
}),
);
});
}, { timeout: 30000 });
});
});
describe("OldTeams - models dropdown options", () => {

4
ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx
View file

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

208
ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx Normal file
View file

@ -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 (
<Form form={form} onFinish={onFinish}>
{children}
<button type="submit">Submit</button>
</Form>
);
};
// ── tests ─────────────────────────────────────────────────────────────────────
describe("OAuthFormFields", () => {
beforeEach(() => {
vi.clearAllMocks();
});
// ── visibility by flow type ─────────────────────────────────────────────────
describe("interactive mode (isM2M=false)", () => {
it("renders Token Validation Rules field", () => {
render(
<WithForm>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
});
it("renders Token Storage TTL field", () => {
render(
<WithForm>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument();
});
it("renders standard interactive fields alongside the new fields", () => {
render(
<WithForm>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
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(
<WithForm>
<OAuthFormFields isM2M={true} />
</WithForm>,
);
expect(screen.queryByText("Token Validation Rules (optional)")).not.toBeInTheDocument();
});
it("does NOT render Token Storage TTL field", () => {
render(
<WithForm>
<OAuthFormFields isM2M={true} />
</WithForm>,
);
expect(screen.queryByText("Token Storage TTL (seconds, optional)")).not.toBeInTheDocument();
});
it("still renders M2M-specific fields", () => {
render(
<WithForm>
<OAuthFormFields isM2M={true} />
</WithForm>,
);
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(
<WithForm onFinish={onFinish}>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
// 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(
<WithForm>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
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(
<WithForm>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
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(
<WithForm>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
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(
<WithForm onFinish={onFinish}>
<OAuthFormFields isM2M={false} />
</WithForm>,
);
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();
});
});
});
});

46
ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx
View file

@ -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<OAuthFormFieldsProps> = ({
>
<TextInput placeholder="https://example.com/oauth/register" className={fieldClassName} />
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Token Validation Rules (optional)"
tooltip='JSON object of key-value rules checked against the OAuth token response before storing. Supports dot-notation for nested fields (e.g. {"organization": "my-org", "team.id": "123"}). Tokens that fail validation are rejected with HTTP 403.'
/>
}
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"));
}
},
},
]}
>
<Input.TextArea
placeholder={'{\n "organization": "my-org",\n "team.id": "123"\n}'}
rows={4}
className="font-mono text-sm rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
/>
</Form.Item>
<Form.Item
label={
<FieldLabel
label="Token Storage TTL (seconds, optional)"
tooltip="How long to cache each user's OAuth access token in Redis before evicting it (regardless of the token's own expires_in). Leave blank to derive the TTL from the token's expires_in, or fall back to the 12-hour default."
/>
}
name="token_storage_ttl_seconds"
>
<InputNumber
min={1}
placeholder="e.g. 3600"
className="w-full rounded-lg"
style={{ width: "100%" }}
/>
</Form.Item>
{oauthFlow && (
<div className="rounded-lg border border-dashed border-gray-300 p-4 space-y-2">
<p className="text-sm text-gray-600">

431
ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx
View file

@ -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(<CreateMCPServer {...defaultProps} />);
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", () => {

20
ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx
View file

@ -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<CreateMCPServerProps> = ({
}
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<CreateMCPServerProps> = ({
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<CreateMCPServerProps> = ({
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<CreateMCPServerProps> = ({
restValues.transport = "http";
}
// Parse token_validation JSON if provided
let tokenValidation: Record<string, any> | 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<string, any> = {
...restValues,
@ -376,6 +389,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
allow_all_keys: Boolean(allowAllKeysRaw),
available_on_public_internet: Boolean(availableOnPublicInternetRaw),
static_headers: staticHeaders,
...(tokenValidation !== null && { token_validation: tokenValidation }),
};
payload.static_headers = staticHeaders;

249
ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx
View file

@ -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: () => <div data-testid="mcp-tool-config" />,
}));
// ── 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(
<MCPServerEdit
mcpServer={interactiveOAuthServer}
accessToken={null}
onCancel={vi.fn()}
onSuccess={vi.fn()}
availableAccessGroups={[]}
/>,
);
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(
<MCPServerEdit
mcpServer={{ ...interactiveOAuthServer, token_validation: tokenValidation }}
accessToken={null}
onCancel={vi.fn()}
onSuccess={vi.fn()}
availableAccessGroups={[]}
/>,
);
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(
<MCPServerEdit
mcpServer={interactiveOAuthServer}
accessToken="access-token"
onCancel={vi.fn()}
onSuccess={onSuccess}
availableAccessGroups={[]}
/>,
);
// 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(
<MCPServerEdit
mcpServer={interactiveOAuthServer}
accessToken="access-token"
onCancel={vi.fn()}
onSuccess={vi.fn()}
availableAccessGroups={[]}
/>,
);
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(
<MCPServerEdit
mcpServer={{ ...interactiveOAuthServer, token_validation: { organization: "old-org" } }}
accessToken="access-token"
onCancel={vi.fn()}
onSuccess={vi.fn()}
availableAccessGroups={[]}
/>,
);
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(
<MCPServerEdit
mcpServer={interactiveOAuthServer}
accessToken="access-token"
onCancel={vi.fn()}
onSuccess={vi.fn()}
availableAccessGroups={[]}
/>,
);
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(
<MCPServerEdit
mcpServer={{ ...interactiveOAuthServer, token_storage_ttl_seconds: 7200 }}
accessToken="access-token"
onCancel={vi.fn()}
onSuccess={vi.fn()}
availableAccessGroups={[]}
/>,
);
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);
});
});

79
ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx
View file

@ -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<MCPServerEditProps> = ({
}
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<MCPServerEditProps> = ({
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<MCPServerEditProps> = ({
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<MCPServerEditProps> = ({
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<MCPServerEditProps> = ({
restValues.transport = "http";
}
// Parse token_validation JSON if provided
let tokenValidation: Record<string, any> | 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<MCPServerEditProps> = ({
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<MCPServerEditProps> = ({
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
/>
</Form.Item>
{!isM2MFlow && (
<>
<Form.Item
label={
<span className="text-sm font-medium text-gray-700 flex items-center">
Token Validation Rules (optional)
<Tooltip title='JSON object of key-value rules checked against the OAuth token response before storing. Supports dot-notation for nested fields (e.g. {"organization": "my-org", "team.id": "123"}). Tokens that fail validation are rejected with HTTP 403.'>
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
</Tooltip>
</span>
}
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"));
}
},
},
]}
>
<Input.TextArea
placeholder={'{\n "organization": "my-org",\n "team.id": "123"\n}'}
rows={4}
className="font-mono text-sm rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
/>
</Form.Item>
<Form.Item
label={
<span className="text-sm font-medium text-gray-700 flex items-center">
Token Storage TTL (seconds, optional)
<Tooltip title="How long to cache each user's OAuth access token in Redis before evicting it (regardless of the token's own expires_in). Leave blank to derive the TTL from the token's expires_in, or fall back to the 12-hour default.">
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
</Tooltip>
</span>
}
name="token_storage_ttl_seconds"
>
<InputNumber
min={1}
placeholder="e.g. 3600"
style={{ width: "100%" }}
className="rounded-lg"
/>
</Form.Item>
</>
)}
<div className="rounded-lg border border-dashed border-gray-300 p-4 space-y-2">
<p className="text-sm text-gray-600">Use OAuth to fetch a fresh access token and temporarily save it in the session as the authentication value.</p>
<Button

3
ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx
View file

@ -20,6 +20,7 @@ import MCPSemanticFilterSettings from "../Settings/AdminSettings/MCPSemanticFilt
import MCPNetworkSettings from "./MCPNetworkSettings";
import MCPDiscovery from "./mcp_discovery";
import { ByokCredentialModal } from "./ByokCredentialModal";
import { getSecureItem } from "@/utils/secureStorage";
const { Text: AntdText, Title: AntdTitle } = Typography;
const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
@ -70,7 +71,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
return;
}
try {
const stored = window.sessionStorage.getItem(EDIT_OAUTH_UI_STATE_KEY);
const stored = getSecureItem(EDIT_OAUTH_UI_STATE_KEY);
if (!stored) {
return;
}

4
ui/litellm-dashboard/src/components/mcp_tools/types.tsx
View file

@ -223,6 +223,10 @@ export interface MCPServer {
submitted_at?: string | null;
reviewed_at?: string | null;
review_notes?: string | null;
/** Per-user OAuth token storage settings (interactive OAuth only) */
token_validation?: Record<string, any> | null;
token_storage_ttl_seconds?: number | null;
}
export interface MCPServerProps {

31
ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
View file

@ -75,6 +75,7 @@ import RealtimePlayground from "./RealtimePlayground";
import { A2ATaskMetadata, MessageType } from "./types";
import { useCodeInterpreter } from "./useCodeInterpreter";
import { useChatHistory } from "./useChatHistory";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
const { TextArea } = Input;
const { Dragger } = Upload;
@ -167,7 +168,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
} = useChatHistory({ simplified });
// codeql[js/clear-text-storage-of-sensitive-data]
const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
const saved = sessionStorage.getItem("apiKeySource");
const saved = getSecureItem("apiKeySource");
if (saved) {
try {
return JSON.parse(saved) as "session" | "custom";
@ -177,8 +178,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
}
return disabledPersonalKeyCreation ? "custom" : "session";
});
// codeql[js/clear-text-storage-of-sensitive-data]
const [apiKey, setApiKey] = useState<string>(() => sessionStorage.getItem("apiKey") || "");
const [apiKey, setApiKey] = useState<string>(() => getSecureItem("apiKey") || "");
const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState<string>(
() => sessionStorage.getItem("customProxyBaseUrl") || "",
);
@ -348,10 +348,12 @@ const ChatUI: React.FC<ChatUIProps> = ({
]);
useEffect(() => {
// codeql[js/clear-text-storage-of-sensitive-data]
sessionStorage.setItem("apiKeySource", JSON.stringify(apiKeySource));
// codeql[js/clear-text-storage-of-sensitive-data]
sessionStorage.setItem("apiKey", apiKey);
try {
setSecureItem("apiKeySource", JSON.stringify(apiKeySource));
setSecureItem("apiKey", apiKey);
} catch {
// Storage full or unavailable — non-critical, skip persisting.
}
sessionStorage.setItem("endpointType", endpointType);
sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags));
sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores));
@ -502,7 +504,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
const handleImageUpload = (file: File) => {
setUploadedImages((prev) => [...prev, file]);
const previewUrl = URL.createObjectURL(file);
const rawPreviewUrl = URL.createObjectURL(file);
// Sanitize: only allow blob: URLs to prevent XSS via img src injection.
const previewUrl = rawPreviewUrl.startsWith("blob:") ? rawPreviewUrl : "";
setImagePreviewUrls((prev) => [...prev, previewUrl]);
return false; // Prevent default upload behavior
};
@ -1827,7 +1831,16 @@ const ChatUI: React.FC<ChatUIProps> = ({
{uploadedImages.map((file, index) => (
<div key={index} className="relative inline-block">
<img
src={imagePreviewUrls[index] || ""}
src={(() => {
const url = imagePreviewUrls[index];
if (!url) return "";
try {
const parsed = new URL(url);
return parsed.protocol === "blob:" ? parsed.href : "";
} catch {
return "";
}
})()}
alt={`Upload preview ${index + 1}`}
className="max-w-32 max-h-32 rounded-md border border-gray-200 object-cover"
/>

1232
ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
View file

File diff suppressed because it is too large Load diff

14
ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx
View file

@ -11,6 +11,7 @@ import {
serverRootPath,
} from "@/components/networking";
import { extractErrorMessage } from "@/utils/errorUtils";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
export type McpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
@ -79,22 +80,13 @@ export const useMcpOAuthFlow = ({
const setStorageItem = (key: string, value: string) => {
if (typeof window === "undefined") return;
try {
// Use sessionStorage only — the flow state may contain client credentials;
// writing them to localStorage would persist across browser sessions and
// make them readable by any injected script (XSS).
// codeql[js/clear-text-storage-of-sensitive-data]
window.sessionStorage.setItem(key, value);
} catch (err) {
console.warn(`Failed to set storage item ${key}`, err);
}
setSecureItem(key, value);
};
const getStorageItem = (key: string): string | null => {
if (typeof window === "undefined") return null;
try {
// Try sessionStorage first, fall back to localStorage
return window.sessionStorage.getItem(key) || window.localStorage.getItem(key);
return getSecureItem(key);
} catch (err) {
console.warn(`Failed to get storage item ${key}`, err);
return null;

16
ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx
View file

@ -23,6 +23,7 @@ import {
} from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { extractErrorMessage } from "@/utils/errorUtils";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
@ -79,22 +80,11 @@ const genChallenge = async (verifier: string) => {
};
const setStorage = (key: string, value: string) => {
try {
// Use sessionStorage only — do not write to localStorage.
// The flow state may contain the LiteLLM access token; writing it to
// localStorage would persist it across browser sessions and make it
// readable by any injected script (XSS).
// codeql[js/clear-text-storage-of-sensitive-data]
window.sessionStorage.setItem(key, value);
} catch (_) {}
setSecureItem(key, value);
};
const getStorage = (key: string): string | null => {
try {
return window.sessionStorage.getItem(key);
} catch (_) {
return null;
}
return getSecureItem(key);
};
const clearStorage = (...keys: string[]) => {

34
ui/litellm-dashboard/src/utils/secureStorage.ts Normal file
View file

@ -0,0 +1,34 @@
function encode(value: string): string {
// btoa cannot handle characters outside Latin-1, so we percent-encode first.
return btoa(
encodeURIComponent(value).replace(
/%([0-9A-F]{2})/g,
(_, p1) => String.fromCharCode(parseInt(p1, 16))
)
);
}
function decode(encoded: string): string {
return decodeURIComponent(
atob(encoded)
.split("")
.map((c) => "%" + c.charCodeAt(0).toString(16).padStart(2, "0"))
.join("")
);
}
export function setSecureItem(key: string, value: string): void {
window.sessionStorage.setItem(key, encode(value));
}
export function getSecureItem(key: string): string | null {
try {
const raw = window.sessionStorage.getItem(key);
if (raw === null) return null;
return decode(raw);
} catch {
// Corrupted or non-encoded legacy value — return null without deleting
// so that in-flight flows (e.g. OAuth) can time out naturally.
return null;
}
}

2
ui/litellm-dashboard/vitest.config.ts
View file

@ -7,7 +7,7 @@ export default defineConfig({
setupFiles: ["tests/setupTests.ts"],
globals: true,
css: true, // lets you import CSS/modules without extra mocks
testTimeout: 10000,
testTimeout: 30000,
coverage: {
provider: "v8",
reporter: ["text", "lcov"],