mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-13 21:06:58 +00:00
@@ -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"]
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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": {
|
||||
"<provider>": {
|
||||
"litellm_credentials": "<credential-name>"
|
||||
}
|
||||
},
|
||||
"<model-name>": {
|
||||
"<provider>": {
|
||||
"litellm_credentials": "<credential-name>"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Description |
|
||||
|---|---|
|
||||
| `defaultconfig` | Fallback credential for any model not explicitly listed |
|
||||
| `<model-name>` | Model-specific override — must match the LiteLLM model group name |
|
||||
| `<provider>` | 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:
|
||||
|
||||
<Tabs>
|
||||
|
||||
<TabItem value="config" label="config.yaml">
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
enable_model_config_credential_overrides: true
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="env" label="Environment Variable">
|
||||
|
||||
```bash
|
||||
export LITELLM_ENABLE_MODEL_CONFIG_CREDENTIAL_OVERRIDES=true
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
</Tabs>
|
||||
|
||||
:::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
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:<digest>
|
||||
```
|
||||
|
||||
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:<digest>
|
||||
```
|
||||
|
||||
### 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
|
||||
@@ -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"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -17,8 +17,15 @@ from litellm.constants import (
|
||||
MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE,
|
||||
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
|
||||
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
|
||||
MCP_PER_USER_TOKEN_DEFAULT_TTL,
|
||||
MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
|
||||
MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -152,6 +159,107 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
||||
mcp_oauth2_token_cache = MCPOAuth2TokenCache()
|
||||
|
||||
|
||||
def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int:
|
||||
"""Compute Redis TTL for a per-user token.
|
||||
|
||||
Uses server.token_storage_ttl_seconds when configured; otherwise derives
|
||||
TTL from expires_in minus the expiry buffer; falls back to the default TTL.
|
||||
"""
|
||||
if server.token_storage_ttl_seconds is not None:
|
||||
return max(server.token_storage_ttl_seconds, 1)
|
||||
if expires_in is not None:
|
||||
return max(
|
||||
expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
|
||||
1,
|
||||
)
|
||||
return MCP_PER_USER_TOKEN_DEFAULT_TTL
|
||||
|
||||
|
||||
class MCPPerUserTokenCache:
|
||||
"""Redis-backed cache for per-user OAuth2 access tokens.
|
||||
|
||||
Uses LiteLLM's existing ``user_api_key_cache`` (DualCache with optional
|
||||
Redis backend). Tokens are NaCl-encrypted with ``encrypt_value_helper``
|
||||
before storage so they are safe at rest in Redis.
|
||||
|
||||
Redis key format: ``mcp:per_user_token:{user_id}:{server_id}``
|
||||
Redis value: ``encrypt_value_helper(access_token)`` — URL-safe base64
|
||||
"""
|
||||
|
||||
def _cache_key(self, user_id: str, server_id: str) -> str:
|
||||
return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}"
|
||||
|
||||
async def get(self, user_id: str, server_id: str) -> Optional[str]:
|
||||
"""Return the plaintext access_token, or None on miss/error."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key = self._cache_key(user_id, server_id)
|
||||
encrypted = await user_api_key_cache.async_get_cache(key)
|
||||
if encrypted is None:
|
||||
return None
|
||||
plaintext = decrypt_value_helper(
|
||||
encrypted,
|
||||
key="mcp_per_user_token",
|
||||
exception_type="debug",
|
||||
)
|
||||
return plaintext or None
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.get failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
async def set(
|
||||
self,
|
||||
user_id: str,
|
||||
server_id: str,
|
||||
access_token: str,
|
||||
ttl: int,
|
||||
) -> None:
|
||||
"""Store NaCl-encrypted access_token in Redis with the given TTL."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key = self._cache_key(user_id, server_id)
|
||||
encrypted = encrypt_value_helper(access_token)
|
||||
await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl)
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds",
|
||||
user_id,
|
||||
server_id,
|
||||
ttl,
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.set failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
async def delete(self, user_id: str, server_id: str) -> None:
|
||||
"""Invalidate the cached token (removes from both in-memory and Redis layers)."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key = self._cache_key(user_id, server_id)
|
||||
await user_api_key_cache.async_delete_cache(key)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.delete failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
|
||||
mcp_per_user_token_cache = MCPPerUserTokenCache()
|
||||
|
||||
|
||||
async def resolve_mcp_auth(
|
||||
server: "MCPServer",
|
||||
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
|
||||
|
||||
@@ -896,11 +896,17 @@ if MCP_AVAILABLE:
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict.
|
||||
"""Look up stored OAuth2 token for (user, server) and return as extra_headers dict.
|
||||
|
||||
Lookup order:
|
||||
1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied
|
||||
2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query
|
||||
3. Auto-refresh when the stored token is expired and a refresh_token exists
|
||||
|
||||
Args:
|
||||
prefetched_creds: Optional dict keyed by server_id with credential payloads.
|
||||
When provided, avoids a per-server DB round-trip.
|
||||
When provided, the Redis and individual DB lookups are
|
||||
skipped in favour of the pre-fetched batch result.
|
||||
"""
|
||||
if server.auth_type != MCPAuth.oauth2:
|
||||
return None
|
||||
@@ -914,8 +920,27 @@ if MCP_AVAILABLE:
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
get_user_oauth_credential,
|
||||
is_oauth_credential_expired,
|
||||
refresh_user_oauth_token,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415
|
||||
_compute_per_user_token_ttl,
|
||||
mcp_per_user_token_cache,
|
||||
)
|
||||
|
||||
# ── Fast path: Redis cache ────────────────────────────────────────
|
||||
# Only used when prefetched_creds is not supplied (individual lookup).
|
||||
if prefetched_creds is None:
|
||||
cached_token = await mcp_per_user_token_cache.get(user_id, server_id)
|
||||
if cached_token is not None:
|
||||
verbose_logger.debug(
|
||||
"_get_user_oauth_extra_headers_from_db: Redis hit for "
|
||||
"user=%s server=%s",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
return {"Authorization": f"Bearer {cached_token}"}
|
||||
|
||||
# ── Slow path: DB lookup ──────────────────────────────────────────
|
||||
if prefetched_creds is not None:
|
||||
cred = prefetched_creds.get(server_id)
|
||||
else:
|
||||
@@ -929,18 +954,83 @@ if MCP_AVAILABLE:
|
||||
cred = await get_user_oauth_credential(
|
||||
prisma_client, user_id, server_id
|
||||
)
|
||||
if cred and cred.get("access_token"):
|
||||
if is_oauth_credential_expired(cred):
|
||||
verbose_logger.debug(
|
||||
f"_get_user_oauth_extra_headers_from_db: token expired for "
|
||||
f"user={user_id} server={server_id}"
|
||||
)
|
||||
|
||||
if not cred or not cred.get("access_token"):
|
||||
return None
|
||||
|
||||
if is_oauth_credential_expired(cred):
|
||||
verbose_logger.debug(
|
||||
"_get_user_oauth_extra_headers_from_db: token expired for "
|
||||
"user=%s server=%s — attempting refresh",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
# Attempt token refresh; requires a DB client (not available from prefetch)
|
||||
if cred.get("refresh_token"):
|
||||
try:
|
||||
from litellm.proxy.utils import ( # noqa: PLC0415
|
||||
get_prisma_client_or_throw,
|
||||
)
|
||||
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Cannot refresh OAuth token."
|
||||
)
|
||||
cred = await refresh_user_oauth_token(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server=server,
|
||||
cred=cred,
|
||||
)
|
||||
except Exception as refresh_exc:
|
||||
verbose_logger.warning(
|
||||
"_get_user_oauth_extra_headers_from_db: refresh failed "
|
||||
"for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
refresh_exc,
|
||||
)
|
||||
cred = None
|
||||
|
||||
if not cred or not cred.get("access_token"):
|
||||
# Clear stale Redis/cache entry so we don't serve it again.
|
||||
# Do this for both the individual and prefetch paths so the
|
||||
# next request doesn't get a stale cache hit.
|
||||
await mcp_per_user_token_cache.delete(user_id, server_id)
|
||||
return None
|
||||
return {"Authorization": f"Bearer {cred['access_token']}"}
|
||||
|
||||
access_token: str = cred["access_token"]
|
||||
|
||||
# Warm (or re-warm) the Redis cache from the DB result.
|
||||
# Always write regardless of whether expires_at is present — tokens
|
||||
# without an expiry are still valid and should be cached using the
|
||||
# server/default TTL so subsequent requests are fast.
|
||||
if prefetched_creds is None:
|
||||
raw_expires = None
|
||||
expires_at = cred.get("expires_at")
|
||||
if expires_at:
|
||||
from datetime import datetime, timezone # noqa: PLC0415
|
||||
|
||||
try:
|
||||
exp_dt = datetime.fromisoformat(expires_at)
|
||||
if exp_dt.tzinfo is None:
|
||||
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
|
||||
remaining = int(
|
||||
(exp_dt - datetime.now(timezone.utc)).total_seconds()
|
||||
)
|
||||
raw_expires = max(remaining, 0) if remaining > 0 else None
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
ttl = _compute_per_user_token_ttl(server, raw_expires)
|
||||
await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl)
|
||||
|
||||
return {"Authorization": f"Bearer {access_token}"}
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
|
||||
f"user={user_id} server={server_id}: {e}"
|
||||
"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
|
||||
"user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
@@ -2504,6 +2594,14 @@ if MCP_AVAILABLE:
|
||||
server_name, client_ip=_client_ip
|
||||
)
|
||||
if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
|
||||
# For servers that store per-user tokens server-side, skip the
|
||||
# pre-emptive 401 — the call_tool / list_tools dispatch will look
|
||||
# up the stored token from Redis / DB and only fail at the MCP
|
||||
# protocol level if none is found, giving the client a proper
|
||||
# tool-execution error rather than an HTTP 401.
|
||||
if server.needs_user_oauth_token:
|
||||
continue
|
||||
|
||||
request = StarletteRequest(scope)
|
||||
base_url = get_request_base_url(request)
|
||||
|
||||
|
||||
+61
-53
@@ -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
|
||||
|
||||
@@ -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`.
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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": [
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
assert "Validation fails" in str(exc.value)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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]}"
|
||||
|
||||
+4
-4
@@ -131,7 +131,7 @@ describe("ModelsAndEndpointsView", () => {
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument();
|
||||
}, 15000);
|
||||
});
|
||||
|
||||
it("should show Missing provider banner by default", async () => {
|
||||
localStorageMock.clear();
|
||||
@@ -149,7 +149,7 @@ describe("ModelsAndEndpointsView", () => {
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument();
|
||||
}, 15000);
|
||||
});
|
||||
|
||||
it("should hide Missing provider banner when dismiss button is clicked and persist to localStorage", async () => {
|
||||
localStorageMock.clear();
|
||||
@@ -180,7 +180,7 @@ describe("ModelsAndEndpointsView", () => {
|
||||
|
||||
// LocalStorage should be updated
|
||||
expect(localStorageMock.getItem("hideMissingProviderBanner")).toBe("true");
|
||||
}, 15000);
|
||||
});
|
||||
|
||||
it("should show compact Request Provider button when banner is dismissed", async () => {
|
||||
// Set localStorage to hide banner
|
||||
@@ -209,7 +209,7 @@ describe("ModelsAndEndpointsView", () => {
|
||||
const requestProviderLinks = document.querySelectorAll('a[href="https://models.litellm.ai/?request=true"]');
|
||||
// There should be a compact button when banner is hidden
|
||||
expect(requestProviderLinks.length).toBeGreaterThan(0);
|
||||
}, 15000);
|
||||
});
|
||||
|
||||
it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => {
|
||||
mockHealthCheckComponent.mockClear();
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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]);
|
||||
|
||||
|
||||
@@ -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]);
|
||||
|
||||
@@ -51,7 +51,7 @@ function renderWithProviders(ui: React.ReactElement) {
|
||||
return render(<QueryClientProvider client={qc}>{ui}</QueryClientProvider>);
|
||||
}
|
||||
|
||||
describe("CreateUserButton", { timeout: 20000 }, () => {
|
||||
describe("CreateUserButton", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockGetProxyUISettings.mockResolvedValue({
|
||||
@@ -62,288 +62,296 @@ describe("CreateUserButton", { timeout: 20000 }, () => {
|
||||
});
|
||||
});
|
||||
|
||||
it("should render the create user form when embedded", () => {
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} isEmbedded />,
|
||||
);
|
||||
expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument();
|
||||
});
|
||||
describe("rendering and visibility", () => {
|
||||
it("should render the create user form when embedded", () => {
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} isEmbedded />,
|
||||
);
|
||||
expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should render the invite user button when not embedded", async () => {
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
it("should render the invite user button when not embedded", async () => {
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("should open the invite modal when invite user button is clicked", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
expect(dialog).toBeInTheDocument();
|
||||
expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should display email invitations info message in embedded mode", () => {
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} isEmbedded />);
|
||||
expect(screen.getByText("Email invitations")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should display user role options when possibleUIRoles is provided", async () => {
|
||||
const possibleUIRoles = {
|
||||
proxy_admin: { ui_label: "Admin", description: "Full access" },
|
||||
proxy_user: { ui_label: "User", description: "Limited access" },
|
||||
};
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={possibleUIRoles} isEmbedded />,
|
||||
);
|
||||
await userEvent.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
expect(screen.getByText("Admin")).toBeInTheDocument();
|
||||
expect(screen.getByText("User")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should close modal when cancel is clicked in standalone mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument();
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.click(within(dialog).getByRole("button", { name: /close/i }));
|
||||
expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("should open the invite modal when invite user button is clicked", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
describe("embedded mode submission", () => {
|
||||
it("should call userCreateCall when form is submitted in embedded mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-1",
|
||||
user_id: "new-user-123",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "test@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({
|
||||
user_email: "test@example.com",
|
||||
user_role: "proxy_user",
|
||||
}));
|
||||
});
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
expect(dialog).toBeInTheDocument();
|
||||
expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should display email invitations info message in embedded mode", () => {
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} isEmbedded />);
|
||||
expect(screen.getByText("Email invitations")).toBeInTheDocument();
|
||||
});
|
||||
it("should call onUserCreated callback when user is created in embedded mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onUserCreated = vi.fn();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } });
|
||||
|
||||
it("should display user role options when possibleUIRoles is provided", async () => {
|
||||
const possibleUIRoles = {
|
||||
proxy_admin: { ui_label: "Admin", description: "Full access" },
|
||||
proxy_user: { ui_label: "User", description: "Limited access" },
|
||||
};
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={possibleUIRoles} isEmbedded />,
|
||||
);
|
||||
await userEvent.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
expect(screen.getByText("Admin")).toBeInTheDocument();
|
||||
expect(screen.getByText("User")).toBeInTheDocument();
|
||||
});
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} onUserCreated={onUserCreated} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
it("should call userCreateCall when form is submitted in embedded mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-1",
|
||||
user_id: "new-user-123",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
await user.type(screen.getByLabelText(/user email/i), "embedded@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
await waitFor(() => {
|
||||
expect(onUserCreated).toHaveBeenCalledWith("new-user-456");
|
||||
});
|
||||
});
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "test@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
it("should show error notification when user creation fails", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } });
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({
|
||||
user_email: "test@example.com",
|
||||
user_role: "proxy_user",
|
||||
}));
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists");
|
||||
});
|
||||
});
|
||||
|
||||
it("should show info notification when making API call", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-3",
|
||||
user_id: "new-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "info@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
it("should call onUserCreated callback when user is created in embedded mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onUserCreated = vi.fn();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } });
|
||||
describe("standalone mode submission", () => {
|
||||
it("should show success notification when user is created successfully in standalone mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-2",
|
||||
user_id: "new-user-789",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} onUserCreated={onUserCreated} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "embedded@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(onUserCreated).toHaveBeenCalledWith("new-user-456");
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created");
|
||||
});
|
||||
});
|
||||
|
||||
it("should show onboarding modal when user is created and SSO is disabled", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-sso",
|
||||
user_id: "sso-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user");
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
it("should show success notification when user is created successfully in standalone mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-2",
|
||||
user_id: "new-user-789",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
describe("organizations", () => {
|
||||
it("should send organizations list in POST body when organizations are selected", async () => {
|
||||
const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations");
|
||||
vi.mocked(useOrganizations).mockReturnValue({
|
||||
data: [{ organization_id: "org-1", organization_alias: "My Org" }],
|
||||
isLoading: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-org",
|
||||
user_id: "org-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
|
||||
// Select org from the dropdown
|
||||
const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i });
|
||||
await user.click(orgSelect);
|
||||
await user.click(screen.getByText("My Org (org-1)"));
|
||||
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({
|
||||
organizations: ["org-1"],
|
||||
}));
|
||||
});
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
it("should not call organizationMemberAddCall after user creation", async () => {
|
||||
const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations");
|
||||
vi.mocked(useOrganizations).mockReturnValue({
|
||||
data: [{ organization_id: "org-1", organization_alias: "My Org" }],
|
||||
isLoading: false,
|
||||
} as any);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created");
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-nma",
|
||||
user_id: "no-member-add-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUserCreateCall).toHaveBeenCalled();
|
||||
});
|
||||
expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
it("should show error notification when user creation fails", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } });
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists");
|
||||
});
|
||||
});
|
||||
|
||||
it("should show info notification when making API call", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-3",
|
||||
user_id: "new-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "info@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call");
|
||||
});
|
||||
});
|
||||
|
||||
it("should close modal when cancel is clicked in standalone mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument();
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.click(within(dialog).getByRole("button", { name: /close/i }));
|
||||
expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show onboarding modal when user is created and SSO is disabled", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-sso",
|
||||
user_id: "sso-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user");
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created");
|
||||
});
|
||||
});
|
||||
|
||||
it("should send organizations list in POST body when organizations are selected", async () => {
|
||||
const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations");
|
||||
vi.mocked(useOrganizations).mockReturnValue({
|
||||
data: [{ organization_id: "org-1", organization_alias: "My Org" }],
|
||||
isLoading: false,
|
||||
} as any);
|
||||
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-org",
|
||||
user_id: "org-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
|
||||
// Select org from the dropdown
|
||||
const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i });
|
||||
await user.click(orgSelect);
|
||||
await user.click(screen.getByText("My Org (org-1)"));
|
||||
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({
|
||||
organizations: ["org-1"],
|
||||
}));
|
||||
});
|
||||
});
|
||||
|
||||
it("should not call organizationMemberAddCall after user creation", async () => {
|
||||
const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations");
|
||||
vi.mocked(useOrganizations).mockReturnValue({
|
||||
data: [{ organization_id: "org-1", organization_alias: "My Org" }],
|
||||
isLoading: false,
|
||||
} as any);
|
||||
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-nma",
|
||||
user_id: "no-member-add-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUserCreateCall).toHaveBeenCalled();
|
||||
});
|
||||
expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -843,7 +843,7 @@ describe("OldTeams - access_group_ids in team create", () => {
|
||||
}),
|
||||
);
|
||||
});
|
||||
}, { timeout: 30000 });
|
||||
});
|
||||
});
|
||||
|
||||
describe("OldTeams - models dropdown options", () => {
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
import React from "react";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { render, screen, waitFor, act, fireEvent } from "@testing-library/react";
|
||||
import { Form } from "antd";
|
||||
import OAuthFormFields from "./OAuthFormFields";
|
||||
|
||||
// ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
/** Minimal Ant Form wrapper so Form.Item registers correctly. */
|
||||
const WithForm: React.FC<{ children: React.ReactNode; onFinish?: (values: any) => void }> = ({
|
||||
children,
|
||||
onFinish,
|
||||
}) => {
|
||||
const [form] = Form.useForm();
|
||||
return (
|
||||
<Form form={form} onFinish={onFinish}>
|
||||
{children}
|
||||
<button type="submit">Submit</button>
|
||||
</Form>
|
||||
);
|
||||
};
|
||||
|
||||
// ── tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
describe("OAuthFormFields", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
// ── visibility by flow type ─────────────────────────────────────────────────
|
||||
|
||||
describe("interactive mode (isM2M=false)", () => {
|
||||
it("renders Token Validation Rules field", () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders Token Storage TTL field", () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders standard interactive fields alongside the new fields", () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
expect(screen.getByText("Authorization URL (optional)")).toBeInTheDocument();
|
||||
expect(screen.getByText("Registration URL (optional)")).toBeInTheDocument();
|
||||
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
|
||||
expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
describe("M2M mode (isM2M=true)", () => {
|
||||
it("does NOT render Token Validation Rules field", () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={true} />
|
||||
</WithForm>,
|
||||
);
|
||||
expect(screen.queryByText("Token Validation Rules (optional)")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does NOT render Token Storage TTL field", () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={true} />
|
||||
</WithForm>,
|
||||
);
|
||||
expect(screen.queryByText("Token Storage TTL (seconds, optional)")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("still renders M2M-specific fields", () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={true} />
|
||||
</WithForm>,
|
||||
);
|
||||
expect(screen.getByText("Client ID")).toBeInTheDocument();
|
||||
expect(screen.getByText("Token URL")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
// ── token_validation_json inline JSON validator ──────────────────────────────
|
||||
|
||||
describe("token_validation_json validation", () => {
|
||||
it("accepts empty value without error", async () => {
|
||||
const onFinish = vi.fn();
|
||||
render(
|
||||
<WithForm onFinish={onFinish}>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
|
||||
// Leave the textarea empty and submit
|
||||
const submitBtn = screen.getByRole("button", { name: "Submit" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitBtn);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("accepts a valid JSON object without error", async () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: '{"organization": "my-org"}' } });
|
||||
});
|
||||
|
||||
const submitBtn = screen.getByRole("button", { name: "Submit" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitBtn);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("shows 'Must be valid JSON' error for malformed JSON", async () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: "not-valid-json{" } });
|
||||
});
|
||||
|
||||
const submitBtn = screen.getByRole("button", { name: "Submit" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitBtn);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Must be valid JSON")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("shows error for a plain string value (not a JSON object)", async () => {
|
||||
render(
|
||||
<WithForm>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
// A bare string is valid JSON but we still want to accept it; only truly
|
||||
// unparseable text should fail. Bare "hello" is actually invalid JSON
|
||||
// (no quotes), so it should fail.
|
||||
fireEvent.change(textarea, { target: { value: "hello" } });
|
||||
});
|
||||
|
||||
const submitBtn = screen.getByRole("button", { name: "Submit" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitBtn);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Must be valid JSON")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("whitespace-only value is treated as empty and passes validation", async () => {
|
||||
const onFinish = vi.fn();
|
||||
render(
|
||||
<WithForm onFinish={onFinish}>
|
||||
<OAuthFormFields isM2M={false} />
|
||||
</WithForm>,
|
||||
);
|
||||
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: " " } });
|
||||
});
|
||||
|
||||
const submitBtn = screen.getByRole("button", { name: "Submit" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitBtn);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,5 +1,5 @@
|
||||
import React from "react";
|
||||
import { Form, Select, Tooltip } from "antd";
|
||||
import { Form, Input, InputNumber, Select, Tooltip } from "antd";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { Button, TextInput } from "@tremor/react";
|
||||
import { OAUTH_FLOW } from "./types";
|
||||
@@ -151,6 +151,50 @@ const OAuthFormFields: React.FC<OAuthFormFieldsProps> = ({
|
||||
>
|
||||
<TextInput placeholder="https://example.com/oauth/register" className={fieldClassName} />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<FieldLabel
|
||||
label="Token Validation Rules (optional)"
|
||||
tooltip='JSON object of key-value rules checked against the OAuth token response before storing. Supports dot-notation for nested fields (e.g. {"organization": "my-org", "team.id": "123"}). Tokens that fail validation are rejected with HTTP 403.'
|
||||
/>
|
||||
}
|
||||
name="token_validation_json"
|
||||
rules={[
|
||||
{
|
||||
validator: (_: any, value: string) => {
|
||||
if (!value || value.trim() === "") return Promise.resolve();
|
||||
try {
|
||||
JSON.parse(value);
|
||||
return Promise.resolve();
|
||||
} catch {
|
||||
return Promise.reject(new Error("Must be valid JSON"));
|
||||
}
|
||||
},
|
||||
},
|
||||
]}
|
||||
>
|
||||
<Input.TextArea
|
||||
placeholder={'{\n "organization": "my-org",\n "team.id": "123"\n}'}
|
||||
rows={4}
|
||||
className="font-mono text-sm rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<FieldLabel
|
||||
label="Token Storage TTL (seconds, optional)"
|
||||
tooltip="How long to cache each user's OAuth access token in Redis before evicting it (regardless of the token's own expires_in). Leave blank to derive the TTL from the token's expires_in, or fall back to the 12-hour default."
|
||||
/>
|
||||
}
|
||||
name="token_storage_ttl_seconds"
|
||||
>
|
||||
<InputNumber
|
||||
min={1}
|
||||
placeholder="e.g. 3600"
|
||||
className="w-full rounded-lg"
|
||||
style={{ width: "100%" }}
|
||||
/>
|
||||
</Form.Item>
|
||||
{oauthFlow && (
|
||||
<div className="rounded-lg border border-dashed border-gray-300 p-4 space-y-2">
|
||||
<p className="text-sm text-gray-600">
|
||||
|
||||
@@ -150,151 +150,139 @@ describe("CreateMCPServer", () => {
|
||||
});
|
||||
});
|
||||
|
||||
it(
|
||||
"should not require auth value when creating a server with API Key auth type",
|
||||
{ timeout: 15000 },
|
||||
async () => {
|
||||
await selectHttpTransport();
|
||||
it("should not require auth value when creating a server with API Key auth type", async () => {
|
||||
await selectHttpTransport();
|
||||
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const user = userEvent.setup({ delay: null });
|
||||
|
||||
// Fill in server name (use id to avoid duplicate placeholder)
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "Test_Server");
|
||||
// Fill in server name (use id to avoid duplicate placeholder)
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "Test_Server");
|
||||
|
||||
// Fill in URL
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
// Fill in URL
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
|
||||
// Select API Key auth type
|
||||
await selectAntOption("Authentication", "API Key");
|
||||
// Select API Key auth type
|
||||
await selectAntOption("Authentication", "API Key");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Authentication Value")).toBeInTheDocument();
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Authentication Value")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Leave auth value empty and submit
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "Test_Server",
|
||||
alias: "Test_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "api_key",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
// Leave auth value empty and submit
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "Test_Server",
|
||||
alias: "Test_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "api_key",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
|
||||
// The form should submit without validation error on auth_value
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
},
|
||||
);
|
||||
// The form should submit without validation error on auth_value
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
it(
|
||||
"should not require auth value when creating a server with Bearer Token auth type",
|
||||
{ timeout: 15000 },
|
||||
async () => {
|
||||
await selectHttpTransport();
|
||||
it("should not require auth value when creating a server with Bearer Token auth type", async () => {
|
||||
await selectHttpTransport();
|
||||
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const user = userEvent.setup({ delay: null });
|
||||
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "Test_Server");
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "Test_Server");
|
||||
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
|
||||
await selectAntOption("Authentication", "Bearer Token");
|
||||
await selectAntOption("Authentication", "Bearer Token");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Authentication Value")).toBeInTheDocument();
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Authentication Value")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Leave auth value empty and submit
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "Test_Server",
|
||||
alias: "Test_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "bearer_token",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
// Leave auth value empty and submit
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "Test_Server",
|
||||
alias: "Test_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "bearer_token",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
},
|
||||
);
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
it(
|
||||
"should successfully create a server when auth value is provided",
|
||||
{ timeout: 15000 },
|
||||
async () => {
|
||||
await selectHttpTransport();
|
||||
it("should successfully create a server when auth value is provided", async () => {
|
||||
await selectHttpTransport();
|
||||
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const user = userEvent.setup({ delay: null });
|
||||
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "My_Server");
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "My_Server");
|
||||
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
|
||||
await selectAntOption("Authentication", "API Key");
|
||||
await selectAntOption("Authentication", "API Key");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Authentication Value")).toBeInTheDocument();
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Authentication Value")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Fill in auth value
|
||||
const authInput = screen.getByPlaceholderText("Enter token or secret");
|
||||
await user.type(authInput, "my-secret-key");
|
||||
// Fill in auth value
|
||||
const authInput = screen.getByPlaceholderText("Enter token or secret");
|
||||
await user.type(authInput, "my-secret-key");
|
||||
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "My_Server",
|
||||
alias: "My_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "api_key",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "My_Server",
|
||||
alias: "My_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "api_key",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [token, payload] = vi.mocked(networking.createMCPServer).mock.calls[0];
|
||||
expect(token).toBe("test-token");
|
||||
expect(payload.credentials).toEqual({ auth_value: "my-secret-key" });
|
||||
},
|
||||
);
|
||||
const [token, payload] = vi.mocked(networking.createMCPServer).mock.calls[0];
|
||||
expect(token).toBe("test-token");
|
||||
expect(payload.credentials).toEqual({ auth_value: "my-secret-key" });
|
||||
});
|
||||
|
||||
it("should not show auth value field when None auth type is selected", async () => {
|
||||
await selectHttpTransport();
|
||||
@@ -307,50 +295,187 @@ describe("CreateMCPServer", () => {
|
||||
});
|
||||
});
|
||||
|
||||
it(
|
||||
"should successfully create a server with no auth",
|
||||
{ timeout: 15000 },
|
||||
async () => {
|
||||
await selectHttpTransport();
|
||||
it("should successfully create a server with no auth", async () => {
|
||||
await selectHttpTransport();
|
||||
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const user = userEvent.setup({ delay: null });
|
||||
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "No_Auth_Server");
|
||||
const nameInput = getServerNameInput();
|
||||
await user.type(nameInput, "No_Auth_Server");
|
||||
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await user.type(urlInput, "https://example.com/mcp");
|
||||
|
||||
await selectAntOption("Authentication", "None");
|
||||
await selectAntOption("Authentication", "None");
|
||||
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "No_Auth_Server",
|
||||
alias: "No_Auth_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "none",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-1",
|
||||
server_name: "No_Auth_Server",
|
||||
alias: "No_Auth_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "none",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0];
|
||||
expect(payload.auth_type).toBe("none");
|
||||
// No credentials should be sent for "none" auth
|
||||
expect(payload.credentials).toBeUndefined();
|
||||
},
|
||||
);
|
||||
const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0];
|
||||
expect(payload.auth_type).toBe("none");
|
||||
// No credentials should be sent for "none" auth
|
||||
expect(payload.credentials).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("when OAuth interactive auth is selected", () => {
|
||||
/** Select HTTP transport + OAuth auth, then wait for the OAuth form to appear. */
|
||||
async function setupOAuthInteractive() {
|
||||
render(<CreateMCPServer {...defaultProps} />);
|
||||
await selectAntOption("Transport Type", "Streamable HTTP");
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
await selectAntOption("Authentication", "OAuth");
|
||||
|
||||
// Wait for OAuthFormFields to render (OAuth Flow Type selector is the sentinel)
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// OAuthFormFields defaults to INTERACTIVE, so the new fields should appear
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
|
||||
expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument();
|
||||
});
|
||||
}
|
||||
|
||||
it("shows Token Validation Rules and Token Storage TTL fields", async () => {
|
||||
await setupOAuthInteractive();
|
||||
// Asserted in setupOAuthInteractive
|
||||
});
|
||||
|
||||
it("includes token_validation in payload when token_validation_json is filled with valid JSON", async () => {
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-oauth",
|
||||
server_name: "OAuth_Server",
|
||||
alias: "OAuth_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "oauth2",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
|
||||
await setupOAuthInteractive();
|
||||
|
||||
// Fill required form fields
|
||||
const nameInput = document.getElementById("server_name") as HTMLInputElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(nameInput, { target: { value: "OAuth_Server" } });
|
||||
});
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await act(async () => {
|
||||
fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
|
||||
});
|
||||
|
||||
// Fill in the token_validation_json textarea
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: '{"organization": "my-org", "team.id": "42"}' } });
|
||||
});
|
||||
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0];
|
||||
expect(payload.token_validation).toEqual({ organization: "my-org", "team.id": "42" });
|
||||
});
|
||||
|
||||
it("omits token_validation from payload when token_validation_json is empty", async () => {
|
||||
vi.mocked(networking.createMCPServer).mockResolvedValue({
|
||||
server_id: "new-server-oauth",
|
||||
server_name: "OAuth_Server",
|
||||
alias: "OAuth_Server",
|
||||
url: "https://example.com/mcp",
|
||||
transport: "http",
|
||||
auth_type: "oauth2",
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
});
|
||||
|
||||
await setupOAuthInteractive();
|
||||
|
||||
const nameInput = document.getElementById("server_name") as HTMLInputElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(nameInput, { target: { value: "OAuth_Server" } });
|
||||
});
|
||||
const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
|
||||
await act(async () => {
|
||||
fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
|
||||
});
|
||||
|
||||
// Leave token_validation_json empty
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.createMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0];
|
||||
expect(payload.token_validation).toBeUndefined();
|
||||
});
|
||||
|
||||
it("does not submit and shows validation error for invalid JSON in token_validation_json", async () => {
|
||||
await setupOAuthInteractive();
|
||||
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: "not-valid-json{" } });
|
||||
});
|
||||
|
||||
const nameInput = document.getElementById("server_name") as HTMLInputElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(nameInput, { target: { value: "OAuth_Server" } });
|
||||
});
|
||||
|
||||
const submitButton = screen.getByRole("button", { name: "Add MCP Server" });
|
||||
await act(async () => {
|
||||
fireEvent.click(submitButton);
|
||||
});
|
||||
|
||||
// Either the inline form validation message or the notification fires —
|
||||
// both indicate the submit was blocked.
|
||||
await waitFor(() => {
|
||||
const inlineError = screen.queryByText("Must be valid JSON");
|
||||
const notCalled = !vi.mocked(networking.createMCPServer).mock.calls.length;
|
||||
expect(inlineError !== null || notCalled).toBe(true);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("when modal is cancelled", () => {
|
||||
|
||||
@@ -17,6 +17,7 @@ import { validateMCPServerUrl, validateMCPServerName } from "./utils";
|
||||
import NotificationsManager from "../molecules/notifications_manager";
|
||||
import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow";
|
||||
import { useTestMCPConnection } from "@/hooks/useTestMCPConnection";
|
||||
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
const asset_logos_folder = "../ui/assets/logos/";
|
||||
export const mcpLogoImg = `${asset_logos_folder}mcp_logo.png`;
|
||||
@@ -94,8 +95,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
}
|
||||
try {
|
||||
const values = form.getFieldsValue(true);
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
window.sessionStorage.setItem(
|
||||
setSecureItem(
|
||||
CREATE_OAUTH_UI_STATE_KEY,
|
||||
JSON.stringify({
|
||||
modalVisible: isModalVisible,
|
||||
@@ -178,7 +178,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
if (typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
const storedState = window.sessionStorage.getItem(CREATE_OAUTH_UI_STATE_KEY);
|
||||
const storedState = getSecureItem(CREATE_OAUTH_UI_STATE_KEY);
|
||||
if (!storedState) {
|
||||
return;
|
||||
}
|
||||
@@ -284,6 +284,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
credentials: credentialValues,
|
||||
allow_all_keys: allowAllKeysRaw,
|
||||
available_on_public_internet: availableOnPublicInternetRaw,
|
||||
token_validation_json: rawTokenValidationJson,
|
||||
...restValues
|
||||
} = values;
|
||||
|
||||
@@ -356,6 +357,18 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
restValues.transport = "http";
|
||||
}
|
||||
|
||||
// Parse token_validation JSON if provided
|
||||
let tokenValidation: Record<string, any> | null = null;
|
||||
if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") {
|
||||
try {
|
||||
tokenValidation = JSON.parse(rawTokenValidationJson);
|
||||
} catch {
|
||||
NotificationsManager.fromBackend("Invalid JSON in Token Validation Rules");
|
||||
setIsLoading(false);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Prepare the payload with cost configuration and allowed tools
|
||||
const payload: Record<string, any> = {
|
||||
...restValues,
|
||||
@@ -376,6 +389,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
|
||||
allow_all_keys: Boolean(allowAllKeysRaw),
|
||||
available_on_public_internet: Boolean(availableOnPublicInternetRaw),
|
||||
static_headers: staticHeaders,
|
||||
...(tokenValidation !== null && { token_validation: tokenValidation }),
|
||||
};
|
||||
|
||||
payload.static_headers = staticHeaders;
|
||||
|
||||
@@ -3,6 +3,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { render, screen, waitFor, fireEvent, act } from "@testing-library/react";
|
||||
import MCPServerEdit from "./mcp_server_edit";
|
||||
import * as networking from "../networking";
|
||||
import NotificationsManager from "../molecules/notifications_manager";
|
||||
|
||||
vi.mock("../networking", () => ({
|
||||
updateMCPServer: vi.fn(),
|
||||
@@ -37,6 +38,29 @@ vi.mock("./mcp_tool_configuration", () => ({
|
||||
default: () => <div data-testid="mcp-tool-config" />,
|
||||
}));
|
||||
|
||||
// ── fixtures ──────────────────────────────────────────────────────────────────
|
||||
|
||||
const interactiveOAuthServer = {
|
||||
server_id: "oauth_server_1",
|
||||
server_name: "OAuthServer",
|
||||
alias: "oauth_server", // underscores: hyphens fail validateMCPServerName
|
||||
description: "Interactive OAuth MCP server",
|
||||
transport: "http",
|
||||
url: "https://example.com/mcp",
|
||||
auth_type: "oauth2",
|
||||
// No token_url → edit form defaults to INTERACTIVE flow
|
||||
token_url: null,
|
||||
authorization_url: null,
|
||||
registration_url: null,
|
||||
created_at: "2024-01-01T00:00:00Z",
|
||||
created_by: "user-1",
|
||||
updated_at: "2024-01-01T00:00:00Z",
|
||||
updated_by: "user-1",
|
||||
mcp_access_groups: [],
|
||||
};
|
||||
|
||||
// ── test suites ───────────────────────────────────────────────────────────────
|
||||
|
||||
describe("MCPServerEdit (stdio)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
@@ -152,3 +176,228 @@ describe("MCPServerEdit (stdio)", () => {
|
||||
expect(payload.env).toEqual({ CIRCLECI_TOKEN: "new-token", CIRCLECI_BASE_URL: "https://circleci.com" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCPServerEdit (interactive OAuth)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("renders Token Validation Rules and Token Storage TTL fields for interactive OAuth server", async () => {
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken={null}
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
|
||||
expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
// Note: The M2M flow hiding logic is tested via OAuthFormFields.test.tsx (isM2M prop directly),
|
||||
// since Form.useWatch doesn't synchronously reflect initialValues in jsdom.
|
||||
|
||||
it("pre-populates token_validation_json from existing server token_validation", async () => {
|
||||
const tokenValidation = { organization: "my-org", "team.id": "123" };
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={{ ...interactiveOAuthServer, token_validation: tokenValidation }}
|
||||
accessToken={null}
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
expect(textarea).not.toBeNull();
|
||||
const parsed = JSON.parse(textarea.value);
|
||||
expect(parsed).toEqual(tokenValidation);
|
||||
});
|
||||
});
|
||||
|
||||
it("includes token_validation in update payload when token_validation_json is filled", async () => {
|
||||
const onSuccess = vi.fn();
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue({
|
||||
...interactiveOAuthServer,
|
||||
token_validation: { organization: "my-org" },
|
||||
});
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken="access-token"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={onSuccess}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
// Wait for the form to mount and the token_validation_json field to appear
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: '{"organization": "my-org"}' } });
|
||||
});
|
||||
|
||||
const saveButtons = screen.getAllByRole("button", { name: "Save Changes" });
|
||||
await act(async () => {
|
||||
fireEvent.click(saveButtons[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0];
|
||||
expect(payload.token_validation).toEqual({ organization: "my-org" });
|
||||
});
|
||||
|
||||
it("does not include token_validation in payload when field is empty and server had none", async () => {
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue(interactiveOAuthServer);
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken="access-token"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Leave token_validation_json empty
|
||||
const saveButtons = screen.getAllByRole("button", { name: "Save Changes" });
|
||||
await act(async () => {
|
||||
fireEvent.click(saveButtons[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0];
|
||||
expect(payload.token_validation).toBeUndefined();
|
||||
});
|
||||
|
||||
it("sends token_validation: null to clear an existing value when textarea is cleared", async () => {
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue({
|
||||
...interactiveOAuthServer,
|
||||
token_validation: null,
|
||||
});
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={{ ...interactiveOAuthServer, token_validation: { organization: "old-org" } }}
|
||||
accessToken="access-token"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
expect(textarea?.value).toContain("old-org");
|
||||
});
|
||||
|
||||
// Clear the textarea
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: "" } });
|
||||
});
|
||||
|
||||
const saveButtons = screen.getAllByRole("button", { name: "Save Changes" });
|
||||
await act(async () => {
|
||||
fireEvent.click(saveButtons[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0];
|
||||
// null signals the backend to clear the existing validation rules
|
||||
expect(payload.token_validation).toBeNull();
|
||||
});
|
||||
|
||||
it("shows inline validation error and does not submit on invalid JSON in token_validation_json", async () => {
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={interactiveOAuthServer}
|
||||
accessToken="access-token"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement;
|
||||
await act(async () => {
|
||||
fireEvent.change(textarea, { target: { value: "{ bad json" } });
|
||||
});
|
||||
|
||||
const saveButtons = screen.getAllByRole("button", { name: "Save Changes" });
|
||||
await act(async () => {
|
||||
fireEvent.click(saveButtons[0]);
|
||||
});
|
||||
|
||||
// The Form.Item inline validator intercepts invalid JSON before handleSave runs,
|
||||
// so the inline error message appears and updateMCPServer is never called.
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Must be valid JSON")).toBeInTheDocument();
|
||||
});
|
||||
expect(networking.updateMCPServer).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("includes token_storage_ttl_seconds in payload when set", async () => {
|
||||
vi.mocked(networking.updateMCPServer).mockResolvedValue({
|
||||
...interactiveOAuthServer,
|
||||
token_storage_ttl_seconds: 7200,
|
||||
});
|
||||
|
||||
render(
|
||||
<MCPServerEdit
|
||||
mcpServer={{ ...interactiveOAuthServer, token_storage_ttl_seconds: 7200 }}
|
||||
accessToken="access-token"
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
availableAccessGroups={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
const saveButtons = screen.getAllByRole("button", { name: "Save Changes" });
|
||||
await act(async () => {
|
||||
fireEvent.click(saveButtons[0]);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(networking.updateMCPServer).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0];
|
||||
expect(payload.token_storage_ttl_seconds).toBe(7200);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import React, { useState, useEffect } from "react";
|
||||
import { Form, Select, Button as AntdButton, Tooltip, Input } from "antd";
|
||||
import { Form, Select, Button as AntdButton, Tooltip, Input, InputNumber } from "antd";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
|
||||
import { AUTH_TYPE, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types";
|
||||
@@ -12,6 +12,7 @@ import MCPLogoSelector from "./MCPLogoSelector";
|
||||
import { validateMCPServerUrl, validateMCPServerName } from "./utils";
|
||||
import NotificationsManager from "../molecules/notifications_manager";
|
||||
import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow";
|
||||
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
interface MCPServerEditProps {
|
||||
mcpServer: MCPServer;
|
||||
@@ -73,8 +74,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
}
|
||||
try {
|
||||
const values = form.getFieldsValue(true);
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
window.sessionStorage.setItem(
|
||||
setSecureItem(
|
||||
EDIT_OAUTH_UI_STATE_KEY,
|
||||
JSON.stringify({
|
||||
serverId: mcpServer.server_id,
|
||||
@@ -190,6 +190,9 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
transport: effectiveTransport,
|
||||
static_headers: initialStaticHeaders,
|
||||
oauth_flow_type: mcpServer.token_url ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE,
|
||||
token_validation_json: mcpServer.token_validation
|
||||
? JSON.stringify(mcpServer.token_validation, null, 2)
|
||||
: undefined,
|
||||
}),
|
||||
[mcpServer, effectiveTransport, initialStaticHeaders, initialEnvJson],
|
||||
);
|
||||
@@ -214,7 +217,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
if (typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
const storedState = window.sessionStorage.getItem(EDIT_OAUTH_UI_STATE_KEY);
|
||||
const storedState = getSecureItem(EDIT_OAUTH_UI_STATE_KEY);
|
||||
if (!storedState) {
|
||||
return;
|
||||
}
|
||||
@@ -400,6 +403,7 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
args: rawArgs,
|
||||
allow_all_keys: allowAllKeysRaw,
|
||||
available_on_public_internet: availableOnPublicInternetRaw,
|
||||
token_validation_json: rawTokenValidationJson,
|
||||
...restValues
|
||||
} = values;
|
||||
|
||||
@@ -522,6 +526,17 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
restValues.transport = "http";
|
||||
}
|
||||
|
||||
// Parse token_validation JSON if provided
|
||||
let tokenValidation: Record<string, any> | null = null;
|
||||
if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") {
|
||||
try {
|
||||
tokenValidation = JSON.parse(rawTokenValidationJson);
|
||||
} catch {
|
||||
NotificationsManager.fromBackend("Invalid JSON in Token Validation Rules");
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Prepare the payload with cost configuration and permission fields
|
||||
const mcpInfoServerName =
|
||||
restValues.server_name ||
|
||||
@@ -556,6 +571,10 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
static_headers: staticHeaders,
|
||||
allow_all_keys: Boolean(allowAllKeysRaw ?? mcpServer.allow_all_keys),
|
||||
available_on_public_internet: Boolean(availableOnPublicInternetRaw ?? mcpServer.available_on_public_internet),
|
||||
// Include token_validation when it is set (non-null) or when clearing an existing value
|
||||
...(tokenValidation !== null || mcpServer.token_validation
|
||||
? { token_validation: tokenValidation }
|
||||
: {}),
|
||||
};
|
||||
|
||||
const includeCredentials = restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type);
|
||||
@@ -863,6 +882,58 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
|
||||
className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
{!isM2MFlow && (
|
||||
<>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
Token Validation Rules (optional)
|
||||
<Tooltip title='JSON object of key-value rules checked against the OAuth token response before storing. Supports dot-notation for nested fields (e.g. {"organization": "my-org", "team.id": "123"}). Tokens that fail validation are rejected with HTTP 403.'>
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name="token_validation_json"
|
||||
rules={[
|
||||
{
|
||||
validator: (_: any, value: string) => {
|
||||
if (!value || value.trim() === "") return Promise.resolve();
|
||||
try {
|
||||
JSON.parse(value);
|
||||
return Promise.resolve();
|
||||
} catch {
|
||||
return Promise.reject(new Error("Must be valid JSON"));
|
||||
}
|
||||
},
|
||||
},
|
||||
]}
|
||||
>
|
||||
<Input.TextArea
|
||||
placeholder={'{\n "organization": "my-org",\n "team.id": "123"\n}'}
|
||||
rows={4}
|
||||
className="font-mono text-sm rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span className="text-sm font-medium text-gray-700 flex items-center">
|
||||
Token Storage TTL (seconds, optional)
|
||||
<Tooltip title="How long to cache each user's OAuth access token in Redis before evicting it (regardless of the token's own expires_in). Leave blank to derive the TTL from the token's expires_in, or fall back to the 12-hour default.">
|
||||
<InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name="token_storage_ttl_seconds"
|
||||
>
|
||||
<InputNumber
|
||||
min={1}
|
||||
placeholder="e.g. 3600"
|
||||
style={{ width: "100%" }}
|
||||
className="rounded-lg"
|
||||
/>
|
||||
</Form.Item>
|
||||
</>
|
||||
)}
|
||||
<div className="rounded-lg border border-dashed border-gray-300 p-4 space-y-2">
|
||||
<p className="text-sm text-gray-600">Use OAuth to fetch a fresh access token and temporarily save it in the session as the authentication value.</p>
|
||||
<Button
|
||||
|
||||
@@ -20,6 +20,7 @@ import MCPSemanticFilterSettings from "../Settings/AdminSettings/MCPSemanticFilt
|
||||
import MCPNetworkSettings from "./MCPNetworkSettings";
|
||||
import MCPDiscovery from "./mcp_discovery";
|
||||
import { ByokCredentialModal } from "./ByokCredentialModal";
|
||||
import { getSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
const { Text: AntdText, Title: AntdTitle } = Typography;
|
||||
const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state";
|
||||
@@ -70,7 +71,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const stored = window.sessionStorage.getItem(EDIT_OAUTH_UI_STATE_KEY);
|
||||
const stored = getSecureItem(EDIT_OAUTH_UI_STATE_KEY);
|
||||
if (!stored) {
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -223,6 +223,10 @@ export interface MCPServer {
|
||||
submitted_at?: string | null;
|
||||
reviewed_at?: string | null;
|
||||
review_notes?: string | null;
|
||||
|
||||
/** Per-user OAuth token storage settings (interactive OAuth only) */
|
||||
token_validation?: Record<string, any> | null;
|
||||
token_storage_ttl_seconds?: number | null;
|
||||
}
|
||||
|
||||
export interface MCPServerProps {
|
||||
|
||||
@@ -75,6 +75,7 @@ import RealtimePlayground from "./RealtimePlayground";
|
||||
import { A2ATaskMetadata, MessageType } from "./types";
|
||||
import { useCodeInterpreter } from "./useCodeInterpreter";
|
||||
import { useChatHistory } from "./useChatHistory";
|
||||
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
const { TextArea } = Input;
|
||||
const { Dragger } = Upload;
|
||||
@@ -167,7 +168,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
||||
} = useChatHistory({ simplified });
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
|
||||
const saved = sessionStorage.getItem("apiKeySource");
|
||||
const saved = getSecureItem("apiKeySource");
|
||||
if (saved) {
|
||||
try {
|
||||
return JSON.parse(saved) as "session" | "custom";
|
||||
@@ -177,8 +178,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
||||
}
|
||||
return disabledPersonalKeyCreation ? "custom" : "session";
|
||||
});
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
const [apiKey, setApiKey] = useState<string>(() => sessionStorage.getItem("apiKey") || "");
|
||||
const [apiKey, setApiKey] = useState<string>(() => getSecureItem("apiKey") || "");
|
||||
const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState<string>(
|
||||
() => sessionStorage.getItem("customProxyBaseUrl") || "",
|
||||
);
|
||||
@@ -348,10 +348,12 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
||||
]);
|
||||
|
||||
useEffect(() => {
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
sessionStorage.setItem("apiKeySource", JSON.stringify(apiKeySource));
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
sessionStorage.setItem("apiKey", apiKey);
|
||||
try {
|
||||
setSecureItem("apiKeySource", JSON.stringify(apiKeySource));
|
||||
setSecureItem("apiKey", apiKey);
|
||||
} catch {
|
||||
// Storage full or unavailable — non-critical, skip persisting.
|
||||
}
|
||||
sessionStorage.setItem("endpointType", endpointType);
|
||||
sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags));
|
||||
sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores));
|
||||
@@ -502,7 +504,9 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
||||
|
||||
const handleImageUpload = (file: File) => {
|
||||
setUploadedImages((prev) => [...prev, file]);
|
||||
const previewUrl = URL.createObjectURL(file);
|
||||
const rawPreviewUrl = URL.createObjectURL(file);
|
||||
// Sanitize: only allow blob: URLs to prevent XSS via img src injection.
|
||||
const previewUrl = rawPreviewUrl.startsWith("blob:") ? rawPreviewUrl : "";
|
||||
setImagePreviewUrls((prev) => [...prev, previewUrl]);
|
||||
return false; // Prevent default upload behavior
|
||||
};
|
||||
@@ -1827,7 +1831,16 @@ const ChatUI: React.FC<ChatUIProps> = ({
|
||||
{uploadedImages.map((file, index) => (
|
||||
<div key={index} className="relative inline-block">
|
||||
<img
|
||||
src={imagePreviewUrls[index] || ""}
|
||||
src={(() => {
|
||||
const url = imagePreviewUrls[index];
|
||||
if (!url) return "";
|
||||
try {
|
||||
const parsed = new URL(url);
|
||||
return parsed.protocol === "blob:" ? parsed.href : "";
|
||||
} catch {
|
||||
return "";
|
||||
}
|
||||
})()}
|
||||
alt={`Upload preview ${index + 1}`}
|
||||
className="max-w-32 max-h-32 rounded-md border border-gray-200 object-cover"
|
||||
/>
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -11,6 +11,7 @@ import {
|
||||
serverRootPath,
|
||||
} from "@/components/networking";
|
||||
import { extractErrorMessage } from "@/utils/errorUtils";
|
||||
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
export type McpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
|
||||
|
||||
@@ -79,22 +80,13 @@ export const useMcpOAuthFlow = ({
|
||||
|
||||
const setStorageItem = (key: string, value: string) => {
|
||||
if (typeof window === "undefined") return;
|
||||
try {
|
||||
// Use sessionStorage only — the flow state may contain client credentials;
|
||||
// writing them to localStorage would persist across browser sessions and
|
||||
// make them readable by any injected script (XSS).
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
window.sessionStorage.setItem(key, value);
|
||||
} catch (err) {
|
||||
console.warn(`Failed to set storage item ${key}`, err);
|
||||
}
|
||||
setSecureItem(key, value);
|
||||
};
|
||||
|
||||
const getStorageItem = (key: string): string | null => {
|
||||
if (typeof window === "undefined") return null;
|
||||
try {
|
||||
// Try sessionStorage first, fall back to localStorage
|
||||
return window.sessionStorage.getItem(key) || window.localStorage.getItem(key);
|
||||
return getSecureItem(key);
|
||||
} catch (err) {
|
||||
console.warn(`Failed to get storage item ${key}`, err);
|
||||
return null;
|
||||
|
||||
@@ -23,6 +23,7 @@ import {
|
||||
} from "@/components/networking";
|
||||
import NotificationsManager from "@/components/molecules/notifications_manager";
|
||||
import { extractErrorMessage } from "@/utils/errorUtils";
|
||||
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
|
||||
|
||||
@@ -79,22 +80,11 @@ const genChallenge = async (verifier: string) => {
|
||||
};
|
||||
|
||||
const setStorage = (key: string, value: string) => {
|
||||
try {
|
||||
// Use sessionStorage only — do not write to localStorage.
|
||||
// The flow state may contain the LiteLLM access token; writing it to
|
||||
// localStorage would persist it across browser sessions and make it
|
||||
// readable by any injected script (XSS).
|
||||
// codeql[js/clear-text-storage-of-sensitive-data]
|
||||
window.sessionStorage.setItem(key, value);
|
||||
} catch (_) {}
|
||||
setSecureItem(key, value);
|
||||
};
|
||||
|
||||
const getStorage = (key: string): string | null => {
|
||||
try {
|
||||
return window.sessionStorage.getItem(key);
|
||||
} catch (_) {
|
||||
return null;
|
||||
}
|
||||
return getSecureItem(key);
|
||||
};
|
||||
|
||||
const clearStorage = (...keys: string[]) => {
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
function encode(value: string): string {
|
||||
// btoa cannot handle characters outside Latin-1, so we percent-encode first.
|
||||
return btoa(
|
||||
encodeURIComponent(value).replace(
|
||||
/%([0-9A-F]{2})/g,
|
||||
(_, p1) => String.fromCharCode(parseInt(p1, 16))
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
function decode(encoded: string): string {
|
||||
return decodeURIComponent(
|
||||
atob(encoded)
|
||||
.split("")
|
||||
.map((c) => "%" + c.charCodeAt(0).toString(16).padStart(2, "0"))
|
||||
.join("")
|
||||
);
|
||||
}
|
||||
|
||||
export function setSecureItem(key: string, value: string): void {
|
||||
window.sessionStorage.setItem(key, encode(value));
|
||||
}
|
||||
|
||||
export function getSecureItem(key: string): string | null {
|
||||
try {
|
||||
const raw = window.sessionStorage.getItem(key);
|
||||
if (raw === null) return null;
|
||||
return decode(raw);
|
||||
} catch {
|
||||
// Corrupted or non-encoded legacy value — return null without deleting
|
||||
// so that in-flight flows (e.g. OAuth) can time out naturally.
|
||||
return null;
|
||||
}
|
||||
}
|
||||
@@ -7,7 +7,7 @@ export default defineConfig({
|
||||
setupFiles: ["tests/setupTests.ts"],
|
||||
globals: true,
|
||||
css: true, // lets you import CSS/modules without extra mocks
|
||||
testTimeout: 10000,
|
||||
testTimeout: 30000,
|
||||
coverage: {
|
||||
provider: "v8",
|
||||
reporter: ["text", "lcov"],
|
||||
|
||||
Reference in New Issue
Block a user