diff --git a/.circleci/config.yml b/.circleci/config.yml index 7f410baf8f..23a62df478 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -4337,7 +4337,8 @@ jobs: name: Check for expected error command: | if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \ - grep -q "ERROR: Application startup failed. Exiting." docker_output.log; then + (grep -q "Database setup failed after multiple retries" docker_output.log || \ + grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then echo "Expected error found. Test passed." else echo "Expected error not found. Test failed." diff --git a/.github/workflows/create_daily_staging_branch.yml b/.github/workflows/create_daily_staging_branch.yml index 9d0093e8b1..08aebd7d04 100644 --- a/.github/workflows/create_daily_staging_branch.yml +++ b/.github/workflows/create_daily_staging_branch.yml @@ -41,3 +41,39 @@ jobs: git push origin $BRANCH_NAME echo "Successfully created and pushed branch: $BRANCH_NAME" fi + + create-internal-dev-branch: + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v3 + with: + fetch-depth: 0 + + - name: Create internal dev branch + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + # Configure Git user + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + + # Generate branch name with MM_DD_YYYY format + BRANCH_NAME="litellm_internal_dev_$(date +'%m_%d_%Y')" + echo "Creating branch: $BRANCH_NAME" + + # Fetch all branches + git fetch --all + + # Check if the branch already exists + if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then + echo "Branch $BRANCH_NAME already exists. Skipping creation." + else + echo "Creating new branch: $BRANCH_NAME" + # Create the new branch from main + git checkout -b $BRANCH_NAME origin/main + # Push the new branch + git push origin $BRANCH_NAME + echo "Successfully created and pushed branch: $BRANCH_NAME" + fi diff --git a/CLAUDE.md b/CLAUDE.md index 0c1caff9b4..5d62d2cdcd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -102,6 +102,22 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: ### UI / Backend Consistency - When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select +### MCP OAuth / OpenAPI Transport Mapping +- `TRANSPORT.OPENAPI` is a UI-only concept. The backend only accepts `"http"`, `"sse"`, or `"stdio"`. Always map it to `"http"` before any API call (including pre-OAuth temp-session calls). +- FastAPI validation errors return `detail` as an array of `{loc, msg, type}` objects. Error extractors must handle: array (map `.msg`), string, nested `{error: string}`, and fallback. +- When an MCP server already has `authorization_url` stored, skip OAuth discovery (`_discovery_metadata`) — the server URL for OpenAPI MCPs is the spec file, not the API base, and fetching it causes timeouts. +- `client_id` should be optional in the `/authorize` endpoint — if the server has a stored `client_id` in credentials, use that. Never require callers to re-supply it. + +### MCP Credential Storage +- OAuth credentials and BYOK credentials share the `litellm_mcpusercredentials` table, distinguished by a `"type"` field in the JSON payload (`"oauth2"` vs plain string). +- When deleting OAuth credentials, check type before deleting to avoid accidentally deleting a BYOK credential for the same `(user_id, server_id)` pair. +- Always pass the raw `expires_at` timestamp to the client — never set it to `None` for expired credentials. Let the frontend compute the "Expired" display state from the timestamp. +- Use `RecordNotFoundError` (not bare `except Exception`) when catching "already deleted" in credential delete endpoints. + +### Browser Storage Safety (UI) +- Never write LiteLLM access tokens or API keys to `localStorage` — use `sessionStorage` only. `localStorage` survives browser close and is readable by any injected script (XSS). +- Shared utility functions (e.g. `extractErrorMessage`) belong in `src/utils/` — never define them inline in hooks or duplicate them across files. + ### Database Migrations - Prisma handles schema migrations - Migration files auto-generated with `prisma migrate dev` diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index ea2c1700ee..90d71fb95c 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -944,7 +944,7 @@ router_settings: | QDRANT_URL | Connection URL for Qdrant database | QDRANT_VECTOR_SIZE | Vector size for Qdrant operations. Default is 1536 | REDIS_CONNECTION_POOL_TIMEOUT | Timeout in seconds for Redis connection pool. Default is 5 -| REDIS_CLUSTER_NODES | JSON-formatted list of Redis cluster startup nodes for Redis Cluster mode. Example: '[{"host": "node1", "port": 6379}]' +| REDIS_CLUSTER_NODES | JSON-formatted list of Redis cluster startup nodes for Redis Cluster mode. Example: `[{"host": "node1", "port": 6379}]` | REDIS_HOST | Hostname for Redis server | REDIS_PASSWORD | Password for Redis service | REDIS_PORT | Port number for Redis server diff --git a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md index e2cb839203..f4411553c6 100644 --- a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md +++ b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md @@ -309,6 +309,10 @@ Response: +## Policy Flow Builder + +For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step pass/fail actions. + ## Config Reference ### `policies` @@ -323,6 +327,7 @@ policies: remove: [...] condition: model: ... + pipeline: ... # optional; see Policy Flow Builder ``` | Field | Type | Description | @@ -332,6 +337,7 @@ policies: | `guardrails.add` | `list[string]` | Guardrails to enable. | | `guardrails.remove` | `list[string]` | Guardrails to disable (useful with inheritance). | | `condition.model` | `string` or `list[string]` | Optional. Only apply when model matches. Supports regex. | +| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions. See [Policy Flow Builder](./policy_flow_builder). | ### `policy_attachments` diff --git a/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md b/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md new file mode 100644 index 0000000000..2a83f3768a --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md @@ -0,0 +1,219 @@ +# Policy Flow Builder + +The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail passes or fails. + +Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors). + +## When to use the Flow Builder + +| Approach | Use case | +|----------|----------| +| **Simple policy** (`guardrails.add`) | All guardrails run in parallel; any failure blocks the request. | +| **Flow Builder** (pipeline) | Guardrails run in sequence; you choose actions per step (next, block, allow, custom response). | + +Use the Flow Builder when you need: + +- **Guardrail fallbacks** — use `on_fail: next` to try a different guardrail when one fails (e.g., fast filter → stricter filter) +- **Retrying the same guardrail** — add the same guardrail as multiple steps; if it fails, `on_fail: next` moves to the next step, which can be the same guardrail again (useful for transient API errors or rate limits) +- **Conditional routing** — e.g., if a fast guardrail fails, run a more advanced one instead of blocking immediately +- **Custom responses** — return a specific message when a guardrail fails instead of a generic block +- **Data chaining** — pass modified data (e.g., PII-masked content) from one step to the next +- **Fine-grained control** — different actions on pass vs. fail per step + +## Concepts + +### Pipeline + +A pipeline has: + +- **Mode**: `pre_call` (before the LLM) or `post_call` (after the LLM) +- **Steps**: Ordered list of guardrail steps + +### Step actions + +Each step defines what happens when the guardrail **passes** and when it **fails**: + +| Action | Description | +|--------|-------------| +| **Next Step** | Continue to the next guardrail in the pipeline | +| **Allow** | Stop the pipeline and allow the request to proceed | +| **Block** | Stop the pipeline and block the request | +| **Custom Response** | Return a custom message instead of the default block | + +### Step options + +| Field | Type | Description | +|-------|------|--------------| +| `guardrail` | `string` | Name of the guardrail to run | +| `on_pass` | `string` | Action when guardrail passes: `next`, `allow`, `block`, `modify_response` | +| `on_fail` | `string` | Action when guardrail fails: `next`, `allow`, `block`, `modify_response` | +| `pass_data` | `boolean` | Forward modified request data (e.g., PII-masked) to the next step | +| `modify_response_message` | `string` | Custom message when using `modify_response` action | + +## Using the Flow Builder (UI) + +1. Go to **Policies** in the LiteLLM Admin UI +2. Click **+ Create New Policy** or **Edit** on an existing policy +3. Select **Flow Builder** (instead of the simple form) +4. Design your flow: + - **Trigger** — Incoming LLM request (runs when the policy matches) + - **Steps** — Add guardrails, set ON PASS and ON FAIL actions per step + - **End** — Request proceeds to the LLM +5. Use the **+** between steps to insert new steps +6. Use the **Test** panel to run sample messages through the pipeline before saving +7. Click **Save** to create or update the policy + +## Config (YAML) + +Define a pipeline in your policy config: + +```yaml showLineNumbers title="config.yaml" +guardrails: + - guardrail_name: pii_masking + litellm_params: + guardrail: presidio + mode: pre_call + + - guardrail_name: prompt_injection + litellm_params: + guardrail: lakera + mode: pre_call + +policies: + my-pipeline-policy: + description: "PII mask first, then check for prompt injection" + guardrails: + add: + - pii_masking + - prompt_injection + pipeline: + mode: pre_call + steps: + - guardrail: pii_masking + on_pass: next + on_fail: block + pass_data: true + - guardrail: prompt_injection + on_pass: allow + on_fail: block + +policy_attachments: + - policy: my-pipeline-policy + scope: "*" +``` + +## Fallbacks and retries + +### Guardrail fallbacks + +Use `on_fail: next` to fall back to another guardrail when one fails. Run a lightweight guardrail first; if it fails, escalate to a stricter or different provider: + +```yaml +policies: + fallback-policy: + guardrails: + add: + - fast_content_filter + - strict_content_filter + pipeline: + mode: pre_call + steps: + - guardrail: fast_content_filter + on_pass: allow + on_fail: next + - guardrail: strict_content_filter + on_pass: allow + on_fail: block +``` + +If `fast_content_filter` passes → allow. If it fails → run `strict_content_filter`; pass → allow, fail → block. + +### Retrying the same guardrail + +Add the same guardrail as multiple steps to retry on failure. Useful for transient errors (API timeouts, rate limits): + +```yaml +policies: + retry-policy: + guardrails: + add: + - lakera_prompt_injection + pipeline: + mode: pre_call + steps: + - guardrail: lakera_prompt_injection + on_pass: allow + on_fail: next + - guardrail: lakera_prompt_injection + on_pass: allow + on_fail: block +``` + +First attempt passes → allow. First attempt fails → retry the same guardrail; second pass → allow, second fail → block. + +## Example: Custom response on fail + +Return a branded message instead of a generic block: + +```yaml +policies: + branded-block-policy: + guardrails: + add: + - pii_detector + pipeline: + mode: pre_call + steps: + - guardrail: pii_detector + on_pass: allow + on_fail: modify_response + modify_response_message: "Your message contains sensitive information. Please remove PII and try again." +``` + +## Test a pipeline (API) + +Test a pipeline with sample messages before attaching it: + +```bash +curl -X POST "http://localhost:4000/policies/test-pipeline" \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{ + "pipeline": { + "mode": "pre_call", + "steps": [ + { + "guardrail": "pii_masking", + "on_pass": "next", + "on_fail": "block", + "pass_data": true + }, + { + "guardrail": "prompt_injection", + "on_pass": "allow", + "on_fail": "block" + } + ] + }, + "test_messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "user", "content": "My SSN is 123-45-6789"} + ] + }' +``` + +Response includes per-step outcomes (pass/fail/error), actions taken, and timing. + +## Pipeline vs simple policy + +When a policy has a `pipeline`, the pipeline defines execution order and actions. The `guardrails.add` list must include all guardrails used in the pipeline steps. + +| Policy type | Execution | +|-------------|-----------| +| Simple (`guardrails.add` only) | All guardrails run; any failure blocks | +| Pipeline (`pipeline` present) | Steps run in order; actions control flow | + +## Related docs + +- [Guardrail Policies](./guardrail_policies) — Policy basics, attachments, inheritance +- [Policy Templates](./policy_templates) — Pre-built policy templates diff --git a/docs/my-website/docs/tutorials/retool_assist.md b/docs/my-website/docs/tutorials/retool_assist.md new file mode 100644 index 0000000000..703ce02ccc --- /dev/null +++ b/docs/my-website/docs/tutorials/retool_assist.md @@ -0,0 +1,143 @@ +import Image from '@theme/IdealImage'; + +# Retool Assist + +This guide walks you through connecting [Retool Assist](https://docs.retool.com/apps/guides/assist/) to LiteLLM Proxy. Retool Assist uses AI to generate and edit apps from within the Retool app IDE. Using LiteLLM with Retool Assist allows you to: + +- Access 100+ LLMs through Retool Assist +- Track spend and usage, set budget limits per virtual key +- Control which models Retool Assist can access +- Use your own LLM providers via a unified OpenAI-compatible API + +
+ +
+ +--- + +:::info +**Hosted Retool requires a public URL.** Retool Cloud runs on Retool's servers, so `localhost` will not work. You must expose your LiteLLM proxy via ngrok, Cloudflare Tunnel, or by deploying to a cloud provider. +::: + +## Quick Reference + +| Setting | Value | +|---------|-------| +| Provider Schema | OpenAI | +| Base URL | Your ngrok URL (e.g. `https://abc123.ngrok-free.app`) or deployed proxy URL | +| API Key | Your LiteLLM Virtual Key | +| Model | Public model name from LiteLLM (e.g. `openai/gpt-4o-mini`, `openai/gpt-5.2-2025-12-11`) | + +--- + +## Prerequisites + +- LiteLLM Proxy running locally or deployed +- [ngrok](https://ngrok.com/download) (or similar tunnel) for local development with hosted Retool +- A [Retool](https://retool.com) account (Cloud or self-hosted) + +## 1. Start LiteLLM Proxy + +Set up LiteLLM Proxy following the [Getting Started Guide](https://docs.litellm.ai/docs/proxy/docker_quick_start). Ensure your proxy is running on port 4000. + +## 2. Expose LiteLLM with a Public URL + + + +Retool Cloud runs on Retool's servers. You must expose your local LiteLLM proxy with a public URL. + +### Using ngrok + +- Install [ngrok](https://ngrok.com/download) +- In a separate terminal, run: + +```bash +ngrok http 4000 +``` +- Copy the generated HTTPS URL (e.g. `https://abc123.ngrok-free.app`). This is your **Base URL** for Retool. + + +### Alternative + +If you deploy LiteLLM to Railway, Render, Fly.io, or another cloud provider, use that public URL as your Base URL. See the [Deploy guide](https://docs.litellm.ai/docs/proxy/deploy) for details. + +## 3. Generate a Virtual Key + + + +Create a virtual key that Retool Assist will use to authenticate with LiteLLM. The key must have access to the models you want to use (e.g. `openai/*` for all OpenAI models). + +### Via LiteLLM UI + +- Navigate to [http://localhost:4000/ui](http://localhost:4000/ui) +- Go to **Virtual Keys** → **+ Create New Key** +- Select the models you need (or `openai/*` for all OpenAI models) +- Copy the key + +## 4. Add LiteLLM as a Custom Provider in Retool + +Inside your Retool dashboard, configure LiteLLM as a custom AI resource: + + + +1. Go to **Resources** + +2. Under the **AI** category, select **Custom Provider** + +3. Fill in the form: + - **Name:** `LiteLLM` + - **Description:** (optional) e.g. `LiteLLM Proxy - 100+ LLMs` + - **Provider Schema:** `OpenAI` + - **Base URL:** Your ngrok-generated URL (e.g. `https://abc123.ngrok-free.app`) or deployed proxy URL—do not add `/v1` unless Retool requires it + - **API Key:** Your LiteLLM virtual key from Step 3 +4. **Add model names** from your LiteLLM proxy (e.g. `openai/gpt-4o-mini`, `openai/gpt-5.2-2025-12-11`). +5. Click **Create Resource** + + + +## 5. Test the Connection + + + +- Open an app in Retool and enable **Assist** (if not already enabled in your organization) +- Use Assist to generate or edit app elements, it will route requests through LiteLLM +- Use the code option from the Sidebar to add a resource query, select the LiteLLM resource, and run it to test the setup. +- Check the LiteLLM **Logs** section to verify requests and track usage + + + +--- + +## Troubleshooting + +### 401 Unauthorized + +- Ensure the **API Key** in Retool matches your LiteLLM virtual key exactly +- Verify the key is not expired or blocked in LiteLLM + +### 401 "key not allowed to access model" + +Your virtual key is restricted to specific models. Generate a new key with `openai/*` or include the model you need (e.g. `openai/gpt-5.2-2025-12-11`) in the key's allowed models list. + +### 500 "api_key client option must be set" + +LiteLLM could not use your OpenAI API key to call the provider. Ensure `OPENAI_API_KEY` is set in your LiteLLM environment (e.g. in `.env` or `docker-compose.yml`) when using `openai/*` models. + +### localhost does not work + +Retool Cloud cannot reach `localhost` it points to Retool's servers. Use ngrok or deploy LiteLLM to a public URL. + +--- + +## Additional Resources + +- [Virtual Keys](https://docs.litellm.ai/docs/proxy/virtual_keys) – Create and manage API keys +- [Deploy LiteLLM](https://docs.litellm.ai/docs/proxy/deploy) – Production deployment options +- [Retool Assist Documentation](https://docs.retool.com/apps/guides/assist/) – Configure Assist and prompting guides diff --git a/docs/my-website/img/litellm_virtual_key.gif b/docs/my-website/img/litellm_virtual_key.gif new file mode 100644 index 0000000000..41daea9ddc Binary files /dev/null and b/docs/my-website/img/litellm_virtual_key.gif differ diff --git a/docs/my-website/img/ngrok_public_url.gif b/docs/my-website/img/ngrok_public_url.gif new file mode 100644 index 0000000000..b6c1079291 Binary files /dev/null and b/docs/my-website/img/ngrok_public_url.gif differ diff --git a/docs/my-website/img/retool_litellm_connection.gif b/docs/my-website/img/retool_litellm_connection.gif new file mode 100644 index 0000000000..13d2250f6e Binary files /dev/null and b/docs/my-website/img/retool_litellm_connection.gif differ diff --git a/docs/my-website/img/retool_litellm_logs.gif b/docs/my-website/img/retool_litellm_logs.gif new file mode 100644 index 0000000000..2055383938 Binary files /dev/null and b/docs/my-website/img/retool_litellm_logs.gif differ diff --git a/docs/my-website/img/retool_llm_setup.gif b/docs/my-website/img/retool_llm_setup.gif new file mode 100644 index 0000000000..c9f46c4936 Binary files /dev/null and b/docs/my-website/img/retool_llm_setup.gif differ diff --git a/docs/my-website/img/retool_resource_setup.gif b/docs/my-website/img/retool_resource_setup.gif new file mode 100644 index 0000000000..e01f32654e Binary files /dev/null and b/docs/my-website/img/retool_resource_setup.gif differ diff --git a/docs/my-website/release_notes/v1.82.0.md b/docs/my-website/release_notes/v1.82.0.md index b249187521..09967d5889 100644 --- a/docs/my-website/release_notes/v1.82.0.md +++ b/docs/my-website/release_notes/v1.82.0.md @@ -1,5 +1,5 @@ --- -title: "[Preview] v1.82.0 - Realtime Guardrails, Projects Management, and 10+ Performance Optimizations" +title: "v1.82.0 - Realtime Guardrails, Projects Management, and 10+ Performance Optimizations" slug: "v1-82-0" date: 2026-02-28T00:00:00 authors: @@ -26,7 +26,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:main-1.82.0 +ghcr.io/berriai/litellm:main-1.82.0-stable ``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 64c8fb291b..8383b826f0 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -100,6 +100,7 @@ const sidebars = { label: "Policies", items: [ "proxy/guardrails/guardrail_policies", + "proxy/guardrails/policy_flow_builder", "proxy/guardrails/policy_templates", "proxy/guardrails/policy_tags", ], @@ -172,7 +173,8 @@ const sidebars = { "tutorials/litellm_gemini_cli", "tutorials/google_genai_sdk", "tutorials/litellm_qwen_code_cli", - "tutorials/openai_codex" + "tutorials/openai_codex", + "tutorials/retool_assist" ] }, { diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl new file mode 100644 index 0000000000..9a5c185de2 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz new file mode 100644 index 0000000000..3e4be95b51 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309115809_add_missing_indexes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309115809_add_missing_indexes/migration.sql deleted file mode 100644 index 7b3e6d089e..0000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309115809_add_missing_indexes/migration.sql +++ /dev/null @@ -1,13 +0,0 @@ --- SkipTransactionBlock - --- Drop invalid indexes left behind by failed CONCURRENTLY builds -DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_VerificationToken_key_alias_idx"; - --- CreateIndex -CREATE INDEX CONCURRENTLY "LiteLLM_VerificationToken_key_alias_idx" ON "LiteLLM_VerificationToken"("key_alias"); - --- Drop invalid indexes left behind by failed CONCURRENTLY builds -DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_SpendLogs_user_startTime_idx"; - --- CreateIndex -CREATE INDEX CONCURRENTLY "LiteLLM_SpendLogs_user_startTime_idx" ON "LiteLLM_SpendLogs"("user", "startTime"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql new file mode 100644 index 0000000000..5ab834695b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql @@ -0,0 +1,11 @@ +-- DropIndex +DROP INDEX "LiteLLM_MCPServerTable_approval_status_idx"; + +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "approval_status", +DROP COLUMN "review_notes", +DROP COLUMN "reviewed_at", +DROP COLUMN "source_url", +DROP COLUMN "submitted_at", +DROP COLUMN "submitted_by"; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d5d17b2bce..ce79c2b3d5 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] @@ -315,6 +316,11 @@ model LiteLLM_MCPServerTable { is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? + approval_status String @default("approved") + submitted_by String? + submitted_at DateTime? + reviewed_at DateTime? + review_notes String? } // Per-user BYOK credentials for MCP servers @@ -388,9 +394,6 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) - - // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC - @@index([key_alias]) } model LiteLLM_JWTKeyMapping { @@ -556,9 +559,6 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) - - // SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ... - @@index([user, startTime]) } // View spend, model, api_key per request diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index ef80f092f1..3191386499 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.53" +version = "0.4.54" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.53" +version = "0.4.54" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 85ed955af4..ee111f3592 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1,6 +1,6 @@ # What is this? ## Helper utilities -from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union, get_args +from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union import httpx diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 828fba66d2..340e1531ff 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1198,6 +1198,9 @@ class AmazonConverseConfig(BaseConfig): + supported_config_params ) inference_params.pop("json_mode", None) # used for handling json_schema + # Anthropic-only key. Bedrock expects `outputConfig` (camelCase) and + # will reject `output_config` if it leaks through pass-through routes. + inference_params.pop("output_config", None) # Extract requestMetadata before processing other parameters request_metadata = inference_params.pop("requestMetadata", None) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d9fc331eed..1e7b1dcffb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7719,6 +7719,536 @@ class BaseLLMHTTPHandler: response=response, ) + async def async_vector_store_retrieve_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ) -> VectorStoreCreateResponse: + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = await async_httpx_client.get( + url=url, headers=headers + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + def vector_store_retrieve_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ) -> Union[ + VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse] + ]: + if _is_async: + return self.async_vector_store_retrieve_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = sync_httpx_client.get(url=url, headers=headers) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + async def async_vector_store_list_handler( + self, + after: Optional[str], + before: Optional[str], + limit: Optional[int], + order: Optional[str], + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ): + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = api_base + + params = {} + if after is not None: + params["after"] = after + if before is not None: + params["before"] = before + if limit is not None: + params["limit"] = limit + if order is not None: + params["order"] = order + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + "params": params, + }, + ) + + try: + response = await async_httpx_client.get( + url=url, headers=headers, params=params + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + + def vector_store_list_handler( + self, + after: Optional[str], + before: Optional[str], + limit: Optional[int], + order: Optional[str], + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ): + if _is_async: + return self.async_vector_store_list_handler( + after=after, + before=before, + limit=limit, + order=order, + vector_store_provider_config=vector_store_provider_config, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = api_base + + params = {} + if after is not None: + params["after"] = after + if before is not None: + params["before"] = before + if limit is not None: + params["limit"] = limit + if order is not None: + params["order"] = order + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + "params": params, + }, + ) + + try: + response = sync_httpx_client.get(url=url, headers=headers, params=params) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + + async def async_vector_store_update_handler( + self, + vector_store_id: str, + vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ) -> VectorStoreCreateResponse: + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + request_body = dict(vector_store_update_optional_params) + + # Clean metadata to only include string values (OpenAI requirement) + if "metadata" in request_body and request_body["metadata"] is not None: + from litellm.utils import add_openai_metadata + request_body["metadata"] = add_openai_metadata(request_body["metadata"]) + + if extra_body: + request_body.update(extra_body) + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": request_body, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = await async_httpx_client.post( + url=url, headers=headers, json=request_body, timeout=timeout + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + def vector_store_update_handler( + self, + vector_store_id: str, + vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ) -> Union[ + VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse] + ]: + if _is_async: + return self.async_vector_store_update_handler( + vector_store_id=vector_store_id, + vector_store_update_optional_params=vector_store_update_optional_params, + vector_store_provider_config=vector_store_provider_config, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + request_body = dict(vector_store_update_optional_params) + + # Clean metadata to only include string values (OpenAI requirement) + if "metadata" in request_body and request_body["metadata"] is not None: + from litellm.utils import add_openai_metadata + request_body["metadata"] = add_openai_metadata(request_body["metadata"]) + + if extra_body: + request_body.update(extra_body) + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": request_body, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = sync_httpx_client.post( + url=url, headers=headers, json=request_body + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return vector_store_provider_config.transform_create_vector_store_response( + response=response, + ) + + async def async_vector_store_delete_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ): + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = await async_httpx_client.delete( + url=url, headers=headers, timeout=timeout + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + + def vector_store_delete_handler( + self, + vector_store_id: str, + vector_store_provider_config: BaseVectorStoreConfig, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + _is_async: bool = False, + ): + if _is_async: + return self.async_vector_store_delete_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + litellm_params=litellm_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + headers = vector_store_provider_config.validate_environment( + headers=extra_headers or {}, litellm_params=litellm_params + ) + + if extra_headers: + headers.update(extra_headers) + + api_base = vector_store_provider_config.get_complete_url( + api_base=litellm_params.api_base, + litellm_params=dict(litellm_params), + ) + + url = f"{api_base}/{vector_store_id}" + + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "api_base": api_base, + "headers": headers, + }, + ) + + try: + response = sync_httpx_client.delete(url=url, headers=headers) + except Exception as e: + raise self._handle_error(e=e, provider_config=vector_store_provider_config) + + return response.json() + ##################################################################### ################ Vector Store Files HANDLERS ######################## ##################################################################### diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index e11cab4138..3e590680a7 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -219,17 +219,32 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): tool_choice: Tool choice in OpenAI format (str or dict) Returns: - Tool choice in Snowflake format (always an object) + Tool choice in Snowflake format (always an object, never a string) - OpenAI format (string): "auto", "required", "none" - OpenAI format (object): {"type": "function", "function": {"name": "get_weather"}} + OpenAI format (string): + "auto", "required", "none" - Snowflake format (string values become objects): {"type": "auto"} - Snowflake format (specific tool): {"type": "tool", "name": ["get_weather"]} + OpenAI format (dict): + {"type": "function", "function": {"name": "get_weather"}} + + Snowflake format: + {"type": "auto"} / {"type": "any"} / {"type": "none"} + {"type": "tool", "name": ["get_weather"]} + + Snowflake's API (like Anthropic) requires tool_choice as an object + with a "type" field, not as a bare string. """ if isinstance(tool_choice, str): - # Snowflake requires object format: {"type": "auto"} not string "auto" - return {"type": tool_choice} + # Snowflake requires object format, not string. + # Map OpenAI string values to Snowflake object format. + # "required" maps to "any" (Snowflake/Anthropic convention). + _type_map = { + "auto": "auto", + "required": "any", + "none": "none", + } + mapped_type = _type_map.get(tool_choice, tool_choice) + return {"type": mapped_type} if isinstance(tool_choice, dict): if tool_choice.get("type") == "function": diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 6772950827..5895a91f3a 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -522,29 +522,6 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False): return parameters -def _build_vertex_schema_for_gemini_2(parameters: dict) -> dict: - """ - Minimal schema builder for Gemini 2.0+ tool parameters. - - Gemini 2.0+ accepts standard JSON Schema natively in tool parameters, - including lowercase types, anyOf with null, and bare {} (TYPE_UNSPECIFIED). - The only transformation needed is resolving $ref/$defs, which Gemini does - NOT support in tool parameters (returns 400). - - This avoids the harmful transforms in _build_vertex_schema that break - JsonValue/Any semantics by coercing {} to {"type": "object"}. - """ - valid_schema_fields = set(get_type_hints(Schema).keys()) - - parameters = dict(parameters) # shallow copy to avoid mutating caller's dict - defs = parameters.pop("$defs", {}) - unpack_defs(parameters, defs) - - parameters = filter_schema_fields(parameters, valid_schema_fields) - - return parameters - - def _build_json_schema(parameters: dict) -> dict: """ Build a JSON Schema for use with Gemini's responseJsonSchema parameter. diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index e23fafcff2..7e686ee279 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -97,7 +97,6 @@ from ..common_utils import ( VertexAIError, _build_json_schema, _build_vertex_schema, - _build_vertex_schema_for_gemini_2, supports_response_json_schema, ) from ..vertex_llm_base import VertexBase @@ -468,7 +467,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return None def _map_function( # noqa: PLR0915 - self, value: List[dict], optional_params: dict, model: str = "" + self, value: List[dict], optional_params: dict ) -> List[Tools]: """ Map OpenAI-style tools/functions to Vertex AI format. @@ -511,21 +510,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "parameters" in _openai_function_object and _openai_function_object["parameters"] is not None and isinstance(_openai_function_object["parameters"], dict) - ): - if supports_response_json_schema(model): - # Gemini 2.0+: minimal transform (resolve $ref only) - _openai_function_object["parameters"] = ( - _build_vertex_schema_for_gemini_2( - _openai_function_object["parameters"] - ) - ) - else: - # Gemini 1.5: full OpenAPI-style transform - _openai_function_object["parameters"] = ( - _build_vertex_schema( - _openai_function_object["parameters"] - ) - ) + ): # OPENAI accepts JSON Schema, Google accepts OpenAPI schema. + _openai_function_object["parameters"] = _build_vertex_schema( + _openai_function_object["parameters"] + ) openai_function_object = _openai_function_object @@ -1060,7 +1048,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ): # Pass optional_params so _map_function can add toolConfig if needed mapped_tools = self._map_function( - value=value, optional_params=optional_params, model=model + value=value, optional_params=optional_params ) optional_params = self._add_tools_to_optional_params( optional_params, mapped_tools diff --git a/litellm/main.py b/litellm/main.py index a6a7cf2b74..0448ceceb8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -63,6 +63,7 @@ from litellm.utils import exception_type, get_litellm_params, get_optional_param # Logging is imported lazily when needed to avoid loading litellm_logging at import time if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import TokenCountResponse from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e21c9be368..9b1d81fee4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7997,6 +7997,80 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "black_forest_labs/flux-kontext-pro": { + "litellm_provider": "black_forest_labs", + "mode": "image_edit", + "output_cost_per_image": 0.04, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/edits", + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-kontext-max": { + "litellm_provider": "black_forest_labs", + "mode": "image_edit", + "output_cost_per_image": 0.08, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/edits", + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-pro-1.0-fill": { + "litellm_provider": "black_forest_labs", + "mode": "image_edit", + "output_cost_per_image": 0.05, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/edits" + ] + }, + "black_forest_labs/flux-pro-1.0-expand": { + "litellm_provider": "black_forest_labs", + "mode": "image_edit", + "output_cost_per_image": 0.05, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/edits" + ] + }, + "black_forest_labs/flux-pro-1.1": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-pro-1.1-ultra": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.06, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-dev": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.025, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "black_forest_labs/flux-pro": { + "litellm_provider": "black_forest_labs", + "mode": "image_generation", + "output_cost_per_image": 0.05, + "source": "https://bfl.ai/pricing", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "cerebras/llama-3.3-70b": { "input_cost_per_token": 8.5e-07, "litellm_provider": "cerebras", @@ -14442,7 +14516,7 @@ "input_cost_per_audio_per_second": 0.00016, "input_cost_per_image": 0.00012, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.0237, + "input_cost_per_video_per_second": 0.00079, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -14453,14 +14527,17 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2-preview": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, "mode": "embedding", "output_cost_per_token": 0, "output_vector_size": 3072, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supports_multimodal": true, "uses_embed_content": true }, @@ -14502,7 +14579,10 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2-preview": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, @@ -14510,7 +14590,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_multimodal": true, "tpm": 10000000 }, diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 289c27059e..8a65936d5b 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,4 +1,6 @@ -from datetime import datetime, timezone +import base64 +import json +from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast from litellm._logging import verbose_proxy_logger @@ -499,7 +501,6 @@ async def store_user_credential( credential: str, ) -> None: """Store a user credential for a BYOK MCP server.""" - import base64 encoded = base64.urlsafe_b64encode(credential.encode()).decode() await prisma_client.db.litellm_mcpusercredentials.upsert( @@ -521,7 +522,6 @@ async def get_user_credential( server_id: str, ) -> Optional[str]: """Return credential for a user+server pair, or None.""" - import base64 row = await prisma_client.db.litellm_mcpusercredentials.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} @@ -563,6 +563,138 @@ async def delete_user_credential( ) +# ── OAuth2 user-credential helpers ──────────────────────────────────────────── + + +async def store_user_oauth_credential( + prisma_client: PrismaClient, + user_id: str, + server_id: str, + access_token: str, + refresh_token: Optional[str] = None, + expires_in: Optional[int] = None, + scopes: Optional[List[str]] = None, +) -> None: + """Persist an OAuth2 access token for a user+server pair. + + The payload is JSON-serialised and stored base64-encoded in the same + ``credential_b64`` column used by BYOK. A ``"type": "oauth2"`` key + differentiates it from plain BYOK API keys. + """ + + expires_at: Optional[str] = None + if expires_in is not None: + expires_at = ( + datetime.now(timezone.utc) + timedelta(seconds=expires_in) + ).isoformat() + + payload: Dict[str, Any] = { + "type": "oauth2", + "access_token": access_token, + "connected_at": datetime.now(timezone.utc).isoformat(), + } + if refresh_token: + payload["refresh_token"] = refresh_token + if expires_at: + payload["expires_at"] = expires_at + if scopes: + payload["scopes"] = scopes + + # 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." + ) + 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( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, + data={ + "create": { + "user_id": user_id, + "server_id": server_id, + "credential_b64": encoded, + }, + "update": {"credential_b64": encoded}, + }, + ) + + +def is_oauth_credential_expired(cred: Dict[str, Any]) -> bool: + """Return True if the OAuth2 credential's access_token has expired. + + Checks the ``expires_at`` ISO-format string stored in the credential payload. + Returns False when ``expires_at`` is absent or unparseable (treat as non-expired). + """ + expires_at = cred.get("expires_at") + if not expires_at: + return False + try: + exp_dt = datetime.fromisoformat(expires_at) + if exp_dt.tzinfo is None: + exp_dt = exp_dt.replace(tzinfo=timezone.utc) + return datetime.now(timezone.utc) > exp_dt + except (ValueError, TypeError): + return False + + +async def get_user_oauth_credential( + prisma_client: PrismaClient, + user_id: str, + server_id: str, +) -> Optional[Dict[str, Any]]: + """Return the decoded OAuth2 payload dict for a user+server pair, or None.""" + + row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} + ) + if row is None: + return None + try: + decoded = base64.urlsafe_b64decode(row.credential_b64).decode() + parsed = json.loads(decoded) + if isinstance(parsed, dict) and parsed.get("type") == "oauth2": + return parsed + # Row exists but is a BYOK (plain string), not an OAuth token + return None + except Exception: + return None + + +async def list_user_oauth_credentials( + prisma_client: PrismaClient, + user_id: str, +) -> List[Dict[str, Any]]: + """Return all OAuth2 credential payloads for a user, tagged with server_id.""" + + rows = await prisma_client.db.litellm_mcpusercredentials.find_many( + where={"user_id": user_id} + ) + results: List[Dict[str, Any]] = [] + for row in rows: + try: + decoded = base64.urlsafe_b64decode(row.credential_b64).decode() + parsed = json.loads(decoded) + if isinstance(parsed, dict) and parsed.get("type") == "oauth2": + parsed["server_id"] = row.server_id + results.append(parsed) + except Exception: + pass # Skip non-OAuth rows (BYOK plain strings) + return results + + async def approve_mcp_server( prisma_client: PrismaClient, server_id: str, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 18db1f22c0..4e68e3f28c 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -312,8 +312,8 @@ async def register_client_with_server( @router.get("/authorize") async def authorize( request: Request, - client_id: str, redirect_uri: str, + client_id: Optional[str] = None, state: str = "", mcp_server_name: Optional[str] = None, code_challenge: Optional[str] = None, @@ -326,19 +326,34 @@ async def authorize( global_mcp_server_manager, ) - lookup_name = mcp_server_name or client_id + lookup_name: Optional[str] = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( - lookup_name, client_ip=client_ip + mcp_server = ( + global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) + if lookup_name + else None ) if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints() if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") + # Use server's stored client_id when caller doesn't supply one. + # Raise a clear error instead of passing an empty string — an empty + # client_id would silently produce a broken authorization URL. + resolved_client_id: str = mcp_server.client_id or client_id or "" + if not resolved_client_id: + raise HTTPException( + status_code=400, + detail={ + "error": "client_id is required but was not supplied and is not " + "stored on the MCP server record. Provide client_id as a query " + "parameter or configure it on the server." + }, + ) return await authorize_with_server( request=request, mcp_server=mcp_server, - client_id=client_id, + client_id=resolved_client_id, redirect_uri=redirect_uri, state=state, code_challenge=code_challenge, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index cdb20da976..164a48ded6 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,6 @@ import importlib -from datetime import datetime -from typing import Any, Awaitable, Callable, Dict, List, Optional, Union +from datetime import datetime, timezone +from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Union from fastapi import APIRouter, Depends, HTTPException, Query, Request @@ -69,6 +69,132 @@ if MCP_AVAILABLE: return server_auth return mcp_auth_header + def _get_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]: + """Return the subset of *allowed_server_ids* whose servers use OAuth2 auth. + + Used as a cheap pre-flight check to skip bulk credential fetching when no + OAuth2 servers are involved in the current request. + """ + return { + sid + for sid in allowed_server_ids + if getattr( + global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None + ) + == MCPAuth.oauth2 + } + + async def _get_user_oauth_extra_headers( + server, + user_api_key_dict: UserAPIKeyAuth, + prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, + ) -> Optional[Dict[str, str]]: + """ + For OAuth2 servers, look up the user's stored access token and return it + as extra_headers {"Authorization": "Bearer "} so that it reaches + the MCP server the same way the admin "Add MCP / Authorize and Fetch" flow does. + Returns None for non-OAuth2 servers or when no credential is stored. + + Args: + prefetched_creds: Optional dict keyed by server_id with credential payloads. + When provided, avoids a per-server DB round-trip. + """ + if getattr(server, "auth_type", None) != MCPAuth.oauth2: + return None + user_id = getattr(user_api_key_dict, "user_id", None) + server_id = getattr(server, "server_id", None) + if not user_id or not server_id: + return None + try: + from litellm.proxy._experimental.mcp_server.db import ( + get_user_oauth_credential, + is_oauth_credential_expired, + ) + + if prefetched_creds is not None: + cred = prefetched_creds.get(server_id) + else: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + 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: token expired for " + f"user={user_id} server={server_id}" + ) + return None + return {"Authorization": f"Bearer {cred['access_token']}"} + except Exception as e: + verbose_logger.warning( + f"_get_user_oauth_extra_headers: failed to retrieve credential for " + f"user={user_id} server={server_id}: {e}" + ) + return None + + async def _prefetch_user_oauth_creds( + user_api_key_dict: UserAPIKeyAuth, + ) -> Dict[str, Dict[str, Any]]: + """Fetch all OAuth2 credentials for the user in a single DB query. + + Returns a dict keyed by server_id. Used to avoid N+1 DB queries when + iterating over multiple OAuth2 MCP servers. + """ + user_id = getattr(user_api_key_dict, "user_id", None) + if not user_id: + return {} + try: + from litellm.proxy._experimental.mcp_server.db import ( + list_user_oauth_credentials, + ) + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + creds = await list_user_oauth_credentials(prisma_client, user_id) + return {c["server_id"]: c for c in creds if "server_id" in c} + except Exception as e: + verbose_logger.warning( + f"_prefetch_user_oauth_creds: failed to prefetch for user={user_id}: {e}" + ) + return {} + + async def _get_bulk_user_oauth_headers( + user_api_key_dict: UserAPIKeyAuth, + ) -> Dict[str, Dict[str, str]]: + """ + Fetch ALL OAuth2 credentials for the current user in a single DB query and + return a mapping of server_id → {"Authorization": "Bearer "}. + + This is the batch alternative to calling _get_user_oauth_extra_headers + per-server inside a loop (N+1 DB queries). + """ + user_id = getattr(user_api_key_dict, "user_id", None) + if not user_id: + return {} + try: + from litellm.proxy._experimental.mcp_server.db import ( + list_user_oauth_credentials, + ) + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + creds = await list_user_oauth_credentials(prisma_client, user_id) + return { + c["server_id"]: {"Authorization": f"Bearer {c['access_token']}"} + for c in creds + if c.get("access_token") and c.get("server_id") + } + except Exception: + verbose_logger.debug("Failed to bulk-fetch OAuth credentials", exc_info=True) + return {} + def _create_tool_response_objects(tools, server_mcp_info): """Helper function to create tool response objects.""" return [ @@ -162,11 +288,13 @@ if MCP_AVAILABLE: server_auth_header, raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, + extra_headers: Optional[Dict[str, str]] = None, ): """Helper function to get tools for a single server.""" tools = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, + extra_headers=extra_headers, add_prefix=False, raw_headers=raw_headers, ) @@ -295,8 +423,23 @@ if MCP_AVAILABLE: # If server_id is specified, only query that specific server if server_id: + # Resolve a server name to its UUID if needed (MCPConnectPicker passes + # server_name strings, but allowed_server_ids_set contains UUIDs). + # _name_resolved is kept so the second check can reuse it for accurate + # IP-filter error reporting if the resolved UUID is not in allowed_server_ids. + _name_resolved = None if server_id not in allowed_server_ids: - _server = global_mcp_server_manager.get_mcp_server_by_id(server_id) + _name_resolved = global_mcp_server_manager.get_mcp_server_by_name(server_id) + if _name_resolved is not None and _name_resolved.server_id in set(allowed_server_ids): + server_id = _name_resolved.server_id + + if server_id not in allowed_server_ids: + # Try UUID lookup first; fall back to the name-resolved server so that + # IP-filter reporting works correctly even when server_id is a name string. + _server = ( + global_mcp_server_manager.get_mcp_server_by_id(server_id) + or _name_resolved + ) if ( _server is not None and _rest_client_ip is not None @@ -334,6 +477,8 @@ if MCP_AVAILABLE: server_auth_header = _get_server_auth_header( server, mcp_server_auth_headers, mcp_auth_header ) + # Single-server request: targeted lookup is more efficient than a bulk fetch. + user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict) try: list_tools_result = await _get_tools_for_single_server( @@ -341,6 +486,7 @@ if MCP_AVAILABLE: server_auth_header, raw_headers_from_request, user_api_key_dict, + extra_headers=user_oauth_extra_headers, ) except Exception as e: verbose_logger.exception( @@ -374,6 +520,14 @@ if MCP_AVAILABLE: }, ) + # Pre-fetch OAuth credentials only when at least one allowed server uses OAuth2, + # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. + prefetched_oauth_creds = ( + await _prefetch_user_oauth_creds(user_api_key_dict) + if _get_oauth2_server_ids(allowed_server_ids) + else {} + ) + # Query all servers the user has access to errors = [] for allowed_server_id in allowed_server_ids: @@ -386,6 +540,9 @@ if MCP_AVAILABLE: server_auth_header = _get_server_auth_header( server, mcp_server_auth_headers, mcp_auth_header ) + user_oauth_extra_headers = await _get_user_oauth_extra_headers( + server, user_api_key_dict, prefetched_creds=prefetched_oauth_creds + ) try: tools_result = await _get_tools_for_single_server( @@ -393,6 +550,7 @@ if MCP_AVAILABLE: server_auth_header, raw_headers_from_request, user_api_key_dict, + extra_headers=user_oauth_extra_headers, ) list_tools_result.extend(tools_result) except Exception as e: @@ -509,6 +667,16 @@ if MCP_AVAILABLE: request, user_api_key_dict, server_id ) + # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). + user_oauth_extra_headers: Optional[Dict[str, str]] = None + target_server = next( + (s for s in allowed_mcp_servers if s.server_id == server_id), None + ) + if target_server is not None: + user_oauth_extra_headers = await _get_user_oauth_extra_headers( + target_server, user_api_key_dict + ) + # Call execute_mcp_tool directly (permission checks already done) result = await execute_mcp_tool( name=tool_name, @@ -518,7 +686,7 @@ if MCP_AVAILABLE: user_api_key_auth=data.get("user_api_key_auth"), mcp_auth_header=data.get("mcp_auth_header"), mcp_server_auth_headers=data.get("mcp_server_auth_headers"), - oauth2_headers=data.get("oauth2_headers"), + oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), raw_headers=data.get("raw_headers"), litellm_logging_obj=data.get("litellm_logging_obj"), ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index fd777b81b2..f7af1127d0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -8,7 +8,7 @@ import contextlib import time import traceback import uuid -from datetime import datetime +from datetime import datetime, timezone from typing import ( Any, AsyncIterator, @@ -877,6 +877,84 @@ if MCP_AVAILABLE: return allowed_mcp_servers + async def _get_user_oauth_extra_headers_from_db( + server: MCPServer, + 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. + + Args: + prefetched_creds: Optional dict keyed by server_id with credential payloads. + When provided, avoids a per-server DB round-trip. + """ + if server.auth_type != MCPAuth.oauth2: + return None + if user_api_key_auth is None: + return None + user_id = getattr(user_api_key_auth, "user_id", None) + server_id = getattr(server, "server_id", None) + if not user_id or not server_id: + return None + try: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + get_user_oauth_credential, + is_oauth_credential_expired, + ) + + if prefetched_creds is not None: + cred = prefetched_creds.get(server_id) + else: + from litellm.proxy.utils import ( # noqa: PLC0415 + get_prisma_client_or_throw, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + 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}" + ) + return None + return {"Authorization": f"Bearer {cred['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}" + ) + return None + + async def _prefetch_oauth_creds_for_user( + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> Dict[str, Dict[str, Any]]: + """Fetch all OAuth2 credentials for the user in one DB query. + + Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. + """ + user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + if not user_id: + return {} + try: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + list_user_oauth_credentials, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + creds = await list_user_oauth_credentials(prisma_client, user_id) + return {c["server_id"]: c for c in creds if "server_id" in c} + except Exception as e: + verbose_logger.warning( + f"_prefetch_oauth_creds_for_user: failed to prefetch for user={user_id}: {e}" + ) + return {} + def _prepare_mcp_server_headers( server: MCPServer, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], @@ -1020,6 +1098,18 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) + # Pre-fetch OAuth credentials only when at least one server uses OAuth2, + # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. + _has_oauth2_server = any( + getattr(s, "auth_type", None) == MCPAuth.oauth2 + for s in allowed_mcp_servers + ) + _prefetched_oauth_creds = ( + await _prefetch_oauth_creds_for_user(user_api_key_auth) + if _has_oauth2_server + else {} + ) + async def _fetch_and_filter_server_tools( server: MCPServer, ) -> List[MCPTool]: @@ -1035,6 +1125,12 @@ if MCP_AVAILABLE: raw_headers=raw_headers, ) + # If no OAuth2 token came from request headers, fall back to pre-fetched creds + if extra_headers is None and server.auth_type == MCPAuth.oauth2: + extra_headers = await _get_user_oauth_extra_headers_from_db( + server, user_api_key_auth, prefetched_creds=_prefetched_oauth_creds + ) + try: tools = await global_mcp_server_manager._get_tools_from_server( server=server, diff --git a/litellm/proxy/_experimental/out/chat.html b/litellm/proxy/_experimental/out/chat/index.html similarity index 100% rename from litellm/proxy/_experimental/out/chat.html rename to litellm/proxy/_experimental/out/chat/index.html diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 50332eeaef..a64ceffda1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,59 +1,40 @@ import enum import json from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + Optional, Union) import httpx -from pydantic import ( - BaseModel, - ConfigDict, - Field, - Json, - field_validator, - model_validator, -) +from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator, + model_validator) from typing_extensions import Required, TypedDict from litellm._uuid import uuid from litellm.types.integrations.slack_alerting import AlertType -from litellm.types.llms.openai import ( - AllMessageValues, - OpenAIFileObject, - ResponsesAPIResponse, -) -from litellm.types.mcp import ( - MCPAuthType, - MCPCredentials, - MCPTransport, - MCPTransportType, -) +from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject, + ResponsesAPIResponse) +from litellm.types.mcp import (MCPAuthType, MCPCredentials, MCPTransport, + MCPTransportType) from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem -from litellm.types.utils import ( - CallTypes, - CostBreakdown, - EmbeddingResponse, - GenericBudgetConfigType, - ImageResponse, - LiteLLMBatch, - LiteLLMFineTuningJob, - LiteLLMPydanticObjectBase, - ModelResponse, - ProviderField, - StandardCallbackDynamicParams, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayloadErrorInformation, - StandardLoggingPayloadStatus, - StandardLoggingVectorStoreRequest, - StandardPassThroughResponseObject, - TextCompletionResponse, -) +from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse, + GenericBudgetConfigType, ImageResponse, + LiteLLMBatch, LiteLLMFineTuningJob, + LiteLLMPydanticObjectBase, ModelResponse, + ProviderField, StandardCallbackDynamicParams, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayloadErrorInformation, + StandardLoggingPayloadStatus, + StandardLoggingVectorStoreRequest, + StandardPassThroughResponseObject, + TextCompletionResponse) from litellm.types.videos.main import VideoObject -from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type +from .types_utils.utils import (get_instance_fn, + validate_custom_validate_return_type) if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -684,6 +665,8 @@ class LiteLLMRoutes(enum.Enum): "/team/daily/activity", "/tag/daily/activity", "/tag/list", + "/audit", + "/audit/{id}", ] + info_routes # All routes accesible by an Org Admin @@ -855,6 +838,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): vector_stores: Optional[List[str]] = None agents: Optional[List[str]] = None agent_access_groups: Optional[List[str]] = None + models: Optional[List[str]] = None class GenerateRequestBase(LiteLLMPydanticObjectBase): @@ -1287,6 +1271,37 @@ class MCPUserCredentialResponse(LiteLLMPydanticObjectBase): has_credential: bool +class MCPOAuthUserCredentialRequest(LiteLLMPydanticObjectBase): + """Stores a user's OAuth2 token for an OpenAPI MCP server.""" + + access_token: str + refresh_token: Optional[str] = None + expires_in: Optional[int] = None # seconds until expiry + scopes: Optional[List[str]] = None + + +class MCPOAuthUserCredentialStatus(LiteLLMPydanticObjectBase): + """Describes whether the calling user has a stored OAuth credential.""" + + server_id: str + has_credential: bool + expires_at: Optional[str] = None # ISO-8601 + is_expired: bool = False + connected_at: Optional[str] = None # ISO-8601 + + +class MCPUserCredentialListItem(LiteLLMPydanticObjectBase): + """One entry in the /user-credentials list.""" + + server_id: str + server_name: Optional[str] = None + alias: Optional[str] = None + credential_type: str # "oauth2" or "byok" + has_credential: bool + expires_at: Optional[str] = None # ISO-8601; None means non-expiring + connected_at: Optional[str] = None # ISO-8601 + + class RejectMCPServerRequest(LiteLLMPydanticObjectBase): review_notes: Optional[str] = None @@ -2430,6 +2445,7 @@ 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 end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None model_config = ConfigDict(arbitrary_types_allowed=True) @@ -2473,7 +2489,8 @@ class UserAPIKeyAuth( This is used to track number of requests/spend for health check calls. """ - from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME return cls( api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, @@ -2505,7 +2522,8 @@ class UserAPIKeyAuth( This is used to track actions performed by automated system jobs. """ - from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME return cls( api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, @@ -2911,7 +2929,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): @model_validator(mode="after") def mask_api_keys(self): - from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + from litellm.litellm_core_utils.sensitive_data_masker import \ + SensitiveDataMasker masker = SensitiveDataMasker(sensitive_patterns={"key"}) diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 955067b486..2a38ceffba 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -18,8 +18,11 @@ async def get_ui_config(): from litellm.proxy.auth.auth_utils import _has_user_setup_sso from litellm.proxy.utils import get_proxy_base_url, get_server_root_path + from litellm.proxy.proxy_server import general_settings + auto_redirect_ui_login_to_sso = ( os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "false").lower() == "true" + or general_settings.get("auto_redirect_ui_login_to_sso", False) is True ) admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index ddeba2100c..b84c74bee4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1122,87 +1122,73 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return response - async def async_post_call_streaming_iterator_hook( + async def _stream_apply_output_masking( self, - user_api_key_dict: UserAPIKeyAuth, response: Any, request_data: dict, ) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]: - """ - Process streaming response chunks to unmask PII tokens when needed. - """ + """Apply Presidio masking to streaming output (apply_to_output=True path).""" from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, ) from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponse - # --- Output masking path (apply_to_output=True) --- - if self.apply_to_output: - all_chunks: List[ModelResponseStream] = [] - try: - async for chunk in response: - if isinstance(chunk, ModelResponseStream): - all_chunks.append(chunk) - elif isinstance(chunk, bytes): - # Anthropic native SSE: pass through as-is - yield chunk # type: ignore[misc] - continue + all_chunks: List[ModelResponseStream] = [] + try: + async for chunk in response: + if isinstance(chunk, ModelResponseStream): + all_chunks.append(chunk) + elif isinstance(chunk, bytes): + yield chunk # type: ignore[misc] + continue - if not all_chunks: - # All chunks were Anthropic native SSE bytes — output - # masking cannot be applied to raw bytes. Log a warning - # so operators know PII masking was skipped for this stream. - verbose_proxy_logger.warning( - "Presidio apply_to_output: streaming response contained only " - "bytes chunks (Anthropic native SSE). Output PII masking was " - "skipped for this response." - ) - return - - assembled_model_response = stream_chunk_builder( - chunks=all_chunks, messages=request_data.get("messages") + if not all_chunks: + verbose_proxy_logger.warning( + "Presidio apply_to_output: streaming response contained only " + "bytes chunks (Anthropic native SSE). Output PII masking was " + "skipped for this response." ) - - if not isinstance(assembled_model_response, ModelResponse): - for chunk in all_chunks: - yield chunk - return - - # Apply Presidio masking on the assembled response - await self._process_response_for_pii( - response=assembled_model_response, - request_data=request_data, - mode="mask", - ) - - mock_response_stream = convert_model_response_to_streaming( - assembled_model_response - ) - yield mock_response_stream return - except Exception as e: - verbose_proxy_logger.error( - f"Error masking streaming PII output: {str(e)}" - ) - # Cannot re-iterate `response` — it's already consumed. - # If we collected chunks before the error, replay those. + assembled_model_response = stream_chunk_builder( + chunks=all_chunks, messages=request_data.get("messages") + ) + + if not isinstance(assembled_model_response, ModelResponse): for chunk in all_chunks: yield chunk return - # --- PII unmasking path (output_parse_pii=True) --- - metadata = (request_data.get("metadata") or {}) if request_data else {} - pii_tokens = metadata.get("pii_tokens", {}) - if not pii_tokens and request_data: - verbose_proxy_logger.debug( - "No pii_tokens in request_data['metadata'] for streaming unmask path" + await self._process_response_for_pii( + response=assembled_model_response, + request_data=request_data, + mode="mask", ) - if not (self.output_parse_pii and pii_tokens): - async for chunk in response: + + mock_response_stream = convert_model_response_to_streaming( + assembled_model_response + ) + yield mock_response_stream + + except Exception as e: + verbose_proxy_logger.error( + f"Error masking streaming PII output: {str(e)}" + ) + for chunk in all_chunks: yield chunk - return + + async def _stream_pii_unmasking( + self, + response: Any, + request_data: dict, + ) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]: + """Apply PII unmasking to streaming output (output_parse_pii=True path).""" + from litellm.llms.base_llm.base_model_iterator import ( + convert_model_response_to_streaming, + ) + from litellm.main import stream_chunk_builder + from litellm.types.utils import ModelResponse remaining_chunks: List[ModelResponseStream] = [] try: @@ -1210,7 +1196,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if isinstance(chunk, ModelResponseStream): remaining_chunks.append(chunk) elif isinstance(chunk, bytes): - # Anthropic native SSE: pass through as-is yield chunk # type: ignore[misc] continue @@ -1226,17 +1211,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk return - # --- PRESERVE USAGE METADATA --- - # stream_chunk_builder might miss usage if it's only in the last chunk - if ( - not getattr(assembled_model_response, "usage", None) - ) and remaining_chunks: - last_chunk = remaining_chunks[-1] - last_chunk_usage = getattr(last_chunk, "usage", None) - if last_chunk_usage: - setattr(assembled_model_response, "usage", last_chunk_usage) + self._preserve_usage_from_last_chunk( + assembled_model_response, remaining_chunks + ) - # Apply PII unmasking to assembled content (unmasking tokens back to original text) await self._process_response_for_pii( response=assembled_model_response, request_data=request_data, @@ -1253,6 +1231,47 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for chunk in remaining_chunks: yield chunk + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict, + ) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]: + """ + Process streaming response chunks to unmask PII tokens when needed. + """ + if self.apply_to_output: + async for chunk in self._stream_apply_output_masking( + response, request_data + ): + yield chunk + return + + metadata = (request_data.get("metadata") or {}) if request_data else {} + pii_tokens = metadata.get("pii_tokens", {}) + if not pii_tokens and request_data: + verbose_proxy_logger.debug( + "No pii_tokens in request_data['metadata'] for streaming unmask path" + ) + if not (self.output_parse_pii and pii_tokens): + async for chunk in response: + yield chunk + return + + async for chunk in self._stream_pii_unmasking(response, request_data): + yield chunk + + @staticmethod + def _preserve_usage_from_last_chunk( + assembled_model_response: Any, + chunks: List[Any], + ) -> None: + """Copy usage metadata from the last chunk when stream_chunk_builder misses it.""" + if not getattr(assembled_model_response, "usage", None) and chunks: + last_chunk_usage = getattr(chunks[-1], "usage", None) + if last_chunk_usage: + setattr(assembled_model_response, "usage", last_chunk_usage) + def get_presidio_settings_from_request_data( self, data: dict ) -> Optional[PresidioPerRequestConfig]: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2af42aa3ab..eb4bb7f884 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1759,6 +1759,34 @@ async def _process_single_key_update( return updated_key_info +async def _validate_mcp_servers_for_key_update( + data: "UpdateKeyRequest", + team_obj: Optional["LiteLLM_TeamTableCachedObj"], + existing_key_row: Any, + prisma_client: Any, + user_api_key_cache: Any, +) -> None: + """Validate MCP servers in object_permission against the effective team.""" + effective_team_obj = team_obj + # If team_id isn't being changed, resolve the existing key's team + if effective_team_obj is None and existing_key_row.team_id: + effective_team_obj = await get_team_object( + team_id=existing_key_row.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + object_permission_dict = ( + data.object_permission.model_dump() + if hasattr(data.object_permission, "model_dump") + else data.object_permission + ) + await validate_key_mcp_servers_against_team( + object_permission=object_permission_dict, + team_obj=effective_team_obj, + ) + + @router.post( "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) @@ -1959,23 +1987,12 @@ async def update_key_fn( # Validate MCP servers in object_permission against the effective team if data.object_permission is not None: - effective_team_obj = team_obj - # If team_id isn't being changed, resolve the existing key's team - if effective_team_obj is None and existing_key_row.team_id: - effective_team_obj = await get_team_object( - team_id=existing_key_row.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - object_permission_dict = ( - data.object_permission.model_dump() - if hasattr(data.object_permission, "model_dump") - else data.object_permission - ) - await validate_key_mcp_servers_against_team( - object_permission=object_permission_dict, - team_obj=effective_team_obj, + await _validate_mcp_servers_for_key_update( + data=data, + team_obj=team_obj, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, ) non_default_values = await prepare_key_update_data( @@ -4587,9 +4604,11 @@ async def _list_key_helper( user_map = {} if expand and "user" in expand: user_ids = [key.user_id for key in keys if key.user_id] - if user_ids: + created_by_ids = [key.created_by for key in keys if key.created_by] + all_ids = list(set(user_ids + created_by_ids)) # Remove duplicates + if all_ids: users = await prisma_client.db.litellm_usertable.find_many( - where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates + where={"user_id": {"in": all_ids}} ) user_map = {user.user_id: user for user in users} @@ -4607,11 +4626,19 @@ async def _list_key_helper( key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) # Include user information if expand includes "user" - if expand and "user" in expand and key.user_id and key.user_id in user_map: - try: - key_dict["user"] = user_map[key.user_id].model_dump() - except Exception: - key_dict["user"] = user_map[key.user_id].dict() + if expand and "user" in expand: + if key.user_id and key.user_id in user_map: + try: + key_dict["user"] = user_map[key.user_id].model_dump() + except Exception: + key_dict["user"] = user_map[key.user_id].dict() + if key.created_by and key.created_by in user_map: + created_by_user = user_map[key.created_by] + key_dict["created_by_user"] = { + "user_id": created_by_user.user_id, + "user_email": created_by_user.user_email, + "user_alias": created_by_user.user_alias, + } if return_full_object is True or (expand and "user" in expand): if use_deleted_table: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index edc2b3048a..77dc1c9724 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -36,6 +36,11 @@ from fastapi import ( ) from fastapi.responses import JSONResponse +try: + from prisma.errors import RecordNotFoundError +except ImportError: + RecordNotFoundError = Exception # type: ignore + import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._uuid import uuid @@ -84,9 +89,13 @@ if MCP_AVAILABLE: delete_user_credential, get_all_mcp_servers_for_user, get_mcp_server, + get_mcp_servers, get_mcp_submissions, + get_user_oauth_credential, + list_user_oauth_credentials, reject_mcp_server, store_user_credential, + store_user_oauth_credential, update_mcp_server, ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -106,7 +115,10 @@ if MCP_AVAILABLE: LitellmUserRoles, MakeMCPServersPublicRequest, MCPApprovalStatus, + MCPOAuthUserCredentialRequest, + MCPOAuthUserCredentialStatus, MCPSubmissionsSummary, + MCPUserCredentialListItem, MCPUserCredentialRequest, MCPUserCredentialResponse, NewMCPServerRequest, @@ -1045,6 +1057,7 @@ if MCP_AVAILABLE: response_model=LiteLLM_MCPServerTable, ) async def fetch_mcp_server( + request: Request, server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -1061,8 +1074,30 @@ if MCP_AVAILABLE: "Database not connected. Connect a database to your proxy" ) - # check to see if server exists for all users + # check to see if server exists (DB first, then registry for config-based servers) mcp_server = await get_mcp_server(prisma_client, server_id) + from_db = mcp_server is not None + + if mcp_server is None: + # Fallback: check registry (config-based servers) - list endpoint uses get_registry() + from litellm.proxy.auth.ip_address_utils import IPAddressUtils + + client_ip = IPAddressUtils.get_mcp_client_ip(request) + registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) + if registry_server is not None and not global_mcp_server_manager._is_server_accessible_from_ip( + registry_server, client_ip + ): + registry_server = None + if registry_server is None: + # Try lookup by server_name or alias (client may use display name in URL) + registry_server = global_mcp_server_manager.get_mcp_server_by_name( + server_id, client_ip=client_ip + ) + if registry_server is not None: + mcp_server = global_mcp_server_manager._build_mcp_server_table( + registry_server + ) + if mcp_server is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1078,10 +1113,17 @@ if MCP_AVAILABLE: if not is_admin_view: # Perform authz check BEFORE any health check (avoid side-effects for # unauthorized callers). - mcp_server_records = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) - exists = does_mcp_server_exist(mcp_server_records, server_id) + if from_db: + mcp_server_records = await get_all_mcp_servers_for_user( + prisma_client, user_api_key_dict + ) + exists = does_mcp_server_exist(mcp_server_records, server_id) + else: + # Registry/config server: use same access logic as list endpoint + allowed_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_dict + ) + exists = mcp_server.server_id in allowed_server_ids if not exists: raise HTTPException( @@ -1095,7 +1137,8 @@ if MCP_AVAILABLE: ) # At this point caller is authorized to view the server. - await global_mcp_server_manager.add_server(mcp_server) + if from_db: + await global_mcp_server_manager.add_server(mcp_server) # Perform health check on the server using server manager try: @@ -1269,10 +1312,21 @@ if MCP_AVAILABLE: def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer: server = get_cached_temporary_mcp_server(server_id) + if server is None: + # Fall back to real DB/config server (e.g. for the user-side OAuth flow + # which calls these endpoints with a real server_id, not a temp session id). + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = ( + global_mcp_server_manager.get_mcp_server_by_id(server_id) + or global_mcp_server_manager.get_mcp_server_by_name(server_id) + ) if server is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"Temporary MCP server {server_id} not found"}, + detail={"error": f"MCP server {server_id} not found"}, ) return server @@ -1283,8 +1337,8 @@ if MCP_AVAILABLE: async def mcp_authorize( request: Request, server_id: str, - client_id: str, - redirect_uri: str, + client_id: Optional[str] = None, + redirect_uri: str = Query(...), state: str = "", code_challenge: Optional[str] = None, code_challenge_method: Optional[str] = None, @@ -1292,10 +1346,23 @@ if MCP_AVAILABLE: scope: Optional[str] = None, ): mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + # Use the server's stored client_id when the caller doesn't supply one + resolved_client_id = mcp_server.client_id or client_id or "" + if not resolved_client_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "missing_client_id", + "message": ( + "No client_id available for this MCP server. " + "Either configure the server with a client_id or supply one in the request." + ), + }, + ) return await authorize_with_server( request=request, mcp_server=mcp_server, - client_id=client_id, + client_id=resolved_client_id, redirect_uri=redirect_uri, state=state, code_challenge=code_challenge, @@ -1314,18 +1381,30 @@ if MCP_AVAILABLE: grant_type: str = Form(...), code: Optional[str] = Form(None), redirect_uri: Optional[str] = Form(None), - client_id: str = Form(...), + client_id: Optional[str] = Form(None), client_secret: Optional[str] = Form(None), code_verifier: Optional[str] = Form(None), ): mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + resolved_client_id = mcp_server.client_id or client_id or "" + if not resolved_client_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "missing_client_id", + "message": ( + "No client_id available for this MCP server. " + "Either configure the server with a client_id or supply one in the request." + ), + }, + ) return await exchange_token_with_server( request=request, mcp_server=mcp_server, grant_type=grant_type, code=code, redirect_uri=redirect_uri, - client_id=client_id, + client_id=resolved_client_id, client_secret=client_secret, code_verifier=code_verifier, ) @@ -1484,7 +1563,7 @@ if MCP_AVAILABLE: ) try: await delete_user_credential(prisma_client, user_id, server_id) - except Exception: + except RecordNotFoundError: pass # Already deleted or didn't exist from litellm.proxy._experimental.mcp_server.server import ( _invalidate_byok_cred_cache, @@ -1493,6 +1572,182 @@ if MCP_AVAILABLE: _invalidate_byok_cred_cache(user_id, server_id) return MCPUserCredentialResponse(server_id=server_id, has_credential=False) + # ── OAuth2 user-credential endpoints ────────────────────────────────────── + + @router.post( + "/server/{server_id}/oauth-user-credential", + description="Store the calling user's OAuth2 token for an OpenAPI MCP server", + dependencies=[Depends(user_api_key_auth)], + response_model=MCPOAuthUserCredentialStatus, + ) + @management_endpoint_wrapper + async def store_mcp_oauth_user_credential( + server_id: str, + payload: MCPOAuthUserCredentialRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """Persist the OAuth2 access token obtained by the calling user.""" + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + mcp_server = await get_mcp_server(prisma_client, server_id) + if mcp_server is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP Server {server_id} not found"}, + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + await store_user_oauth_credential( + prisma_client, + user_id, + server_id, + payload.access_token, + refresh_token=payload.refresh_token, + expires_in=payload.expires_in, + scopes=payload.scopes, + ) + # Read back the persisted record so the response reflects the stored + # expires_at rather than recomputing it here (which could diverge by + # milliseconds or if the storage logic ever adds a grace period). + stored = await get_user_oauth_credential(prisma_client, user_id, server_id) + expires_at: Optional[str] = stored.get("expires_at") if stored else None + return MCPOAuthUserCredentialStatus( + server_id=server_id, + has_credential=True, + expires_at=expires_at, + is_expired=False, + ) + + @router.delete( + "/server/{server_id}/oauth-user-credential", + description="Revoke the calling user's stored OAuth2 token for an MCP server", + dependencies=[Depends(user_api_key_auth)], + response_model=MCPOAuthUserCredentialStatus, + ) + @management_endpoint_wrapper + async def delete_mcp_oauth_user_credential( + server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """Revoke/delete the user's OAuth2 credential.""" + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + # Only delete if the stored credential is actually an OAuth2 token. + # This prevents accidentally deleting a BYOK credential if one exists + # for the same (user_id, server_id) pair. + cred_to_delete = await get_user_oauth_credential(prisma_client, user_id, server_id) + if cred_to_delete is not None: + try: + await delete_user_credential(prisma_client, user_id, server_id) + except RecordNotFoundError: + pass # Already gone — treat as a successful delete + return MCPOAuthUserCredentialStatus( + server_id=server_id, + has_credential=False, + is_expired=False, + ) + + @router.get( + "/server/{server_id}/oauth-user-credential/status", + description="Check whether the calling user has a stored OAuth2 credential for this MCP server", + dependencies=[Depends(user_api_key_auth)], + response_model=MCPOAuthUserCredentialStatus, + ) + @management_endpoint_wrapper + async def get_mcp_oauth_user_credential_status( + server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """Return credential status (has_credential, expiry) without exposing the token.""" + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + cred = await get_user_oauth_credential(prisma_client, user_id, server_id) + if cred is None: + return MCPOAuthUserCredentialStatus( + server_id=server_id, has_credential=False, is_expired=False + ) + expires_at: Optional[str] = cred.get("expires_at") + is_expired = False + if expires_at: + try: + exp = datetime.fromisoformat(expires_at) + is_expired = exp < datetime.now(timezone.utc) + except Exception: + pass + return MCPOAuthUserCredentialStatus( + server_id=server_id, + has_credential=True, + expires_at=expires_at, + is_expired=is_expired, + connected_at=cred.get("connected_at"), + ) + + @router.get( + "/user-credentials", + description="List all OAuth2 MCP credentials stored for the calling user", + dependencies=[Depends(user_api_key_auth)], + response_model=List[MCPUserCredentialListItem], + ) + @management_endpoint_wrapper + async def list_mcp_user_credentials( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """Return all servers the calling user has connected via OAuth2.""" + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + user_id = user_api_key_dict.user_id or "" + if not user_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "User ID not found in token"}, + ) + oauth_creds = await list_user_oauth_credentials(prisma_client, user_id) + if not oauth_creds: + return [] + # Fetch server metadata for display names — single batch query instead of N+1. + server_ids = [c["server_id"] for c in oauth_creds] + servers = { + srv.server_id: srv + for srv in await get_mcp_servers(prisma_client, server_ids) + } + items: List[MCPUserCredentialListItem] = [] + for cred in oauth_creds: + sid = cred["server_id"] + srv = servers.get(sid) + expires_at: Optional[str] = cred.get("expires_at") + items.append( + MCPUserCredentialListItem( + server_id=sid, + server_name=getattr(srv, "server_name", None) if srv else None, + alias=getattr(srv, "alias", None) if srv else None, + credential_type="oauth2", + has_credential=True, + expires_at=expires_at, # always pass the raw timestamp; client computes expiry state + connected_at=cred.get("connected_at"), + ) + ) + return items + @router.put( "/server", description="Allows deleting mcp serves in the db", diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 7fdd3475c0..19ca2c9f6b 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -26,7 +26,6 @@ from litellm.types.tool_management import ( ToolDetailResponse, ToolInputPolicy, ToolListResponse, - ToolOutputPolicy, ToolPolicyOption, ToolPolicyOptionsResponse, ToolPolicyUpdateRequest, diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 319a0b5eb7..0f426bf604 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -4,7 +4,7 @@ organizations, teams, and keys. """ import json -from typing import Dict, List, Optional, Set, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union from fastapi import HTTPException, status @@ -12,6 +12,12 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient + +if TYPE_CHECKING: + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTableCachedObj, + ) diff --git a/litellm/proxy/mcp_registry.json b/litellm/proxy/mcp_registry.json index 2e1e8f64ea..7c5b21dc39 100644 --- a/litellm/proxy/mcp_registry.json +++ b/litellm/proxy/mcp_registry.json @@ -33,8 +33,8 @@ "icon_url": "https://cdn.simpleicons.org/atlassian", "category": "Developer Tools", "registry_url": "https://registry.modelcontextprotocol.io/servers/com.atlassian%2Fatlassian-mcp-server", - "transport": "sse", - "url": "https://mcp.atlassian.com/v1/sse", + "transport": "http", + "url": "https://mcp.atlassian.com/v1/mcp", "env_vars": [] }, { diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a51d5e82b0..fc4d1bffa5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1133,6 +1133,7 @@ def create_pass_through_route( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), subpath: str = "", # captures sub-paths when include_subpath=True + custom_body: Optional[dict] = None, # caller-supplied body takes precedence over request-parsed body ): from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, @@ -1208,9 +1209,11 @@ def create_pass_through_route( ) if query_params: final_query_params.update(query_params) - # Use the body parsed from the raw request + # Caller-supplied custom_body takes precedence over the request-parsed body final_custom_body: Optional[dict] = None - if isinstance(custom_body_data, dict): + if custom_body is not None: + final_custom_body = custom_body + elif isinstance(custom_body_data, dict): final_custom_body = custom_body_data return await pass_through_request( # type: ignore diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3af72d65b5..b68872e2ed 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] @@ -397,9 +398,6 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) - - // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC - @@index([key_alias]) } model LiteLLM_JWTKeyMapping { @@ -565,9 +563,6 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) - - // SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ... - @@index([user, startTime]) } // View spend, model, api_key per request diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b2832b62d7..3eacc19a6d 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -52,8 +52,8 @@ def _get_max_string_length_prompt_in_db() -> int: return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB -def _is_master_key(api_key: str, _master_key: Optional[str]) -> bool: - if _master_key is None: +def _is_master_key(api_key: Optional[str], _master_key: Optional[str]) -> bool: + if _master_key is None or api_key is None: return False ## string comparison diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 7dfbed5660..8df215d998 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -451,6 +451,7 @@ async def _update_litellm_setting( # Update the in-memory settings in_memory_var = settings.model_dump(exclude_none=True) + setattr(litellm, settings_key, in_memory_var) # Load existing config config = await proxy_config.get_config() @@ -459,7 +460,7 @@ async def _update_litellm_setting( if "litellm_settings" not in config: config["litellm_settings"] = {} - config["litellm_settings"][settings_key] = settings.model_dump(exclude_none=True) + config["litellm_settings"][settings_key] = in_memory_var # Save the updated config await proxy_config.save_config(new_config=config) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index cbfe9464b4..0b68807774 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -288,6 +288,268 @@ async def vector_store_create( ) +@router.get("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.get("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +async def vector_store_retrieve( + request: Request, + vector_store_id: str, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Retrieve a vector store. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/retrieve + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = {"vector_store_id": vector_store_id} + + data = _update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + ) + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_retrieve", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + +@router.get("/v1/vector_stores", dependencies=[Depends(user_api_key_auth)]) +@router.get("/vector_stores", dependencies=[Depends(user_api_key_auth)]) +async def vector_store_list( + request: Request, + fastapi_response: Response, + after: Optional[str] = None, + before: Optional[str] = None, + limit: Optional[int] = 20, + order: Optional[str] = "desc", + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + List vector stores. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/list + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = {} + if after is not None: + data["after"] = after + if before is not None: + data["before"] = before + if limit is not None: + data["limit"] = limit + if order is not None: + data["order"] = order + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_list", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + +@router.post("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.post("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +async def vector_store_update( + request: Request, + vector_store_id: str, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update a vector store. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/modify + """ + from litellm.proxy.proxy_server import ( + _read_request_body, + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = await _read_request_body(request=request) + if "vector_store_id" not in data: + data["vector_store_id"] = vector_store_id + + data = _update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + ) + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_update", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + +@router.delete("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.delete("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +async def vector_store_delete( + request: Request, + vector_store_id: str, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete a vector store. + + API Reference: + https://platform.openai.com/docs/api-reference/vector-stores/delete + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = {"vector_store_id": vector_store_id} + + data = _update_request_data_with_litellm_managed_vector_store_registry( + data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict + ) + + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="avector_store_delete", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + + @router.post( "/v1/indexes", dependencies=[Depends(user_api_key_auth)], diff --git a/litellm/responses/main.py b/litellm/responses/main.py index ffe3d7b742..e4c4df50b6 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -24,6 +24,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i LiteLLMResponsesTransformationHandler, ) from litellm.constants import request_timeout +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_responses_input_with_model_file_ids, @@ -655,34 +656,37 @@ def responses( # Native MCP Responses API ######################################################### if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): - return aresponses_api_with_mcp( - input=input, - model=model, - include=include, - instructions=instructions, - max_output_tokens=max_output_tokens, - prompt=prompt, - metadata=metadata, - parallel_tool_calls=parallel_tool_calls, - previous_response_id=previous_response_id, - reasoning=reasoning, - store=store, - background=background, - stream=stream, - temperature=temperature, - text=text, - tool_choice=tool_choice, - tools=tools, - top_p=top_p, - truncation=truncation, - user=user, - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - custom_llm_provider=custom_llm_provider, + mcp_call_kwargs = { + "input": input, + "model": model, + "include": include, + "instructions": instructions, + "max_output_tokens": max_output_tokens, + "prompt": prompt, + "metadata": metadata, + "parallel_tool_calls": parallel_tool_calls, + "previous_response_id": previous_response_id, + "reasoning": reasoning, + "store": store, + "background": background, + "stream": stream, + "temperature": temperature, + "text": text, + "tool_choice": tool_choice, + "tools": tools, + "top_p": top_p, + "truncation": truncation, + "user": user, + "extra_headers": extra_headers, + "extra_query": extra_query, + "extra_body": extra_body, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, **kwargs, - ) + } + if _is_async: + return aresponses_api_with_mcp(**mcp_call_kwargs) + return run_async_function(aresponses_api_with_mcp, **mcp_call_kwargs) # get provider config responses_api_provider_config: Optional[ diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 5776ef95ac..7a3934ffda 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,3 +1,4 @@ +import re import traceback from datetime import datetime from typing import ( @@ -43,6 +44,13 @@ ToolParam = Any LITELLM_PROXY_MCP_SERVER_URL = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" +# Matches any URL whose path ends with /mcp/ — covers both root-path +# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments. +# A false-positive match (e.g. an external URL that happens to end with /mcp/) results +# in a "server not found" error from the internal gateway, not a silent failure or data leak, +# so this broad pattern is intentional and preferred over anchoring to localhost only. +_PROXY_MCP_PATH_RE = re.compile(r"^https?://.+/mcp/([^/]+)$") + class LiteLLM_Proxy_MCP_Handler: """ @@ -54,7 +62,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _should_use_litellm_mcp_gateway(tools: Optional[Iterable[ToolParam]]) -> bool: """ - Returns True if the user passed a MCP tool with server_url="litellm_proxy" + Returns True if any MCP tool should be handled via the litellm proxy MCP gateway. + This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/. """ if tools: for tool in tools: @@ -64,6 +73,10 @@ class LiteLLM_Proxy_MCP_Handler: LITELLM_PROXY_MCP_SERVER_URL ): return True + if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match( + server_url + ): + return True return False @staticmethod @@ -87,6 +100,18 @@ class LiteLLM_Proxy_MCP_Handler: LITELLM_PROXY_MCP_SERVER_URL ): mcp_tools_with_litellm_proxy.append(tool) + elif isinstance(server_url, str): + # Also intercept URLs like http://localhost:4000/mcp/atlassian_test + # by rewriting them to the internal litellm_proxy format. + m = _PROXY_MCP_PATH_RE.match(server_url) + if m: + rewritten = { + **tool, + "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}", + } + mcp_tools_with_litellm_proxy.append(rewritten) + else: + other_tools.append(tool) else: other_tools.append(tool) else: diff --git a/litellm/router.py b/litellm/router.py index 06def6ceb4..47de15655a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -164,7 +164,11 @@ from litellm.types.utils import ( ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage +from litellm.types.utils import ( + ModelResponseStream, + StandardLoggingPayload, + Usage, +) from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -913,7 +917,19 @@ class Router: def _initialize_vector_store_endpoints(self): """Initialize vector store endpoints.""" - from litellm.vector_stores.main import asearch, create, search + from litellm.vector_stores.main import ( + adelete, + alist, + aretrieve, + asearch, + aupdate, + create, + delete, + list, + retrieve, + search, + update, + ) self.avector_store_search = self.factory_function( asearch, call_type="avector_store_search" @@ -924,6 +940,30 @@ class Router: self.vector_store_create = self.factory_function( create, call_type="vector_store_create" ) + self.avector_store_retrieve = self.factory_function( + aretrieve, call_type="avector_store_retrieve" + ) + self.vector_store_retrieve = self.factory_function( + retrieve, call_type="vector_store_retrieve" + ) + self.avector_store_list = self.factory_function( + alist, call_type="avector_store_list" + ) + self.vector_store_list = self.factory_function( + list, call_type="vector_store_list" + ) + self.avector_store_update = self.factory_function( + aupdate, call_type="avector_store_update" + ) + self.vector_store_update = self.factory_function( + update, call_type="vector_store_update" + ) + self.avector_store_delete = self.factory_function( + adelete, call_type="avector_store_delete" + ) + self.vector_store_delete = self.factory_function( + delete, call_type="vector_store_delete" + ) def _initialize_vector_store_file_endpoints(self): """Initialize vector store file endpoints.""" @@ -4725,6 +4765,10 @@ class Router: "generate_content_stream", "avector_store_search", "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", "avector_store_file_create", "avector_store_file_list", "avector_store_file_retrieve", @@ -4733,6 +4777,10 @@ class Router: "avector_store_file_delete", "vector_store_search", "vector_store_create", + "vector_store_retrieve", + "vector_store_list", + "vector_store_update", + "vector_store_delete", "vector_store_file_create", "vector_store_file_list", "vector_store_file_retrieve", @@ -4798,6 +4846,10 @@ class Router: "generate_content_stream", "vector_store_search", "vector_store_create", + "vector_store_retrieve", + "vector_store_list", + "vector_store_update", + "vector_store_delete", "ocr", "search", "video_generation", @@ -4946,6 +4998,10 @@ class Router: elif call_type in ( "avector_store_search", "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", ): return await self._init_vector_store_api_endpoints( original_function=original_function, diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 951fbfcabd..efb2e73bfb 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -172,6 +172,8 @@ class AgentObjectPermission(TypedDict, total=False): mcp_servers: Optional[List[str]] mcp_access_groups: Optional[List[str]] mcp_tool_permissions: Optional[Dict[str, List[str]]] + models: Optional[List[str]] + agents: Optional[List[str]] class AgentConfig(TypedDict, total=False): diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 2b4d1aaa46..36799b4a9d 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -479,3 +479,588 @@ def search( completion_kwargs=local_vars, extra_kwargs=kwargs, ) + + +@client +async def aretrieve( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> VectorStoreCreateResponse: + """ + Async: Retrieve a vector store. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["aretrieve"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + retrieve, + vector_store_id=vector_store_id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def retrieve( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + """ + Retrieve a vector store. + + Args: + vector_store_id: The ID of the vector store to retrieve. + + Returns: + VectorStoreCreateResponse containing the vector store details. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("aretrieve", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store retrieve is not supported for {custom_llm_provider}" + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={"vector_store_id": vector_store_id}, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_retrieve_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +async def alist( + after: Optional[str] = None, + before: Optional[str] = None, + limit: Optional[int] = 20, + order: Optional[str] = "desc", + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + Async: List vector stores. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["alist"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + list, + after=after, + before=before, + limit=limit, + order=order, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def list( + after: Optional[str] = None, + before: Optional[str] = None, + limit: Optional[int] = 20, + order: Optional[str] = "desc", + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + List vector stores. + + Args: + after: A cursor for use in pagination. + before: A cursor for use in pagination. + limit: A limit on the number of objects to be returned. + order: Sort order by the created_at timestamp. + + Returns: + List of vector stores. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("alist", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store list is not supported for {custom_llm_provider}" + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={ + "after": after, + "before": before, + "limit": limit, + "order": order, + }, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_list_handler( + after=after, + before=before, + limit=limit, + order=order, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +async def aupdate( + vector_store_id: str, + name: Optional[str] = None, + expires_after: Optional[Dict] = None, + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> VectorStoreCreateResponse: + """ + Async: Update a vector store. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["aupdate"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + update, + vector_store_id=vector_store_id, + name=name, + expires_after=expires_after, + metadata=metadata, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def update( + vector_store_id: str, + name: Optional[str] = None, + expires_after: Optional[Dict] = None, + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]: + """ + Update a vector store. + + Args: + vector_store_id: The ID of the vector store to update. + name: The name of the vector store. + expires_after: The expiration policy for the vector store. + metadata: Set of 16 key-value pairs that can be attached to an object. + + Returns: + VectorStoreCreateResponse containing the updated vector store details. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("aupdate", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store update is not supported for {custom_llm_provider}" + ) + + local_vars.update(kwargs) + + vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams = ( + VectorStoreRequestUtils.get_requested_vector_store_create_optional_param( + local_vars + ) + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={ + "vector_store_id": vector_store_id, + "name": name, + **vector_store_update_optional_params, + }, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_update_handler( + vector_store_id=vector_store_id, + vector_store_update_optional_params=vector_store_update_optional_params, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +async def adelete( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + Async: Delete a vector store. + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["adelete"] = True + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + func = partial( + delete, + vector_store_id=vector_store_id, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def delete( + vector_store_id: str, + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + **kwargs, +): + """ + Delete a vector store. + + Args: + vector_store_id: The ID of the vector store to delete. + + Returns: + Deletion confirmation response. + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("adelete", False) is True + + litellm_params = GenericLiteLLMParams(**kwargs) + + if custom_llm_provider is None: + custom_llm_provider = "openai" + + if "/" in custom_llm_provider: + api_type, custom_llm_provider, _, _ = get_llm_provider( + model=custom_llm_provider, + custom_llm_provider=None, + litellm_params=None, + ) + else: + api_type = None + custom_llm_provider = custom_llm_provider + + vector_store_provider_config = ( + ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + ) + + if vector_store_provider_config is None: + raise ValueError( + f"Vector store delete is not supported for {custom_llm_provider}" + ) + + litellm_logging_obj.update_environment_variables( + model=None, + optional_params={"vector_store_id": vector_store_id}, + litellm_params={"litellm_call_id": litellm_call_id}, + custom_llm_provider=custom_llm_provider, + ) + + response = base_llm_http_handler.vector_store_delete_handler( + vector_store_id=vector_store_id, + vector_store_provider_config=vector_store_provider_config, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=None, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f99cf4c0a2..9b1d81fee4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14516,7 +14516,7 @@ "input_cost_per_audio_per_second": 0.00016, "input_cost_per_image": 0.00012, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.0237, + "input_cost_per_video_per_second": 0.00079, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -14527,14 +14527,17 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2-preview": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, "mode": "embedding", "output_cost_per_token": 0, "output_vector_size": 3072, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supports_multimodal": true, "uses_embed_content": true }, @@ -14576,7 +14579,10 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2-preview": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_audio_per_second": 0.00016, + "input_cost_per_image": 0.00012, + "input_cost_per_token": 2e-07, + "input_cost_per_video_per_second": 0.00079, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, @@ -14584,7 +14590,7 @@ "output_cost_per_token": 0, "output_vector_size": 3072, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_multimodal": true, "tpm": 10000000 }, diff --git a/poetry.lock b/poetry.lock index c63d0df793..23f4fad175 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3222,15 +3222,15 @@ files = [ [[package]] name = "litellm-proxy-extras" -version = "0.4.53" +version = "0.4.54" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." optional = true python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" groups = ["main"] markers = "extra == \"proxy\"" files = [ - {file = "litellm_proxy_extras-0.4.53-py3-none-any.whl", hash = "sha256:9224c667144774b6119e4de9b4b2d52fafc58442e6db317785c43b2d833665d6"}, - {file = "litellm_proxy_extras-0.4.53.tar.gz", hash = "sha256:22c53fa8890d93d4a0d24171726e4e2bba8be6fef4838317cb74284fa9d27f70"}, + {file = "litellm_proxy_extras-0.4.54-py3-none-any.whl", hash = "sha256:6621cf529f7f3647eb2dd0d2c417d91db8c7a05c3c592bef251887a122928837"}, + {file = "litellm_proxy_extras-0.4.54.tar.gz", hash = "sha256:2c777ecdf39901c4007ade4466eb6398985ed4000afe3fc2cac997e1169e8cee"}, ] [[package]] @@ -8002,4 +8002,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "3036cfcdc06fb4293e248a2edd9c32a7afe6846920167527e247b2aefd74cfa6" +content-hash = "5ed0af4e3644bc7b5a02b8bfc8b3eda15c014b43aa6da7a9a97a9b070fba5366" diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index b1d4d5a116..64942636a9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2518,6 +2518,39 @@ "search": true, "a2a": false } + }, + "black_forest_labs": { + "display_name": "Black Forest Labs (`black_forest_labs`)", + "url": "https://docs.litellm.ai/docs/providers/black_forest_labs", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": true, + "image_edits": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } + }, + "charity_engine": { + "display_name": "Charity Engine (`charity_engine`)", + "url": "https://docs.litellm.ai/docs/providers/charity_engine", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } } }, "endpoints": { diff --git a/pyproject.toml b/pyproject.toml index dd8747b664..1ddbeac099 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ boto3 = { version = "^1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "^0.4.53", optional = true} +litellm-proxy-extras = {version = "^0.4.54", optional = true} rich = {version = "^13.7.1", optional = true} litellm-enterprise = {version = "^0.1.33", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index ccbfa281d9..1243f46401 100644 --- a/requirements.txt +++ b/requirements.txt @@ -57,7 +57,7 @@ grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.53 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.54 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/schema.prisma b/schema.prisma index d5d17b2bce..939f1eb0f4 100644 --- a/schema.prisma +++ b/schema.prisma @@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable { vector_stores String[] @default([]) agents String[] @default([]) agent_access_groups String[] @default([]) + models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] @@ -388,9 +389,6 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) - - // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC - @@index([key_alias]) } model LiteLLM_JWTKeyMapping { @@ -556,9 +554,6 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) - - // SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ... - @@index([user, startTime]) } // View spend, model, api_key per request diff --git a/tests/image_gen_tests/base_image_generation_test.py b/tests/image_gen_tests/base_image_generation_test.py index e3c7d79b3b..ab46bd36fe 100644 --- a/tests/image_gen_tests/base_image_generation_test.py +++ b/tests/image_gen_tests/base_image_generation_test.py @@ -93,6 +93,8 @@ class BaseImageGenTest(ABC): except Exception as e: if "Your task failed as a result of our safety system." in str(e): pass + elif "ModelDeprecated" in str(e): + pass # Azure model deployment has been deprecated - skip else: pytest.fail(f"An exception occurred - {str(e)}") diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 779798702c..e6af29e0e8 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -271,7 +271,7 @@ def test_trimming_should_not_change_original_messages(): assert messages == messages_copy -@pytest.mark.parametrize("model", ["gpt-4-0125-preview", "claude-3-opus-20240229"]) +@pytest.mark.parametrize("model", ["gpt-4-0125-preview", "claude-sonnet-4-6"]) def test_trimming_with_model_cost_max_input_tokens(model): messages = [ {"role": "system", "content": "This is a normal system message"}, @@ -521,7 +521,7 @@ def test_function_to_dict(): ("gpt-3.5-turbo", True), ("azure/gpt-4-1106-preview", True), ("groq/gemma-7b-it", True), - ("gemini/gemini-1.5-flash", True), + ("gemini/gemini-2.5-flash", True), ], ) def test_supports_function_calling(model, expected_bool): @@ -1062,8 +1062,8 @@ def test_parse_content_for_reasoning(content, expected_reasoning, expected_conte @pytest.mark.parametrize( "model, expected_bool", [ - ("vertex_ai/gemini-1.5-pro", True), - ("gemini/gemini-1.5-pro", True), + ("vertex_ai/gemini-2.5-pro", True), + ("gemini/gemini-2.5-pro", True), ("predibase/llama3-8b-instruct", True), ("databricks/databricks-meta-llama-3-1-70b-instruct", True), ("gpt-3.5-turbo", False), @@ -1074,7 +1074,7 @@ def test_supports_response_schema(model, expected_bool): """ Unit tests for 'supports_response_schema' helper function. - Should be true for gemini-1.5-pro on google ai studio / vertex ai AND predibase models + Should be true for gemini-2.5-pro on google ai studio / vertex ai AND predibase models Should be false otherwise """ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -1093,7 +1093,7 @@ def test_supports_response_schema(model, expected_bool): ("gpt-3.5-turbo", True), ("gpt-4", True), ("command-nightly", False), - ("gemini-pro", True), + ("gemini-2.5-pro", True), ], ) def test_supports_function_calling_v2(model, expected_bool): @@ -1109,10 +1109,10 @@ def test_supports_function_calling_v2(model, expected_bool): @pytest.mark.parametrize( "model, expected_bool", [ - ("gpt-4-vision-preview", True), + ("gpt-4o", True), ("gpt-3.5-turbo", False), - ("claude-3-opus-20240229", True), - ("gemini-pro-vision", True), + ("claude-sonnet-4-6", True), + ("gemini-2.5-flash", True), ("command-nightly", False), ], ) @@ -1727,7 +1727,7 @@ def test_supports_vision_gemini(): litellm.model_cost = litellm.get_model_cost_map(url="") from litellm.utils import supports_vision - assert supports_vision("gemini-1.5-pro") is True + assert supports_vision("gemini-2.5-pro") is True def test_pick_cheapest_chat_model_from_llm_provider(): diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index c9ee362539..796b35b436 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -271,7 +271,7 @@ def test_gemini_context_caching_separate_messages(): def test_gemini_image_generation(): # litellm._turn_on_debug() response = completion( - model="gemini/gemini-2.0-flash-exp-image-generation", + model="gemini/gemini-2.5-flash-image-preview", messages=[{"role": "user", "content": "Generate an image of a cat"}], modalities=["image", "text"], ) diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index 2f78f27361..dd060a56d2 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -1012,8 +1012,8 @@ def test_completion_cost_azure_common_deployment_name(): @pytest.mark.parametrize( "model, custom_llm_provider", [ - ("claude-3-5-sonnet-20240620", "anthropic"), - ("gemini/gemini-1.5-flash-001", "gemini"), + ("claude-sonnet-4-6", "anthropic"), + ("claude-haiku-4-5", "anthropic"), ], ) def test_completion_cost_prompt_caching(model, custom_llm_provider): diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index fcdfcfe6e7..a28151d47a 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -1085,10 +1085,15 @@ def test_standard_logging_payload(model, turn_off_message_logging): if turn_off_message_logging: print("checks redacted-by-litellm") assert "redacted-by-litellm" == slobject["messages"][0]["content"] - # response is a full ModelResponse dict (choices format) since d84e5e381acf response = slobject["response"] - assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" - assert response["choices"][0]["message"].get("audio") is None + if "choices" in response: + assert ( + response["choices"][0]["message"]["content"] + == "redacted-by-litellm" + ) + assert response["choices"][0]["message"].get("audio") is None + else: + assert response["text"] == "redacted-by-litellm" @pytest.mark.parametrize( @@ -1188,10 +1193,15 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream): if turn_off_message_logging: print("checks redacted-by-litellm") assert "redacted-by-litellm" == slobject["messages"][0]["content"] - # response is a full ModelResponse dict (choices format) since d84e5e381acf response = slobject["response"] - assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" - assert response["choices"][0]["message"].get("audio") is None + if "choices" in response: + assert ( + response["choices"][0]["message"]["content"] + == "redacted-by-litellm" + ) + assert response["choices"][0]["message"].get("audio") is None + else: + assert response["text"] == "redacted-by-litellm" @pytest.mark.skip(reason="Works locally. Flaky on ci/cd") diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 4cc2723ace..2c950d7906 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -927,7 +927,7 @@ def test_anthropic_tool_calling_exception(): ] try: litellm.completion( - model="claude-3-5-sonnet-20240620", + model="claude-haiku-4-5-20251001", messages=[{"role": "user", "content": "Hey, how's it going?"}], tools=tools, ) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 128ca517ba..2544e06598 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1456,6 +1456,7 @@ async def test_add_update_server_with_alias(): mock_mcp_server.authorization_url = None mock_mcp_server.registration_url = None mock_mcp_server.token_url = None + mock_mcp_server.oauth2_flow = None # Additional fields used by build_mcp_server_from_table mock_mcp_server.extra_headers = None mock_mcp_server.allow_all_keys = False @@ -1511,6 +1512,7 @@ async def test_add_update_server_without_alias(): mock_mcp_server.authorization_url = None mock_mcp_server.registration_url = None mock_mcp_server.token_url = None + mock_mcp_server.oauth2_flow = None # Additional fields used by build_mcp_server_from_table mock_mcp_server.extra_headers = None mock_mcp_server.allow_all_keys = False @@ -1566,6 +1568,7 @@ async def test_add_update_server_fallback_to_server_id(): mock_mcp_server.authorization_url = None mock_mcp_server.registration_url = None mock_mcp_server.token_url = None + mock_mcp_server.oauth2_flow = None # Additional fields used by build_mcp_server_from_table - set explicitly # to avoid MagicMock objects being passed to Pydantic MCPServer constructor mock_mcp_server.extra_headers = None @@ -1823,6 +1826,7 @@ async def test_get_tools_for_single_server(): mock_manager._get_tools_from_server.assert_called_once_with( server=mock_server, mcp_auth_header="Bearer test_token", + extra_headers=None, add_prefix=False, raw_headers=None, ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 00c751c6fd..e907e92e66 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -42,7 +42,9 @@ from litellm.types.utils import CacheCreationTokenDetails, Usage def test_reasoning_tokens_no_price_set(): - model = "o1-mini" + # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics + # (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token) + model = "o1" custom_llm_provider = "openai" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -266,7 +268,8 @@ def test_image_tokens_fallback_to_base_cost(): def test_generic_cost_per_token_above_200k_tokens(): - model = "gemini-2.5-pro-exp-03-25" + # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing + model = "gemini-2.5-pro" custom_llm_provider = "vertex_ai" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 9eb8ae542e..a3bd0274dd 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -121,17 +121,20 @@ def test_get_cost_for_built_in_tools_file_search(): def test_get_cost_for_anthropic_web_search(): """ - Test that the cost for a web search is 0.00 when no response object is provided + Test that Anthropic web search cost is tracked when usage.server_tool_use.web_search_requests + is set. Use claude-3-7-sonnet-20250219 (has search_context_cost_per_query) and + custom_llm_provider=anthropic so get_cost_for_anthropic_web_search is invoked. """ from litellm.types.utils import ServerToolUse, Usage - model = "claude-3-7-sonnet-latest" + model = "claude-3-7-sonnet-20250219" usage = Usage(server_tool_use=ServerToolUse(web_search_requests=1)) cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( model=model, usage=usage, response_object=None, standard_built_in_tools_params=None, + custom_llm_provider="anthropic", ) assert cost > 0.0 diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index acb55a9739..3dede83032 100644 --- a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -219,6 +219,7 @@ class TestAgentCoreStreamingJsonFallback: messages=[{"role": "user", "content": "test"}], stream=True, client=client, + api_key="test-jwt-token", ) # Collect content across all chunks @@ -257,6 +258,7 @@ class TestAgentCoreStreamingJsonFallback: messages=[{"role": "user", "content": "test"}], stream=True, client=client, + api_key="test-jwt-token", ) # Collect content across all chunks @@ -289,6 +291,7 @@ class TestAgentCoreStreamingJsonFallback: messages=[{"role": "user", "content": "test"}], stream=True, client=client, + api_key="test-jwt-token", ) async def test_async_streaming_malformed_json_raises_error(self): @@ -316,4 +319,5 @@ class TestAgentCoreStreamingJsonFallback: messages=[{"role": "user", "content": "test"}], stream=True, client=client, + api_key="test-jwt-token", ) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 345f3ae7c5..317faa5457 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2616,11 +2616,11 @@ def test_empty_assistant_message_handling(): empty or whitespace-only content with a placeholder to prevent AWS Bedrock Converse API 400 Bad Request errors. """ + # Import the litellm module that factory.py uses to ensure we patch the correct reference + import litellm.litellm_core_utils.prompt_templates.factory as factory_module from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, ) - # Import the litellm module that factory.py uses to ensure we patch the correct reference - import litellm.litellm_core_utils.prompt_templates.factory as factory_module # Test case 1: Empty string content - test with modify_params=True to prevent merging messages = [ @@ -3135,7 +3135,12 @@ def test_native_structured_output_no_fake_stream(): def test_transform_request_with_output_config(): """Test that outputConfig flows through _transform_request_helper into the final request.""" - from litellm.types.llms.bedrock import OutputConfigBlock, OutputFormat, OutputFormatStructure, JsonSchemaDefinition + from litellm.types.llms.bedrock import ( + JsonSchemaDefinition, + OutputConfigBlock, + OutputFormat, + OutputFormatStructure, + ) config = AmazonConverseConfig() @@ -3170,6 +3175,29 @@ def test_transform_request_with_output_config(): assert result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "TestSchema" +def test_transform_request_strips_anthropic_output_config(): + """ + output_config is Anthropic-specific and must never be forwarded to Bedrock. + """ + config = AmazonConverseConfig() + messages = [{"role": "user", "content": "hello"}] + + result = config._transform_request( + model="us.amazon.nova-pro-v1:0", + messages=messages, + optional_params={ + "maxTokens": 64, + "output_config": {"effort": "low"}, + }, + litellm_params={}, + headers={}, + ) + + assert "outputConfig" not in result + additional_fields = result.get("additionalModelRequestFields", {}) + assert "output_config" not in additional_fields + + def test_transform_response_native_structured_output(): """Test response handling when model returns JSON as text content (native structured output).""" response_json = { diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 3c1fa52cb0..5bb4942dde 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -107,12 +107,19 @@ class TestSnowflakeToolTransformation: """ Test that string tool_choice values are transformed to Snowflake object format. - Snowflake requires tool_choice to be an object, not a string. - Ref: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/reference/cortex-inference#post--api-v2-cortex-inference-complete-req-body-schema + Snowflake's API (like Anthropic) requires tool_choice as an object + with a "type" field, not as a bare string. OpenAI's "required" maps + to Snowflake's "any". """ config = SnowflakeConfig() - for value in ["auto", "required", "none"]: + expected_mappings = { + "auto": {"type": "auto"}, + "required": {"type": "any"}, + "none": {"type": "none"}, + } + + for value, expected in expected_mappings.items(): optional_params = {"tool_choice": value} transformed_request = config.transform_request( @@ -123,8 +130,10 @@ class TestSnowflakeToolTransformation: headers={}, ) - # Snowflake requires object format: {"type": "auto"} not string "auto" - assert transformed_request["tool_choice"] == {"type": value} + assert transformed_request["tool_choice"] == expected, ( + f"tool_choice='{value}' should be transformed to {expected}, " + f"got {transformed_request['tool_choice']}" + ) def test_transform_response_with_tool_calls(self): """ diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 525b6b2ce7..d483a81a34 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -11,7 +11,6 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm.llms.vertex_ai.common_utils import ( - _build_vertex_schema_for_gemini_2, _get_vertex_url, convert_anyof_null_to_nullable, get_vertex_location_from_url, @@ -1383,93 +1382,3 @@ def test_add_object_type_does_not_add_type_when_anyof_present(): # Verify type was not added (anyOf handles the type) assert "type" not in input_schema, "type should not be added when anyOf is present" - - -class TestBuildVertexSchemaForGemini2: - """Tests for _build_vertex_schema_for_gemini_2 — minimal transform for Gemini 2.0+ tools.""" - - def test_jsonvalue_standalone_preserved(self): - """JsonValue (bare {}) should NOT be coerced to {"type": "object"}.""" - schema = { - "type": "object", - "properties": { - "name": {"type": "string"}, - "value": {}, - }, - "required": ["name", "value"], - } - result = _build_vertex_schema_for_gemini_2(schema) - assert result["properties"]["value"] == {} - - def test_optional_jsonvalue_anyof_preserved(self): - """Optional[JsonValue] anyOf with null should be preserved, not converted to nullable.""" - schema = { - "type": "object", - "properties": { - "name": {"type": "string"}, - "value": { - "anyOf": [ - {"type": "array", "items": {}}, - {}, - {"type": "null"}, - ] - }, - }, - "required": ["name"], - } - result = _build_vertex_schema_for_gemini_2(schema) - value_schema = result["properties"]["value"] - assert "anyOf" in value_schema - assert len(value_schema["anyOf"]) == 3 - assert {"type": "null"} in value_schema["anyOf"] - assert {} in value_schema["anyOf"] - - def test_ref_defs_resolved(self): - """$ref/$defs should be resolved since Gemini doesn't support them in tool params.""" - schema = { - "type": "object", - "properties": { - "value": {"$ref": "#/$defs/JsonValue"}, - }, - "$defs": {"JsonValue": {}}, - } - result = _build_vertex_schema_for_gemini_2(schema) - assert "$ref" not in result["properties"]["value"] - assert "$defs" not in result - assert result["properties"]["value"] == {} - - def test_unsupported_fields_stripped(self): - """Fields not in Vertex Schema TypedDict should be removed.""" - schema = { - "type": "object", - "properties": { - "name": {"type": "string", "additionalProperties": False}, - }, - "additionalProperties": False, - "$schema": "http://json-schema.org/draft-07/schema#", - } - result = _build_vertex_schema_for_gemini_2(schema) - assert "additionalProperties" not in result - assert "$schema" not in result - - def test_no_type_coercion(self): - """Schemas without type should NOT have type: object added.""" - schema = { - "type": "object", - "properties": { - "data": {"description": "Any data"}, - }, - } - result = _build_vertex_schema_for_gemini_2(schema) - assert "type" not in result["properties"]["data"] - - def test_items_empty_preserved(self): - """items: {} should NOT be coerced to items: {"type": "object"}.""" - schema = { - "type": "object", - "properties": { - "values": {"type": "array", "items": {}}, - }, - } - result = _build_vertex_schema_for_gemini_2(schema) - assert result["properties"]["values"]["items"] == {} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index de2ec13b4a..314eed9598 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2093,3 +2093,83 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab assert spend_meta["tool_count_total"] == 1 assert spend_meta["allowed_server_count"] == 1 assert spend_meta["per_server_tool_counts"]["server_a"] == 1 + + +@pytest.mark.asyncio +async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token(): + """ + When _get_tools_from_mcp_servers is called for an OAuth2 MCP server and no + oauth2_headers are provided in the request (e.g. a /responses API call from a + chat UI), the per-user stored token must be fetched from the DB and passed as + extra_headers to _get_tools_from_server. + + The implementation pre-fetches all user credentials in a single bulk query + (_prefetch_oauth_creds_for_user) to avoid N+1 queries in the gather loop. + + This covers the bug where OAuth2 MCP tools were always empty in the /responses + API because the stored credential was never injected. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp import MCPAuth + except ImportError: + pytest.skip("MCP server not available") + + STORED_TOKEN = "atlassian-oauth-access-token-xyz" + SERVER_ID = "srv-oauth2-id" + USER_ID = "user-123" + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id=USER_ID) + + oauth2_server = MagicMock(name="atlassian_server") + oauth2_server.name = "atlassian_test" + oauth2_server.alias = "atlassian_test" + oauth2_server.server_name = "atlassian_test" + oauth2_server.server_id = SERVER_ID + oauth2_server.auth_type = MCPAuth.oauth2 + oauth2_server.extra_headers = None + + # Simulate the DB returning a valid credential for this user+server + prefetched_creds = {SERVER_ID: {"access_token": STORED_TOKEN, "server_id": SERVER_ID}} + + tool_1 = MagicMock() + tool_1.name = "atlassian_test-search" + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[oauth2_server]), + ), patch( + # Patch the bulk prefetch so no real DB connection is needed + "litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user", + new=AsyncMock(return_value=prefetched_creds), + ) as mock_prefetch, patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + ) as mock_manager, patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + side_effect=lambda tools, _server: tools, + ), patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + new=AsyncMock(side_effect=lambda tools, **_: tools), + ): + mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) + + tools = await _get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["atlassian_test"], + mcp_server_auth_headers=None, + oauth2_headers=None, # No token from request — must fall back to DB + ) + + # Bulk credential prefetch was called once (not once per server) + mock_prefetch.assert_awaited_once_with(user_auth) + + # The stored token was forwarded to the MCP transport layer as extra_headers + mock_manager._get_tools_from_server.assert_awaited_once() + call_kwargs = mock_manager._get_tools_from_server.await_args.kwargs + assert call_kwargs["extra_headers"] == {"Authorization": f"Bearer {STORED_TOKEN}"} + + assert tools == [tool_1] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 4f93270c16..1d296f0440 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -484,7 +484,7 @@ class TestListToolsRestAPI: captured = {"called": False} async def fake_get_tools( - server, server_auth_header, raw_headers=None, user_api_key_auth=None + server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None ): captured["called"] = True captured["server"] = server @@ -529,6 +529,175 @@ class TestListToolsRestAPI: assert result["error"] is None assert result["message"] == "Successfully retrieved tools" + async def test_name_resolution_finds_server_by_uuid(self, monkeypatch): + """When server_id is a name string, it should be resolved to its UUID + and used for the tools lookup when the UUID is in allowed_server_ids.""" + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + stub_server = MCPServer( + server_id="uuid-abc-123", + name="my-server", + transport=MCPTransport.sse, + ) + stub_server.alias = "my-server" + stub_server.server_name = "my-server" + stub_server.available_on_public_internet = True + stub_server.allowed_tools = None + stub_server.mcp_info = {"server_name": "my-server"} + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + # Allowed list contains the UUID, not the name + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["uuid-abc-123"] + + captured = {"called": False, "server_arg": None} + + async def fake_get_tools(server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None): + captured["called"] = True + captured["server_arg"] = server + return ["tool-x"] + + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_name", + lambda name: stub_server if name == "my-server" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id", + lambda sid: stub_server if sid == "uuid-abc-123" else None, + raising=False, + ) + monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="my-server", # pass name, not UUID + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert captured["called"] is True + assert captured["server_arg"] is stub_server + assert result["tools"] == ["tool-x"] + assert result["error"] is None + + async def test_name_not_in_allowed_returns_access_denied(self, monkeypatch): + """When name resolves to a server whose UUID is NOT in allowed_server_ids, + the result should be an access_denied error (not a crash or silent pass).""" + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + stub_server = MCPServer( + server_id="uuid-xyz-999", + name="restricted-server", + transport=MCPTransport.sse, + ) + stub_server.available_on_public_internet = True + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + # No allowed servers for this key + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return [] + + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_name", + lambda name: stub_server if name == "restricted-server" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id", + lambda sid: stub_server if sid == "uuid-xyz-999" else None, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="restricted-server", + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert result["tools"] == [] + assert result["error"] == "unexpected_error" + assert "access_denied" in result["message"] + + async def test_oauth2_user_token_injected_for_single_server(self, monkeypatch): + """For a single-server OAuth2 request, _get_user_oauth_extra_headers is called + and the returned headers are forwarded to _get_tools_for_single_server.""" + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + stub_server = MCPServer( + server_id="oauth-server-id", + name="oauth-server", + transport=MCPTransport.sse, + ) + stub_server.alias = "oauth-server" + stub_server.server_name = "oauth-server" + stub_server.available_on_public_internet = True + stub_server.allowed_tools = None + stub_server.mcp_info = {"server_name": "oauth-server"} + stub_server.auth_type = MCPAuth.oauth2 + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["oauth-server-id"] + + oauth_headers = {"Authorization": "Bearer user-oauth-token"} + + async def fake_get_user_oauth_extra_headers(server, user_api_key_dict, prefetched_creds=None): + return oauth_headers + + captured = {} + + async def fake_get_tools(server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None): + captured["server"] = server + captured["auth_header"] = server_auth_header + return ["oauth-tool"] + + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id", + lambda sid: stub_server if sid == "oauth-server-id" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, "_get_user_oauth_extra_headers", + fake_get_user_oauth_extra_headers, raising=False, + ) + monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="oauth-server-id", + user_api_key_dict=UserAPIKeyAuth(user_id="user-123"), + ) + + assert result["tools"] == ["oauth-tool"] + assert result["error"] is None + class TestCallToolRestAPI: pytestmark = pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index c16ee78379..f20c14aa61 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -976,6 +976,43 @@ def test_proxy_admin_viewer_can_access_global_spend_tags(): ) +@pytest.mark.parametrize("route", ["/audit", "/audit/some-log-id"]) +def test_proxy_admin_viewer_can_access_audit_logs(route): + """ + Test that proxy_admin_viewer can access /audit endpoints. + + Admin viewers should be able to view audit logs since these are read-only. + """ + + user_obj = LiteLLM_UserTable( + user_id="viewer_user", + user_email="viewer@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + + valid_token = UserAPIKeyAuth( + user_id="viewer_user", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + + request = MagicMock(spec=Request) + request.query_params = {} + + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as e: + pytest.fail( + f"proxy_admin_viewer should be able to access {route} route. Got error: {str(e)}" + ) + + class TestModelsRouteExemptFromDisableLLMEndpoints: """ Test that /models and /v1/models are exempt from DISABLE_LLM_API_ENDPOINTS. diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index 9d0c771e1d..f15960a607 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -175,6 +175,46 @@ def test_ui_discovery_endpoints_both_routes_return_same_data(): assert response1.json() == response2.json() +def test_ui_discovery_endpoints_with_auto_redirect_via_general_settings(): + """When auto_redirect_ui_login_to_sso is set in general_settings (config.yaml), it should be honored.""" + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \ + patch("litellm.proxy.proxy_server.general_settings", {"auto_redirect_ui_login_to_sso": True}), \ + patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False): + os.environ.pop("AUTO_REDIRECT_UI_LOGIN_TO_SSO", None) + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["auto_redirect_to_sso"] is True + assert data["sso_configured"] is True + + +def test_ui_discovery_endpoints_with_auto_redirect_env_var_overrides_general_settings(): + """Env var and general_settings should both work — either being true enables the feature.""" + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \ + patch("litellm.proxy.proxy_server.general_settings", {"auto_redirect_ui_login_to_sso": False}), \ + patch.dict(os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true", "DISABLE_ADMIN_UI": "false"}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["auto_redirect_to_sso"] is True + + def test_ui_discovery_endpoints_with_admin_ui_disabled(): app = FastAPI() app.include_router(router) diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index f8affda25d..1355ca0abb 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -155,6 +155,103 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN +@pytest.mark.asyncio +async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeypatch): + """ + Default user role set to 'Internal User' via UI, + but SCIM-created users get 'Internal Viewer' instead. + + The UI saves the setting via _update_litellm_setting, which should update + litellm.default_internal_user_params in memory. Then SCIM create_user + should read that in-memory value and assign the correct role. + + This test simulates the full flow: + 1. Start with default_internal_user_params = None + 2. Call update_internal_user_settings (the UI endpoint) to set role to INTERNAL_USER + 3. Create a user via SCIM + 4. Assert the user gets INTERNAL_USER (not INTERNAL_USER_VIEW_ONLY) + """ + from litellm.proxy._types import DefaultInternalUserParams + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _update_litellm_setting, + ) + + # Step 1: Start with no default params (fresh proxy state) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + # Step 2: Simulate the UI saving "Internal User (Create/Delete/View)" as default role + # Mock the proxy_config and store_model_in_db that _update_litellm_setting needs + mock_proxy_config = mocker.MagicMock() + mock_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {}}) + mock_proxy_config.save_config = AsyncMock() + + mocker.patch( + "litellm.proxy.proxy_server.proxy_config", + mock_proxy_config, + ) + mocker.patch( + "litellm.proxy.proxy_server.store_model_in_db", + True, + ) + + import litellm + + settings = DefaultInternalUserParams( + user_role=LitellmUserRoles.INTERNAL_USER, + ) + await _update_litellm_setting( + settings=settings, + settings_key="default_internal_user_params", + in_memory_var=litellm.default_internal_user_params, + success_message="ok", + ) + + # Verify the in-memory variable was actually updated + assert litellm.default_internal_user_params is not None, ( + "BUG: _update_litellm_setting did not update litellm.default_internal_user_params in memory. " + "The local variable reassignment (in_memory_var = ...) doesn't propagate back." + ) + assert litellm.default_internal_user_params.get("user_role") == LitellmUserRoles.INTERNAL_USER + + # Step 3: Create a user via SCIM + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="idontexist@krakentest.tech", + emails=[SCIMUserEmail(value="idontexist@krakentest.tech")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + + new_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id="idontexist@krakentest.tech")), + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await create_user(user=scim_user) + + # Step 4: Verify the user got INTERNAL_USER, not INTERNAL_USER_VIEW_ONLY + called_args = new_user_mock.call_args.kwargs["data"] + assert called_args.user_role == LitellmUserRoles.INTERNAL_USER, ( + f"BUG: SCIM created user with role {called_args.user_role} instead of " + f"{LitellmUserRoles.INTERNAL_USER}. The default_internal_user_params " + f"in-memory variable was not updated by _update_litellm_setting." + ) + + @pytest.mark.asyncio async def test_handle_existing_user_by_email_no_email(mocker): """Should return None when new_user_request has no email""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 09dfdb81cb..de7e865fa3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -4003,6 +4003,7 @@ async def test_list_keys_with_expand_user(): mock_key1 = MagicMock() mock_key1.token = "token1" mock_key1.user_id = "user123" + mock_key1.created_by = None # Set up model_dump() to raise AttributeError so it falls back to dict() mock_key1.model_dump = MagicMock(side_effect=AttributeError("model_dump not available")) mock_key1.dict = MagicMock(return_value=key1_dict) @@ -4016,6 +4017,7 @@ async def test_list_keys_with_expand_user(): mock_key2 = MagicMock() mock_key2.token = "token2" mock_key2.user_id = "user456" + mock_key2.created_by = None # Set up model_dump() to raise AttributeError so it falls back to dict() mock_key2.model_dump = MagicMock(side_effect=AttributeError("model_dump not available")) mock_key2.dict = MagicMock(return_value=key2_dict) @@ -4120,6 +4122,98 @@ async def test_list_keys_with_expand_user(): } +@pytest.mark.asyncio +async def test_list_keys_with_expand_user_includes_created_by_user(): + """ + Test that expand=user also resolves created_by to a user object. + """ + mock_prisma_client = AsyncMock() + + # Key created by user789 but owned by user123 + key1_dict = { + "token": "token1", + "user_id": "user123", + "created_by": "user789", + "key_alias": "key1", + "models": ["gpt-4"], + } + mock_key1 = MagicMock() + mock_key1.token = "token1" + mock_key1.user_id = "user123" + mock_key1.created_by = "user789" + mock_key1.model_dump = MagicMock(return_value=key1_dict) + + mock_find_many_keys = AsyncMock(return_value=[mock_key1]) + mock_count_keys = AsyncMock(return_value=1) + + # Create mock users for both user_id and created_by + mock_user_owner = MagicMock() + mock_user_owner.user_id = "user123" + mock_user_owner.user_email = "owner@example.com" + mock_user_owner.user_alias = "Owner" + mock_user_owner.model_dump = MagicMock(return_value={ + "user_id": "user123", + "user_email": "owner@example.com", + "user_alias": "Owner", + }) + + mock_user_creator = MagicMock() + mock_user_creator.user_id = "user789" + mock_user_creator.user_email = "creator@example.com" + mock_user_creator.user_alias = "Creator" + + mock_find_many_users = AsyncMock(return_value=[mock_user_owner, mock_user_creator]) + + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys + mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users + + async def mock_attach_object_permission(d, _): + return d + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.attach_object_permission_to_dict", + side_effect=mock_attach_object_permission, + ): + args = { + "prisma_client": mock_prisma_client, + "page": 1, + "size": 50, + "user_id": None, + "team_id": None, + "organization_id": None, + "key_alias": None, + "key_hash": None, + "exclude_team_id": None, + "return_full_object": False, + "admin_team_ids": None, + "include_created_by_keys": False, + "expand": ["user"], + } + + result = await _list_key_helper(**args) + + # Verify that the user lookup included both user_id and created_by + call_args = mock_find_many_users.call_args + user_ids_in_query = set(call_args.kwargs["where"]["user_id"]["in"]) + assert user_ids_in_query == {"user123", "user789"} + + # Verify created_by_user is attached + key_result = result["keys"][0] + assert key_result.created_by_user == { + "user_id": "user789", + "user_email": "creator@example.com", + "user_alias": "Creator", + } + + # Verify user (owner) is also still attached + assert key_result.user == { + "user_id": "user123", + "user_email": "owner@example.com", + "user_alias": "Owner", + } + + @pytest.mark.asyncio async def test_list_keys_with_status_deleted(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index ea51965ebf..eeaeb49832 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -76,6 +76,15 @@ def generate_mock_mcp_server_config_record( ) +def _make_mock_request(ip: str = "127.0.0.1"): + """Create a mock Request for fetch_mcp_server tests (IP used for access control).""" + req = MagicMock() + req.client = MagicMock() + req.client.host = ip + req.headers = {} + return req + + def generate_mock_user_api_key_auth( user_role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN, user_id: str = "test_user_id", @@ -735,7 +744,9 @@ class TestListMCPServers: ) result = await fetch_mcp_server( - server_id="server-1", user_api_key_dict=mock_user_auth + request=_make_mock_request(), + server_id="server-1", + user_api_key_dict=mock_user_auth, ) assert result.server_id == "server-1" @@ -788,7 +799,9 @@ class TestListMCPServers: ) result = await fetch_mcp_server( - server_id="server-2", user_api_key_dict=mock_user_auth + request=_make_mock_request(), + server_id="server-2", + user_api_key_dict=mock_user_auth, ) assert result.server_id == "server-2" @@ -796,6 +809,290 @@ class TestListMCPServers: assert not hasattr(result, "credentials") assert result.status == "healthy" + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_from_registry_config_based(self): + """ + Test that fetch_mcp_server finds config-based servers when not in DB. + Config servers appear in list via get_registry() but were 404 on fetch. + """ + config_server = generate_mock_mcp_server_config_record( + server_id="serper_custom_dev", + name="Serper MCP", + url="https://serper.example.com/mcp", + transport="http", + ) + + mock_health_result = generate_mock_mcp_server_db_record( + server_id="serper_custom_dev", alias="Serper MCP" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id = MagicMock( + side_effect=lambda sid: config_server if sid == "serper_custom_dev" else None + ) + mock_manager.get_mcp_server_by_name = MagicMock(return_value=None) + mock_manager._build_mcp_server_table = MagicMock( + return_value=generate_mock_mcp_server_db_record( + server_id="serper_custom_dev", + alias="Serper MCP", + url="https://serper.example.com/mcp", + transport="http", + ) + ) + mock_manager.get_allowed_mcp_servers = AsyncMock( + return_value=["serper_custom_dev"] + ) + mock_manager.health_check_server = AsyncMock(return_value=mock_health_result) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(), + server_id="serper_custom_dev", + user_api_key_dict=mock_user_auth, + ) + + assert result.server_id == "serper_custom_dev" + assert result.status == "healthy" + mock_manager.get_mcp_server_by_id.assert_called_with("serper_custom_dev") + mock_manager._build_mcp_server_table.assert_called_once() + + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_from_registry_by_name_passes_client_ip(self): + """ + When lookup by server_id fails, fallback to get_mcp_server_by_name. + Verify client_ip is passed for IP-based access control (security). + """ + config_server = generate_mock_mcp_server_config_record( + server_id="serper_custom_dev", + name="Serper MCP", + url="https://serper.example.com/mcp", + transport="http", + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id = MagicMock(return_value=None) + mock_manager.get_mcp_server_by_name = MagicMock(return_value=config_server) + mock_manager._build_mcp_server_table = MagicMock( + return_value=generate_mock_mcp_server_db_record( + server_id="serper_custom_dev", + alias="Serper MCP", + url="https://serper.example.com/mcp", + transport="http", + ) + ) + mock_manager.get_allowed_mcp_servers = AsyncMock( + return_value=["serper_custom_dev"] + ) + mock_manager.health_check_server = AsyncMock( + return_value=generate_mock_mcp_server_db_record( + server_id="serper_custom_dev", alias="Serper MCP" + ) + ) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(ip="192.168.1.100"), + server_id="Serper MCP", + user_api_key_dict=mock_user_auth, + ) + + assert result.server_id == "serper_custom_dev" + mock_manager.get_mcp_server_by_id.assert_called_with("Serper MCP") + mock_manager.get_mcp_server_by_name.assert_called_once_with( + "Serper MCP", client_ip="192.168.1.100" + ) + + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_from_registry_non_admin_denied(self): + """ + Non-admin user: config server NOT in allowed_server_ids -> 403. + """ + config_server = generate_mock_mcp_server_config_record( + server_id="restricted_server", + name="Restricted MCP", + url="https://restricted.example.com/mcp", + transport="http", + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id = MagicMock( + side_effect=lambda sid: config_server if sid == "restricted_server" else None + ) + mock_manager.get_mcp_server_by_name = MagicMock(return_value=None) + mock_manager._build_mcp_server_table = MagicMock( + return_value=generate_mock_mcp_server_db_record( + server_id="restricted_server", + alias="Restricted MCP", + url="https://restricted.example.com/mcp", + transport="http", + ) + ) + mock_manager.get_allowed_mcp_servers = AsyncMock( + return_value=["other_server"] # restricted_server NOT in list + ) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=False, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + with pytest.raises(HTTPException) as exc_info: + await fetch_mcp_server( + request=_make_mock_request(), + server_id="restricted_server", + user_api_key_dict=mock_user_auth, + ) + + assert exc_info.value.status_code == 403 + mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth) + + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_from_registry_non_admin_granted(self): + """ + Non-admin user: config server IS in allowed_server_ids -> 200. + """ + config_server = generate_mock_mcp_server_config_record( + server_id="allowed_config_server", + name="Allowed MCP", + url="https://allowed.example.com/mcp", + transport="http", + ) + + mock_health_result = generate_mock_mcp_server_db_record( + server_id="allowed_config_server", alias="Allowed MCP" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id = MagicMock( + side_effect=lambda sid: config_server if sid == "allowed_config_server" else None + ) + mock_manager.get_mcp_server_by_name = MagicMock(return_value=None) + mock_manager._build_mcp_server_table = MagicMock( + return_value=generate_mock_mcp_server_db_record( + server_id="allowed_config_server", + alias="Allowed MCP", + url="https://allowed.example.com/mcp", + transport="http", + ) + ) + mock_manager.get_allowed_mcp_servers = AsyncMock( + return_value=["allowed_config_server"] + ) + mock_manager.health_check_server = AsyncMock(return_value=mock_health_result) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=False, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(), + server_id="allowed_config_server", + user_api_key_dict=mock_user_auth, + ) + + assert result.server_id == "allowed_config_server" + assert result.status == "healthy" + mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth) + class TestTeamScopedMCPServerAccess: """Tests for cross-team information disclosure and restricted key bypass fixes.""" @@ -961,6 +1258,11 @@ class TestTemporaryMCPSessionEndpoints: existing_server.client_id = "client-123" existing_server.client_secret = "secret-xyz" existing_server.scopes = ["scope:a", "scope:b"] + existing_server.aws_access_key_id = None + existing_server.aws_secret_access_key = None + existing_server.aws_session_token = None + existing_server.aws_region_name = None + existing_server.aws_service_name = None mock_manager = MagicMock() mock_manager.get_mcp_server_by_id.return_value = existing_server @@ -1068,6 +1370,11 @@ class TestTemporaryMCPSessionEndpoints: client_id="client-id", client_secret="client-secret", scopes=["scope1"], + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_region_name=None, + aws_service_name=None, ) built_server = generate_mock_mcp_server_config_record(server_id="temp-server") mock_manager = MagicMock() @@ -2005,3 +2312,183 @@ class TestValidateMCPRequiredFields: _validate_mcp_required_fields(payload) assert exc_info.value.status_code == 500 assert "source_Url" in str(exc_info.value.detail) + + +# ── OAuth user credential endpoint unit tests ────────────────────────────────── + + +def _make_user_auth(user_id: str = "user-abc") -> "UserAPIKeyAuth": + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def _make_prisma_client(): + """Return a minimal mock PrismaClient accepted by get_prisma_client_or_throw.""" + client = MagicMock() + client.db = MagicMock() + return client + + +@pytest.mark.asyncio +async def test_store_mcp_oauth_user_credential_returns_status(): + """store_mcp_oauth_user_credential persists the token and echoes back status.""" + from litellm.proxy._types import ( + MCPOAuthUserCredentialRequest, + MCPOAuthUserCredentialStatus, + ) + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + store_mcp_oauth_user_credential, + ) + + server_id = "srv-1" + user_id = "user-123" + stored_payload = { + "type": "oauth2", + "access_token": "tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "connected_at": "2026-01-01T00:00:00+00:00", + "server_id": server_id, + } + + mock_prisma = _make_prisma_client() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential", + new=AsyncMock(return_value=stored_payload), + ), + ): + result = await store_mcp_oauth_user_credential( + server_id=server_id, + payload=MCPOAuthUserCredentialRequest( + access_token="tok", + expires_in=3600, + ), + user_api_key_dict=_make_user_auth(user_id), + ) + + assert isinstance(result, MCPOAuthUserCredentialStatus) + assert result.has_credential is True + assert result.server_id == server_id + # expires_at should come from the stored record, not be recomputed + assert result.expires_at == "2099-01-01T00:00:00+00:00" + + +@pytest.mark.asyncio +async def test_delete_mcp_oauth_user_credential_only_deletes_oauth(): + """delete_mcp_oauth_user_credential only deletes OAuth2 credentials, not BYOK.""" + from litellm.proxy._types import MCPOAuthUserCredentialStatus + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + delete_mcp_oauth_user_credential, + ) + + server_id = "srv-2" + user_id = "user-456" + delete_mock = AsyncMock(return_value=None) + + # When get_user_oauth_credential returns None (no OAuth cred), delete should NOT be called. + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential", + new=delete_mock, + ), + ): + result = await delete_mcp_oauth_user_credential( + server_id=server_id, + user_api_key_dict=_make_user_auth(user_id), + ) + + delete_mock.assert_not_called() + assert isinstance(result, MCPOAuthUserCredentialStatus) + assert result.has_credential is False + + +@pytest.mark.asyncio +async def test_list_mcp_user_credentials_batch_server_fetch(): + """list_mcp_user_credentials uses a single batch DB call, not N+1 queries.""" + from litellm.proxy._types import MCPUserCredentialListItem + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + list_mcp_user_credentials, + ) + + user_id = "user-789" + server_id = "srv-3" + stored_creds = [ + { + "type": "oauth2", + "access_token": "tok", + "expires_at": "2099-01-01T00:00:00+00:00", + "connected_at": "2026-01-01T00:00:00+00:00", + "server_id": server_id, + } + ] + mock_server = generate_mock_mcp_server_db_record(server_id=server_id, alias="My Server") + # get_mcp_servers (batch) should be called once; get_mcp_server (single) must not be called. + batch_mock = AsyncMock(return_value=[mock_server]) + single_mock = AsyncMock(return_value=mock_server) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_user_oauth_credentials", + new=AsyncMock(return_value=stored_creds), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_servers", + new=batch_mock, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=single_mock, + ), + ): + result = await list_mcp_user_credentials( + user_api_key_dict=_make_user_auth(user_id), + ) + + batch_mock.assert_called_once() + single_mock.assert_not_called() + assert len(result) == 1 + assert isinstance(result[0], MCPUserCredentialListItem) + assert result[0].server_id == server_id + assert result[0].alias == "My Server" + # expires_at should always be the raw timestamp (not set to None when expired) + assert result[0].expires_at == "2099-01-01T00:00:00+00:00" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py index 2c3bbc0e6e..be38d08327 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py @@ -242,7 +242,11 @@ class TestGeminiPassthroughLoggingHandler: assert "kwargs" in result @pytest.mark.asyncio - async def test_pass_through_success_handler_gemini_routing(self): + @patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler.litellm.completion_cost", + return_value=0.000050, + ) + async def test_pass_through_success_handler_gemini_routing(self, mock_completion_cost): """Test that the success handler correctly routes Gemini requests to the Gemini handler""" handler = PassThroughEndpointLogging() @@ -263,7 +267,7 @@ class TestGeminiPassthroughLoggingHandler: httpx_response=mock_response, response_body=self.mock_gemini_response, logging_obj=mock_logging_obj, - url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent", result="", start_time=self.start_time, end_time=self.end_time, @@ -277,15 +281,15 @@ class TestGeminiPassthroughLoggingHandler: assert result is None # Verify that the logging object has the cost set (from Gemini handler) - assert mock_logging_obj.model_call_details["response_cost"] is not None - assert mock_logging_obj.model_call_details["model"] == "gemini-1.5-flash" + assert mock_logging_obj.model_call_details["response_cost"] == 0.000050 + assert mock_logging_obj.model_call_details["model"] == "gemini-2.0-flash" assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini" # Verify that _handle_logging was called with the correct kwargs handler._handle_logging.assert_called_once() call_kwargs = handler._handle_logging.call_args[1] - assert call_kwargs["response_cost"] is not None - assert call_kwargs["model"] == "gemini-1.5-flash" + assert call_kwargs["response_cost"] == 0.000050 + assert call_kwargs["model"] == "gemini-2.0-flash" assert call_kwargs["custom_llm_provider"] == "gemini" @patch("litellm.completion_cost") diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index c2f6d3fd53..41a573689f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -243,7 +243,7 @@ class TestVertexAIBatchPassthroughHandler: ] total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( - vertex_ai_batch_responses, model_name="gemini-1.5-flash-001" + vertex_ai_batch_responses, model_name="gemini-2.0-flash-001" ) assert usage.total_tokens == 15 diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9a64e641b5..7174253538 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -29,6 +29,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_response_for_spend_logs_payload, _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, + _is_master_key, _sanitize_request_body_for_spend_logs_payload, _should_store_prompts_and_responses_in_spend_logs, get_logging_payload, @@ -1440,3 +1441,30 @@ def test_get_logging_payload_includes_request_duration_ms(): ) assert payload["request_duration_ms"] == 3000 + + +class TestIsMasterKey: + """Tests for _is_master_key handling None inputs without raising TypeError.""" + + def test_none_api_key_returns_false(self): + """Regression: _is_master_key(None, 'sk-master') should return False, not raise TypeError.""" + assert _is_master_key(api_key=None, _master_key="sk-master-key") is False + + def test_none_master_key_returns_false(self): + assert _is_master_key(api_key="sk-some-key", _master_key=None) is False + + def test_both_none_returns_false(self): + assert _is_master_key(api_key=None, _master_key=None) is False + + def test_matching_key_returns_true(self): + assert _is_master_key(api_key="sk-master", _master_key="sk-master") is True + + def test_non_matching_key_returns_false(self): + assert _is_master_key(api_key="sk-other", _master_key="sk-master") is False + + def test_hashed_key_returns_true(self): + from litellm.proxy.utils import hash_token + + master = "sk-master-key-123" + hashed = hash_token(master) + assert _is_master_key(api_key=hashed, _master_key=master) is True diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index b991dcaf4e..1204b11934 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -970,12 +970,12 @@ def test_cost_discount_vertex_ai(): # Save original config original_discount_config = litellm.cost_discount_config.copy() - # Create mock response + # Create mock response (use a model that exists in model_prices_and_context_window.json) response = ModelResponse( id="test-id", choices=[], created=1234567890, - model="gemini-pro", + model="gemini-3-pro-preview", object="chat.completion", usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), ) @@ -984,7 +984,7 @@ def test_cost_discount_vertex_ai(): litellm.cost_discount_config = {} cost_without_discount = completion_cost( completion_response=response, - model="vertex_ai/gemini-pro", + model="vertex_ai/gemini-3-pro-preview", custom_llm_provider="vertex_ai", ) @@ -994,7 +994,7 @@ def test_cost_discount_vertex_ai(): # Calculate cost with discount cost_with_discount = completion_cost( completion_response=response, - model="vertex_ai/gemini-pro", + model="vertex_ai/gemini-3-pro-preview", custom_llm_provider="vertex_ai", ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 44b7ffb30d..64488e2fb6 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -90,7 +90,7 @@ def test_supports_function_calling_github_openai_alias(): def test_supports_function_calling_github_anthropic_alias(): assert ( litellm.utils.supports_function_calling( - model="github/claude-3-5-sonnet-latest" + model="github/claude-3-7-sonnet-20250219" ) is True ) @@ -384,14 +384,14 @@ def test_all_model_configs(): assert ( "max_completion_tokens" in VertexAIAnthropicConfig().get_supported_openai_params( - model="claude-3-5-sonnet-20240620" + model="claude-sonnet-4-6" ) ) assert VertexAIAnthropicConfig().map_openai_params( non_default_params={"max_completion_tokens": 10}, optional_params={}, - model="claude-3-5-sonnet-20240620", + model="claude-sonnet-4-6", drop_params=False, ) == {"max_tokens": 10} @@ -619,6 +619,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_image_above_128k_tokens": {"type": "number"}, "input_cost_per_image_token": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, + "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, "cache_read_input_token_cost_flex": {"type": "number"}, "cache_read_input_token_cost_priority": {"type": "number"}, @@ -700,6 +701,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token": {"type": "number"}, "output_cost_per_token_above_128k_tokens": {"type": "number"}, "output_cost_per_token_above_200k_tokens": {"type": "number"}, + "output_cost_per_token_above_256k_tokens": {"type": "number"}, "output_cost_per_token_above_272k_tokens": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": { @@ -738,6 +740,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, "supports_url_context": {"type": "boolean"}, + "supports_multimodal": {"type": "boolean"}, + "uses_embed_content": {"type": "boolean"}, "supports_reasoning": {"type": "boolean"}, "supports_none_reasoning_effort": {"type": "boolean"}, "supports_xhigh_reasoning_effort": {"type": "boolean"}, @@ -1136,7 +1140,7 @@ def test_get_model_info_shows_supports_computer_use(): [ ("gpt-3.5-turbo", "openai"), ("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"), - ("gemini-1.5-pro", "vertex_ai"), + ("gemini-2.5-pro", "vertex_ai"), ], ) def test_pre_process_non_default_params(model, custom_llm_provider): @@ -1225,14 +1229,14 @@ class TestProxyFunctionCalling: ), # Anthropic models (Claude supports function calling) ( - "claude-3-5-sonnet-20240620", - "litellm_proxy/claude-3-5-sonnet-20240620", + "claude-sonnet-4-6", + "litellm_proxy/claude-sonnet-4-6", True, ), # Google models - ("gemini-pro", "litellm_proxy/gemini-pro", True), - ("gemini/gemini-1.5-pro", "litellm_proxy/gemini/gemini-1.5-pro", True), - ("gemini/gemini-1.5-flash", "litellm_proxy/gemini/gemini-1.5-flash", True), + ("gemini-2.5-pro", "litellm_proxy/gemini-2.5-pro", True), + ("gemini/gemini-2.5-pro", "litellm_proxy/gemini/gemini-2.5-pro", True), + ("gemini/gemini-2.5-flash", "litellm_proxy/gemini/gemini-2.5-flash", True), # Groq models (mixed support) ("groq/gemma-7b-it", "litellm_proxy/groq/gemma-7b-it", True), ( @@ -1437,8 +1441,8 @@ class TestProxyFunctionCalling: ("litellm_proxy/gpt-3.5-turbo", True), ("litellm_proxy/gpt-4", True), ("litellm_proxy/gpt-4o", True), - ("litellm_proxy/claude-3-5-sonnet-20240620", True), - ("litellm_proxy/gemini/gemini-1.5-pro", True), + ("litellm_proxy/claude-sonnet-4-6", True), + ("litellm_proxy/gemini/gemini-2.5-pro", True), # Test proxy models that should not support function calling ("litellm_proxy/command-nightly", False), ("litellm_proxy/anthropic.claude-instant-v1", False), @@ -1483,8 +1487,8 @@ class TestProxyFunctionCalling: [ "litellm_proxy/gpt-3.5-turbo", "litellm_proxy/gpt-4", - "litellm_proxy/claude-3-5-sonnet-20240620", - "litellm_proxy/gemini/gemini-1.5-pro", + "litellm_proxy/claude-sonnet-4-6", + "litellm_proxy/gemini/gemini-2.5-pro", ], ) def test_proxy_model_with_custom_llm_provider_none(self, model_name): diff --git a/tests/test_new_vector_store_endpoints.py b/tests/test_new_vector_store_endpoints.py new file mode 100644 index 0000000000..05774c3667 --- /dev/null +++ b/tests/test_new_vector_store_endpoints.py @@ -0,0 +1,364 @@ +""" +Comprehensive test for new vector store endpoints: retrieve, list, update, delete +Tests both basic functionality and complex scenarios including target_model_names +""" +import asyncio +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_vector_store_retrieve_basic(): + """Test basic vector store retrieve functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_test123", + "object": "vector_store", + "created_at": 1699061776, + "name": "Test Vector Store", + "file_counts": { + "in_progress": 0, + "completed": 5, + "failed": 0, + "cancelled": 0, + "total": 5, + }, + "status": "completed", + "usage_bytes": 12345, + } + + with patch( + "litellm.vector_stores.main.aretrieve", + new=AsyncMock(return_value=mock_response), + ) as mock_retrieve: + result = await router.avector_store_retrieve( + vector_store_id="vs_test123", + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_test123" + assert result["object"] == "vector_store" + assert result["status"] == "completed" + mock_retrieve.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_list_basic(): + """Test basic vector store list functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "object": "list", + "data": [ + { + "id": "vs_test1", + "object": "vector_store", + "created_at": 1699061776, + "name": "Store 1", + }, + { + "id": "vs_test2", + "object": "vector_store", + "created_at": 1699061777, + "name": "Store 2", + }, + ], + "first_id": "vs_test1", + "last_id": "vs_test2", + "has_more": False, + } + + with patch( + "litellm.vector_stores.main.alist", + new=AsyncMock(return_value=mock_response), + ) as mock_list: + result = await router.avector_store_list( + limit=20, + order="desc", + custom_llm_provider="openai", + ) + + assert result["object"] == "list" + assert len(result["data"]) == 2 + assert result["data"][0]["id"] == "vs_test1" + mock_list.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_update_basic(): + """Test basic vector store update functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_test123", + "object": "vector_store", + "created_at": 1699061776, + "name": "Updated Name", + "metadata": {"key": "value"}, + "status": "completed", + } + + with patch( + "litellm.vector_stores.main.aupdate", + new=AsyncMock(return_value=mock_response), + ) as mock_update: + result = await router.avector_store_update( + vector_store_id="vs_test123", + name="Updated Name", + metadata={"key": "value"}, + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_test123" + assert result["name"] == "Updated Name" + assert result["metadata"]["key"] == "value" + mock_update.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_delete_basic(): + """Test basic vector store delete functionality.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_test123", + "object": "vector_store.deleted", + "deleted": True, + } + + with patch( + "litellm.vector_stores.main.adelete", + new=AsyncMock(return_value=mock_response), + ) as mock_delete: + result = await router.avector_store_delete( + vector_store_id="vs_test123", + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_test123" + assert result["deleted"] is True + assert result["object"] == "vector_store.deleted" + mock_delete.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_retrieve(): + """Test async vector store retrieve.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_async123", + "object": "vector_store", + "name": "Async Test Store", + } + + with patch( + "litellm.vector_stores.main.aretrieve", + new=AsyncMock(return_value=mock_response), + ) as mock_aretrieve: + result = await router.avector_store_retrieve( + vector_store_id="vs_async123", + custom_llm_provider="openai", + ) + + assert result["id"] == "vs_async123" + mock_aretrieve.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_list(): + """Test async vector store list.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "object": "list", + "data": [{"id": "vs_1"}, {"id": "vs_2"}], + } + + with patch( + "litellm.vector_stores.main.alist", + new=AsyncMock(return_value=mock_response), + ) as mock_alist: + result = await router.avector_store_list( + limit=10, + custom_llm_provider="openai", + ) + + assert len(result["data"]) == 2 + mock_alist.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_update(): + """Test async vector store update.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_async123", + "name": "Updated Async Name", + } + + with patch( + "litellm.vector_stores.main.aupdate", + new=AsyncMock(return_value=mock_response), + ) as mock_aupdate: + result = await router.avector_store_update( + vector_store_id="vs_async123", + name="Updated Async Name", + custom_llm_provider="openai", + ) + + assert result["name"] == "Updated Async Name" + mock_aupdate.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_store_delete(): + """Test async vector store delete.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "id": "vs_async123", + "deleted": True, + } + + with patch( + "litellm.vector_stores.main.adelete", + new=AsyncMock(return_value=mock_response), + ) as mock_adelete: + result = await router.avector_store_delete( + vector_store_id="vs_async123", + custom_llm_provider="openai", + ) + + assert result["deleted"] is True + mock_adelete.assert_called_once() + + +@pytest.mark.asyncio +async def test_vector_store_list_with_pagination(): + """Test vector store list with pagination parameters.""" + router = litellm.Router(model_list=[]) + + mock_response = { + "object": "list", + "data": [{"id": f"vs_{i}"} for i in range(5)], + "has_more": True, + "first_id": "vs_0", + "last_id": "vs_4", + } + + with patch( + "litellm.vector_stores.main.list", + return_value=mock_response, + ) as mock_list: + result = router.vector_store_list( + limit=5, + after="vs_previous", + order="asc", + custom_llm_provider="openai", + ) + + assert result["has_more"] is True + assert len(result["data"]) == 5 + + # Verify pagination params were passed + call_kwargs = mock_list.call_args.kwargs + assert call_kwargs["limit"] == 5 + assert call_kwargs["after"] == "vs_previous" + assert call_kwargs["order"] == "asc" + + +@pytest.mark.asyncio +async def test_vector_store_update_with_expires_after(): + """Test vector store update with expiration policy.""" + router = litellm.Router(model_list=[]) + + expires_after = { + "anchor": "last_active_at", + "days": 7, + } + + mock_response = { + "id": "vs_test123", + "expires_after": expires_after, + "expires_at": 1699668576, + } + + with patch( + "litellm.vector_stores.main.update", + return_value=mock_response, + ) as mock_update: + result = router.vector_store_update( + vector_store_id="vs_test123", + expires_after=expires_after, + custom_llm_provider="openai", + ) + + assert result["expires_after"]["days"] == 7 + assert result["expires_at"] is not None + + call_kwargs = mock_update.call_args.kwargs + assert call_kwargs["expires_after"] == expires_after + + +def test_router_initializes_new_endpoints(): + """Test that router properly initializes the new vector store endpoints.""" + router = litellm.Router(model_list=[]) + + # Verify all new endpoints are initialized + assert hasattr(router, "vector_store_retrieve") + assert hasattr(router, "avector_store_retrieve") + assert hasattr(router, "vector_store_list") + assert hasattr(router, "avector_store_list") + assert hasattr(router, "vector_store_update") + assert hasattr(router, "avector_store_update") + assert hasattr(router, "vector_store_delete") + assert hasattr(router, "avector_store_delete") + + # Verify they are callable + assert callable(router.vector_store_retrieve) + assert callable(router.avector_store_retrieve) + assert callable(router.vector_store_list) + assert callable(router.avector_store_list) + assert callable(router.vector_store_update) + assert callable(router.avector_store_update) + assert callable(router.vector_store_delete) + assert callable(router.avector_store_delete) + + +if __name__ == "__main__": + # Run basic smoke tests + print("Running smoke tests for new vector store endpoints...") + + # Test router initialization + print("✓ Testing router initialization...") + test_router_initializes_new_endpoints() + print("✓ Router initialization successful") + + # Test basic sync operations + print("✓ Testing basic sync operations...") + asyncio.run(test_vector_store_retrieve_basic()) + asyncio.run(test_vector_store_list_basic()) + asyncio.run(test_vector_store_update_basic()) + asyncio.run(test_vector_store_delete_basic()) + print("✓ Basic sync operations successful") + + # Test async operations + print("✓ Testing async operations...") + asyncio.run(test_async_vector_store_retrieve()) + asyncio.run(test_async_vector_store_list()) + asyncio.run(test_async_vector_store_update()) + asyncio.run(test_async_vector_store_delete()) + print("✓ Async operations successful") + + print("\n✅ All smoke tests passed!") diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index d4f60ff0d7..c062356ebb 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -46,7 +46,7 @@ "@testing-library/user-event": "^14.6.1", "@types/babel__traverse": "^7.28.0", "@types/lodash": "^4.17.15", - "@types/node": "^20", + "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "^5.0.7", "@types/react-dom": "^18", @@ -65,7 +65,7 @@ "postcss": "^8.4.33", "prettier": "3.2.5", "tailwindcss": "^3.4.1", - "typescript": "^5.3.3", + "typescript": "5.9.3", "vite": "^7.1.11", "vitest": "^3.2.4" }, @@ -3394,9 +3394,9 @@ "license": "MIT" }, "node_modules/@types/node": { - "version": "20.19.30", - "resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.30.tgz", - "integrity": "sha512-WJtwWJu7UdlvzEAUm484QNg5eAoq5QR08KDNx7g45Usrs2NtOPiX8ugDqmKdXkyL03rBqU5dYNYVQetEpBHq2g==", + "version": "20.19.37", + "resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.37.tgz", + "integrity": "sha512-8kzdPJ3FsNsVIurqBs7oodNnCEVbni9yUEkaHbgptDACOPW04jimGagZ51E6+lXUwJjgnBw+hyko/lkFWCldqw==", "license": "MIT", "dependencies": { "undici-types": "~6.21.0" @@ -12555,9 +12555,9 @@ } }, "node_modules/typescript": { - "version": "5.3.3", - "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.3.3.tgz", - "integrity": "sha512-pXWcraxM0uxAS+tN0AG/BF2TyqmHO014Z070UsJ+pFvYuRSq8KH8DmWpnbXe0pEPDHXZV3FcAbJkijJ5oNEnWw==", + "version": "5.9.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", + "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", "devOptional": true, "license": "Apache-2.0", "bin": { @@ -13279,21 +13279,6 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } - }, - "node_modules/@next/swc-win32-ia32-msvc": { - "version": "14.2.33", - "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz", - "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==", - "cpu": [ - "ia32" - ], - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } } } } diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 06aff85e3f..5cbe1ead88 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -58,7 +58,7 @@ "@testing-library/user-event": "^14.6.1", "@types/babel__traverse": "^7.28.0", "@types/lodash": "^4.17.15", - "@types/node": "^20", + "@types/node": "20.19.37", "@types/react": "18.2.48", "@types/react-copy-to-clipboard": "^5.0.7", "@types/react-dom": "^18", @@ -77,7 +77,7 @@ "postcss": "^8.4.33", "prettier": "3.2.5", "tailwindcss": "^3.4.1", - "typescript": "^5.3.3", + "typescript": "5.9.3", "vite": "^7.1.11", "vitest": "^3.2.4" }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts index 3681ffc747..ee03a0ab7c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.test.ts @@ -74,7 +74,7 @@ describe("useMCPServers", () => { expect(result.current.isSuccess).toBe(true); }); - expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken); + expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken, undefined); expect(result.current.data).toEqual(mockServers); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index b7d4db2618..32ab83ea75 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -20,11 +20,15 @@ vi.mock("@/components/molecules/notifications_manager", () => ({ // Mock react-query const mockInvalidateQueries = vi.fn(); -vi.mock("@tanstack/react-query", () => ({ - useQueryClient: () => ({ - invalidateQueries: mockInvalidateQueries, - }), -})); +vi.mock("@tanstack/react-query", async (importOriginal) => { + const actual = await importOriginal() as any; + return { + ...actual, + useQueryClient: () => ({ + invalidateQueries: mockInvalidateQueries, + }), + }; +}); // Mock the useModelsInfo hook const mockUseModelsInfo = vi.fn(() => ({ @@ -553,7 +557,7 @@ describe("AllModelsTab", () => { mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); - render(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("gpt-4-delete-test")).toBeInTheDocument(); @@ -597,7 +601,7 @@ describe("AllModelsTab", () => { mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); - render(); + renderWithProviders(); await waitFor(() => { expect(screen.getByText("gpt-4-clickable")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx index 73ff8c51ba..0c4cad8cb0 100644 --- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx +++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx @@ -3,7 +3,11 @@ import { Suspense, useEffect, useMemo } from "react"; import { useSearchParams } from "next/navigation"; -const RESULT_STORAGE_KEY = "litellm-mcp-oauth-result"; +// Written to sessionStorage so both the admin hook (useMcpOAuthFlow) and the +// user hook (useUserMcpOAuthFlow) can pick up the result. Each hook reads +// its own namespace to avoid cross-flow collisions. +const ADMIN_RESULT_KEY = "litellm-mcp-oauth-result"; +const USER_RESULT_KEY = "litellm-user-mcp-oauth-result"; const RETURN_URL_STORAGE_KEY = "litellm-mcp-oauth-return-url"; const resolveDefaultRedirect = () => { @@ -32,6 +36,10 @@ const McpOAuthCallbackContent = () => { type: "litellm-mcp-oauth", code: searchParams.get("code"), state: searchParams.get("state"), + // Forward OAuth provider error params so the hook can surface the real + // reason (e.g. "access_denied") instead of a generic "code missing" error. + error: searchParams.get("error"), + error_description: searchParams.get("error_description"), }; }, [searchParams]); @@ -41,16 +49,16 @@ const McpOAuthCallbackContent = () => { } try { - // Store in both sessionStorage and localStorage for redundancy - window.sessionStorage.setItem(RESULT_STORAGE_KEY, JSON.stringify(payload)); - window.localStorage.setItem(RESULT_STORAGE_KEY, JSON.stringify(payload)); + // 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); } catch (err) { // Silently ignore storage errors } - // Check both sessionStorage and localStorage for return URL - const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY) || - window.localStorage.getItem(RETURN_URL_STORAGE_KEY); + const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY); const destination = returnUrl || resolveDefaultRedirect(); window.location.replace(destination); }, [payload]); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index bddcb0ab59..5f2921203f 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -547,7 +547,7 @@ function CreateKeyPageContent() { ) : page == "policies" ? ( ) : page == "agents" ? ( - + ) : page == "prompts" ? ( ) : page == "transform-request" ? ( diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx index c29ade5d65..5c23cf71ab 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx @@ -125,6 +125,208 @@ describe("EntityUsage", () => { }, }; + const mockAgentSpendData = { + results: [ + { + date: "2025-01-01", + metrics: { + spend: 245.8, + api_requests: 3200, + successful_requests: 3100, + failed_requests: 100, + total_tokens: 1250000, + prompt_tokens: 850000, + completion_tokens: 400000, + cache_read_input_tokens: 50000, + cache_creation_input_tokens: 10000, + }, + breakdown: { + entities: { + "agent-code-review": { + metrics: { + spend: 120.4, + api_requests: 1500, + successful_requests: 1450, + failed_requests: 50, + total_tokens: 620000, + prompt_tokens: 420000, + completion_tokens: 200000, + cache_read_input_tokens: 30000, + cache_creation_input_tokens: 5000, + }, + metadata: { agent_name: "Code Review Agent" }, + api_key_breakdown: {}, + }, + "agent-customer-support": { + metrics: { + spend: 85.2, + api_requests: 1200, + successful_requests: 1170, + failed_requests: 30, + total_tokens: 430000, + prompt_tokens: 290000, + completion_tokens: 140000, + cache_read_input_tokens: 15000, + cache_creation_input_tokens: 3000, + }, + metadata: { agent_name: "Customer Support Agent" }, + api_key_breakdown: {}, + }, + "agent-data-analyst": { + metrics: { + spend: 40.2, + api_requests: 500, + successful_requests: 480, + failed_requests: 20, + total_tokens: 200000, + prompt_tokens: 140000, + completion_tokens: 60000, + cache_read_input_tokens: 5000, + cache_creation_input_tokens: 2000, + }, + metadata: { agent_name: "Data Analyst Agent" }, + api_key_breakdown: {}, + }, + }, + models: { + "gpt-4o": { + metrics: { + spend: 180.0, + api_requests: 2000, + successful_requests: 1950, + failed_requests: 50, + total_tokens: 900000, + prompt_tokens: 600000, + completion_tokens: 300000, + cache_read_input_tokens: 40000, + cache_creation_input_tokens: 8000, + }, + metadata: {}, + api_key_breakdown: {}, + }, + "claude-sonnet-4-20250514": { + metrics: { + spend: 65.8, + api_requests: 1200, + successful_requests: 1150, + failed_requests: 50, + total_tokens: 350000, + prompt_tokens: 250000, + completion_tokens: 100000, + cache_read_input_tokens: 10000, + cache_creation_input_tokens: 2000, + }, + metadata: {}, + api_key_breakdown: {}, + }, + }, + api_keys: {}, + providers: { + openai: { + metrics: { + spend: 180.0, + api_requests: 2000, + successful_requests: 1950, + failed_requests: 50, + total_tokens: 900000, + prompt_tokens: 600000, + completion_tokens: 300000, + cache_read_input_tokens: 40000, + cache_creation_input_tokens: 8000, + }, + }, + anthropic: { + metrics: { + spend: 65.8, + api_requests: 1200, + successful_requests: 1150, + failed_requests: 50, + total_tokens: 350000, + prompt_tokens: 250000, + completion_tokens: 100000, + cache_read_input_tokens: 10000, + cache_creation_input_tokens: 2000, + }, + }, + }, + }, + }, + { + date: "2025-01-02", + metrics: { + spend: 198.5, + api_requests: 2800, + successful_requests: 2720, + failed_requests: 80, + total_tokens: 980000, + prompt_tokens: 670000, + completion_tokens: 310000, + cache_read_input_tokens: 42000, + cache_creation_input_tokens: 9000, + }, + breakdown: { + entities: { + "agent-code-review": { + metrics: { + spend: 95.3, + api_requests: 1300, + successful_requests: 1270, + failed_requests: 30, + total_tokens: 510000, + prompt_tokens: 350000, + completion_tokens: 160000, + cache_read_input_tokens: 25000, + cache_creation_input_tokens: 4000, + }, + metadata: { agent_name: "Code Review Agent" }, + api_key_breakdown: {}, + }, + "agent-customer-support": { + metrics: { + spend: 68.7, + api_requests: 1000, + successful_requests: 970, + failed_requests: 30, + total_tokens: 320000, + prompt_tokens: 220000, + completion_tokens: 100000, + cache_read_input_tokens: 12000, + cache_creation_input_tokens: 3000, + }, + metadata: { agent_name: "Customer Support Agent" }, + api_key_breakdown: {}, + }, + "agent-data-analyst": { + metrics: { + spend: 34.5, + api_requests: 500, + successful_requests: 480, + failed_requests: 20, + total_tokens: 150000, + prompt_tokens: 100000, + completion_tokens: 50000, + cache_read_input_tokens: 5000, + cache_creation_input_tokens: 2000, + }, + metadata: { agent_name: "Data Analyst Agent" }, + api_key_breakdown: {}, + }, + }, + models: {}, + api_keys: {}, + providers: {}, + }, + }, + ], + metadata: { + total_spend: 444.3, + total_api_requests: 6000, + total_successful_requests: 5820, + total_failed_requests: 180, + total_tokens: 2230000, + }, + }; + const defaultProps = { accessToken: "test-token", entityType: "tag" as const, @@ -153,7 +355,7 @@ describe("EntityUsage", () => { mockTeamDailyActivityCall.mockResolvedValue(mockSpendData); mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData); mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData); - mockAgentDailyActivityCall.mockResolvedValue(mockSpendData); + mockAgentDailyActivityCall.mockResolvedValue(mockAgentSpendData); mockUserDailyActivityCall.mockResolvedValue(mockSpendData); }); @@ -231,7 +433,7 @@ describe("EntityUsage", () => { expect(screen.getByText("Agent Spend Overview")).toBeInTheDocument(); await waitFor(() => { - const spendElements = screen.getAllByText("$100.50"); + const spendElements = screen.getAllByText("$444.30"); expect(spendElements.length).toBeGreaterThan(0); }); }); @@ -385,6 +587,87 @@ describe("EntityUsage", () => { }); }); + it("should display Agent Activity tab for team entity type", async () => { + render(); + + await waitFor(() => { + expect(mockTeamDailyActivityCall).toHaveBeenCalled(); + }); + + expect(screen.getByText("Agent Activity")).toBeInTheDocument(); + }); + + it("should not display Agent Activity tab for non-team entity types", async () => { + render(); + + await waitFor(() => { + expect(mockTagDailyActivityCall).toHaveBeenCalled(); + }); + + expect(screen.queryByText("Agent Activity")).not.toBeInTheDocument(); + }); + + it("should display Top Agents Driving Spend card for team entity type", async () => { + render(); + + await waitFor(() => { + expect(mockTeamDailyActivityCall).toHaveBeenCalled(); + }); + + expect(screen.getByText("Top Agents Driving Spend")).toBeInTheDocument(); + }); + + it("should not display Top Agents Driving Spend card for non-team entity types", async () => { + render(); + + await waitFor(() => { + expect(mockTagDailyActivityCall).toHaveBeenCalled(); + }); + + expect(screen.queryByText("Top Agents Driving Spend")).not.toBeInTheDocument(); + }); + + it("should fetch agent activity data when entity type is team", async () => { + render(); + + await waitFor(() => { + expect(mockAgentDailyActivityCall).toHaveBeenCalledWith( + "test-token", + expect.any(Date), + expect.any(Date), + 1, + null, + ); + }); + }); + + it("should not fetch agent activity data for non-team entity types", async () => { + render(); + + await waitFor(() => { + expect(mockTagDailyActivityCall).toHaveBeenCalled(); + }); + + expect(mockAgentDailyActivityCall).not.toHaveBeenCalled(); + }); + + it("should switch to Agent Activity tab for team entity type", async () => { + render(); + + await waitFor(() => { + expect(mockTeamDailyActivityCall).toHaveBeenCalled(); + }); + + const agentActivityTab = screen.getByText("Agent Activity"); + act(() => { + fireEvent.click(agentActivityTab); + }); + + await waitFor(() => { + expect(screen.getAllByText("Activity Metrics").length).toBeGreaterThan(0); + }); + }); + it("should fallback to entity value when no entityList and no team_alias", async () => { const spendDataWithoutAlias = { ...mockSpendData, diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx index a106910cff..3e0343fdf6 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx @@ -100,11 +100,24 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti }); const { teams } = useTeams(); + const [agentSpendData, setAgentSpendData] = useState({ + results: [], + metadata: { + total_spend: 0, + total_api_requests: 0, + total_successful_requests: 0, + total_failed_requests: 0, + total_tokens: 0, + }, + }); + const modelMetrics = processActivityData(spendData, "models", teams || []); const keyMetrics = processActivityData(spendData, "api_keys", teams || []); + const agentMetrics = entityType === "team" ? processActivityData(agentSpendData, "entities", teams || []) : {}; const [selectedTags, setSelectedTags] = useState([]); const [topKeysLimit, setTopKeysLimit] = useState(5); const [topModelsLimit, setTopModelsLimit] = useState(5); + const [topAgentsLimit, setTopAgentsLimit] = useState(5); const fetchSpendData = async () => { if (!accessToken || !dateValue.from || !dateValue.to) return; @@ -171,8 +184,21 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti } }; + const fetchAgentSpendData = async () => { + if (!accessToken || !dateValue.from || !dateValue.to || entityType !== "team") return; + const startTime = new Date(dateValue.from); + const endTime = new Date(dateValue.to); + try { + const data = await agentDailyActivityCall(accessToken, startTime, endTime, 1, null); + setAgentSpendData(data); + } catch (e) { + console.error("Failed to fetch agent activity data:", e); + } + }; + useEffect(() => { fetchSpendData(); + fetchAgentSpendData(); }, [accessToken, dateValue, entityId, selectedTags]); const getTopModels = () => { @@ -209,6 +235,37 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti .slice(0, topModelsLimit); }; + const getTopAgents = () => { + const agentSpend: { [key: string]: any } = {}; + agentSpendData.results.forEach((day) => { + Object.entries(day.breakdown.entities || {}).forEach(([agentId, data]) => { + if (!agentSpend[agentId]) { + agentSpend[agentId] = { + spend: 0, + requests: 0, + successful_requests: 0, + failed_requests: 0, + tokens: 0, + agent_name: (data.metadata as any)?.agent_name || agentId, + }; + } + agentSpend[agentId].spend += data.metrics.spend; + agentSpend[agentId].requests += data.metrics.api_requests; + agentSpend[agentId].successful_requests += data.metrics.successful_requests; + agentSpend[agentId].failed_requests += data.metrics.failed_requests; + agentSpend[agentId].tokens += data.metrics.total_tokens; + }); + }); + + return Object.entries(agentSpend) + .map(([agentId, metrics]) => ({ + key: metrics.agent_name, + ...metrics, + })) + .sort((a, b) => b.spend - a.spend) + .slice(0, topAgentsLimit); + }; + const getTopAPIKeys = () => { console.log("debugTags", { spendData }); const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; @@ -408,6 +465,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti Cost {entityType === "agent" ? "Request / Token Consumption" : "Model Activity"} + {entityType === "team" && Agent Activity} Key Activity Endpoint Activity @@ -621,6 +679,20 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti + {/* Top Agents - only for team entity type */} + {entityType === "team" && ( + + + Top Agents Driving Spend + + + + )} + {/* Spend by Provider */} @@ -696,6 +768,11 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti + {entityType === "team" && ( + + + + )} diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx index 410d751017..b9fe1687e6 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx @@ -249,6 +249,8 @@ vi.mock("@ant-design/icons", async () => { CalendarOutlined: Icon, InfoCircleOutlined: Icon, UserOutlined: Icon, + DownOutlined: Icon, + RightOutlined: Icon, LoadingOutlined, }; }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 8b0d5ffac0..e81e4ceaa3 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -6,7 +6,7 @@ * Works at 1m+ spend logs, by querying an aggregate table instead. */ -import { InfoCircleOutlined, LoadingOutlined } from "@ant-design/icons"; +import { DownOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons"; import { BarChart, Card, @@ -148,6 +148,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const [showCredentialBanner, setShowCredentialBanner] = useState(true); const [topKeysLimit, setTopKeysLimit] = useState(5); const [topModelsLimit, setTopModelsLimit] = useState(5); + const [showTokenBreakdown, setShowTokenBreakdown] = useState(false); const getAllTags = async () => { if (!accessToken) { return; @@ -176,7 +177,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const totalSpend = userSpendData.metadata?.total_spend || 0; // Calculate top models from the breakdown data - const getTopModels = (limit: number = 5) => { + const topModels = useMemo(() => { const modelSpend: { [key: string]: MetricWithMetadata } = {}; userSpendData.results.forEach((day) => { Object.entries(day.breakdown.models || {}).forEach(([model, metrics]) => { @@ -219,10 +220,10 @@ const UsagePage: React.FC = ({ teams, organizations }) => { tokens: metrics.metrics.total_tokens, })) .sort((a, b) => b.spend - a.spend) - .slice(0, limit); - }; + .slice(0, topModelsLimit); + }, [userSpendData.results, topModelsLimit]); - const getTopModelGroups = (limit: number = 5) => { + const topModelGroups = useMemo(() => { const modelGroupSpend: { [key: string]: MetricWithMetadata } = {}; userSpendData.results.forEach((day) => { Object.entries(day.breakdown.model_groups || {}).forEach(([modelGroup, metrics]) => { @@ -266,16 +267,16 @@ const UsagePage: React.FC = ({ teams, organizations }) => { tokens: metrics.metrics.total_tokens, })) .sort((a, b) => b.spend - a.spend) - .slice(0, limit); - }; + .slice(0, topModelsLimit); + }, [userSpendData.results, topModelsLimit]); // Calculate provider spend from the breakdown data - const getProviderSpend = () => { - const providerSpend: { [key: string]: MetricWithMetadata } = {}; + const providerSpend = useMemo(() => { + const providerSpendMap: { [key: string]: MetricWithMetadata } = {}; userSpendData.results.forEach((day) => { Object.entries(day.breakdown.providers || {}).forEach(([provider, metrics]) => { - if (!providerSpend[provider]) { - providerSpend[provider] = { + if (!providerSpendMap[provider]) { + providerSpendMap[provider] = { metrics: { spend: 0, prompt_tokens: 0, @@ -291,19 +292,19 @@ const UsagePage: React.FC = ({ teams, organizations }) => { api_key_breakdown: {}, }; } - providerSpend[provider].metrics.spend += metrics.metrics.spend; - providerSpend[provider].metrics.prompt_tokens += metrics.metrics.prompt_tokens; - providerSpend[provider].metrics.completion_tokens += metrics.metrics.completion_tokens; - providerSpend[provider].metrics.total_tokens += metrics.metrics.total_tokens; - providerSpend[provider].metrics.api_requests += metrics.metrics.api_requests; - providerSpend[provider].metrics.successful_requests += metrics.metrics.successful_requests || 0; - providerSpend[provider].metrics.failed_requests += metrics.metrics.failed_requests || 0; - providerSpend[provider].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0; - providerSpend[provider].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0; + providerSpendMap[provider].metrics.spend += metrics.metrics.spend; + providerSpendMap[provider].metrics.prompt_tokens += metrics.metrics.prompt_tokens; + providerSpendMap[provider].metrics.completion_tokens += metrics.metrics.completion_tokens; + providerSpendMap[provider].metrics.total_tokens += metrics.metrics.total_tokens; + providerSpendMap[provider].metrics.api_requests += metrics.metrics.api_requests; + providerSpendMap[provider].metrics.successful_requests += metrics.metrics.successful_requests || 0; + providerSpendMap[provider].metrics.failed_requests += metrics.metrics.failed_requests || 0; + providerSpendMap[provider].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0; + providerSpendMap[provider].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0; }); }); - return Object.entries(providerSpend).map(([provider, metrics]) => ({ + return Object.entries(providerSpendMap).map(([provider, metrics]) => ({ provider, spend: metrics.metrics.spend, requests: metrics.metrics.api_requests, @@ -311,10 +312,10 @@ const UsagePage: React.FC = ({ teams, organizations }) => { failed_requests: metrics.metrics.failed_requests, tokens: metrics.metrics.total_tokens, })); - }; + }, [userSpendData.results]); // Calculate top API keys from the breakdown data - const getTopKeys = (limit: number = 5) => { + const topKeys = useMemo(() => { const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; userSpendData.results.forEach((day) => { Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => { @@ -334,7 +335,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { metadata: { key_alias: metrics.metadata.key_alias, team_id: null, - tags: metrics.metadata.tags || [], // This gets key-level tags + tags: metrics.metadata.tags || [], }, }; } @@ -350,18 +351,16 @@ const UsagePage: React.FC = ({ teams, organizations }) => { }); }); - console.log("debugTags", { keySpend, userSpendData }); - return Object.entries(keySpend) .map(([api_key, metrics]) => ({ api_key, - key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias - tags: metrics.metadata.tags || [], // This will show key-level tags + key_alias: metrics.metadata.key_alias || "-", + tags: metrics.metadata.tags || [], spend: metrics.metrics.spend, })) .sort((a, b) => b.spend - a.spend) - .slice(0, limit); - }; + .slice(0, topKeysLimit); + }, [userSpendData.results, topKeysLimit]); const fetchUserSpendData = useCallback(async () => { if (!accessToken || !dateValue.from || !dateValue.to) return; @@ -399,11 +398,15 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const pageData = await userDailyActivityCall(accessToken, startTime, endTime, page, effectiveUserId); allResults.push(...pageData.results); if (pageData.metadata) { - aggregatedMetadata.total_spend += pageData.metadata.total_spend || 0; - aggregatedMetadata.total_api_requests += pageData.metadata.total_api_requests || 0; - aggregatedMetadata.total_successful_requests += pageData.metadata.total_successful_requests || 0; - aggregatedMetadata.total_failed_requests += pageData.metadata.total_failed_requests || 0; - aggregatedMetadata.total_tokens += pageData.metadata.total_tokens || 0; + aggregatedMetadata.total_spend = (aggregatedMetadata.total_spend || 0) + (pageData.metadata.total_spend || 0); + aggregatedMetadata.total_api_requests = (aggregatedMetadata.total_api_requests || 0) + (pageData.metadata.total_api_requests || 0); + aggregatedMetadata.total_successful_requests = (aggregatedMetadata.total_successful_requests || 0) + (pageData.metadata.total_successful_requests || 0); + aggregatedMetadata.total_failed_requests = (aggregatedMetadata.total_failed_requests || 0) + (pageData.metadata.total_failed_requests || 0); + aggregatedMetadata.total_tokens = (aggregatedMetadata.total_tokens || 0) + (pageData.metadata.total_tokens || 0); + aggregatedMetadata.total_prompt_tokens = (aggregatedMetadata.total_prompt_tokens || 0) + (pageData.metadata.total_prompt_tokens || 0); + aggregatedMetadata.total_completion_tokens = (aggregatedMetadata.total_completion_tokens || 0) + (pageData.metadata.total_completion_tokens || 0); + aggregatedMetadata.total_cache_read_input_tokens = (aggregatedMetadata.total_cache_read_input_tokens || 0) + (pageData.metadata.total_cache_read_input_tokens || 0); + aggregatedMetadata.total_cache_creation_input_tokens = (aggregatedMetadata.total_cache_creation_input_tokens || 0) + (pageData.metadata.total_cache_creation_input_tokens || 0); } } @@ -440,9 +443,13 @@ const UsagePage: React.FC = ({ teams, organizations }) => { return () => clearTimeout(timeoutId); }, [fetchUserSpendData]); - const modelMetrics = processActivityData(userSpendData, "models", teams); - const keyMetrics = processActivityData(userSpendData, "api_keys", teams); - const mcpServerMetrics = processActivityData(userSpendData, "mcp_servers", teams); + const sortedDailyResults = useMemo( + () => [...userSpendData.results].sort((a, b) => new Date(a.date).getTime() - new Date(b.date).getTime()), + [userSpendData.results], + ); + const modelMetrics = useMemo(() => processActivityData(userSpendData, "models", teams), [userSpendData, teams]); + const keyMetrics = useMemo(() => processActivityData(userSpendData, "api_keys", teams), [userSpendData, teams]); + const mcpServerMetrics = useMemo(() => processActivityData(userSpendData, "mcp_servers", teams), [userSpendData, teams]); return (
@@ -626,12 +633,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { {userSpendData.metadata?.total_failed_requests?.toLocaleString() || 0} - - Total Tokens - - {userSpendData.metadata?.total_tokens?.toLocaleString() || 0} - - Average Cost per Request @@ -642,7 +643,51 @@ const UsagePage: React.FC = ({ teams, organizations }) => { )} + setShowTokenBreakdown(!showTokenBreakdown)} + > +
+ Total Tokens + {showTokenBreakdown ? ( + + ) : ( + + )} +
+ + {userSpendData.metadata?.total_tokens?.toLocaleString() || 0} + + + {showTokenBreakdown && ( + + + Input Tokens + + {userSpendData.metadata?.total_prompt_tokens?.toLocaleString() || 0} + + + + Output Tokens + + {userSpendData.metadata?.total_completion_tokens?.toLocaleString() || 0} + + + + Cache Read Tokens + + {userSpendData.metadata?.total_cache_read_input_tokens?.toLocaleString() || 0} + + + + Cache Write Tokens + + {userSpendData.metadata?.total_cache_creation_input_tokens?.toLocaleString() || 0} + + + + )} @@ -654,9 +699,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { ) : ( new Date(a.date).getTime() - new Date(b.date).getTime(), - )} + data={sortedDailyResults} index="date" categories={["metrics.spend"]} colors={["cyan"]} @@ -688,7 +731,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { Top Virtual Keys = ({ teams, organizations }) => { {(() => { const modelData = modelViewType === "groups" - ? getTopModelGroups(topModelsLimit) - : getTopModels(topModelsLimit); + ? topModelGroups + : topModels; return ( = ({ teams, organizations }) => { diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 7d7bb924ac..d63154f0ce 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -555,6 +555,94 @@ it("should display 'Default Proxy Admin' for created_by when value is 'default_u }); +it("should display created_by_user email in 'Created By' column when available", async () => { + const keyWithCreatedByUser = { + ...mockKey, + created_by: "some-uuid-1234", + created_by_user: { + user_id: "some-uuid-1234", + user_email: "creator@example.com", + user_alias: null, + }, + }; + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [keyWithCreatedByUser], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("creator@example.com")).toBeInTheDocument(); + }); +}); + +it("should display created_by_user alias over email when both available", async () => { + const keyWithCreatedByUser = { + ...mockKey, + created_by: "some-uuid-1234", + created_by_user: { + user_id: "some-uuid-1234", + user_email: "creator@example.com", + user_alias: "The Creator", + }, + }; + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [keyWithCreatedByUser], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("The Creator")).toBeInTheDocument(); + }); +}); + it("should render table without crashing when models is null", async () => { const keyWithNullModels = { ...mockKey, diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index fd0cd4dd50..6091794170 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -311,25 +311,40 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo cell: (info) => { const userId = info.getValue() as string | null; if (!userId) return "-"; + const key = info.row.original; + const createdByUser = key.created_by_user; + const userAlias = createdByUser?.user_alias ?? null; + const userEmail = createdByUser?.user_email ?? null; const isDefaultAdmin = userId === "default_user_id"; + const displayValue = userAlias || userEmail || userId; const width = 160; const popoverContent = (
-
- User ID - - {userId} - -
+ {[ + { label: "User Alias", value: userAlias }, + { label: "User Email", value: userEmail }, + { label: "User ID", value: userId }, + ].map(({ label, value }) => ( +
+ {label} + {value ? ( + + {value} + + ) : ( + - + )} +
+ ))}
); - if (isDefaultAdmin) { + if (isDefaultAdmin && !userAlias && !userEmail) { return ( @@ -345,7 +360,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo className="font-mono text-xs truncate block cursor-default" style={{ maxWidth: width, overflow: "hidden" }} > - {userId} + {displayValue} ); diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index 2337d41d8d..e805eed8c8 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -362,7 +362,7 @@ export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string, // Process data function export const processActivityData = ( dailyActivity: { results: DailyData[] }, - key: "models" | "api_keys" | "mcp_servers", + key: "models" | "api_keys" | "mcp_servers" | "entities", teams: Team[] = [], ): Record => { const modelMetrics: Record = {}; @@ -371,7 +371,11 @@ export const processActivityData = ( Object.entries(day.breakdown[key] || {}).forEach(([model, modelData]) => { if (!modelMetrics[model]) { modelMetrics[model] = { - label: key === "api_keys" ? formatKeyLabel(modelData as KeyMetricWithMetadata, model, teams) : model, + label: key === "api_keys" + ? formatKeyLabel(modelData as KeyMetricWithMetadata, model, teams) + : key === "entities" + ? ((modelData as any).metadata?.agent_name || (modelData as any).metadata?.team_alias || model) + : model, total_requests: 0, total_successful_requests: 0, total_failed_requests: 0, diff --git a/ui/litellm-dashboard/src/components/agents.tsx b/ui/litellm-dashboard/src/components/agents.tsx index 169017ec86..542d890e12 100644 --- a/ui/litellm-dashboard/src/components/agents.tsx +++ b/ui/litellm-dashboard/src/components/agents.tsx @@ -19,19 +19,21 @@ import { isAdminRole } from "@/utils/roles"; import AgentInfoView from "./agents/agent_info"; import NotificationsManager from "./molecules/notifications_manager"; import { Agent, AgentKeyInfo } from "./agents/types"; +import { Team } from "./key_team_helpers/key_list"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; interface AgentsPanelProps { accessToken: string | null; userRole?: string; + teams?: Team[] | null; } interface AgentsResponse { agents: Agent[]; } -const AgentsPanel: React.FC = ({ accessToken, userRole }) => { +const AgentsPanel: React.FC = ({ accessToken, userRole, teams }) => { const [agentsList, setAgentsList] = useState([]); const [keyInfoMap, setKeyInfoMap] = useState>({}); const [isAddModalVisible, setIsAddModalVisible] = useState(false); @@ -282,6 +284,7 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { onClose={handleCloseModal} accessToken={accessToken} onSuccess={handleSuccess} + teams={teams} /> {agentToDelete && ( diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index 0cec0331f4..c5518596b8 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -6,6 +6,7 @@ import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; import { createAgentCall, getAgentCreateMetadata, + getAgentsList, keyCreateForAgentCall, keyListCall, keyUpdateCall, @@ -14,11 +15,14 @@ import { } from "../networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { Team } from "../key_team_helpers/key_list"; +import TeamDropdown from "../common_components/team_dropdown"; import AgentFormFields from "./agent_form_fields"; import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields"; import { getDefaultFormValues, buildAgentDataFromForm } from "./agent_config"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; +import GuardrailSelector from "../guardrails/GuardrailSelector"; const { Step } = Steps; @@ -29,6 +33,7 @@ interface AddAgentFormProps { onClose: () => void; accessToken: string | null; onSuccess: () => void; + teams?: Team[] | null; } const AddAgentForm: React.FC = ({ @@ -36,6 +41,7 @@ const AddAgentForm: React.FC = ({ onClose, accessToken, onSuccess, + teams, }) => { const { userId, userRole } = useAuthorized(); const [form] = Form.useForm(); @@ -45,7 +51,7 @@ const AddAgentForm: React.FC = ({ const [agentTypeMetadata, setAgentTypeMetadata] = useState([]); const [loadingMetadata, setLoadingMetadata] = useState(false); - // Step 1: key assignment state + // Step 3: key assignment state const [keyAssignOption, setKeyAssignOption] = useState<"create_new" | "existing_key" | "skip">("create_new"); const [newKeyName, setNewKeyName] = useState(""); const [newKeyModels, setNewKeyModels] = useState([]); @@ -54,8 +60,10 @@ const AddAgentForm: React.FC = ({ const [loadingKeys, setLoadingKeys] = useState(false); const [availableModels, setAvailableModels] = useState([]); const [loadingModels, setLoadingModels] = useState(false); + const [availableAgents, setAvailableAgents] = useState<{agent_id: string; agent_name: string}[]>([]); + const [loadingAgents, setLoadingAgents] = useState(false); - // Step 2: results + // Step 4: results const [createdAgentName, setCreatedAgentName] = useState(""); const [createdKeyValue, setCreatedKeyValue] = useState(null); const [assignedKeyAlias, setAssignedKeyAlias] = useState(null); @@ -82,9 +90,9 @@ const AddAgentForm: React.FC = ({ fetchMetadata(); }, []); - // Fetch existing keys when assign key step becomes active (step 2) + // Fetch existing keys when Agent Management step becomes active (step 3) useEffect(() => { - if (currentStep === 2 && accessToken && existingKeys.length === 0) { + if (currentStep === 3 && accessToken && existingKeys.length === 0) { const fetchKeys = async () => { setLoadingKeys(true); try { @@ -100,9 +108,9 @@ const AddAgentForm: React.FC = ({ } }, [currentStep, accessToken]); - // Fetch available models when Assign Key step is active (same list as key generation) + // Fetch available models when Agent Management step is active (same list as key generation) useEffect(() => { - if (currentStep !== 2 || !accessToken || !userId || !userRole) return; + if ((currentStep !== 1 && currentStep !== 3) || !accessToken || !userId || !userRole) return; let cancelled = false; setLoadingModels(true); modelAvailableCall(accessToken, userId, userRole) @@ -125,6 +133,25 @@ const AddAgentForm: React.FC = ({ }; }, [currentStep, accessToken, userId, userRole]); + useEffect(() => { + if (currentStep !== 1 || !accessToken) return; + let cancelled = false; + setLoadingAgents(true); + getAgentsList(accessToken) + .then((response) => { + if (cancelled) return; + const agents = response?.agents ?? []; + setAvailableAgents(agents.map((a: any) => ({ agent_id: a.agent_id, agent_name: a.agent_name }))); + }) + .catch((error) => { + if (!cancelled) console.error("Error fetching agents:", error); + }) + .finally(() => { + if (!cancelled) setLoadingAgents(false); + }); + return () => { cancelled = true; }; + }, [currentStep, accessToken]); + const selectedAgentTypeInfo = agentTypeMetadata.find( (info) => info.agent_type === agentType ); @@ -207,11 +234,14 @@ const AddAgentForm: React.FC = ({ // Build object_permission from MCP Tools step (allowed_mcp_servers_and_groups, mcp_tool_permissions) const mcpServersAndGroups = values.allowed_mcp_servers_and_groups; const mcpToolPermissions = values.mcp_tool_permissions || {}; - if ( - mcpServersAndGroups && - (mcpServersAndGroups.servers?.length > 0 || mcpServersAndGroups.accessGroups?.length > 0) || - Object.keys(mcpToolPermissions).length > 0 - ) { + const entitlementModels = values.entitlement_models || []; + const entitlementAgents = values.entitlement_agents || []; + const hasObjectPermission = + (mcpServersAndGroups?.servers?.length > 0 || mcpServersAndGroups?.accessGroups?.length > 0) || + Object.keys(mcpToolPermissions).length > 0 || + entitlementModels.length > 0 || + entitlementAgents.length > 0; + if (hasObjectPermission) { agentData.object_permission = {}; if (mcpServersAndGroups?.servers?.length > 0) { agentData.object_permission.mcp_servers = mcpServersAndGroups.servers; @@ -222,6 +252,12 @@ const AddAgentForm: React.FC = ({ if (Object.keys(mcpToolPermissions).length > 0) { agentData.object_permission.mcp_tool_permissions = mcpToolPermissions; } + if (entitlementModels.length > 0) { + agentData.object_permission.models = entitlementModels; + } + if (entitlementAgents.length > 0) { + agentData.object_permission.agents = entitlementAgents; + } } // Wire trace-id flags and budget controls into agent litellm_params (before create call) @@ -237,6 +273,17 @@ const AddAgentForm: React.FC = ({ } } + const selectedGuardrails = values.guardrails || []; + if (selectedGuardrails.length > 0) { + if (!agentData.litellm_params) agentData.litellm_params = {}; + agentData.litellm_params.guardrails = selectedGuardrails; + } + + const selectedTeamId = values.team_id || null; + if (selectedTeamId) { + agentData.team_id = selectedTeamId; + } + const agentResponse = await createAgentCall(accessToken, agentData); const agentId: string = agentResponse.agent_id; const agentName: string = agentResponse.agent_name || values.agent_name || agentId; @@ -248,6 +295,8 @@ const AddAgentForm: React.FC = ({ agentId, newKeyName, newKeyModels, + undefined, + selectedTeamId, ); setCreatedKeyValue(keyResponse.key || null); } else if (keyAssignOption === "existing_key") { @@ -264,7 +313,7 @@ const AddAgentForm: React.FC = ({ setAssignedKeyAlias(keyInfo?.key_alias || selectedExistingKey.slice(0, 12) + "…"); } - setCurrentStep(3); + setCurrentStep(4); onSuccess(); } catch (error) { console.error("Error creating agent:", error); @@ -293,11 +342,54 @@ const AddAgentForm: React.FC = ({ onClose(); }; - const renderMCPToolsStep = () => ( + const renderEntitlementsStep = () => (

- Optionally restrict which MCP servers and tools this agent can use. Leave empty to allow all (subject to key/team permissions). + Configure which models, agents, and MCP tools this agent is allowed to use. Leave fields empty to allow all (subject to key/team permissions).

+ + Allowed Models} + name="entitlement_models" + tooltip="Restrict which models this agent can call. Leave empty to allow all." + > + + (option?.label as string ?? "").toLowerCase().includes(input.toLowerCase()) + } + options={availableAgents.map((a) => ({ + label: a.agent_name, + value: a.agent_id, + }))} + /> + + + + @@ -338,122 +430,137 @@ const AddAgentForm: React.FC = ({
)} +
+ ); - Tracing, - children: ( -
-
-
- - Require x-litellm-trace-id on calls TO this agent - -

- Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent). -

-
- -
- -
-
- - Require x-litellm-trace-id on calls BY this agent - -

- Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking. -

-
- { - setRequireTraceIdOutbound(checked); - if (!checked) { - setMaxIterations(null); - setMaxBudgetPerSession(null); - } - }} - /> -
-
- ), - }, - { - key: "budgets_and_rate_limits", - label: Budgets & Rate Limits, - children: ( -
- {!requireTraceIdOutbound && ( -
- Enable "Require x-litellm-trace-id on calls BY this agent" in Tracing to configure budgets and rate limits. -
- )} - -
Session Budgets
-
-
- - setMaxIterations(val)} - /> -

Hard cap on LLM calls per session

-
-
- - setMaxBudgetPerSession(val)} - /> -

Max spend per trace before returning 429

-
-
- - - -
Agent Rate Limits
-

- Global rate limits applied across all callers of this agent. + const renderObservabilityStep = () => ( +

+
+

Tracing

+
+
+
+ + Require x-litellm-trace-id on calls TO this agent + +

+ Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent).

-
- - - - - - -
- -
Per-Session Rate Limits
-

- Rate limits per session (x-litellm-trace-id). Each session gets its own counters. -

-
- - - - - - -
- ), - }, - ]} /> + +
+ +
+
+ + Require x-litellm-trace-id on calls BY this agent + +

+ Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking. +

+
+ { + setRequireTraceIdOutbound(checked); + if (!checked) { + setMaxIterations(null); + setMaxBudgetPerSession(null); + } + }} + /> +
+
+
+ + + +
+

Budgets & Rate Limits

+
+ {!requireTraceIdOutbound && ( +
+ Enable "Require x-litellm-trace-id on calls BY this agent" in Tracing to configure budgets and rate limits. +
+ )} + +
Session Budgets
+
+
+ + setMaxIterations(val)} + /> +

Hard cap on LLM calls per session

+
+
+ + setMaxBudgetPerSession(val)} + /> +

Max spend per trace before returning 429

+
+
+ + + +
Agent Rate Limits
+

+ Global rate limits applied across all callers of this agent. +

+
+ + + + + + +
+ +
Per-Session Rate Limits
+

+ Rate limits per session (x-litellm-trace-id). Each session gets its own counters. +

+
+ + + + + + +
+
+
+ + + +
+

Guardrails

+

+ Apply guardrails to this agent. Selected guardrails will run on all calls made by this agent. +

+ + form.setFieldsValue({ guardrails: selected })} + /> + +
); @@ -610,6 +717,19 @@ const AddAgentForm: React.FC = ({
+ Assign to Team} + name="team_id" + tooltip="Optionally assign this agent to a team. The agent and its key will belong to the selected team." + > + + + + +
{/* Option: Create new key */}
= ({ placeholder="e.g. my-agent-key" />
-
- - } @@ -447,14 +462,26 @@ const MCPToolConfiguration: React.FC = ({ size="large" /> - {/* Tool list with checkboxes */} - {filteredTools.length === 0 ? ( -
- - No tools found matching "{toolSearchTerm}" -
- ) : ( -
+ {/* CRUD grouped view */} + {viewMode === "crud" && ( + onAllowedToolsChange(allowed)} + /> + )} + + {/* Flat list view */} + {viewMode === "flat" && ( + <> + {filteredTools.length === 0 ? ( +
+ + No tools found matching "{toolSearchTerm}" +
+ ) : ( +
{pinnedFiltered.length > 0 && ( <>
@@ -533,7 +560,9 @@ const MCPToolConfiguration: React.FC = ({ ))} )} -
+
+ )} + )}
)} diff --git a/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx b/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx index 0a8dacbadd..b20c85ecb3 100644 --- a/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx +++ b/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx @@ -612,18 +612,19 @@ describe("columns", () => { it("should allow Admin to delete DB models", async () => { const user = userEvent.setup(); - const setSelectedModelId = vi.fn(); + const onDeleteClick = vi.fn(); const cols = columns( "Admin", "admin-user", defaultProps.premiumUser, - setSelectedModelId, + defaultProps.setSelectedModelId, defaultProps.setSelectedTeamId, defaultProps.getDisplayModelName, defaultProps.handleEditClick, defaultProps.handleRefreshClick, defaultProps.expandedRows, defaultProps.setExpandedRows, + onDeleteClick, ); const model = createMockModel({ @@ -639,23 +640,24 @@ describe("columns", () => { expect(deleteButton).toBeInTheDocument(); await user.click(deleteButton); - expect(setSelectedModelId).toHaveBeenCalledWith("deletable-model"); + expect(onDeleteClick).toHaveBeenCalledWith("deletable-model"); }); it("should allow model creator to delete their own DB models", async () => { const user = userEvent.setup(); - const setSelectedModelId = vi.fn(); + const onDeleteClick = vi.fn(); const cols = columns( "User", "model-creator", defaultProps.premiumUser, - setSelectedModelId, + defaultProps.setSelectedModelId, defaultProps.setSelectedTeamId, defaultProps.getDisplayModelName, defaultProps.handleEditClick, defaultProps.handleRefreshClick, defaultProps.expandedRows, defaultProps.setExpandedRows, + onDeleteClick, ); const model = createMockModel({ @@ -672,7 +674,7 @@ describe("columns", () => { expect(deleteButton).toBeInTheDocument(); await user.click(deleteButton); - expect(setSelectedModelId).toHaveBeenCalledWith("user-model"); + expect(onDeleteClick).toHaveBeenCalledWith("user-model"); }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index ea3d5a1622..6b5ec5139e 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -907,6 +907,7 @@ export const keyCreateForAgentCall = async ( keyAlias: string, models: string[], metadata?: Record, + teamId?: string | null, ) => { const url = proxyBaseUrl ? `${proxyBaseUrl}/key/generate` : `/key/generate`; const body: Record = { @@ -914,6 +915,9 @@ export const keyCreateForAgentCall = async ( key_alias: keyAlias, models: models.length > 0 ? models : [], }; + if (teamId) { + body.team_id = teamId; + } if (metadata && Object.keys(metadata).length > 0) { body.metadata = metadata; } @@ -9603,7 +9607,17 @@ export const storeMCPOAuthUserCredential = async ( }); if (!response.ok) { const err = await response.json().catch(() => ({})); - throw new Error((err as { detail?: { error?: string } })?.detail?.error || "Failed to store OAuth credential"); + const errObj = err as { detail?: unknown }; + const detail = errObj?.detail; + const detailMsg = + Array.isArray(detail) + ? detail.map((d: unknown) => (d && typeof d === "object" ? (d as Record).msg ?? JSON.stringify(d) : String(d))).join("; ") + : typeof detail === "string" + ? detail + : detail && typeof (detail as Record).error === "string" + ? (detail as Record).error as string + : undefined; + throw new Error(detailMsg || "Failed to store OAuth credential"); } return response.json(); }; @@ -9621,7 +9635,17 @@ export const deleteMCPOAuthUserCredential = async ( }); if (!response.ok) { const err = await response.json().catch(() => ({})); - throw new Error((err as { detail?: { error?: string } })?.detail?.error || "Failed to revoke OAuth credential"); + const errObj = err as { detail?: unknown }; + const detail = errObj?.detail; + const detailMsg = + Array.isArray(detail) + ? detail.map((d: unknown) => (d && typeof d === "object" ? (d as Record).msg ?? JSON.stringify(d) : String(d))).join("; ") + : typeof detail === "string" + ? detail + : detail && typeof (detail as Record).error === "string" + ? (detail as Record).error as string + : undefined; + throw new Error(detailMsg || "Failed to revoke OAuth credential"); } return response.json(); }; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index 9936f34452..eba9d37669 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -60,7 +60,8 @@ import CodeInterpreterOutput from "./CodeInterpreterOutput"; import CodeInterpreterTool from "./CodeInterpreterTool"; import { generateCodeSnippet } from "./CodeSnippets"; import EndpointSelector from "./EndpointSelector"; -import MCPEventsDisplay, { MCPEvent } from "./MCPEventsDisplay"; +import MCPEventsDisplay from "./MCPEventsDisplay"; +import type { MCPEvent } from "../../mcp_tools/types"; import { EndpointType, getEndpointType } from "./mode_endpoint_mapping"; import ReasoningContent from "./ReasoningContent"; import ResponseMetrics, { TokenUsage } from "./ResponseMetrics"; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx index 048ea9bfa1..3197c9409c 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx @@ -3,8 +3,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; -import { MCPServer } from "../../mcp_tools/types"; -import { MCPEvent } from "../chat_ui/MCPEventsDisplay"; +import { MCPServer, type MCPEvent } from "../../mcp_tools/types"; export async function makeOpenAIChatCompletionRequest( chatHistory: { role: string; content: string | any[] }[], diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx index f14ac52b84..77ff5fd00b 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx @@ -127,15 +127,15 @@ describe("responses_api", () => { expect(callArgs.tools).toEqual([ { type: "mcp", - server_label: "litellm", - server_url: "litellm_proxy/mcp/alpha", + server_label: "Alpha", + server_url: "https://example.com/mcp/Alpha", require_approval: "never", allowed_tools: ["toolA"], }, { type: "mcp", - server_label: "litellm", - server_url: "litellm_proxy/mcp/Beta", + server_label: "Beta", + server_url: "https://example.com/mcp/Beta", require_approval: "never", allowed_tools: ["toolB", "toolC"], }, diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx new file mode 100644 index 0000000000..013e222c6b --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx @@ -0,0 +1,71 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { render, screen, act } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import CreatedKeyDisplay from "./CreatedKeyDisplay"; + +vi.mock("antd", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + message: { success: vi.fn() }, + }; +}); + +import { message } from "antd"; + +describe("CreatedKeyDisplay", () => { + beforeEach(() => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("should render", () => { + render(); + expect(screen.getByText("sk-test-123")).toBeInTheDocument(); + }); + + it("should display the security warning", () => { + render(); + expect(screen.getByText(/you will not be able to view it again/i)).toBeInTheDocument(); + }); + + it("should show the copy button with initial label", () => { + render(); + expect(screen.getByRole("button", { name: /copy virtual key/i })).toBeInTheDocument(); + }); + + it("should change button text to Copied after clicking copy", async () => { + const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); + render(); + + await user.click(screen.getByRole("button", { name: /copy virtual key/i })); + + expect(screen.getByRole("button", { name: /copied/i })).toBeInTheDocument(); + }); + + it("should show a success message when the key is copied", async () => { + const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); + render(); + + await user.click(screen.getByRole("button", { name: /copy virtual key/i })); + + expect(message.success).toHaveBeenCalledWith("Key copied to clipboard"); + }); + + it("should revert button text back after 2 seconds", async () => { + const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); + render(); + + await user.click(screen.getByRole("button", { name: /copy virtual key/i })); + expect(screen.getByRole("button", { name: /copied/i })).toBeInTheDocument(); + + act(() => { + vi.advanceTimersByTime(2000); + }); + + expect(screen.getByRole("button", { name: /copy virtual key/i })).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/chart_loader.test.tsx b/ui/litellm-dashboard/src/components/shared/chart_loader.test.tsx new file mode 100644 index 0000000000..8c924931a2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/chart_loader.test.tsx @@ -0,0 +1,27 @@ +import { describe, it, expect } from "vitest"; +import { render, screen } from "@testing-library/react"; +import { ChartLoader } from "./chart_loader"; + +describe("ChartLoader", () => { + it("should render", () => { + render(); + expect(screen.getByText("Loading chart data...")).toBeInTheDocument(); + }); + + it("should show loading text by default", () => { + render(); + expect(screen.getByText("Fetching your data")).toBeInTheDocument(); + }); + + it("should show date-changing text when isDateChanging is true", () => { + render(); + expect(screen.getByText("Processing date selection...")).toBeInTheDocument(); + expect(screen.getByText("This will only take a moment")).toBeInTheDocument(); + }); + + it("should not show default loading text when isDateChanging is true", () => { + render(); + expect(screen.queryByText("Loading chart data...")).not.toBeInTheDocument(); + expect(screen.queryByText("Fetching your data")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index 97d48510c2..67ace5db40 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -10,20 +10,10 @@ import { registerMcpOAuthClient, serverRootPath, } from "@/components/networking"; +import { extractErrorMessage } from "@/utils/errorUtils"; export type McpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; -function extractErrorMessage(err: unknown): string { - if (err instanceof Error) return err.message; - if (err && typeof err === "object") { - const e = err as Record; - if (typeof e.detail === "string") return e.detail; - if (typeof e.message === "string") return e.message; - return JSON.stringify(err); - } - return String(err); -} - interface UseMcpOAuthFlowOptions { accessToken: string | null; getCredentials: () => { @@ -90,9 +80,10 @@ export const useMcpOAuthFlow = ({ const setStorageItem = (key: string, value: string) => { if (typeof window === "undefined") return; try { - // Store in both sessionStorage and localStorage for redundancy + // 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). window.sessionStorage.setItem(key, value); - window.localStorage.setItem(key, value); } catch (err) { console.warn(`Failed to set storage item ${key}`, err); } diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx new file mode 100644 index 0000000000..f8c0db2689 --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx @@ -0,0 +1,285 @@ +"use client"; + +/** + * OAuth2 PKCE flow for the *user* connect path. + * + * Unlike useMcpOAuthFlow (used in the admin create-server form), this hook + * targets a server that already exists in the database. It therefore skips + * the temp-session cache step and calls /server/oauth/{serverId}/authorize + * directly with the real server_id. + * + * On success it calls storeMCPOAuthUserCredential to persist the token for + * the user and then invokes onSuccess so the caller can refresh UI state. + */ + +import { useCallback, useEffect, useRef, useState } from "react"; +import { + buildMcpOAuthAuthorizeUrl, + exchangeMcpOAuthToken, + getProxyBaseUrl, + registerMcpOAuthClient, + serverRootPath, + storeMCPOAuthUserCredential, +} from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { extractErrorMessage } from "@/utils/errorUtils"; + +export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; + +interface UseUserMcpOAuthFlowOptions { + accessToken: string; + serverId: string; + serverAlias?: string | null; + /** Scopes to request, e.g. ["repo", "read:user"] */ + scopes?: string[]; + /** Pre-configured client_id if the MCP server record has one. */ + clientId?: string | null; + onSuccess: () => void; +} + +interface UseUserMcpOAuthFlowResult { + startOAuthFlow: () => Promise; + status: UserMcpOAuthStatus; + error: string | null; +} + +const FLOW_STATE_KEY = "litellm-user-mcp-oauth-flow-state"; +// Use a user-flow-specific key to avoid collisions with the admin OAuth flow +// (useMcpOAuthFlow) which uses "litellm-mcp-oauth-result". +const RESULT_KEY = "litellm-user-mcp-oauth-result"; +const RETURN_URL_KEY = "litellm-mcp-oauth-return-url"; + +type StoredFlowState = { + state: string; + codeVerifier: string; + serverId: string; + redirectUri: string; + clientId?: string; + clientSecret?: string; + scopes?: string[]; +}; + +const b64url = (buf: ArrayBuffer) => { + const bytes = new Uint8Array(buf); + let s = ""; + bytes.forEach((b) => (s += String.fromCharCode(b))); + return btoa(s).replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, ""); +}; + +const genVerifier = () => { + const arr = new Uint8Array(32); + window.crypto.getRandomValues(arr); + return b64url(arr.buffer); +}; + +const genChallenge = async (verifier: string) => { + const data = new TextEncoder().encode(verifier); + const digest = await window.crypto.subtle.digest("SHA-256", data); + return b64url(digest); +}; + +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). + window.sessionStorage.setItem(key, value); + } catch (_) {} +}; + +const getStorage = (key: string): string | null => { + try { + return window.sessionStorage.getItem(key); + } catch (_) { + return null; + } +}; + +const clearStorage = (...keys: string[]) => { + keys.forEach((k) => { + try { + window.sessionStorage.removeItem(k); + } catch (_) {} + }); +}; + +const buildCallbackUrl = (): string => { + if (typeof window !== "undefined") { + const path = window.location.pathname || ""; + const idx = path.indexOf("/ui"); + const prefix = idx >= 0 ? path.slice(0, idx + 3).replace(/\/+$/, "") : ""; + return `${window.location.origin}${prefix}/mcp/oauth/callback`; + } + const base = (getProxyBaseUrl() || "").replace(/\/+$/, ""); + const root = serverRootPath && serverRootPath !== "/" ? serverRootPath : ""; + return `${base}${root}/ui/mcp/oauth/callback`; +}; + +export const useUserMcpOAuthFlow = ({ + accessToken, + serverId, + serverAlias, + scopes, + clientId: preClientId, + onSuccess, +}: UseUserMcpOAuthFlowOptions): UseUserMcpOAuthFlowResult => { + const [status, setStatus] = useState("idle"); + const [error, setError] = useState(null); + const processingRef = useRef(false); + + const startOAuthFlow = useCallback(async () => { + if (typeof window === "undefined") return; + try { + setStatus("authorizing"); + setError(null); + + let clientId: string | undefined = preClientId ?? undefined; + let clientSecret: string | undefined; + + if (!clientId) { + // Attempt dynamic client registration against the server's registration endpoint. + try { + const reg = await registerMcpOAuthClient(accessToken, serverId, { + client_name: serverAlias || serverId, + grant_types: ["authorization_code", "refresh_token"], + response_types: ["code"], + token_endpoint_auth_method: "none", + }); + clientId = reg?.client_id; + clientSecret = reg?.client_secret; + } catch (_) { + // Registration is optional; proceed without client_id + } + } + + const verifier = genVerifier(); + const challenge = await genChallenge(verifier); + const state = crypto.randomUUID(); + const redirectUri = buildCallbackUrl(); + const scopeString = scopes?.filter((s) => s.trim()).join(" "); + + const authorizeUrl = buildMcpOAuthAuthorizeUrl({ + serverId, + clientId, + redirectUri, + state, + codeChallenge: challenge, + scope: scopeString, + }); + + const flowState: StoredFlowState = { + state, + codeVerifier: verifier, + serverId, + redirectUri, + clientId, + clientSecret, + scopes, + }; + + setStorage(FLOW_STATE_KEY, JSON.stringify(flowState)); + const returnUrl = new URL(window.location.href); + returnUrl.searchParams.set("mcpOauthReturn", "apps"); + setStorage(RETURN_URL_KEY, returnUrl.toString()); + + window.location.href = authorizeUrl; + } catch (err) { + const msg = extractErrorMessage(err); + setError(msg); + setStatus("error"); + NotificationsManager.error(msg); + } + }, [accessToken, serverId, serverAlias, scopes, preClientId]); + + const resumeOAuthFlow = useCallback(async () => { + if (typeof window === "undefined" || processingRef.current) return; + + const storedResult = getStorage(RESULT_KEY); + if (!storedResult) return; + + // When multiple OAuth2ConnectButton components are mounted (one per server + // card), each holds its own hook instance. All run resumeOAuthFlow() on + // mount and would compete for the same RESULT_KEY. Peek at the stored + // flow state first: only the hook instance whose serverId matches the one + // that initiated the OAuth flow should consume the result. + const rawFlowState = getStorage(FLOW_STATE_KEY); + if (rawFlowState) { + try { + const peeked = JSON.parse(rawFlowState) as StoredFlowState; + if (peeked.serverId && peeked.serverId !== serverId) return; + } catch (_) {} + } + + processingRef.current = true; + clearStorage(RESULT_KEY); + + let payload: Record | null = null; + let flowState: StoredFlowState | null = null; + + try { + payload = JSON.parse(storedResult); + const raw = getStorage(FLOW_STATE_KEY); + flowState = raw ? JSON.parse(raw) : null; + } catch (_) { + setError("Failed to resume OAuth flow. Please retry."); + setStatus("error"); + processingRef.current = false; + clearStorage(FLOW_STATE_KEY); + return; + } + + try { + if (!flowState?.state || !flowState.codeVerifier || !flowState.serverId) { + throw new Error("OAuth session state was lost. Please retry."); + } + if (!payload?.state || payload.state !== flowState.state) { + throw new Error("OAuth state mismatch. Please retry."); + } + if (payload.error) { + throw new Error((payload.error_description as string) || (payload.error as string)); + } + if (!payload.code) { + throw new Error("Authorization code missing in callback."); + } + + setStatus("exchanging"); + const token = await exchangeMcpOAuthToken({ + serverId: flowState.serverId, + code: payload.code as string, + clientId: flowState.clientId, + clientSecret: flowState.clientSecret, + codeVerifier: flowState.codeVerifier, + redirectUri: flowState.redirectUri, + }); + + // Persist the token for this user via the backend. + // accessToken comes from props — it is never stored in sessionStorage. + await storeMCPOAuthUserCredential(accessToken, flowState.serverId, { + access_token: token.access_token, + refresh_token: token.refresh_token, + expires_in: token.expires_in, + scopes: flowState.scopes, + }); + + setStatus("success"); + setError(null); + NotificationsManager.success("Connected successfully"); + onSuccess(); + } catch (err) { + const msg = extractErrorMessage(err); + setError(msg); + setStatus("error"); + NotificationsManager.error(msg); + } finally { + clearStorage(FLOW_STATE_KEY); + setTimeout(() => { processingRef.current = false; }, 1000); + } + }, [accessToken, serverId, onSuccess]); + + useEffect(() => { + resumeOAuthFlow(); + }, [resumeOAuthFlow]); + + return { startOAuthFlow, status, error }; +}; diff --git a/ui/litellm-dashboard/src/utils/errorUtils.ts b/ui/litellm-dashboard/src/utils/errorUtils.ts new file mode 100644 index 0000000000..665e4e9cb7 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/errorUtils.ts @@ -0,0 +1,35 @@ +/** + * Shared error-message extraction utility. + * + * Handles the common shapes returned by LiteLLM / FastAPI: + * - Error instances (err.message) + * - { detail: "string" } + * - { detail: [{ msg, loc, type }] } (FastAPI 422) + * - { detail: { error: "string" } } + * - { message: "string" } + * - anything else → JSON.stringify / String() + */ +export function extractErrorMessage(err: unknown): string { + if (err instanceof Error) return err.message; + if (err && typeof err === "object") { + const e = err as Record; + const detail = e.detail; + if (typeof detail === "string") return detail; + if (Array.isArray(detail)) { + return detail.map((d: unknown) => { + if (d && typeof d === "object") { + const item = d as Record; + return typeof item.msg === "string" ? item.msg : JSON.stringify(d); + } + return String(d); + }).join("; "); + } + if (detail && typeof detail === "object") { + const detailObj = detail as Record; + if (typeof detailObj.error === "string") return detailObj.error; + } + if (typeof e.message === "string") return e.message; + return JSON.stringify(err); + } + return String(err); +} diff --git a/ui/litellm-dashboard/src/utils/mcpToolCrudClassification.ts b/ui/litellm-dashboard/src/utils/mcpToolCrudClassification.ts new file mode 100644 index 0000000000..523995cce4 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/mcpToolCrudClassification.ts @@ -0,0 +1,86 @@ +export type CrudOp = "read" | "create" | "update" | "delete" | "unknown"; + +const DELETE_RE = /\b(delete|remove|destroy|purge|drop|erase|unlink)\b/i; +const CREATE_RE = /\b(create|add|insert|new|post|submit|register|make|generate|write|upload)\b/i; +const UPDATE_RE = /\b(update|edit|modify|change|patch|put|set|rename|move|transform)\b/i; +const READ_RE = /\b(get|read|list|fetch|search|find|query|retrieve|show|view|check|describe|info)\b/i; + +export interface MCPToolEntry { + name: string; + description?: string; +} + +/** + * Classifies a tool by its name first; falls back to description only when + * the name alone yields no match. This prevents incidental phrasing in + * free-form descriptions (e.g. "removes noise from…") from promoting a safe + * tool into a high-risk bucket. + * + * READ is checked before DELETE/UPDATE so that tools like `get_removed_entries` + * or `list_deleted_items` — where the primary verb is a read operation — are + * not silently blocked by the delete-by-default policy for new servers. + */ +export function classifyToolOp(name: string, description = ""): CrudOp { + const nameLower = name.toLowerCase(); + if (READ_RE.test(nameLower)) return "read"; + if (DELETE_RE.test(nameLower)) return "delete"; + if (UPDATE_RE.test(nameLower)) return "update"; + if (CREATE_RE.test(nameLower)) return "create"; + + // Only consult description when the name is unrecognised. + if (description) { + const descLower = description.toLowerCase(); + if (READ_RE.test(descLower)) return "read"; + if (DELETE_RE.test(descLower)) return "delete"; + if (UPDATE_RE.test(descLower)) return "update"; + if (CREATE_RE.test(descLower)) return "create"; + } + + return "unknown"; +} + +export function groupToolsByCrud(tools: MCPToolEntry[]): Record { + const groups: Record = { + read: [], + create: [], + update: [], + delete: [], + unknown: [], + }; + for (const tool of tools) { + const op = classifyToolOp(tool.name, tool.description); + groups[op].push(tool); + } + return groups; +} + +export const CRUD_GROUP_META: Record< + CrudOp, + { label: string; description: string; risk: "low" | "medium" | "high" | "unknown" } +> = { + read: { + label: "Read", + description: "Safe operations — fetch, list, search. No side effects.", + risk: "low", + }, + create: { + label: "Create", + description: "Add new resources — insert, upload, register.", + risk: "medium", + }, + update: { + label: "Update", + description: "Modify existing resources — edit, patch, rename.", + risk: "medium", + }, + delete: { + label: "Delete", + description: "Destructive operations — remove, purge, destroy.", + risk: "high", + }, + unknown: { + label: "Other", + description: "Operations that could not be automatically classified.", + risk: "unknown", + }, +};