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