diff --git a/docker/Dockerfile.custom_ui b/docker/Dockerfile.custom_ui index f836190a49..cc44893bf9 100644 --- a/docker/Dockerfile.custom_ui +++ b/docker/Dockerfile.custom_ui @@ -71,8 +71,16 @@ WORKDIR /app RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh +# Run as non-root user +RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser \ + && chown -R appuser:appuser /app +USER appuser + # Expose the necessary port EXPOSE 4000/tcp +HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \ + CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:4000/health')"] + # Override the CMD instruction with your desired command and arguments CMD ["--port", "4000", "--config", "config.yaml", "--detailed_debug"] \ No newline at end of file diff --git a/docker/Dockerfile.health_check b/docker/Dockerfile.health_check index fb9cc201d2..e968bec340 100644 --- a/docker/Dockerfile.health_check +++ b/docker/Dockerfile.health_check @@ -13,12 +13,12 @@ RUN pip install --no-cache-dir -r requirements.txt RUN chmod +x /app/health_check_client.py # Run as non-root user -RUN adduser --disabled-password --gecos "" --uid 1001 healthcheck -USER healthcheck +RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser +USER appuser # Health check -HEALTHCHECK --interval=30s --timeout=5s --retries=3 \ - CMD python /app/health_check_client.py --help || exit 1 +HEALTHCHECK --interval=30s --timeout=5s --start-period=5s --retries=3 \ + CMD ["python", "/app/health_check_client.py", "--help"] # Set entrypoint ENTRYPOINT ["python", "/app/health_check_client.py"] diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 88cbcac52c..c64d475fda 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -602,6 +602,8 @@ router_settings: | MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE | Maximum number of entries in MCP OAuth2 token cache. Default is 200 | MCP_OAUTH2_TOKEN_CACHE_MIN_TTL | Minimum TTL in seconds for MCP OAuth2 token cache. Default is 10 | MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from token expiry when computing cache TTL. Default is 60 +| MCP_PER_USER_TOKEN_DEFAULT_TTL | Default TTL in seconds for per-user MCP OAuth tokens stored in Redis. Default is 43200 (12 hours) +| MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from per-user MCP OAuth token expiry when computing Redis TTL. Default is 60 | DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20 | DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10 | DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602 diff --git a/docs/my-website/docs/proxy/credential_routing.md b/docs/my-website/docs/proxy/credential_routing.md new file mode 100644 index 0000000000..2af57c6b49 --- /dev/null +++ b/docs/my-website/docs/proxy/credential_routing.md @@ -0,0 +1,274 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Per-Team/Project Credential Routing + +Route the same model to different LLM provider endpoints (e.g. different Azure instances) based on which team or project makes the request. + +## Overview + +In multi-tenant deployments, different teams often need the same model name (e.g., `gpt-4`) to hit different provider endpoints — for example, separate Azure OpenAI instances per business unit for cost isolation, data residency, or rate limit separation. + +**Credential routing** lets you configure this in team/project metadata using the existing [credentials table](./ui_credentials.md), without duplicating model definitions or creating separate model groups per team. + +``` +Hotel Team → gpt-4 → https://hotel-eastus.openai.azure.com/ +Flight Team → gpt-4 → https://flight-centralus.openai.azure.com/ +``` + +### Precedence Chain + +When a request comes in, the system walks this precedence chain (first match wins): + +1. **Clientside credentials** — `api_base`/`api_key` passed in the request body ([docs](./clientside_auth.md)) +2. **Project model-specific** — override for this exact model in the project's `model_config` +3. **Project default** — `defaultconfig` in the project's `model_config` +4. **Team model-specific** — override for this exact model in the team's `model_config` +5. **Team default** — `defaultconfig` in the team's `model_config` +6. **Deployment default** — the model's `litellm_params` as configured in `config.yaml` + +## Quick Start + +### Step 1: Create Credentials + +Store your Azure endpoint credentials in the credentials table. You can do this via the [UI](./ui_credentials.md) or API: + +```bash showLineNumbers +# Create credential for Hotel team's Azure endpoint +curl -X POST 'http://0.0.0.0:4000/credentials' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "credential_name": "hotel-azure-eastus", + "credential_values": { + "api_base": "https://hotel-eastus.openai.azure.com/", + "api_key": "sk-azure-hotel-key-xxx" + } +}' +``` + +```bash showLineNumbers +# Create credential for Flight team's Azure endpoint +curl -X POST 'http://0.0.0.0:4000/credentials' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "credential_name": "flight-azure-centralus", + "credential_values": { + "api_base": "https://flight-centralus.openai.azure.com/", + "api_key": "sk-azure-flight-key-xxx" + } +}' +``` + +### Step 2: Set `model_config` on Teams + +Add a `model_config` key to the team's metadata referencing the credential by name: + +```bash showLineNumbers +# Hotel team — default Azure endpoint for all models +curl -X PATCH 'http://0.0.0.0:4000/team/update' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "team_id": "hotel-team-id", + "metadata": { + "model_config": { + "defaultconfig": { + "azure": { + "litellm_credentials": "hotel-azure-eastus" + } + } + } + } +}' +``` + +```bash showLineNumbers +# Flight team — default Azure endpoint for all models +curl -X PATCH 'http://0.0.0.0:4000/team/update' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "team_id": "flight-team-id", + "metadata": { + "model_config": { + "defaultconfig": { + "azure": { + "litellm_credentials": "flight-azure-centralus" + } + } + } + } +}' +``` + +### Step 3: Make Requests + +Requests are automatically routed to the correct Azure endpoint based on the API key's team: + +```bash showLineNumbers +# Request using Hotel team's API key → routes to hotel-eastus.openai.azure.com +curl http://localhost:4000/v1/chat/completions \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-hotel-team-key' \ +-d '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' + +# Request using Flight team's API key → routes to flight-centralus.openai.azure.com +curl http://localhost:4000/v1/chat/completions \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-flight-team-key' \ +-d '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' +``` + +## Per-Model Overrides + +You can set different credentials for specific models while keeping a default for everything else: + +```bash showLineNumbers +curl -X PATCH 'http://0.0.0.0:4000/team/update' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "team_id": "hotel-team-id", + "metadata": { + "model_config": { + "defaultconfig": { + "azure": { + "litellm_credentials": "hotel-azure-eastus" + } + }, + "gpt-4": { + "azure": { + "litellm_credentials": "hotel-azure-westus" + } + } + } + } +}' +``` + +With this config: +- `gpt-4` requests → `hotel-azure-westus` credential (model-specific) +- All other models → `hotel-azure-eastus` credential (default) + +## Project-Level Overrides + +Projects inherit their team's `model_config` but can override at the project level. Project overrides take precedence over team overrides. + +```bash showLineNumbers +# Project overrides the team default for all models +curl -X PATCH 'http://0.0.0.0:4000/project/update' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{ + "project_id": "hotel-rec-app-id", + "metadata": { + "model_config": { + "defaultconfig": { + "azure": { + "litellm_credentials": "hotel-rec-azure" + } + }, + "gpt-4-vision": { + "azure": { + "litellm_credentials": "hotel-rec-vision" + } + } + } + } +}' +``` + +### Full Example: Hotel Team with Two Projects + +**Setup:** +- **Hotel Team**: default `hotel-azure-eastus`, GPT-4 override to `hotel-azure-westus` +- **Hotel Rec App** (project): default `hotel-rec-azure`, GPT-4-Vision override to `hotel-rec-vision` +- **Hotel Review App** (project): no overrides — inherits team config + +**Resolution:** + +| Request | Resolved Credential | Why | +|---|---|---| +| Hotel Rec App → `gpt-4` | `hotel-rec-azure` | Project default (no project model-specific match for gpt-4) | +| Hotel Rec App → `gpt-4-vision` | `hotel-rec-vision` | Project model-specific | +| Hotel Review App → `gpt-3.5` | `hotel-azure-eastus` | Team default (no project config) | +| Hotel Review App → `gpt-4` | `hotel-azure-westus` | Team model-specific | + +## `model_config` Schema + +The `model_config` key is a JSON object in team/project `metadata`: + +```json +{ + "model_config": { + "defaultconfig": { + "": { + "litellm_credentials": "" + } + }, + "": { + "": { + "litellm_credentials": "" + } + } + } +} +``` + +| Field | Description | +|---|---| +| `defaultconfig` | Fallback credential for any model not explicitly listed | +| `` | Model-specific override — must match the LiteLLM model group name | +| `` | Provider key (e.g. `azure`, `openai`, `bedrock`). When the model name includes a provider prefix (e.g. `azure/gpt-4`), the system prefers the matching provider key | +| `litellm_credentials` | Name of a credential in the [credentials table](./ui_credentials.md) | + +### Credential Values + +The referenced credential can contain any combination of: + +| Key | Description | +|---|---| +| `api_base` | Provider endpoint URL | +| `api_key` | API key for the provider | +| `api_version` | API version (e.g. for Azure) | + +Only keys present in the credential are applied. Keys already in the request (e.g. clientside `api_version`) are never overwritten. + +## Enabling the Feature + +This feature is **disabled by default** and must be explicitly enabled. To enable it: + + + + + +```yaml +litellm_settings: + enable_model_config_credential_overrides: true +``` + + + + + +```bash +export LITELLM_ENABLE_MODEL_CONFIG_CREDENTIAL_OVERRIDES=true +``` + + + + + +:::info +The feature flag must be enabled before `model_config` entries in team/project metadata take effect. Without it, credential routing is completely inert — no metadata is read, no credentials are resolved. +::: + +## Related Documentation + +- [Adding LLM Credentials](./ui_credentials.md) — Create and manage reusable credentials +- [Project Management](./project_management.md) — Project hierarchy and API +- [Team Budgets](./team_budgets.md) — Team-level budget management +- [Clientside LLM Credentials](./clientside_auth.md) — Passing credentials in the request body +- [Credential Usage Tracking](./credential_usage_tracking.md) — Track spend by credential diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 4b087afd84..f7833fc02c 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -99,7 +99,7 @@ The following checks were performed on each of these signatures: - The signatures were verified against the specified public key ``` -Learn more about LiteLLM's release signing in the [CI/CD v2 announcement](https://docs.litellm.ai/blog/ci-cd-v2-improvements#verify-docker-image-signatures). +Learn more about LiteLLM's release signing in the [CI/CD v2 announcement](https://docs.litellm.ai/blog/ci-cd-v2-improvements#verify-docker-image-signatures). For a complete guide covering all image variants, CI/CD enforcement, and deployment best practices, see the [Docker Image Security Guide](./docker_image_security.md). ### Docker Run diff --git a/docs/my-website/docs/proxy/docker_image_security.md b/docs/my-website/docs/proxy/docker_image_security.md new file mode 100644 index 0000000000..41ace2174b --- /dev/null +++ b/docs/my-website/docs/proxy/docker_image_security.md @@ -0,0 +1,189 @@ +# Docker Image Security Guide + +LiteLLM signs every Docker image published to GHCR with [cosign](https://docs.sigstore.dev/cosign/overview/) starting from **v1.83.0**. This page covers how to verify signatures, enforce verification in CI/CD, and follow recommended deployment patterns. + +## Signed images + +All image variants published to `ghcr.io/berriai/` are signed with the same cosign key: + +| Image | Description | +|---|---| +| `ghcr.io/berriai/litellm` | Core proxy | +| `ghcr.io/berriai/litellm-database` | Proxy with Postgres dependencies | +| `ghcr.io/berriai/litellm-non_root` | Non-root variant | +| `ghcr.io/berriai/litellm-spend_logs` | Spend-logs sidecar | + +The signing key was introduced in [commit `0112e53`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0) and the public key is checked into the repository at [`cosign.pub`](https://github.com/BerriAI/litellm/blob/main/cosign.pub). + +:::info Enterprise images +Enterprise images (`litellm-ee`) follow the same signing process. Contact [support@berri.ai](mailto:support@berri.ai) to confirm coverage for your specific enterprise image tag. +::: + +## Verify image signatures + +Install cosign following the [official instructions](https://docs.sigstore.dev/cosign/system_config/installation/). + +### Verify with the pinned commit hash (recommended) + +A commit hash is cryptographically immutable, making this the strongest verification method: + +```bash +cosign verify \ + --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \ + ghcr.io/berriai/litellm:v1.83.0-stable +``` + +Replace the image reference with any signed variant: + +```bash +# litellm-database +cosign verify \ + --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \ + ghcr.io/berriai/litellm-database:v1.83.0-stable + +# litellm-non_root +cosign verify \ + --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \ + ghcr.io/berriai/litellm-non_root:v1.83.0-stable +``` + +### Verify with a release tag (convenience) + +Tags are protected in this repository and resolve to the same key: + +```bash +cosign verify \ + --key https://raw.githubusercontent.com/BerriAI/litellm/v1.83.0-stable/cosign.pub \ + ghcr.io/berriai/litellm-database:v1.83.0-stable +``` + +### Expected output + +``` +The following checks were performed on each of these signatures: + - The cosign claims were validated + - The signatures were verified against the specified public key +``` + +## Enforce verification in CI/CD + +### Kubernetes — Sigstore Policy Controller + +The [Sigstore Policy Controller](https://docs.sigstore.dev/policy-controller/overview/) rejects pods whose images fail cosign verification. + +1. Install the controller: + +```bash +helm repo add sigstore https://sigstore.github.io/helm-charts +helm install policy-controller sigstore/policy-controller \ + -n cosign-system --create-namespace +``` + +2. Create a `ClusterImagePolicy` with the LiteLLM public key: + +```yaml +apiVersion: policy.sigstore.dev/v1beta1 +kind: ClusterImagePolicy +metadata: + name: litellm-signed-images +spec: + images: + - glob: "ghcr.io/berriai/litellm*" + authorities: + - key: + data: | + -----BEGIN PUBLIC KEY----- + MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEKi4ivqGpE231OGH50PKbqy1Y1Kkb + POJC8+i2Wko82gBOUCe3M0Vw86H/4rhUhfoYEti4gdJ9wZbYmK0I2EE96g== + -----END PUBLIC KEY----- +``` + +3. Label the namespace to enable enforcement: + +```bash +kubectl label namespace litellm policy.sigstore.dev/include=true +``` + +Any pod in that namespace using an unsigned `ghcr.io/berriai/litellm*` image will be rejected at admission. + +### GCP — Binary Authorization + +[Binary Authorization](https://cloud.google.com/binary-authorization/docs) can enforce cosign signatures on Cloud Run and GKE. + +1. Create a cosign-based attestor using the LiteLLM public key: + +```bash +# Import the public key into a Cloud KMS keyring or use a PGP/PKIX attestor. +# See: https://cloud.google.com/binary-authorization/docs/creating-attestors-console +``` + +2. Configure a Binary Authorization policy that requires the attestor for `ghcr.io/berriai/litellm*` images. + +3. Enable the policy on your Cloud Run service or GKE cluster. + +Refer to the [GCP Binary Authorization docs](https://cloud.google.com/binary-authorization/docs/setting-up) for full setup steps. + +### AWS — ECS / ECR + +AWS does not natively verify cosign signatures at deploy time. Common approaches: + +- **CI/CD gate**: Run `cosign verify` in your deployment pipeline before pushing to ECR or updating the ECS task definition. Fail the pipeline if verification fails. +- **OPA/Gatekeeper on EKS**: If running on EKS, use the Sigstore Policy Controller (same as the Kubernetes approach above). + +### GitHub Actions gate + +Add a verification step before any deployment job: + +```yaml +- name: Verify LiteLLM image signature + run: | + cosign verify \ + --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \ + ghcr.io/berriai/litellm-database:${{ env.LITELLM_VERSION }} +``` + +## Recommended deployment patterns + +### Pin by digest + +Digest pinning guarantees the exact image content regardless of tag mutations: + +```yaml +image: ghcr.io/berriai/litellm-database@sha256: +``` + +Get the digest after pulling: + +```bash +docker inspect --format='{{index .RepoDigests 0}}' \ + ghcr.io/berriai/litellm-database:v1.83.0-stable +``` + +Cosign verification works with digests too: + +```bash +cosign verify \ + --key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \ + ghcr.io/berriai/litellm-database@sha256: +``` + +### Use stable release tags + +If digest pinning is too rigid for your workflow, use `-stable` release tags (e.g. `v1.83.0-stable`). These are immutable release tags that will not be overwritten. + +Avoid `main-latest` or `main-stable` in production — these rolling tags point to the most recent build and can change between deployments. + +### Safe upgrade checklist + +1. **Verify the new image** — run `cosign verify` against the new release tag or digest. +2. **Test in staging** — deploy the verified image to a non-production environment. +3. **Update your pinned reference** — change the digest or tag in your deployment manifest. +4. **Deploy to production** — roll out using your standard deployment process. +5. **Monitor `/health`** — confirm the proxy is healthy after the upgrade. + +## Further reading + +- [CI/CD v2 announcement](https://docs.litellm.ai/blog/ci-cd-v2-improvements) — background on LiteLLM's signing infrastructure +- [Docker deployment guide](./deploy.md) — full Docker, Helm, and Terraform setup +- [cosign documentation](https://docs.sigstore.dev/cosign/overview/) — cosign usage and key management +- [Sigstore Policy Controller](https://docs.sigstore.dev/policy-controller/overview/) — Kubernetes admission control diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ab8f257c7d..54581dceb9 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -349,6 +349,7 @@ const sidebars = { "proxy/debugging", "proxy/error_diagnosis", "proxy/deploy", + "proxy/docker_image_security", "proxy/health", "proxy/master_key_rotations", "proxy/model_management", @@ -563,7 +564,8 @@ const sidebars = { "proxy/model_access", "proxy/model_access_groups", "proxy/access_groups", - "proxy/team_model_add" + "proxy/team_model_add", + "proxy/credential_routing" ] }, { diff --git a/litellm/__init__.py b/litellm/__init__.py index d4418c661a..64c60ca337 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -318,6 +318,7 @@ return_response_headers: bool = ( False # get response headers from LLM Api providers - example x-remaining-requests, ) enable_json_schema_validation: bool = False +enable_model_config_credential_overrides: bool = False enable_key_alias_format_validation: bool = ( False # opt-in validation of key_alias format on /key/generate and /key/update ) diff --git a/litellm/constants.py b/litellm/constants.py index a7d86ddb16..337cb1243f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -135,6 +135,15 @@ MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int( MCP_NPM_CACHE_DIR = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache") MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10")) +# Per-user OAuth token Redis cache (for server-side token storage) +MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX = "mcp:per_user_token" +MCP_PER_USER_TOKEN_DEFAULT_TTL = int( + os.getenv("MCP_PER_USER_TOKEN_DEFAULT_TTL", "43200") # 12 hours +) +MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int( + os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60") +) + # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0")) MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c9abaad90..7a8820a878 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,5 +1,6 @@ import json import ssl +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse from typing import ( TYPE_CHECKING, Any, @@ -5027,6 +5028,16 @@ class BaseLLMHTTPHandler: litellm_params={}, ) ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + # OpenAI's WebSocket responses endpoint requires ?model= in the URL, + # matching the Realtime API convention (wss://.../v1/realtime?model=...). + # Use urllib.parse so existing query params (e.g. api-version) are preserved. + _parsed = urlparse(ws_url) + _qs = parse_qs(_parsed.query) + if "model" not in _qs: + _qs["model"] = [model] + ws_url = urlunparse( + _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) + ) try: ssl_context = get_shared_realtime_ssl_context() diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index fbef33c32e..e9bd41bb95 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -21,7 +21,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.utils import PrismaClient +from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPCredentials @@ -576,6 +578,7 @@ async def store_user_oauth_credential( refresh_token: Optional[str] = None, expires_in: Optional[int] = None, scopes: Optional[List[str]] = None, + skip_byok_guard: bool = False, ) -> None: """Persist an OAuth2 access token for a user+server pair. @@ -604,21 +607,26 @@ async def store_user_oauth_credential( # Guard against silently overwriting a BYOK credential with an OAuth token. # BYOK credentials lack a "type" field (or use a non-"oauth2" type). - existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} - ) - if existing is not None: - _byok_error = ValueError( - f"A non-OAuth2 credential already exists for user {user_id} " - f"and server {server_id}. Refusing to overwrite." + # Skip the guard when the caller knows the row is already an OAuth2 credential + # (e.g. during token refresh), saving an extra DB round-trip. + if not skip_byok_guard: + existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - try: - raw = json.loads(base64.urlsafe_b64decode(existing.credential_b64).decode()) - except Exception: - # Credential is not base64+JSON — it's a plain-text BYOK key. - raise _byok_error - if raw.get("type") != "oauth2": - raise _byok_error + if existing is not None: + _byok_error = ValueError( + f"A non-OAuth2 credential already exists for user {user_id} " + f"and server {server_id}. Refusing to overwrite." + ) + try: + raw = json.loads( + base64.urlsafe_b64decode(existing.credential_b64).decode() + ) + except Exception: + # Credential is not base64+JSON — it's a plain-text BYOK key. + raise _byok_error + if raw.get("type") != "oauth2": + raise _byok_error encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() await prisma_client.db.litellm_mcpusercredentials.upsert( @@ -697,6 +705,115 @@ async def list_user_oauth_credentials( return results +async def refresh_user_oauth_token( + prisma_client: PrismaClient, + user_id: str, + server: Any, + cred: Dict[str, Any], +) -> Optional[Dict[str, Any]]: + """Attempt to refresh a per-user OAuth2 token using its stored refresh_token. + + POSTs to ``server.token_url`` with ``grant_type=refresh_token``. + + On success: persists the new credential via ``store_user_oauth_credential`` + and returns the updated payload dict. + On failure (network error, invalid_grant, missing refresh_token, …): logs a + warning and returns ``None`` — the caller is responsible for clearing the + stale credential and triggering re-authentication. + """ + refresh_token: Optional[str] = cred.get("refresh_token") + token_url: Optional[str] = getattr(server, "token_url", None) + server_id: str = getattr(server, "server_id", "") + client_id: Optional[str] = getattr(server, "client_id", None) + client_secret: Optional[str] = getattr(server, "client_secret", None) + + if not refresh_token: + verbose_proxy_logger.debug( + "refresh_user_oauth_token: no refresh_token stored for user=%s server=%s", + user_id, + server_id, + ) + return None + if not token_url: + verbose_proxy_logger.debug( + "refresh_user_oauth_token: server=%s has no token_url configured", + server_id, + ) + return None + + token_data: Dict[str, str] = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + } + if client_id: + token_data["client_id"] = client_id + if client_secret: + token_data["client_secret"] = client_secret + + try: + async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.Oauth2Check + ) + response = await async_client.post( + token_url, + headers={"Accept": "application/json"}, + data=token_data, + ) + response.raise_for_status() + body: Dict[str, Any] = response.json() + except Exception as exc: + verbose_proxy_logger.warning( + "refresh_user_oauth_token: refresh request failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + return None + + access_token: Optional[str] = body.get("access_token") + if not access_token: + verbose_proxy_logger.warning( + "refresh_user_oauth_token: token response missing access_token for " + "user=%s server=%s", + user_id, + server_id, + ) + return None + + expires_in: Optional[int] = None + raw_expires = body.get("expires_in") + try: + expires_in = int(raw_expires) if raw_expires is not None else None + except (TypeError, ValueError): + pass + + # Rotate refresh token when the provider returns a new one + new_refresh_token: Optional[str] = body.get("refresh_token") or refresh_token + + raw_scope = body.get("scope") + scopes: Optional[List[str]] = ( + raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + ) or cred.get("scopes") + + await store_user_oauth_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=server_id, + access_token=access_token, + refresh_token=new_refresh_token, + expires_in=expires_in, + scopes=scopes, + skip_byok_guard=True, # Row is already OAuth2; skip the extra find_unique check + ) + + verbose_proxy_logger.info( + "refresh_user_oauth_token: refreshed token for user=%s server=%s", + user_id, + server_id, + ) + return await get_user_oauth_credential(prisma_client, user_id, server_id) + + async def approve_mcp_server( prisma_client: PrismaClient, server_id: str, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 07309eb57f..d0d6198632 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,10 +1,11 @@ import json -from typing import Optional +from typing import Any, Dict, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -147,6 +148,160 @@ def _resolve_oauth2_server_for_root_endpoints( return None +def _validate_token_response( + token_response: Dict[str, Any], + validation_rules: Dict[str, Any], + server_id: str, +) -> None: + """Raise HTTPException 403 if any validation rule doesn't match the token response. + + Supports dot-notation for nested fields (e.g. ``"team.enterprise_id"`` checks + ``token_response["team"]["enterprise_id"]``). Top-level keys are tried first, + then dot-split traversal. All comparisons are string-coerced so that numeric + values in the response (e.g. ``"org_id": 12345``) match string rules + (``"org_id": "12345"``). + """ + for key, expected in validation_rules.items(): + actual: Any = token_response.get(key) + # Try dot-notation traversal when top-level lookup returns None + if actual is None and "." in key: + obj: Any = token_response + for part in key.split("."): + if isinstance(obj, dict): + obj = obj.get(part) + else: + obj = None + break + actual = obj + # Treat absent fields as a distinct failure from a mismatched value + if actual is None: + raise HTTPException( + status_code=403, + detail={ + "error": "token_validation_failed", + "server_id": server_id, + "field": key, + "message": ( + f"OAuth token rejected: required field '{key}' is absent" + ), + }, + ) + if str(actual) != str(expected): + raise HTTPException( + status_code=403, + detail={ + "error": "token_validation_failed", + "server_id": server_id, + "field": key, + "message": ( + f"OAuth token rejected: '{key}' = '{actual}', " + f"expected '{expected}'" + ), + }, + ) + + +async def _extract_user_id_from_request(request: Request) -> Optional[str]: + """Best-effort extraction of LiteLLM user_id from the request's Authorization header. + + Called at the OAuth token endpoint so that per-user tokens can be stored + server-side. Uses a read-only cache lookup to avoid re-running the full + auth pipeline (which has side effects such as rate-limit increments and + spend logging). Returns ``None`` if no cached credential is found. + """ + auth_header = request.headers.get("Authorization") or request.headers.get( + "authorization" + ) + if not auth_header: + return None + lower = auth_header.lower() + if not lower.startswith("bearer "): + return None + token = auth_header[7:].strip() + try: + from litellm.proxy._types import hash_token # noqa: PLC0415 + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + cached = await user_api_key_cache.async_get_cache(hash_token(token)) + return getattr(cached, "user_id", None) + except Exception: + return None + + +async def _store_per_user_token_server_side( + server: MCPServer, + user_id: str, + token_response: Dict[str, Any], +) -> None: + """Persist the OAuth token server-side and warm the Redis cache. + + Called from the token endpoint after a successful code exchange or refresh. + Errors are logged but NOT re-raised — the token is always returned to the + client even when server-side storage fails. + """ + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415 + _compute_per_user_token_ttl, + mcp_per_user_token_cache, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + access_token: Optional[str] = token_response.get("access_token") + if not access_token: + return + + raw_expires = token_response.get("expires_in") + try: + expires_in: Optional[int] = int(raw_expires) if raw_expires is not None else None + except (TypeError, ValueError): + expires_in = None + + refresh_token: Optional[str] = token_response.get("refresh_token") or None + raw_scope = token_response.get("scope") + scopes: Optional[list] = ( + raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + ) + + try: + prisma_client = get_prisma_client_or_throw( + "Database not connected. Cannot store per-user OAuth token." + ) + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + store_user_oauth_credential, + ) + + await store_user_oauth_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=server.server_id, + access_token=access_token, + refresh_token=refresh_token, + expires_in=expires_in, + scopes=scopes, + ) + verbose_logger.info( + "_store_per_user_token_server_side: stored token for user=%s server=%s", + user_id, + server.server_id, + ) + except Exception as exc: + verbose_logger.warning( + "_store_per_user_token_server_side: DB storage failed for user=%s server=%s: %s", + user_id, + server.server_id, + exc, + ) + return # Don't warm Redis if DB write failed + + # Warm the Redis cache so the first subsequent MCP call is a cache hit + ttl = _compute_per_user_token_ttl(server, expires_in) + await mcp_per_user_token_cache.set( + user_id=user_id, + server_id=server.server_id, + access_token=access_token, + ttl=ttl, + ) + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -266,6 +421,44 @@ async def exchange_token_with_server( token_response = response.json() access_token = token_response["access_token"] + # Validate token response against server-configured rules before any storage. + # This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc. + if mcp_server.token_validation and isinstance(mcp_server.token_validation, dict): + _validate_token_response( + token_response=token_response, + validation_rules=mcp_server.token_validation, + server_id=mcp_server.server_id, + ) + + # Store server-side when the server is configured for per-user OAuth and + # the calling client has provided a valid LiteLLM identity. + # Errors are non-fatal: the token is still returned to the client. + if mcp_server.needs_user_oauth_token: + user_id = await _extract_user_id_from_request(request) + if user_id: + try: + await _store_per_user_token_server_side( + server=mcp_server, + user_id=user_id, + token_response=token_response, + ) + except Exception as exc: + verbose_logger.warning( + "exchange_token_with_server: server-side storage failed " + "for user=%s server=%s: %s", + user_id, + mcp_server.server_id, + exc, + ) + else: + verbose_logger.debug( + "exchange_token_with_server: no LiteLLM user_id found in request; " + "per-user token for server=%s will not be stored server-side. " + "The client should call POST /mcp/server/{id}/oauth-user-credential " + "to store it manually.", + mcp_server.server_id, + ) + result = { "access_token": access_token, "token_type": token_response.get("token_type", "Bearer"), diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 402e12d935..8d3831e75f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2455,6 +2455,37 @@ class MCPServerManager: ) tasks.append(during_hook_task) + # For per-user OAuth servers: if the client didn't supply a token in + # oauth2_headers, look up the stored token from Redis / DB. This is the + # call_tool equivalent of _get_user_oauth_extra_headers_from_db used in + # list_tools. + if ( + mcp_server.needs_user_oauth_token + and not oauth2_headers + and user_api_key_auth is not None + ): + user_id = getattr(user_api_key_auth, "user_id", None) + if user_id: + try: + from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + _get_user_oauth_extra_headers_from_db, + ) + + stored_headers = await _get_user_oauth_extra_headers_from_db( + server=mcp_server, + user_api_key_auth=user_api_key_auth, + ) + if stored_headers: + oauth2_headers = stored_headers + except Exception as _lookup_exc: + verbose_logger.debug( + "call_tool: per-user token lookup failed for " + "user=%s server=%s: %s", + user_id, + mcp_server.server_id, + _lookup_exc, + ) + # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: verbose_logger.debug( diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 84a2e94467..476e215666 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -17,8 +17,15 @@ from litellm.constants import ( MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_PER_USER_TOKEN_DEFAULT_TTL, + MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -152,6 +159,107 @@ class MCPOAuth2TokenCache(InMemoryCache): mcp_oauth2_token_cache = MCPOAuth2TokenCache() +def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int: + """Compute Redis TTL for a per-user token. + + Uses server.token_storage_ttl_seconds when configured; otherwise derives + TTL from expires_in minus the expiry buffer; falls back to the default TTL. + """ + if server.token_storage_ttl_seconds is not None: + return max(server.token_storage_ttl_seconds, 1) + if expires_in is not None: + return max( + expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, + 1, + ) + return MCP_PER_USER_TOKEN_DEFAULT_TTL + + +class MCPPerUserTokenCache: + """Redis-backed cache for per-user OAuth2 access tokens. + + Uses LiteLLM's existing ``user_api_key_cache`` (DualCache with optional + Redis backend). Tokens are NaCl-encrypted with ``encrypt_value_helper`` + before storage so they are safe at rest in Redis. + + Redis key format: ``mcp:per_user_token:{user_id}:{server_id}`` + Redis value: ``encrypt_value_helper(access_token)`` — URL-safe base64 + """ + + def _cache_key(self, user_id: str, server_id: str) -> str: + return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + + async def get(self, user_id: str, server_id: str) -> Optional[str]: + """Return the plaintext access_token, or None on miss/error.""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + encrypted = await user_api_key_cache.async_get_cache(key) + if encrypted is None: + return None + plaintext = decrypt_value_helper( + encrypted, + key="mcp_per_user_token", + exception_type="debug", + ) + return plaintext or None + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.get failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + return None + + async def set( + self, + user_id: str, + server_id: str, + access_token: str, + ttl: int, + ) -> None: + """Store NaCl-encrypted access_token in Redis with the given TTL.""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + encrypted = encrypt_value_helper(access_token) + await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl) + verbose_logger.debug( + "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds", + user_id, + server_id, + ttl, + ) + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.set failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + + async def delete(self, user_id: str, server_id: str) -> None: + """Invalidate the cached token (removes from both in-memory and Redis layers).""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + await user_api_key_cache.async_delete_cache(key) + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.delete failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + + +mcp_per_user_token_cache = MCPPerUserTokenCache() + + async def resolve_mcp_auth( server: "MCPServer", mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7fc28b68e9..99578d006e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -896,11 +896,17 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, ) -> Optional[Dict[str, str]]: - """Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict. + """Look up stored OAuth2 token for (user, server) and return as extra_headers dict. + + Lookup order: + 1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied + 2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query + 3. Auto-refresh when the stored token is expired and a refresh_token exists Args: prefetched_creds: Optional dict keyed by server_id with credential payloads. - When provided, avoids a per-server DB round-trip. + When provided, the Redis and individual DB lookups are + skipped in favour of the pre-fetched batch result. """ if server.auth_type != MCPAuth.oauth2: return None @@ -914,8 +920,27 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 get_user_oauth_credential, is_oauth_credential_expired, + refresh_user_oauth_token, + ) + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415 + _compute_per_user_token_ttl, + mcp_per_user_token_cache, ) + # ── Fast path: Redis cache ──────────────────────────────────────── + # Only used when prefetched_creds is not supplied (individual lookup). + if prefetched_creds is None: + cached_token = await mcp_per_user_token_cache.get(user_id, server_id) + if cached_token is not None: + verbose_logger.debug( + "_get_user_oauth_extra_headers_from_db: Redis hit for " + "user=%s server=%s", + user_id, + server_id, + ) + return {"Authorization": f"Bearer {cached_token}"} + + # ── Slow path: DB lookup ────────────────────────────────────────── if prefetched_creds is not None: cred = prefetched_creds.get(server_id) else: @@ -929,18 +954,83 @@ if MCP_AVAILABLE: cred = await get_user_oauth_credential( prisma_client, user_id, server_id ) - if cred and cred.get("access_token"): - if is_oauth_credential_expired(cred): - verbose_logger.debug( - f"_get_user_oauth_extra_headers_from_db: token expired for " - f"user={user_id} server={server_id}" - ) + + if not cred or not cred.get("access_token"): + return None + + if is_oauth_credential_expired(cred): + verbose_logger.debug( + "_get_user_oauth_extra_headers_from_db: token expired for " + "user=%s server=%s — attempting refresh", + user_id, + server_id, + ) + # Attempt token refresh; requires a DB client (not available from prefetch) + if cred.get("refresh_token"): + try: + from litellm.proxy.utils import ( # noqa: PLC0415 + get_prisma_client_or_throw, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Cannot refresh OAuth token." + ) + cred = await refresh_user_oauth_token( + prisma_client=prisma_client, + user_id=user_id, + server=server, + cred=cred, + ) + except Exception as refresh_exc: + verbose_logger.warning( + "_get_user_oauth_extra_headers_from_db: refresh failed " + "for user=%s server=%s: %s", + user_id, + server_id, + refresh_exc, + ) + cred = None + + if not cred or not cred.get("access_token"): + # Clear stale Redis/cache entry so we don't serve it again. + # Do this for both the individual and prefetch paths so the + # next request doesn't get a stale cache hit. + await mcp_per_user_token_cache.delete(user_id, server_id) return None - return {"Authorization": f"Bearer {cred['access_token']}"} + + access_token: str = cred["access_token"] + + # Warm (or re-warm) the Redis cache from the DB result. + # Always write regardless of whether expires_at is present — tokens + # without an expiry are still valid and should be cached using the + # server/default TTL so subsequent requests are fast. + if prefetched_creds is None: + raw_expires = None + expires_at = cred.get("expires_at") + if expires_at: + from datetime import datetime, timezone # noqa: PLC0415 + + try: + exp_dt = datetime.fromisoformat(expires_at) + if exp_dt.tzinfo is None: + exp_dt = exp_dt.replace(tzinfo=timezone.utc) + remaining = int( + (exp_dt - datetime.now(timezone.utc)).total_seconds() + ) + raw_expires = max(remaining, 0) if remaining > 0 else None + except (ValueError, TypeError): + pass + ttl = _compute_per_user_token_ttl(server, raw_expires) + await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl) + + return {"Authorization": f"Bearer {access_token}"} except Exception as e: verbose_logger.warning( - f"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " - f"user={user_id} server={server_id}: {e}" + "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " + "user=%s server=%s: %s", + user_id, + server_id, + e, ) return None @@ -2504,6 +2594,14 @@ if MCP_AVAILABLE: server_name, client_ip=_client_ip ) if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: + # For servers that store per-user tokens server-side, skip the + # pre-emptive 401 — the call_tool / list_tools dispatch will look + # up the stored token from Redis / DB and only fail at the MCP + # protocol level if none is found, giving the client a proper + # tool-execution error rather than an HTTP 401. + if server.needs_user_oauth_token: + continue + request = StarletteRequest(scope) base_url = get_request_base_url(request) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 793742891f..96e221a9ac 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -494,10 +494,12 @@ class LiteLLMRoutes(enum.Enum): "/v2/key/info", "/model_group/info", "/health", + "/health/services", "/key/list", "/user/filter/ui", "/models", "/v1/models", + "/sso/get/ui_settings", ] # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend @@ -566,6 +568,8 @@ class LiteLLMRoutes(enum.Enum): "/spend/tags", "/spend/calculate", "/spend/logs", + "/spend/logs/ui", + "/spend/logs/session/ui", "/cost/estimate", ] @@ -581,6 +585,7 @@ class LiteLLMRoutes(enum.Enum): "/global/spend/report", "/global/spend/provider", "/global/spend/tags", + "/global/spend/all_tag_names", ] public_routes = set( @@ -602,6 +607,9 @@ class LiteLLMRoutes(enum.Enum): ] ) + # Retained for backwards compatibility with JWT auth configs that reference + # "ui_routes" in admin_allowed_routes. Not used by the proxy's own route + # authorization — UI tokens now go through the same RBAC path as API tokens. ui_routes = [ "/sso", "/sso/get/ui_settings", @@ -627,19 +635,16 @@ class LiteLLMRoutes(enum.Enum): internal_user_routes = ( [ - "/global/spend/tags", - "/global/spend/keys", - "/global/spend/models", - "/global/spend/provider", - "/global/spend/end_users", "/global/activity", "/global/activity/model", + "/global/activity/cache_hits", "/v1/models/{model_id}", "/models/{model_id}", "/guardrails/list", "/v2/guardrails/list", ] + spend_tracking_routes + + global_spend_tracking_routes + key_management_routes ) @@ -694,6 +699,9 @@ class LiteLLMRoutes(enum.Enum): "/tag/list", "/audit", "/audit/{id}", + "/global/activity", + "/global/activity/model", + "/global/activity/cache_hits", ] + info_routes # All routes accesible by an Org Admin @@ -892,9 +900,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1036,9 +1044,9 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None - grace_period: Optional[str] = ( - None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke - ) + grace_period: Optional[ + str + ] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke class ResetSpendRequest(LiteLLMPydanticObjectBase): @@ -1562,12 +1570,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1590,12 +1598,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -1685,15 +1693,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1790,9 +1798,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -2134,9 +2142,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2495,9 +2503,9 @@ class UserAPIKeyAuth( user_max_budget: Optional[float] = None request_route: Optional[str] = None user: Optional[Any] = None # Expanded user object when expand=user is used - created_by_user: Optional[Any] = ( - None # Expanded created_by user when expand=user is used - ) + created_by_user: Optional[ + Any + ] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None # Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery # and forwarded into outbound tokens by guardrails such as MCPJWTSigner. @@ -2636,9 +2644,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available litellm_budget_table: Optional[LiteLLM_BudgetTable] = None user_email: Optional[str] = None @@ -3793,9 +3801,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -4050,9 +4058,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -4214,9 +4222,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 68bde8434a..56958a88f6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -196,9 +196,7 @@ def _is_model_cost_zero( return True -def _is_cost_explicitly_configured( - model: str, llm_router: "Router" -) -> bool: +def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: """ Check if any deployment in the model group has cost fields explicitly set in its litellm.model_cost entry. @@ -215,10 +213,7 @@ def _is_cost_explicitly_configured( if model_id is None: continue raw_entry = litellm.model_cost.get(model_id, {}) - if ( - "input_cost_per_token" in raw_entry - or "output_cost_per_token" in raw_entry - ): + if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: return True return False @@ -596,17 +591,12 @@ async def common_checks( # noqa: PLR0915 user_object=user_object, route=route, request_body=request_body ) - token_team = getattr(valid_token, "team_id", None) - token_type: Literal["ui", "api"] = ( - "ui" if token_team is not None and token_team == "litellm-dashboard" else "api" - ) - _is_route_allowed = _is_allowed_route( + _is_route_allowed = _is_api_route_allowed( route=route, - token_type=token_type, - user_obj=user_object, request=request, request_data=request_body, valid_token=valid_token, + user_obj=user_object, ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store @@ -629,31 +619,6 @@ async def common_checks( # noqa: PLR0915 return True -def _is_ui_route( - route: str, - user_obj: Optional[LiteLLM_UserTable] = None, -) -> bool: - """ - - Check if the route is a UI used route - """ - # this token is only used for managing the ui - allowed_routes = LiteLLMRoutes.ui_routes.value - # check if the current route startswith any of the allowed routes - if ( - route is not None - and isinstance(route, str) - and any(route.startswith(allowed_route) for allowed_route in allowed_routes) - ): - # Do something if the current route starts with any of the allowed routes - return True - elif any( - RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) - for allowed_route in allowed_routes - ): - return True - return False - - def _get_user_role( user_obj: Optional[LiteLLM_UserTable], ) -> Optional[LitellmUserRoles]: @@ -717,30 +682,6 @@ def _is_user_proxy_admin(user_obj: Optional[LiteLLM_UserTable]): return False -def _is_allowed_route( - route: str, - token_type: Literal["ui", "api"], - request: Request, - request_data: dict, - valid_token: Optional[UserAPIKeyAuth], - user_obj: Optional[LiteLLM_UserTable] = None, -) -> bool: - """ - - Route b/w ui token check and normal token check - """ - - if token_type == "ui" and _is_ui_route(route=route, user_obj=user_obj): - return True - else: - return _is_api_route_allowed( - route=route, - request=request, - request_data=request_data, - valid_token=valid_token, - user_obj=user_obj, - ) - - def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: """ Return if a user is allowed to access route. Helper function for `allowed_routes_check`. diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ece72b1060..2b8c16ed12 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -10,6 +10,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( AddTeamCallback, @@ -1264,6 +1265,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 user_api_key_dict=user_api_key_dict, ) + # Save pre-alias model name for credential override lookup + _pre_alias_model = data.get("model") + # Team Model Aliases _update_model_if_team_alias_exists( data=data, @@ -1280,6 +1284,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 "[PROXY] returned data from litellm_pre_call_utils: %s", data ) + # Team/Project credential overrides from model_config + # Placed after the debug log to avoid leaking credential secrets in logs + _apply_credential_overrides_from_model_config( + data=data, + user_api_key_dict=user_api_key_dict, + pre_alias_model_name=_pre_alias_model, + ) + ## ENFORCED PARAMS CHECK # loop through each enforced param # example enforced_params ['user', 'metadata', 'metadata.generation_name'] @@ -1407,6 +1419,175 @@ def _update_model_if_key_alias_exists( return +def _apply_credential_overrides_from_model_config( + data: dict, + user_api_key_dict: UserAPIKeyAuth, + pre_alias_model_name: Optional[str] = None, +) -> None: + """ + Walk the model_config precedence chain in team/project metadata. + If a matching credential is found, set api_base/api_key/api_version on data + so they override deployment defaults in the router. + + Precedence (highest to lowest): + 1. Clientside credentials (already in data — skip if present) + 2. Project model-specific override + 3. Project default override (defaultconfig) + 4. Team model-specific override + 5. Team default override (defaultconfig) + 6. Deployment default (no action needed) + """ + # Feature flag gate — disabled by default, opt in with litellm.enable_model_config_credential_overrides = True + if not litellm.enable_model_config_credential_overrides: + return + + # Respect clientside credentials — highest precedence + if data.get("api_base") is not None or data.get("api_key") is not None: + return + + model_name = data.get("model") + if not model_name: + return + + project_metadata = user_api_key_dict.project_metadata or {} + team_metadata = user_api_key_dict.team_metadata or {} + + project_model_config = project_metadata.get("model_config") + team_model_config = team_metadata.get("model_config") + + if not project_model_config and not team_model_config: + return + + # Extract provider hint from model name (e.g. "azure/gpt-4" -> "azure") + provider: Optional[str] = None + if "/" in model_name: + provider = model_name.split("/", 1)[0] + + credential_name = _resolve_credential_from_model_config( + model_name=model_name, + project_model_config=project_model_config, + team_model_config=team_model_config, + pre_alias_model_name=pre_alias_model_name, + provider=provider, + ) + + if not credential_name: + return + + credential_values = CredentialAccessor.get_credential_values(credential_name) + if not credential_values: + _safe_cred = str(credential_name).replace("\n", "").replace("\r", "") + verbose_proxy_logger.warning( + "model_config references credential '%s' but it was not found or has no values", + _safe_cred, + ) + return + + # Apply credential overrides only for keys not already in the request + for key in ("api_base", "api_key", "api_version"): + if key in credential_values and key not in data: + data[key] = credential_values[key] + + _safe_model = str(model_name).replace("\n", "").replace("\r", "") + _safe_cred = str(credential_name).replace("\n", "").replace("\r", "") + verbose_proxy_logger.debug( + "Applied credential override '%s' for model '%s'", + _safe_cred, + _safe_model, + ) + + +def _resolve_credential_from_model_config( + model_name: str, + project_model_config: Optional[dict], + team_model_config: Optional[dict], + pre_alias_model_name: Optional[str] = None, + provider: Optional[str] = None, +) -> Optional[str]: + """ + Walk the precedence chain and return the first matching credential name. + + Checks (in order): + 1. project_model_config[model_name][provider] — project model-specific + 2. project_model_config[pre_alias_model_name][provider] — project pre-alias + 3. project_model_config["defaultconfig"][provider] — project default + 4. team_model_config[model_name][provider] — team model-specific + 5. team_model_config[pre_alias_model_name][provider] — team pre-alias + 6. team_model_config["defaultconfig"][provider] — team default + + When a model-specific entry exists but contains no litellm_credentials, + the function falls through to defaultconfig. This is intentional — + an entry without litellm_credentials is treated as incomplete config, + not as an explicit "no override" signal. + """ + # Build the list of model names to try (post-alias first, then pre-alias) + model_names_to_try = [model_name] + if pre_alias_model_name and pre_alias_model_name != model_name: + model_names_to_try.append(pre_alias_model_name) + + for model_config in (project_model_config, team_model_config): + if not model_config or not isinstance(model_config, dict): + continue + + # Model-specific check (try resolved name, then pre-alias name) + for name in model_names_to_try: + model_entry = model_config.get(name) + if model_entry: + credential_name = _extract_credential_from_entry( + model_entry, provider=provider + ) + if credential_name: + return credential_name + _safe_name = str(name).replace("\n", "").replace("\r", "") + verbose_proxy_logger.debug( + "model_config entry '%s' found but has no litellm_credentials, " + "trying next candidate", + _safe_name, + ) + + # Default check + default_entry = model_config.get("defaultconfig") + if default_entry: + credential_name = _extract_credential_from_entry( + default_entry, provider=provider + ) + if credential_name: + return credential_name + + return None + + +def _extract_credential_from_entry( + entry: dict, provider: Optional[str] = None +) -> Optional[str]: + """ + Extract litellm_credentials from a model_config entry. + + Entry structure: {"azure": {"litellm_credentials": "name"}, ...} + + When provider is given (e.g. "azure"), tries an exact provider match first. + Falls back to the first credential found across all provider keys. + """ + if not isinstance(entry, dict): + return None + + # Prefer exact provider match when provider hint is available + if provider and provider in entry: + provider_config = entry[provider] + if isinstance(provider_config, dict): + credential_name = provider_config.get("litellm_credentials") + if credential_name: + return credential_name + + # Fall back to first available provider + for provider_config in entry.values(): + if isinstance(provider_config, dict): + credential_name = provider_config.get("litellm_credentials") + if credential_name: + return credential_name + return None + + def _get_enforced_params( general_settings: Optional[dict], user_api_key_dict: UserAPIKeyAuth ) -> Optional[list]: diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index db7657a017..a7d0968c0e 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -71,6 +71,15 @@ class MCPServer(BaseModel): # OAuth2 flow type. Defaults to None (interactive / authorization_code). # Set to "client_credentials" to enable M2M token fetching. oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None + # Per-user OAuth server-side storage config. + # token_validation: key-value pairs that must match fields in the OAuth token + # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). + # Tokens that fail validation are rejected before storage. + token_validation: Optional[Dict[str, Any]] = None + # Optional TTL override (seconds) for the Redis per-user token cache. + # Defaults to the token's expires_in minus the expiry buffer, or + # MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent. + token_storage_ttl_seconds: Optional[int] = None model_config = ConfigDict(arbitrary_types_allowed=True) @property diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index a1684e2376..001b946400 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -178,7 +178,6 @@ async def test_get_response(): async def test_aavertex_ai_anthropic_async(): # load_vertex_ai_credentials() try: - model = "claude-3-5-sonnet@20240620" vertex_ai_project = "pathrise-convert-1606954137718" @@ -351,7 +350,6 @@ def test_avertex_ai_stream(): @pytest.mark.flaky(retries=3, delay=1) @pytest.mark.asyncio async def test_async_vertexai_response_basic(): - load_vertex_ai_credentials() try: user_message = "Hello, how are you?" @@ -1382,7 +1380,6 @@ async def test_gemini_pro_json_schema_args_sent_httpx( ] ) elif resp is not None: - assert resp.model == model.split("/")[1] @@ -2291,6 +2288,8 @@ def test_prompt_factory_nested(): async def test_completion_fine_tuned_model(): load_vertex_ai_credentials() mock_response = AsyncMock() + mock_response.headers = {} + mock_response.status_code = 200 def return_val(): return { @@ -2326,7 +2325,6 @@ async def test_completion_fine_tuned_model(): } mock_response.json = return_val - mock_response.status_code = 200 expected_payload = { "contents": [ diff --git a/tests/mcp_tests/test_per_user_oauth_cache.py b/tests/mcp_tests/test_per_user_oauth_cache.py new file mode 100644 index 0000000000..36c26a5a50 --- /dev/null +++ b/tests/mcp_tests/test_per_user_oauth_cache.py @@ -0,0 +1,527 @@ +""" +Unit tests for per-user MCP OAuth token storage: +- MCPPerUserTokenCache (NaCl-encrypted Redis cache) +- _validate_token_response (token validation rules) +- _compute_per_user_token_ttl (TTL computation) +- refresh_user_oauth_token (token refresh flow) +""" + +import sys +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# Stub out modules that aren't available in the unit-test environment +# so we can import the targets without a full proxy stack. +for _mod in ("orjson",): + if _mod not in sys.modules: + sys.modules[_mod] = MagicMock() + +from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: E402 + MCPPerUserTokenCache, + _compute_per_user_token_ttl, + mcp_per_user_token_cache, +) +from litellm.types.mcp import MCPAuth, MCPTransport # noqa: E402 +from litellm.types.mcp_server.mcp_server_manager import MCPServer # noqa: E402 + + +def _import_validate(): + """Lazy import to avoid pulling orjson at collection time.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _validate_token_response, + ) + + return _validate_token_response + + +# ── Fixtures ───────────────────────────────────────────────────────────────── + + +def _make_server(**kwargs) -> MCPServer: + defaults: Dict[str, Any] = { + "server_id": "slack-test", + "name": "Slack", + "server_name": "slack", + "url": "https://slack-mcp.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "client_id": "SLACK_CLIENT_ID", + "client_secret": "SLACK_CLIENT_SECRET", + "token_url": "https://slack.com/api/oauth.v2.access", + "authorization_url": "https://slack.com/oauth/v2/authorize", + } + defaults.update(kwargs) + return MCPServer(**defaults) + + +# ── _validate_token_response ────────────────────────────────────────────────── + + +class TestValidateTokenResponse: + def test_passes_when_all_rules_match(self): + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "enterprise_id": "E04XXXXXX", + "team": {"id": "T123", "name": "Acme"}, + } + # Should not raise + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + + def test_raises_on_mismatch(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "xoxb-123", "enterprise_id": "E99999999"} + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + detail = exc_info.value.detail + assert detail["error"] == "token_validation_failed" + assert detail["field"] == "enterprise_id" + + def test_raises_when_field_absent(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "xoxb-123"} + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + # Absent field should produce a distinct "absent" message, not str(None) + assert "absent" in exc_info.value.detail["message"] + + def test_absent_field_does_not_match_string_none(self): + """str(None)='None' must NOT match the string rule value 'None'.""" + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "tok"} # enterprise_id absent + # Even if admin writes validation_rules={"enterprise_id": "None"}, absent + # field should raise, not pass. + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "None"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + assert "absent" in exc_info.value.detail["message"] + + def test_dot_notation_nested_field(self): + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "team": {"enterprise_id": "E04XXXXXX"}, + } + # Should not raise — dot-notation traverses nested dict + _validate_token_response( + token_response=token_response, + validation_rules={"team.enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + + def test_dot_notation_mismatch(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "team": {"enterprise_id": "WRONG"}, + } + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"team.enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["field"] == "team.enterprise_id" + + def test_numeric_value_string_coercion(self): + """Numeric values in token response should match string rules.""" + _validate_token_response = _import_validate() + token_response = {"access_token": "tok", "org_id": 12345} + # Should not raise — str(12345) == "12345" + _validate_token_response( + token_response=token_response, + validation_rules={"org_id": "12345"}, + server_id="test", + ) + + def test_multiple_rules_all_must_match(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = { + "access_token": "tok", + "enterprise_id": "E04XXXXXX", + "cloud_id": "WRONG_CLOUD", + } + with pytest.raises(HTTPException): + _validate_token_response( + token_response=token_response, + validation_rules={ + "enterprise_id": "E04XXXXXX", + "cloud_id": "abc-123", + }, + server_id="atlassian", + ) + + +# ── _compute_per_user_token_ttl ────────────────────────────────────────────── + + +class TestComputePerUserTokenTtl: + def test_uses_server_override_when_set(self): + server = _make_server(token_storage_ttl_seconds=7200) + assert _compute_per_user_token_ttl(server, expires_in=99999) == 7200 + + def test_uses_expires_in_minus_buffer(self): + server = _make_server() + # Default buffer is 60s + ttl = _compute_per_user_token_ttl(server, expires_in=3600) + assert ttl == 3600 - 60 + + def test_minimum_ttl_is_1(self): + server = _make_server() + # expires_in smaller than buffer → clamp to 1 + ttl = _compute_per_user_token_ttl(server, expires_in=30) + assert ttl == 1 + + def test_default_ttl_when_expires_in_none(self): + from litellm.constants import MCP_PER_USER_TOKEN_DEFAULT_TTL + + server = _make_server() + ttl = _compute_per_user_token_ttl(server, expires_in=None) + assert ttl == MCP_PER_USER_TOKEN_DEFAULT_TTL + + +# ── MCPPerUserTokenCache ────────────────────────────────────────────────────── + + +class TestMCPPerUserTokenCache: + """Tests for Redis-backed per-user token cache. + + Patches ``user_api_key_cache`` to avoid needing a real Redis instance. + Patches ``encrypt_value_helper`` / ``decrypt_value_helper`` to verify + encryption is applied before Redis writes and decryption after reads. + """ + + @pytest.fixture + def cache(self): + return MCPPerUserTokenCache() + + @pytest.fixture + def mock_dual_cache(self): + dc = MagicMock() + dc.async_get_cache = AsyncMock(return_value=None) + dc.async_set_cache = AsyncMock() + return dc + + @pytest.mark.asyncio + async def test_get_returns_none_on_miss(self, cache, mock_dual_cache): + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper" + ) as mock_decrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = None + result = await cache.get("alice", "slack-test") + assert result is None + mock_decrypt.assert_not_called() + + @pytest.mark.asyncio + async def test_get_decrypts_cached_value(self, cache, mock_dual_cache): + fake_encrypted = "encrypted_blob_abc123" + fake_plaintext = "xoxb-slack-token" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper", + return_value=fake_plaintext, + ) as mock_decrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = fake_encrypted + result = await cache.get("alice", "slack-test") + + assert result == fake_plaintext + mock_decrypt.assert_called_once_with( + fake_encrypted, + key="mcp_per_user_token", + exception_type="debug", + ) + + @pytest.mark.asyncio + async def test_set_encrypts_before_storing(self, cache, mock_dual_cache): + fake_encrypted = "encrypted_blob_xyz" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value=fake_encrypted, + ) as mock_encrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.set("alice", "slack-test", "xoxb-token", ttl=3540) + + mock_encrypt.assert_called_once_with("xoxb-token") + mock_dual_cache.async_set_cache.assert_called_once() + call_kwargs = mock_dual_cache.async_set_cache.call_args + assert call_kwargs[0][1] == fake_encrypted # encrypted value stored + assert call_kwargs[1]["ttl"] == 3540 + + @pytest.mark.asyncio + async def test_set_uses_correct_cache_key(self, cache, mock_dual_cache): + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value="enc", + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.set("bob", "github-server", "ghp_token", ttl=3600) + + key_used = mock_dual_cache.async_set_cache.call_args[0][0] + assert key_used == "mcp:per_user_token:bob:github-server" + + @pytest.mark.asyncio + async def test_delete_calls_async_delete_cache(self, cache, mock_dual_cache): + mock_dual_cache.async_delete_cache = AsyncMock() + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.delete("alice", "slack-test") + + mock_dual_cache.async_delete_cache.assert_called_once_with( + "mcp:per_user_token:alice:slack-test" + ) + mock_dual_cache.async_set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_get_returns_none_on_decrypt_failure(self, cache, mock_dual_cache): + """Cache misses and decrypt errors should both return None without raising.""" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper", + return_value=None, # decrypt returns None on failure + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = "bad_encrypted_data" + result = await cache.get("alice", "slack-test") + + assert result is None + + @pytest.mark.asyncio + async def test_set_is_noop_on_cache_error(self, cache, mock_dual_cache): + """Errors in the cache layer must not propagate to the caller.""" + mock_dual_cache.async_set_cache.side_effect = RuntimeError("Redis down") + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value="enc", + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + # Should not raise + await cache.set("alice", "slack-test", "token", ttl=3600) + + +# ── refresh_user_oauth_token ────────────────────────────────────────────────── + + +class TestRefreshUserOauthToken: + """Tests for the DB-level token refresh helper.""" + + @pytest.fixture + def server(self): + return _make_server() + + @pytest.fixture + def cred(self): + return { + "type": "oauth2", + "access_token": "OLD_TOKEN", + "refresh_token": "REFRESH_TOKEN_123", + "expires_at": ( + datetime.now(timezone.utc) - timedelta(hours=1) + ).isoformat(), + } + + @pytest.mark.asyncio + async def test_returns_none_when_no_refresh_token(self, server): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + cred = {"type": "oauth2", "access_token": "OLD"} # no refresh_token + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_returns_none_when_no_token_url(self, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + server = _make_server(token_url=None) + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_returns_none_on_http_error(self, server, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + mock_client = AsyncMock() + mock_client.post.side_effect = Exception("Connection refused") + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ): + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_stores_and_returns_new_credential(self, server, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + new_token_response = MagicMock() + new_token_response.json.return_value = { + "access_token": "NEW_TOKEN", + "expires_in": 3600, + "refresh_token": "NEW_REFRESH", + "scope": "channels:read chat:write", + } + new_token_response.raise_for_status = MagicMock() + + mock_client = AsyncMock() + mock_client.post.return_value = new_token_response + + stored_cred = { + "type": "oauth2", + "access_token": "NEW_TOKEN", + "refresh_token": "NEW_REFRESH", + } + mock_prisma = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ), patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new_callable=AsyncMock, + ) as mock_store, patch( + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new_callable=AsyncMock, + return_value=stored_cred, + ): + result = await refresh_user_oauth_token( + prisma_client=mock_prisma, + user_id="alice", + server=server, + cred=cred, + ) + + assert result == stored_cred + mock_store.assert_called_once() + call_kwargs = mock_store.call_args[1] + assert call_kwargs["access_token"] == "NEW_TOKEN" + assert call_kwargs["refresh_token"] == "NEW_REFRESH" + assert call_kwargs["expires_in"] == 3600 + assert call_kwargs["scopes"] == ["channels:read", "chat:write"] + # Refresh path must skip the BYOK guard (row is already OAuth2) + assert call_kwargs.get("skip_byok_guard") is True + + @pytest.mark.asyncio + async def test_falls_back_to_old_refresh_token_when_not_rotated( + self, server, cred + ): + """When provider doesn't return a new refresh_token, keep the old one.""" + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + new_token_response = MagicMock() + new_token_response.json.return_value = { + "access_token": "NEW_TOKEN", + "expires_in": 3600, + # No refresh_token in response + } + new_token_response.raise_for_status = MagicMock() + + mock_client = AsyncMock() + mock_client.post.return_value = new_token_response + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ), patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new_callable=AsyncMock, + ) as mock_store, patch( + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new_callable=AsyncMock, + return_value={"type": "oauth2", "access_token": "NEW_TOKEN"}, + ): + await refresh_user_oauth_token( + prisma_client=AsyncMock(), + user_id="alice", + server=server, + cred=cred, + ) + + call_kwargs = mock_store.call_args[1] + # Old refresh_token preserved when provider doesn't rotate + assert call_kwargs["refresh_token"] == "REFRESH_TOKEN_123" + + +# ── MCPServer new fields ────────────────────────────────────────────────────── + + +class TestMCPServerNewFields: + def test_token_validation_default_none(self): + server = _make_server() + assert server.token_validation is None + + def test_token_validation_set(self): + server = _make_server(token_validation={"enterprise_id": "E04XXXXXX"}) + assert server.token_validation == {"enterprise_id": "E04XXXXXX"} + + def test_token_storage_ttl_default_none(self): + server = _make_server() + assert server.token_storage_ttl_seconds is None + + def test_token_storage_ttl_set(self): + server = _make_server(token_storage_ttl_seconds=7200) + assert server.token_storage_ttl_seconds == 7200 + + def test_needs_user_oauth_token_true_for_oauth2_without_m2m(self): + server = _make_server(auth_type=MCPAuth.oauth2) + assert server.needs_user_oauth_token is True + + def test_needs_user_oauth_token_false_for_m2m(self): + server = _make_server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + ) + assert server.needs_user_oauth_token is False diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index a5be1a3a42..9a8d6d3702 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -934,10 +934,7 @@ async def mock_user_object(*args, **kwargs): user_id = kwargs.get("user_id") user_email = kwargs.get("user_email") return LiteLLM_UserTable( - spend=0, - user_id=user_id, - max_budget=None, - user_email=user_email + spend=0, user_id=user_id, max_budget=None, user_email=user_email ) @@ -1170,15 +1167,13 @@ async def test_end_user_jwt_auth(monkeypatch): # use generated key to auth in from litellm import Router from litellm.types.router import RouterGeneralSettings - + # Create a router with pass_through_all_models enabled router = Router( model_list=[], - router_general_settings=RouterGeneralSettings( - pass_through_all_models=True - ), + router_general_settings=RouterGeneralSettings(pass_through_all_models=True), ) - + setattr(litellm.proxy.proxy_server, "premium_user", True) setattr( litellm.proxy.proxy_server, @@ -1196,7 +1191,7 @@ async def test_end_user_jwt_auth(monkeypatch): cost_tracking() result = await user_api_key_auth(request=request, api_key=bearer_token) - + # Assert that end_user_id is correctly extracted from JWT token's 'sub' field assert result.end_user_id == "81b3e52a-67a6-4efb-9645-70527e101479" @@ -1228,7 +1223,9 @@ async def test_end_user_jwt_auth(monkeypatch): ), ) - with patch("litellm.acompletion", new=AsyncMock(return_value=mock_response)) as mock_completion: + with patch( + "litellm.acompletion", new=AsyncMock(return_value=mock_response) + ) as mock_completion: resp = await chat_completion( request=request, fastapi_response=temp_response, @@ -1243,10 +1240,13 @@ async def test_end_user_jwt_auth(monkeypatch): # Verify the completion was called with correct end_user_id mock_completion.assert_called_once() call_kwargs = mock_completion.call_args.kwargs - + # end_user_id is passed in metadata as 'user_api_key_end_user_id' metadata = call_kwargs.get("metadata", {}) - assert metadata.get("user_api_key_end_user_id") == "81b3e52a-67a6-4efb-9645-70527e101479" + assert ( + metadata.get("user_api_key_end_user_id") + == "81b3e52a-67a6-4efb-9645-70527e101479" + ) def test_can_rbac_role_call_route(): @@ -1278,13 +1278,13 @@ def test_user_api_key_auth_jwt_hashing(): """ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.handle_jwt import JWTHandler - + # Test with a JWT token (3 parts separated by dots) jwt_token = "test-jwt-token-header.payload.signature" - + # Create UserAPIKeyAuth instance with JWT user_auth = UserAPIKeyAuth(api_key=jwt_token) - + # Verify that the API key is hashed with "hashed-jwt-" prefix # critical - the raw JWT token should not be in the api_key or token assert user_auth.api_key.startswith("hashed-jwt-") @@ -1292,19 +1292,18 @@ def test_user_api_key_auth_jwt_hashing(): assert jwt_token not in user_auth.api_key assert jwt_token not in user_auth.token - # Test with a regular API key (should not be hashed) regular_api_key = "sk-1234567890abcdef" user_auth_regular = UserAPIKeyAuth(api_key=regular_api_key) - + # Verify that regular API key is hashed normally (without "hashed-jwt-" prefix) assert not user_auth_regular.api_key.startswith("hashed-jwt-") assert not user_auth_regular.token.startswith("hashed-jwt-") - + # Test with a non-JWT, non-sk string (should not be hashed) non_jwt_key = "some-random-key" user_auth_non_jwt = UserAPIKeyAuth(api_key=non_jwt_key) - + # Verify that non-JWT key is not hashed assert user_auth_non_jwt.api_key == non_jwt_key assert user_auth_non_jwt.token == non_jwt_key @@ -1315,19 +1314,19 @@ def test_jwt_handler_is_jwt_static_method(): Test that JWTHandler.is_jwt is a static method and works correctly """ from litellm.proxy.auth.handle_jwt import JWTHandler - + # Test with valid JWT format valid_jwt = "test-jwt-token-header.payload.signature" assert JWTHandler.is_jwt(valid_jwt) == True - + # Test with invalid JWT format (only 2 parts) invalid_jwt = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ" assert JWTHandler.is_jwt(invalid_jwt) == False - + # Test with regular API key regular_key = "sk-1234567890abcdef" assert JWTHandler.is_jwt(regular_key) == False - + # Test with empty string assert JWTHandler.is_jwt("") == False @@ -1461,7 +1460,13 @@ async def test_auth_jwt_es256_jwk_path(monkeypatch): now = int(time.time()) token = jwt.encode( - {"sub": "alice", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + { + "sub": "alice", + "aud": "litellm-proxy", + "iss": "http://example", + "iat": now, + "exp": now + 300, + }, ec_priv_pem, algorithm="ES256", headers={"kid": "ec1"}, @@ -1508,7 +1513,13 @@ async def test_auth_jwt_rs256_regression(monkeypatch): now = int(time.time()) token = jwt.encode( - {"sub": "bob", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + { + "sub": "bob", + "aud": "litellm-proxy", + "iss": "http://example", + "iat": now, + "exp": now + 300, + }, rsa_priv_pem, algorithm="RS256", headers={"kid": "rsa1"}, @@ -1540,7 +1551,13 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch): ) now = int(time.time()) token = jwt.encode( - {"sub": "mallory", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + { + "sub": "mallory", + "aud": "litellm-proxy", + "iss": "http://example", + "iat": now, + "exp": now + 300, + }, ec_priv_pem, algorithm="ES256", headers={"kid": "ec1"}, @@ -1566,4 +1583,4 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch): with patch.object(h, "get_public_key", new=AsyncMock(return_value=rsa_jwk)): with pytest.raises(Exception) as exc: await h.auth_jwt(token) - assert "Validation fails" in str(exc.value) \ No newline at end of file + assert "Validation fails" in str(exc.value) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 1a6e2eda9a..75f0d5e319 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -359,27 +359,38 @@ async def test_auth_with_allowed_routes(route, should_raise_error): @pytest.mark.parametrize( - "route, user_role, expected_result", + "route, user_role, should_be_allowed", [ - # Proxy Admin checks + # Admin can access everything + ("/config/update", "proxy_admin", True), ("/global/spend/logs", "proxy_admin", True), - ("/key/delete", "proxy_admin", False), - ("/key/generate", "proxy_admin", False), - ("/key/regenerate", "proxy_admin", False), - # Internal User checks - allowed routes + ("/global/activity/cache_hits", "proxy_admin", True), + # Internal User - allowed read-only routes ("/global/spend/logs", "internal_user", True), - ("/key/delete", "internal_user", False), - ("/key/generate", "internal_user", False), - ("/key/82akk800000000jjsk/regenerate", "internal_user", False), - # Internal User Viewer - ("/key/generate", "internal_user_viewer", False), - # Internal User checks - disallowed routes + ("/spend/logs/ui", "internal_user", True), + ("/global/activity/cache_hits", "internal_user", True), + ("/health/services", "internal_user", True), + # Internal User - BLOCKED from admin routes (security fix) + ("/config/update", "internal_user", False), + ("/config/pass_through_endpoint", "internal_user", False), + ("/config/field/update", "internal_user", False), ("/organization/member_add", "internal_user", False), + # Internal User Viewer - allowed spend routes only + ("/spend/logs/ui", "internal_user_viewer", True), + ("/global/spend/all_tag_names", "internal_user_viewer", True), + # Internal User Viewer - blocked from admin routes + ("/config/update", "internal_user_viewer", False), + ("/key/generate", "internal_user_viewer", False), ], ) -def test_is_ui_route_allowed(route, user_role, expected_result): - from litellm.proxy.auth.auth_checks import _is_ui_route - from litellm.proxy._types import LiteLLM_UserTable +def test_ui_token_route_access(route, user_role, should_be_allowed): + """ + Verify that UI tokens (team_id=litellm-dashboard) go through the same + RBAC checks as API tokens. Non-admin dashboard users must not be able + to access admin-only routes like /config/update. + """ + from litellm.proxy.auth.auth_checks import _is_api_route_allowed + from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth user_obj = LiteLLM_UserTable( user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", @@ -395,18 +406,36 @@ def test_is_ui_route_allowed(route, user_role, expected_result): organization_memberships=[], ) - received_args: dict = { - "route": route, - "user_obj": user_obj, - } - try: - assert _is_ui_route(**received_args) == expected_result - except Exception as e: - # If expected result is False, we expect an error - if expected_result is False: - pass - else: - raise e + valid_token = UserAPIKeyAuth( + user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", + team_id="litellm-dashboard", + user_role=user_role, + ) + + from starlette.datastructures import URL + from fastapi import Request + + request = Request(scope={"type": "http"}) + request._url = URL(url=route) + + if should_be_allowed: + result = _is_api_route_allowed( + route=route, + request=request, + request_data={}, + valid_token=valid_token, + user_obj=user_obj, + ) + assert result is True + else: + with pytest.raises(Exception): + _is_api_route_allowed( + route=route, + request=request, + request_data={}, + valid_token=valid_token, + user_obj=user_obj, + ) @pytest.mark.parametrize( @@ -684,7 +713,7 @@ async def test_soft_budget_alert(): def test_is_allowed_route(): - from litellm.proxy.auth.auth_checks import _is_allowed_route + from litellm.proxy.auth.auth_checks import _is_api_route_allowed from litellm.proxy._types import UserAPIKeyAuth import datetime @@ -692,7 +721,6 @@ def test_is_allowed_route(): args = { "route": "/embeddings", - "token_type": "api", "request": request, "request_data": {"input": ["hello world"], "model": "embedding-small"}, "valid_token": UserAPIKeyAuth( @@ -752,7 +780,7 @@ def test_is_allowed_route(): "user_obj": None, } - assert _is_allowed_route(**args) + assert _is_api_route_allowed(**args) @pytest.mark.parametrize( @@ -836,7 +864,6 @@ async def test_user_api_key_auth_websocket(): with patch( "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True ) as mock_user_api_key_auth: - # Make the call to the WebSocket function await user_api_key_auth_websocket(mock_websocket) @@ -845,10 +872,14 @@ async def test_user_api_key_auth_websocket(): # Get the request object that was passed to user_api_key_auth request_arg = mock_user_api_key_auth.call_args.kwargs["request"] - + # Verify that the request has headers set - assert hasattr(request_arg, "headers"), "Request object should have headers attribute" - assert "authorization" in request_arg.headers, "Request headers should contain authorization" + assert hasattr( + request_arg, "headers" + ), "Request object should have headers attribute" + assert ( + "authorization" in request_arg.headers + ), "Request headers should contain authorization" assert request_arg.headers["authorization"] == "Bearer some_api_key" assert ( @@ -1036,7 +1067,10 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): # Create request request = Request( - scope={"type": "http", "headers": [(b"authorization", b"Bearer fake.jwt.token")]} + scope={ + "type": "http", + "headers": [(b"authorization", b"Bearer fake.jwt.token")], + } ) request._url = URL(url="/team/new") @@ -1101,14 +1135,14 @@ async def test_x_litellm_api_key(): ignored_key = "aj12445" # Create request with headers as bytes - request = Request( - scope={ - "type": "http" - } - ) + request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - valid_token = await user_api_key_auth(request=request, api_key="Bearer " + ignored_key, custom_litellm_key_header=master_key) + valid_token = await user_api_key_auth( + request=request, + api_key="Bearer " + ignored_key, + custom_litellm_key_header=master_key, + ) assert valid_token.token == hash_token(master_key) @@ -1123,7 +1157,9 @@ async def test_user_api_key_from_query_param(): from litellm.proxy.proxy_server import hash_token, user_api_key_cache user_key = "sk-query-1234" - user_api_key_cache.set_cache(key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key))) + user_api_key_cache.set_cache( + key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key)) + ) setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -1136,7 +1172,9 @@ async def test_user_api_key_from_query_param(): "query_string": f"alt=sse&key={user_key}".encode(), } ) - request._url = URL(url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}") + request._url = URL( + url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}" + ) async def return_body(): return b"{}" @@ -1145,4 +1183,3 @@ async def test_user_api_key_from_query_param(): valid_token = await user_api_key_auth(request=request, api_key="") assert valid_token.token == hash_token(user_key) - diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 04af5cd008..cf7e71b14d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -13,14 +13,18 @@ from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup, + _apply_credential_overrides_from_model_config, + _extract_credential_from_entry, _get_dynamic_logging_metadata, _get_enforced_params, _get_metadata_variable_name, + _resolve_credential_from_model_config, _update_model_if_key_alias_exists, add_guardrails_from_policy_engine, add_litellm_data_to_request, check_if_token_is_service_account, ) +from litellm.types.utils import CredentialItem sys.path.insert( 0, os.path.abspath("../../..") @@ -1912,3 +1916,542 @@ async def test_bearer_token_not_in_debug_logs(): f"Bearer token leaked in debug logs. " f"Found token in log output:\n{log_output[:500]}" ) + + +# ============================================================================ +# Tests for credential overrides from model_config (team/project metadata) +# ============================================================================ + + +@pytest.fixture() +def setup_test_credentials(): + """Populate litellm.credential_list with test credentials and enable feature flag, clean up after.""" + original = litellm.credential_list[:] + original_flag = litellm.enable_model_config_credential_overrides + litellm.enable_model_config_credential_overrides = True + litellm.credential_list.extend( + [ + CredentialItem( + credential_name="hotel-azure-eastus", + credential_info={}, + credential_values={ + "api_base": "https://hotel-eastus.openai.azure.com/", + "api_key": "key-hotel-eastus", + }, + ), + CredentialItem( + credential_name="hotel-azure-westus", + credential_info={}, + credential_values={ + "api_base": "https://hotel-westus.openai.azure.com/", + "api_key": "key-hotel-westus", + }, + ), + CredentialItem( + credential_name="hotel-rec-azure", + credential_info={}, + credential_values={ + "api_base": "https://hotel-rec-app.openai.azure.com/", + "api_key": "key-hotel-rec", + }, + ), + CredentialItem( + credential_name="hotel-rec-vision", + credential_info={}, + credential_values={ + "api_base": "https://hotel-rec-vision.openai.azure.com/", + "api_key": "key-hotel-rec-vision", + "api_version": "2024-06-01", + }, + ), + CredentialItem( + credential_name="flight-azure-centralus", + credential_info={}, + credential_values={ + "api_base": "https://flight-centralus.openai.azure.com/", + "api_key": "key-flight-centralus", + }, + ), + ] + ) + yield + litellm.credential_list[:] = original + litellm.enable_model_config_credential_overrides = original_flag + + +# --- Unit tests for _extract_credential_from_entry --- + + +def test_extract_credential_from_entry_azure(): + entry = {"azure": {"litellm_credentials": "my-cred"}} + assert _extract_credential_from_entry(entry) == "my-cred" + + +def test_extract_credential_from_entry_no_credential(): + entry = {"azure": {"some_other_key": "value"}} + assert _extract_credential_from_entry(entry) is None + + +def test_extract_credential_from_entry_empty(): + assert _extract_credential_from_entry({}) is None + + +def test_extract_credential_from_entry_non_dict_value(): + entry = {"azure": "not-a-dict"} + assert _extract_credential_from_entry(entry) is None + + +def test_extract_credential_from_entry_non_dict_entry(): + """Non-dict entry (e.g. string) should return None, not crash.""" + assert _extract_credential_from_entry("my-cred-name") is None + assert _extract_credential_from_entry(["a", "list"]) is None + assert _extract_credential_from_entry(42) is None + + +# --- Unit tests for _resolve_credential_from_model_config --- + + +def test_resolve_project_model_specific_wins(): + project_config = { + "gpt-4": {"azure": {"litellm_credentials": "proj-gpt4"}}, + "defaultconfig": {"azure": {"litellm_credentials": "proj-default"}}, + } + team_config = { + "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, + "defaultconfig": {"azure": {"litellm_credentials": "team-default"}}, + } + result = _resolve_credential_from_model_config( + "gpt-4", project_config, team_config + ) + assert result == "proj-gpt4" + + +def test_resolve_project_default_wins_over_team(): + project_config = { + "defaultconfig": {"azure": {"litellm_credentials": "proj-default"}}, + } + team_config = { + "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, + "defaultconfig": {"azure": {"litellm_credentials": "team-default"}}, + } + result = _resolve_credential_from_model_config( + "gpt-4", project_config, team_config + ) + assert result == "proj-default" + + +def test_resolve_team_model_specific_wins_over_team_default(): + team_config = { + "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, + "defaultconfig": {"azure": {"litellm_credentials": "team-default"}}, + } + result = _resolve_credential_from_model_config("gpt-4", None, team_config) + assert result == "team-gpt4" + + +def test_resolve_team_default_used_as_fallback(): + team_config = { + "defaultconfig": {"azure": {"litellm_credentials": "team-default"}}, + } + result = _resolve_credential_from_model_config("gpt-3.5", None, team_config) + assert result == "team-default" + + +def test_resolve_no_match_returns_none(): + result = _resolve_credential_from_model_config("gpt-4", None, None) + assert result is None + + +def test_resolve_empty_configs_returns_none(): + result = _resolve_credential_from_model_config("gpt-4", {}, {}) + assert result is None + + +def test_resolve_model_not_in_any_config(): + project_config = {"gpt-4": {"azure": {"litellm_credentials": "x"}}} + result = _resolve_credential_from_model_config("gpt-3.5", project_config, None) + assert result is None + + +# --- Integration tests for _apply_credential_overrides_from_model_config --- + + +def test_apply_overrides_project_model_specific(setup_test_credentials): + """Scenario 2: Hotel Rec App -> gpt-4-vision -> project model-specific.""" + data = {"model": "gpt-4-vision"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-azure-eastus"} + }, + "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, + } + }, + project_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-rec-azure"} + }, + "gpt-4-vision": { + "azure": {"litellm_credentials": "hotel-rec-vision"} + }, + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" + assert data["api_key"] == "key-hotel-rec-vision" + assert data["api_version"] == "2024-06-01" + + +def test_apply_overrides_project_default(setup_test_credentials): + """Scenario 1: Hotel Rec App -> gpt-4 -> project default.""" + data = {"model": "gpt-4"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-azure-eastus"} + }, + "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, + } + }, + project_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-rec-azure"} + }, + "gpt-4-vision": { + "azure": {"litellm_credentials": "hotel-rec-vision"} + }, + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert data["api_base"] == "https://hotel-rec-app.openai.azure.com/" + assert data["api_key"] == "key-hotel-rec" + + +def test_apply_overrides_team_model_specific(setup_test_credentials): + """Scenario 4: Hotel Review App -> gpt-4 -> team model-specific.""" + data = {"model": "gpt-4"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-azure-eastus"} + }, + "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, + } + }, + project_metadata={}, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert data["api_base"] == "https://hotel-westus.openai.azure.com/" + assert data["api_key"] == "key-hotel-westus" + + +def test_apply_overrides_team_default(setup_test_credentials): + """Scenario 3: Hotel Review App -> gpt-3.5 -> team default.""" + data = {"model": "gpt-3.5"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-azure-eastus"} + }, + "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, + } + }, + project_metadata={}, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" + assert data["api_key"] == "key-hotel-eastus" + + +def test_apply_overrides_no_config(setup_test_credentials): + """Scenario 6: No model_config anywhere -> data unchanged.""" + data = {"model": "gpt-4"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={}, + project_metadata={}, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert "api_base" not in data + assert "api_key" not in data + + +def test_apply_overrides_clientside_credentials_take_precedence( + setup_test_credentials, +): + """Clientside api_base/api_key in data should block model_config override.""" + data = { + "model": "gpt-4", + "api_base": "https://my-custom-endpoint.openai.azure.com/", + "api_key": "my-custom-key", + } + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-azure-eastus"} + } + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert data["api_base"] == "https://my-custom-endpoint.openai.azure.com/" + assert data["api_key"] == "my-custom-key" + + +def test_apply_overrides_missing_credential_name(setup_test_credentials): + """model_config references a credential that doesn't exist -> no override.""" + data = {"model": "gpt-4"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "gpt-4": { + "azure": {"litellm_credentials": "nonexistent-credential"} + } + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert "api_base" not in data + assert "api_key" not in data + + +def test_apply_overrides_api_version_only_if_present(setup_test_credentials): + """api_version should only be set if the credential contains it.""" + data = {"model": "gpt-3.5"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "hotel-azure-eastus"} + } + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" + assert data["api_key"] == "key-hotel-eastus" + assert "api_version" not in data + + +def test_apply_overrides_no_model_in_data(setup_test_credentials): + """No model in request data -> skip override.""" + data = {"messages": [{"role": "user", "content": "hello"}]} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "defaultconfig": { + "azure": {"litellm_credentials": "some-cred"} + } + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert "api_base" not in data + + +def test_apply_overrides_none_metadata(setup_test_credentials): + """None metadata on both team and project -> skip override.""" + data = {"model": "gpt-4"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata=None, + project_metadata=None, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert "api_base" not in data + + +def test_apply_overrides_clientside_api_version_preserved(setup_test_credentials): + """Clientside api_version should not be overwritten by credential.""" + data = {"model": "gpt-4-vision", "api_version": "2025-01-01"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "gpt-4-vision": { + "azure": {"litellm_credentials": "hotel-rec-vision"} + } + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + # api_base and api_key should be set from credential + assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" + assert data["api_key"] == "key-hotel-rec-vision" + # api_version should be preserved from the request, not overwritten + assert data["api_version"] == "2025-01-01" + + +def test_resolve_non_dict_model_config_ignored(): + """Non-dict model_config (e.g. string) should be safely skipped.""" + result = _resolve_credential_from_model_config("gpt-4", "not-a-dict", None) + assert result is None + + result = _resolve_credential_from_model_config( + "gpt-4", None, ["also", "not", "a", "dict"] + ) + assert result is None + + # Valid config still works alongside invalid one + result = _resolve_credential_from_model_config( + "gpt-4", + "invalid", + {"gpt-4": {"azure": {"litellm_credentials": "valid-cred"}}}, + ) + assert result == "valid-cred" + + +def test_resolve_pre_alias_model_name_fallback(): + """model_config keyed on pre-alias name should match after alias resolution.""" + team_config = { + "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, + } + # Post-alias name doesn't match, but pre-alias does (team scope) + result = _resolve_credential_from_model_config( + "azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4" + ) + assert result == "team-gpt4" + + # Same test for project scope + project_config = { + "gpt-4": {"azure": {"litellm_credentials": "proj-gpt4"}}, + } + result = _resolve_credential_from_model_config( + "azure/gpt-4-0613", project_config, None, pre_alias_model_name="gpt-4" + ) + assert result == "proj-gpt4" + + +def test_resolve_post_alias_name_takes_priority(): + """Post-alias (resolved) name should be tried before pre-alias name.""" + team_config = { + "gpt-4": {"azure": {"litellm_credentials": "pre-alias-cred"}}, + "gpt-4o-team-1": {"azure": {"litellm_credentials": "post-alias-cred"}}, + } + # Team scope + result = _resolve_credential_from_model_config( + "gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4" + ) + assert result == "post-alias-cred" + + # Project scope + result = _resolve_credential_from_model_config( + "gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4" + ) + assert result == "post-alias-cred" + + +def test_apply_overrides_with_alias(setup_test_credentials): + """Credential override should work when model name was changed by alias.""" + # Simulate: user called "my-gpt4", alias resolved to "azure/gpt-4-custom" + # model_config is keyed on "my-gpt4" (the pre-alias name) + data = {"model": "azure/gpt-4-custom"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "my-gpt4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, + user_api_key_dict=user_api_key_dict, + pre_alias_model_name="my-gpt4", + ) + assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" + assert data["api_key"] == "key-hotel-eastus" + + +def test_apply_overrides_feature_flag_disabled_by_default(): + """Feature flag defaults to False — credential overrides are inert until explicitly enabled.""" + assert litellm.enable_model_config_credential_overrides is False + data = {"model": "gpt-4"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_metadata={ + "model_config": { + "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}} + } + }, + ) + _apply_credential_overrides_from_model_config( + data=data, user_api_key_dict=user_api_key_dict + ) + assert "api_base" not in data + assert "api_key" not in data + + +def test_extract_credential_provider_hint_prefers_exact_match(): + """Provider hint selects the correct provider in a multi-provider entry.""" + entry = { + "openai": {"litellm_credentials": "openai-cred"}, + "azure": {"litellm_credentials": "azure-cred"}, + } + # With provider hint, should pick the exact match + assert _extract_credential_from_entry(entry, provider="azure") == "azure-cred" + assert _extract_credential_from_entry(entry, provider="openai") == "openai-cred" + + # Without provider hint, falls back to first key (insertion order) + result = _extract_credential_from_entry(entry) + assert result in ("openai-cred", "azure-cred") + + # Unknown provider falls back to first available + result = _extract_credential_from_entry(entry, provider="bedrock") + assert result in ("openai-cred", "azure-cred") + + +def test_resolve_provider_hint_from_model_name(): + """Provider prefix in model name (e.g. azure/gpt-4) threads through to entry extraction.""" + config = { + "gpt-4": { + "openai": {"litellm_credentials": "openai-cred"}, + "azure": {"litellm_credentials": "azure-cred"}, + }, + } + # Model name "azure/gpt-4" -> provider="azure" -> should prefer azure-cred + # But _resolve_credential_from_model_config tries "azure/gpt-4" first (no match), + # then falls to defaultconfig (no match). So we need to use pre_alias_model_name. + result = _resolve_credential_from_model_config( + "azure/gpt-4", config, None, pre_alias_model_name="gpt-4", provider="azure" + ) + assert result == "azure-cred" diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 0d83b9f88d..efec841cc4 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -971,3 +971,106 @@ class TestWebSocketChunkTypes: ) assert len(messages) == 1 assert messages[0]["content"][0]["text"] == "Part 1Part 2" + + +class TestNativeWebSocketUrlConstruction: + """Test that native WebSocket URLs include the model query parameter. + + These tests mock websockets.connect so they exercise the actual URL-building + code inside BaseLLMHTTPHandler.async_responses_websocket rather than + reimplementing the logic themselves. + """ + + @pytest.mark.asyncio + async def test_openai_ws_url_includes_model(self): + """Handler must pass ?model= in the URL to the backend WebSocket.""" + from unittest.mock import AsyncMock, MagicMock, patch + + captured_urls = [] + + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) + + async def __aenter__(self): + raise Exception("stop") + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = "https://api.openai.com/v1/responses" + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o-mini", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o-mini"], f"Expected model in URL, got: {captured_urls[0]}" + + @pytest.mark.asyncio + async def test_ws_url_preserves_existing_params_and_adds_model(self): + """When api_base already has query params, model is added alongside them.""" + from unittest.mock import AsyncMock, MagicMock, patch + + captured_urls = [] + + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) + + async def __aenter__(self): + raise Exception("stop") + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = ( + "https://custom.example.com/v1/responses?api-version=2024-05-01" + ) + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o"], f"model missing from URL: {captured_urls[0]}" + assert qs.get("api-version") == ["2024-05-01"], f"existing param lost: {captured_urls[0]}" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx index e1b3b35830..3c5101fc2d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx @@ -131,7 +131,7 @@ describe("ModelsAndEndpointsView", () => { , ); expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument(); - }, 15000); + }); it("should show Missing provider banner by default", async () => { localStorageMock.clear(); @@ -149,7 +149,7 @@ describe("ModelsAndEndpointsView", () => { , ); expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument(); - }, 15000); + }); it("should hide Missing provider banner when dismiss button is clicked and persist to localStorage", async () => { localStorageMock.clear(); @@ -180,7 +180,7 @@ describe("ModelsAndEndpointsView", () => { // LocalStorage should be updated expect(localStorageMock.getItem("hideMissingProviderBanner")).toBe("true"); - }, 15000); + }); it("should show compact Request Provider button when banner is dismissed", async () => { // Set localStorage to hide banner @@ -209,7 +209,7 @@ describe("ModelsAndEndpointsView", () => { const requestProviderLinks = document.querySelectorAll('a[href="https://models.litellm.ai/?request=true"]'); // There should be a compact button when banner is hidden expect(requestProviderLinks.length).toBeGreaterThan(0); - }, 15000); + }); it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => { mockHealthCheckComponent.mockClear(); diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 202820a11a..7ad3e32ef5 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -46,10 +46,17 @@ function LoginPageContent() { // Cross-origin SSO: worker redirected back with a single-use code. // Exchange it for the JWT via the worker's /v3/login/exchange endpoint. const params = new URLSearchParams(window.location.search); - const ssoCode = params.get("code"); + const rawSsoCode = params.get("code"); + // Validate the SSO code is a plausible OAuth authorization code (alphanumeric + // plus common URL-safe chars) so that arbitrary user input cannot trigger the + // exchange endpoint. + const ssoCode = + rawSsoCode && /^[a-zA-Z0-9._~+/=-]+$/.test(rawSsoCode) ? rawSsoCode : null; if (ssoCode) { - // codeql[js/user-controlled-bypass] - const workerUrl = localStorage.getItem("litellm_worker_url"); + const rawWorkerUrl = localStorage.getItem("litellm_worker_url"); + // Validate the stored worker URL: only allow http(s) URLs. + const workerUrl = + rawWorkerUrl && /^https?:\/\/.+/.test(rawWorkerUrl) ? rawWorkerUrl : null; exchangeLoginCode(ssoCode, workerUrl).then(() => { params.delete("code"); const cleanSearch = params.toString(); diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx index 0c4cad8cb0..0539d6d8f1 100644 --- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx +++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx @@ -2,6 +2,7 @@ import { Suspense, useEffect, useMemo } from "react"; import { useSearchParams } from "next/navigation"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; // Written to sessionStorage so both the admin hook (useMcpOAuthFlow) and the // user hook (useUserMcpOAuthFlow) can pick up the result. Each hook reads @@ -52,14 +53,24 @@ const McpOAuthCallbackContent = () => { // Write to both namespace keys (admin and user) so whichever hook is // active can consume the result. sessionStorage only — no localStorage. const serialized = JSON.stringify(payload); - window.sessionStorage.setItem(ADMIN_RESULT_KEY, serialized); - window.sessionStorage.setItem(USER_RESULT_KEY, serialized); + setSecureItem(ADMIN_RESULT_KEY, serialized); + setSecureItem(USER_RESULT_KEY, serialized); } catch (err) { // Silently ignore storage errors } - const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY); - const destination = returnUrl || resolveDefaultRedirect(); + const returnUrl = getSecureItem(RETURN_URL_STORAGE_KEY); + let destination = resolveDefaultRedirect(); + if (returnUrl) { + try { + const parsed = new URL(returnUrl, window.location.origin); + if (parsed.origin === window.location.origin) { + destination = parsed.href; + } + } catch { + // invalid URL — fall through to default + } + } window.location.replace(destination); }, [payload]); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 44df1b5bd4..d3fab5cf5b 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -277,13 +277,18 @@ function CreateKeyPageContent() { // Check for a stored return URL const returnUrl = consumeReturnUrl(); if (returnUrl && isValidReturnUrl(returnUrl)) { + // Inline origin check: only redirect to same-origin URLs to prevent open redirect. + const safeUrl = new URL(returnUrl, window.location.origin); + if (safeUrl.origin !== window.location.origin) { + return; + } const currentUrl = window.location.href; const normalizedReturnUrl = normalizeUrlForCompare(returnUrl); const normalizedCurrentUrl = normalizeUrlForCompare(currentUrl); // Only redirect if the return URL is different from the current URL // This prevents infinite redirect loops if (normalizedReturnUrl !== normalizedCurrentUrl) { - window.location.replace(returnUrl); + window.location.replace(safeUrl.href); } } }, [authLoading, token]); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 9a4659da9d..03d982ca7c 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -51,7 +51,7 @@ function renderWithProviders(ui: React.ReactElement) { return render({ui}); } -describe("CreateUserButton", { timeout: 20000 }, () => { +describe("CreateUserButton", () => { beforeEach(() => { vi.clearAllMocks(); mockGetProxyUISettings.mockResolvedValue({ @@ -62,288 +62,296 @@ describe("CreateUserButton", { timeout: 20000 }, () => { }); }); - it("should render the create user form when embedded", () => { - renderWithProviders( - , - ); - expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument(); - }); + describe("rendering and visibility", () => { + it("should render the create user form when embedded", () => { + renderWithProviders( + , + ); + expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument(); + }); - it("should render the invite user button when not embedded", async () => { - renderWithProviders(); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + it("should render the invite user button when not embedded", async () => { + renderWithProviders(); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + }); + + it("should open the invite modal when invite user button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + expect(dialog).toBeInTheDocument(); + expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument(); + }); + + it("should display email invitations info message in embedded mode", () => { + renderWithProviders(); + expect(screen.getByText("Email invitations")).toBeInTheDocument(); + }); + + it("should display user role options when possibleUIRoles is provided", async () => { + const possibleUIRoles = { + proxy_admin: { ui_label: "Admin", description: "Full access" }, + proxy_user: { ui_label: "User", description: "Limited access" }, + }; + renderWithProviders( + , + ); + await userEvent.click(screen.getByRole("combobox", { name: /user role/i })); + expect(screen.getByText("Admin")).toBeInTheDocument(); + expect(screen.getByText("User")).toBeInTheDocument(); + }); + + it("should close modal when cancel is clicked in standalone mode", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument(); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.click(within(dialog).getByRole("button", { name: /close/i })); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); }); }); - it("should open the invite modal when invite user button is clicked", async () => { - const user = userEvent.setup(); - renderWithProviders(); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + describe("embedded mode submission", () => { + it("should call userCreateCall when form is submitted in embedded mode", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-1", + user_id: "new-user-123", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "test@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ + user_email: "test@example.com", + user_role: "proxy_user", + })); + }); }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - expect(dialog).toBeInTheDocument(); - expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument(); - }); - it("should display email invitations info message in embedded mode", () => { - renderWithProviders(); - expect(screen.getByText("Email invitations")).toBeInTheDocument(); - }); + it("should call onUserCreated callback when user is created in embedded mode", async () => { + const user = userEvent.setup(); + const onUserCreated = vi.fn(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } }); - it("should display user role options when possibleUIRoles is provided", async () => { - const possibleUIRoles = { - proxy_admin: { ui_label: "Admin", description: "Full access" }, - proxy_user: { ui_label: "User", description: "Limited access" }, - }; - renderWithProviders( - , - ); - await userEvent.click(screen.getByRole("combobox", { name: /user role/i })); - expect(screen.getByText("Admin")).toBeInTheDocument(); - expect(screen.getByText("User")).toBeInTheDocument(); - }); + renderWithProviders( + , + ); - it("should call userCreateCall when form is submitted in embedded mode", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-1", - user_id: "new-user-123", - has_user_setup_sso: false, - } as any); + await user.type(screen.getByLabelText(/user email/i), "embedded@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); - renderWithProviders( - , - ); + await waitFor(() => { + expect(onUserCreated).toHaveBeenCalledWith("new-user-456"); + }); + }); - await user.type(screen.getByLabelText(/user email/i), "test@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); + it("should show error notification when user creation fails", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } }); - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ - user_email: "test@example.com", - user_role: "proxy_user", - })); + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists"); + }); + }); + + it("should show info notification when making API call", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-3", + user_id: "new-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "info@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call"); + }); }); }); - it("should call onUserCreated callback when user is created in embedded mode", async () => { - const user = userEvent.setup(); - const onUserCreated = vi.fn(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } }); + describe("standalone mode submission", () => { + it("should show success notification when user is created successfully in standalone mode", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-2", + user_id: "new-user-789", + has_user_setup_sso: false, + } as any); - renderWithProviders( - , - ); + renderWithProviders( + , + ); - await user.type(screen.getByLabelText(/user email/i), "embedded@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - await waitFor(() => { - expect(onUserCreated).toHaveBeenCalledWith("new-user-456"); + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + }); + }); + + it("should show onboarding modal when user is created and SSO is disabled", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-sso", + user_id: "sso-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user"); + }); + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + }); }); }); - it("should show success notification when user is created successfully in standalone mode", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-2", - user_id: "new-user-789", - has_user_setup_sso: false, - } as any); + describe("organizations", () => { + it("should send organizations list in POST body when organizations are selected", async () => { + const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); + vi.mocked(useOrganizations).mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "My Org" }], + isLoading: false, + } as any); - renderWithProviders( - , - ); + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-org", + user_id: "org-user", + has_user_setup_sso: false, + } as any); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + + // Select org from the dropdown + const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i }); + await user.click(orgSelect); + await user.click(screen.getByText("My Org (org-1)")); + + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ + organizations: ["org-1"], + })); + }); }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + it("should not call organizationMemberAddCall after user creation", async () => { + const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); + vi.mocked(useOrganizations).mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "My Org" }], + isLoading: false, + } as any); - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-nma", + user_id: "no-member-add-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalled(); + }); + expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled(); }); }); - - it("should show error notification when user creation fails", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } }); - - renderWithProviders( - , - ); - - await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); - - await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists"); - }); - }); - - it("should show info notification when making API call", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-3", - user_id: "new-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await user.type(screen.getByLabelText(/user email/i), "info@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); - - await waitFor(() => { - expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call"); - }); - }); - - it("should close modal when cancel is clicked in standalone mode", async () => { - const user = userEvent.setup(); - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument(); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.click(within(dialog).getByRole("button", { name: /close/i })); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - }); - - it("should show onboarding modal when user is created and SSO is disabled", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-sso", - user_id: "sso-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user"); - }); - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); - }); - }); - - it("should send organizations list in POST body when organizations are selected", async () => { - const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); - vi.mocked(useOrganizations).mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "My Org" }], - isLoading: false, - } as any); - - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-org", - user_id: "org-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - - // Select org from the dropdown - const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i }); - await user.click(orgSelect); - await user.click(screen.getByText("My Org (org-1)")); - - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ - organizations: ["org-1"], - })); - }); - }); - - it("should not call organizationMemberAddCall after user creation", async () => { - const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); - vi.mocked(useOrganizations).mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "My Org" }], - isLoading: false, - } as any); - - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-nma", - user_id: "no-member-add-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalled(); - }); - expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled(); - }); }); diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index 651d1495c6..4b89820bad 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -843,7 +843,7 @@ describe("OldTeams - access_group_ids in team create", () => { }), ); }); - }, { timeout: 30000 }); + }); }); describe("OldTeams - models dropdown options", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 197bcd6569..0a8d2d9124 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -175,7 +175,7 @@ describe("Add Model Tab", () => { ); expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument(); - }, 10000); // This test is flaky, adding a timeout until we find a better solution + }); it("should display both Add Model and Add Auto Router tabs", async () => { const props = createTestProps(); @@ -269,7 +269,7 @@ describe("Add Model Tab", () => { }, { timeout: 10000 }, ); - }, 15000); // 15 second timeout to allow waitFor to complete + }); it("should show team selection when team-only switch is enabled", async () => { const props = createTestProps(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx new file mode 100644 index 0000000000..888f306625 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx @@ -0,0 +1,208 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen, waitFor, act, fireEvent } from "@testing-library/react"; +import { Form } from "antd"; +import OAuthFormFields from "./OAuthFormFields"; + +// ── helpers ────────────────────────────────────────────────────────────────── + +/** Minimal Ant Form wrapper so Form.Item registers correctly. */ +const WithForm: React.FC<{ children: React.ReactNode; onFinish?: (values: any) => void }> = ({ + children, + onFinish, +}) => { + const [form] = Form.useForm(); + return ( +
+ {children} + +
+ ); +}; + +// ── tests ───────────────────────────────────────────────────────────────────── + +describe("OAuthFormFields", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + // ── visibility by flow type ───────────────────────────────────────────────── + + describe("interactive mode (isM2M=false)", () => { + it("renders Token Validation Rules field", () => { + render( + + + , + ); + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + it("renders Token Storage TTL field", () => { + render( + + + , + ); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + + it("renders standard interactive fields alongside the new fields", () => { + render( + + + , + ); + expect(screen.getByText("Authorization URL (optional)")).toBeInTheDocument(); + expect(screen.getByText("Registration URL (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + }); + + describe("M2M mode (isM2M=true)", () => { + it("does NOT render Token Validation Rules field", () => { + render( + + + , + ); + expect(screen.queryByText("Token Validation Rules (optional)")).not.toBeInTheDocument(); + }); + + it("does NOT render Token Storage TTL field", () => { + render( + + + , + ); + expect(screen.queryByText("Token Storage TTL (seconds, optional)")).not.toBeInTheDocument(); + }); + + it("still renders M2M-specific fields", () => { + render( + + + , + ); + expect(screen.getByText("Client ID")).toBeInTheDocument(); + expect(screen.getByText("Token URL")).toBeInTheDocument(); + }); + }); + + // ── token_validation_json inline JSON validator ────────────────────────────── + + describe("token_validation_json validation", () => { + it("accepts empty value without error", async () => { + const onFinish = vi.fn(); + render( + + + , + ); + + // Leave the textarea empty and submit + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + + it("accepts a valid JSON object without error", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: '{"organization": "my-org"}' } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + + it("shows 'Must be valid JSON' error for malformed JSON", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "not-valid-json{" } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + }); + + it("shows error for a plain string value (not a JSON object)", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + // A bare string is valid JSON but we still want to accept it; only truly + // unparseable text should fail. Bare "hello" is actually invalid JSON + // (no quotes), so it should fail. + fireEvent.change(textarea, { target: { value: "hello" } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + }); + + it("whitespace-only value is treated as empty and passes validation", async () => { + const onFinish = vi.fn(); + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: " " } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx index 85487a8a47..4a808ca489 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Form, Select, Tooltip } from "antd"; +import { Form, Input, InputNumber, Select, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { OAUTH_FLOW } from "./types"; @@ -151,6 +151,50 @@ const OAuthFormFields: React.FC = ({ > + + } + name="token_validation_json" + rules={[ + { + validator: (_: any, value: string) => { + if (!value || value.trim() === "") return Promise.resolve(); + try { + JSON.parse(value); + return Promise.resolve(); + } catch { + return Promise.reject(new Error("Must be valid JSON")); + } + }, + }, + ]} + > + + + + } + name="token_storage_ttl_seconds" + > + + {oauthFlow && (

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

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

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