mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge pull request #25460 from BerriAI/staging/mcp-per-user-oauth
feat(mcp): configurable token validation rules for per-user OAuth credentials
This commit is contained in:
commit
478bbeb53b
24 changed files with 2081 additions and 30 deletions
76
docs/my-website/docs/mcp_token_validation.md
Normal file
76
docs/my-website/docs/mcp_token_validation.md
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
# MCP Token Validation
|
||||
|
||||
Enforce that per-user OAuth tokens belong to your corporate workspace before LiteLLM stores them.
|
||||
|
||||
**Use case:** Your company uses Slack. You want to ensure employees can only connect their `@yourco.com` Slack account — not a personal one. Configure `token_validation` rules on the MCP server, and LiteLLM rejects any token whose claims don't match.
|
||||
|
||||
## How it works
|
||||
|
||||
When a user completes OAuth and submits their token via `POST /server/{id}/oauth-user-credential`, LiteLLM checks the `token_response_metadata` fields against the `token_validation` rules you've configured. If any field is missing or doesn't match, the request is rejected with a `403` before the credential is stored.
|
||||
|
||||
## Configuration
|
||||
|
||||
### Via the UI
|
||||
|
||||
Open **MCP Servers → Add New MCP Server**, set auth type to **OAuth**, and scroll to **Required Token Claims**:
|
||||
|
||||

|
||||
|
||||
Click a quick-add chip (Slack enterprise, Jira cloud, GitHub Enterprise) or type your own key. Fill in the required value.
|
||||
|
||||
### Via API
|
||||
|
||||
```bash
|
||||
curl -X POST https://your-proxy/v1/mcp/server \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"name": "slack-corp",
|
||||
"url": "https://slack.com/api/mcp",
|
||||
"transport": "http",
|
||||
"auth_type": "oauth2",
|
||||
"token_validation": {
|
||||
"enterprise_id": "E04COREWEAVE"
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
### Via config.yaml
|
||||
|
||||
```yaml title="config.yaml"
|
||||
mcp_servers:
|
||||
slack_corp:
|
||||
url: "https://slack.com/api/mcp"
|
||||
auth_type: oauth2
|
||||
token_validation:
|
||||
enterprise_id: "E04COREWEAVE"
|
||||
```
|
||||
|
||||
## Expected behavior
|
||||
|
||||
When a user submits their OAuth credential:
|
||||
|
||||
| Scenario | Result |
|
||||
|----------|--------|
|
||||
| `token_response_metadata: {"enterprise_id": "E04COREWEAVE"}` | ✅ 200 — credential stored |
|
||||
| `token_response_metadata: {"enterprise_id": "E99PERSONAL"}` | ❌ 403 — `'enterprise_id' = 'E99PERSONAL', expected 'E04COREWEAVE'` |
|
||||
| `token_response_metadata: {}` (field absent) | ❌ 403 — `required field 'enterprise_id' is absent` |
|
||||
|
||||
## Common claim keys
|
||||
|
||||
| Provider | Key | Example value |
|
||||
|----------|-----|---------------|
|
||||
| Slack | `enterprise_id` | `E04XXXXXXX` |
|
||||
| Jira / Confluence | `cloud_id` | `abc-123-def` |
|
||||
| GitHub Enterprise | `enterprise_id` | `your-org` |
|
||||
|
||||
## Token Storage TTL
|
||||
|
||||
Optionally set how long credentials are cached (in seconds):
|
||||
|
||||
```json
|
||||
{
|
||||
"token_validation": {"enterprise_id": "E04COREWEAVE"},
|
||||
"token_storage_ttl_seconds": 3600
|
||||
}
|
||||
```
|
||||
BIN
docs/my-website/img/mcp_token_validation.png
Normal file
BIN
docs/my-website/img/mcp_token_validation.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 78 KiB |
|
|
@ -318,6 +318,7 @@ const sidebars = {
|
|||
"mcp_usage",
|
||||
"mcp_openapi",
|
||||
"mcp_oauth",
|
||||
"mcp_token_validation",
|
||||
"mcp_aws_sigv4",
|
||||
"mcp_zero_trust",
|
||||
"mcp_public_internet",
|
||||
|
|
@ -685,6 +686,7 @@ const sidebars = {
|
|||
"mcp_usage",
|
||||
"mcp_openapi",
|
||||
"mcp_oauth",
|
||||
"mcp_token_validation",
|
||||
"mcp_aws_sigv4",
|
||||
"mcp_zero_trust",
|
||||
"mcp_public_internet",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,5 @@
|
|||
-- Add OAuth token validation fields to LiteLLM_MCPServerTable
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable"
|
||||
ADD COLUMN IF NOT EXISTS "token_validation" JSONB,
|
||||
ADD COLUMN IF NOT EXISTS "token_storage_ttl_seconds" INTEGER;
|
||||
|
|
@ -328,6 +328,10 @@ model LiteLLM_MCPServerTable {
|
|||
submitted_at DateTime?
|
||||
reviewed_at DateTime?
|
||||
review_notes String?
|
||||
// Per-user OAuth token validation rules: key-value pairs that must match fields
|
||||
// in the OAuth token response. Tokens that don't match are rejected with HTTP 403.
|
||||
token_validation Json?
|
||||
token_storage_ttl_seconds Int?
|
||||
|
||||
@@index([approval_status])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -82,6 +84,11 @@ def _prepare_mcp_server_data(
|
|||
# but be explicit to ensure a False value is always written to the DB).
|
||||
data_dict["is_byok"] = getattr(data, "is_byok", False)
|
||||
|
||||
# token_validation is a JSON dict — serialize it
|
||||
token_validation = getattr(data, "token_validation", None)
|
||||
if token_validation is not None:
|
||||
data_dict["token_validation"] = safe_dumps(token_validation)
|
||||
|
||||
return data_dict
|
||||
|
||||
|
||||
|
|
@ -576,6 +583,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 +612,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 +710,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,
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -693,6 +693,13 @@ class MCPServerManager:
|
|||
aws_service_name=aws_creds.get("aws_service_name"),
|
||||
aws_role_name=aws_creds.get("aws_role_name"),
|
||||
aws_session_name=aws_creds.get("aws_session_name"),
|
||||
# OAuth token validation rules persisted in the DB
|
||||
token_validation=_deserialize_json_dict(
|
||||
getattr(mcp_server, "token_validation", None)
|
||||
),
|
||||
token_storage_ttl_seconds=getattr(
|
||||
mcp_server, "token_storage_ttl_seconds", None
|
||||
),
|
||||
)
|
||||
return new_server
|
||||
|
||||
|
|
@ -2455,6 +2462,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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -1139,6 +1139,13 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
byok_description: List[str] = Field(default_factory=list)
|
||||
byok_api_key_help_url: Optional[str] = None
|
||||
source_url: Optional[str] = None
|
||||
# OAuth token validation rules: key-value pairs that must match fields in the
|
||||
# OAuth token response. Tokens that don't match are rejected with HTTP 403.
|
||||
# Example: {"enterprise_id": "E04COREWEAVE"}
|
||||
token_validation: Optional[Dict[str, Any]] = None
|
||||
# How long to cache per-user OAuth tokens in Redis (seconds).
|
||||
# Defaults to the token's own expires_in, or 12 hours if absent.
|
||||
token_storage_ttl_seconds: Optional[int] = None
|
||||
# BYOM submission fields — set by the endpoint, not by the caller.
|
||||
# Any caller-provided values are silently overridden before persistence.
|
||||
approval_status: Optional[str] = Field(
|
||||
|
|
@ -1220,6 +1227,8 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
byok_description: List[str] = Field(default_factory=list)
|
||||
byok_api_key_help_url: Optional[str] = None
|
||||
source_url: Optional[str] = None
|
||||
token_validation: Optional[Dict[str, Any]] = None
|
||||
token_storage_ttl_seconds: Optional[int] = None
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
|
|
@ -1323,6 +1332,9 @@ class MCPOAuthUserCredentialRequest(LiteLLMPydanticObjectBase):
|
|||
refresh_token: Optional[str] = None
|
||||
expires_in: Optional[int] = None # seconds until expiry
|
||||
scopes: Optional[List[str]] = None
|
||||
# Extra fields from the token response (e.g. Slack's enterprise_id, team_id).
|
||||
# These are checked against the server's token_validation rules.
|
||||
token_response_metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class MCPOAuthUserCredentialStatus(LiteLLMPydanticObjectBase):
|
||||
|
|
|
|||
|
|
@ -1619,6 +1619,22 @@ if MCP_AVAILABLE:
|
|||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"error": f"MCP Server {server_id} not found"},
|
||||
)
|
||||
|
||||
# Enforce token validation rules configured on the server.
|
||||
# The caller must pass token_response_metadata with any fields
|
||||
# that the server requires (e.g. enterprise_id from Slack's token response).
|
||||
token_validation = getattr(mcp_server, "token_validation", None)
|
||||
if token_validation and isinstance(token_validation, dict) and token_validation:
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
_validate_token_response,
|
||||
)
|
||||
|
||||
_validate_token_response(
|
||||
token_response=payload.token_response_metadata or {},
|
||||
validation_rules=token_validation,
|
||||
server_id=server_id,
|
||||
)
|
||||
|
||||
user_id = user_api_key_dict.user_id or ""
|
||||
if not user_id:
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -328,6 +328,10 @@ model LiteLLM_MCPServerTable {
|
|||
submitted_at DateTime?
|
||||
reviewed_at DateTime?
|
||||
review_notes String?
|
||||
// Per-user OAuth token validation rules: key-value pairs that must match fields
|
||||
// in the OAuth token response. Tokens that don't match are rejected with HTTP 403.
|
||||
token_validation Json?
|
||||
token_storage_ttl_seconds Int?
|
||||
|
||||
@@index([approval_status])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -328,6 +328,10 @@ model LiteLLM_MCPServerTable {
|
|||
submitted_at DateTime?
|
||||
reviewed_at DateTime?
|
||||
review_notes String?
|
||||
// Per-user OAuth token validation rules: key-value pairs that must match fields
|
||||
// in the OAuth token response. Tokens that don't match are rejected with HTTP 403.
|
||||
token_validation Json?
|
||||
token_storage_ttl_seconds Int?
|
||||
|
||||
@@index([approval_status])
|
||||
}
|
||||
|
|
|
|||
527
tests/mcp_tests/test_per_user_oauth_cache.py
Normal file
527
tests/mcp_tests/test_per_user_oauth_cache.py
Normal 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
|
||||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,9 +1,119 @@
|
|||
import React from "react";
|
||||
import { Form, Select, Tooltip } from "antd";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import React, { useState } from "react";
|
||||
import { Form, Input, InputNumber, Select, Tooltip } from "antd";
|
||||
import { InfoCircleOutlined, PlusOutlined, CloseOutlined } from "@ant-design/icons";
|
||||
import { Button, TextInput } from "@tremor/react";
|
||||
import { OAUTH_FLOW } from "./types";
|
||||
|
||||
const CLAIM_PRESETS = [
|
||||
{ label: "Slack enterprise", field: "enterprise_id", placeholder: "e.g. E04XXXXXXX" },
|
||||
{ label: "Jira cloud", field: "cloud_id", placeholder: "e.g. abc-123" },
|
||||
{ label: "GitHub Enterprise", field: "enterprise", placeholder: "e.g. my-org" },
|
||||
];
|
||||
|
||||
interface ClaimRow { field: string; value: string }
|
||||
|
||||
interface TokenClaimsBuilderProps {
|
||||
value?: string; // JSON string from antd Form
|
||||
onChange?: (v: string) => void;
|
||||
}
|
||||
|
||||
const TokenClaimsBuilder: React.FC<TokenClaimsBuilderProps> = ({ value, onChange }) => {
|
||||
const parse = (v?: string): ClaimRow[] => {
|
||||
if (!v || v.trim() === "") return [];
|
||||
try {
|
||||
const obj = JSON.parse(v);
|
||||
return Object.entries(obj).map(([field, val]) => ({ field, value: String(val) }));
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
};
|
||||
|
||||
const [rows, setRows] = useState<ClaimRow[]>(() => parse(value));
|
||||
|
||||
const emit = (next: ClaimRow[]) => {
|
||||
setRows(next);
|
||||
const nonempty = next.filter(r => r.field.trim() !== "");
|
||||
if (nonempty.length === 0) { onChange?.(""); return; }
|
||||
const obj: Record<string, string> = {};
|
||||
nonempty.forEach(r => { obj[r.field.trim()] = r.value; });
|
||||
onChange?.(JSON.stringify(obj));
|
||||
};
|
||||
|
||||
const update = (i: number, key: keyof ClaimRow, val: string) =>
|
||||
emit(rows.map((r, idx) => idx === i ? { ...r, [key]: val } : r));
|
||||
|
||||
const remove = (i: number) => emit(rows.filter((_, idx) => idx !== i));
|
||||
|
||||
const addPreset = (field: string) => {
|
||||
if (rows.some(r => r.field === field)) return;
|
||||
emit([...rows, { field, value: "" }]);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="space-y-3">
|
||||
<div className="flex flex-wrap gap-2">
|
||||
<span className="text-xs text-gray-400 self-center">Quick add:</span>
|
||||
{CLAIM_PRESETS.map(p => (
|
||||
<button
|
||||
key={p.field}
|
||||
type="button"
|
||||
onClick={() => addPreset(p.field)}
|
||||
disabled={rows.some(r => r.field === p.field)}
|
||||
className="px-2 py-0.5 text-xs rounded-full border border-blue-300 text-blue-600 hover:bg-blue-50 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
|
||||
>
|
||||
{p.label}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
|
||||
{rows.length > 0 && (
|
||||
<div className="rounded-lg border border-gray-200 overflow-hidden">
|
||||
<div className="grid grid-cols-[1fr_1fr_32px] bg-gray-50 border-b border-gray-200 px-3 py-1.5">
|
||||
<span className="text-xs font-medium text-gray-500">Token field</span>
|
||||
<span className="text-xs font-medium text-gray-500">Required value</span>
|
||||
<span />
|
||||
</div>
|
||||
{rows.map((row, i) => {
|
||||
const preset = CLAIM_PRESETS.find(p => p.field === row.field);
|
||||
return (
|
||||
<div key={i} className="grid grid-cols-[1fr_1fr_32px] items-center px-3 py-2 gap-2 border-b border-gray-100 last:border-b-0 hover:bg-gray-50">
|
||||
<input
|
||||
className="w-full text-sm border border-gray-200 rounded px-2 py-1 focus:outline-none focus:border-blue-400"
|
||||
placeholder="e.g. enterprise_id"
|
||||
value={row.field}
|
||||
onChange={e => update(i, "field", e.target.value)}
|
||||
/>
|
||||
<input
|
||||
className="w-full text-sm border border-gray-200 rounded px-2 py-1 focus:outline-none focus:border-blue-400"
|
||||
placeholder={preset?.placeholder ?? "required value"}
|
||||
value={row.value}
|
||||
onChange={e => update(i, "value", e.target.value)}
|
||||
/>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => remove(i)}
|
||||
className="flex items-center justify-center w-6 h-6 rounded hover:bg-red-50 text-gray-400 hover:text-red-500 transition-colors"
|
||||
>
|
||||
<CloseOutlined style={{ fontSize: 11 }} />
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => emit([...rows, { field: "", value: "" }])}
|
||||
className="flex items-center gap-1 text-xs text-blue-600 hover:text-blue-800 transition-colors"
|
||||
>
|
||||
<PlusOutlined style={{ fontSize: 11 }} />
|
||||
Add claim
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
interface OAuthFlowStatus {
|
||||
startOAuthFlow: () => void;
|
||||
status: string;
|
||||
|
|
@ -151,6 +261,33 @@ const OAuthFormFields: React.FC<OAuthFormFieldsProps> = ({
|
|||
>
|
||||
<TextInput placeholder="https://example.com/oauth/register" className={fieldClassName} />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<FieldLabel
|
||||
label="Required Token Claims (optional)"
|
||||
tooltip="Block personal accounts by requiring specific fields in the OAuth token response. For example, require enterprise_id to match your Slack workspace. Tokens that don't match are rejected with HTTP 403."
|
||||
/>
|
||||
}
|
||||
name="token_validation_json"
|
||||
>
|
||||
<TokenClaimsBuilder />
|
||||
</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">
|
||||
|
|
|
|||
|
|
@ -353,6 +353,147 @@ describe("CreateMCPServer", () => {
|
|||
);
|
||||
});
|
||||
|
||||
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", () => {
|
||||
it("should call setModalVisible(false) when cancel is clicked", async () => {
|
||||
render(<CreateMCPServer {...defaultProps} />);
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
@ -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],
|
||||
);
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue