diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 3d11345e850..0b7cce2e4be 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -34,8 +34,6 @@ jobs: build-mode: none - language: python build-mode: none - - language: ruby - build-mode: none steps: - name: Checkout repository diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index df483ab927d..51af22b7a46 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -13,6 +13,10 @@ spec: {{- if and (not .Values.keda.enabled) (not .Values.autoscaling.enabled) }} replicas: {{ .Values.replicaCount }} {{- end }} + {{- with .Values.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} selector: matchLabels: {{- include "litellm.selectorLabels" . | nindent 6 }} diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index bab909c8954..f8944bddd53 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -35,6 +35,14 @@ deploymentLabels: {} podAnnotations: {} podLabels: {} +# -- Deployment strategy configuration +# Example: +# type: RollingUpdate +# rollingUpdate: +# maxUnavailable: 0 +# maxSurge: 1 +strategy: {} + terminationGracePeriodSeconds: 90 topologySpreadConstraints: [] diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 57bc1d57ffd..600f69547d4 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -217,6 +217,7 @@ mcp_servers: | `bearer_token` | `Authorization: Bearer ` | | `basic` | `Authorization: Basic ` | | `authorization` | `Authorization: ` | + | `aws_sigv4` | Per-request AWS SigV4 signature ([details](./mcp_aws_sigv4.md)) | - **Extra Headers**: Optional list of additional header names that should be forwarded from client to the MCP server - **Static Headers**: Optional map of header key/value pairs to include every request to the MCP server. @@ -257,6 +258,16 @@ mcp_servers: auth_type: "authorization" auth_value: "Token example123" # headers={"Authorization": "Token example123"} + # AWS SigV4 for Bedrock AgentCore MCP servers + agentcore_mcp: + url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes//invocations" + transport: "http" + auth_type: "aws_sigv4" + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: us-east-1 + aws_service_name: bedrock-agentcore + # Example with extra headers forwarding github_mcp: url: "https://api.githubcopilot.com/mcp" diff --git a/docs/my-website/docs/mcp_aws_sigv4.md b/docs/my-website/docs/mcp_aws_sigv4.md new file mode 100644 index 00000000000..e00cee4fd52 --- /dev/null +++ b/docs/my-website/docs/mcp_aws_sigv4.md @@ -0,0 +1,144 @@ +# MCP - AWS SigV4 Auth + +Use AWS SigV4 authentication to connect LiteLLM to MCP servers hosted on [AWS Bedrock AgentCore](https://docs.aws.amazon.com/bedrock/latest/userguide/agentcore.html). + +## Why SigV4? + +AWS services authenticate requests using [Signature Version 4](https://docs.aws.amazon.com/general/latest/gr/signature-version-4.html) — a per-request signing protocol that includes the request body in the cryptographic signature. This is fundamentally different from static-header auth types (`api_key`, `bearer_token`, etc.) which send the same header on every request. + +LiteLLM's `aws_sigv4` auth type handles this automatically: every outgoing MCP request is signed with your AWS credentials before it's sent. + +## Quick Start + +### 1. Set AWS credentials + +```bash +export AWS_ACCESS_KEY_ID="AKIA..." +export AWS_SECRET_ACCESS_KEY="..." +export AWS_REGION_NAME="us-east-1" +``` + +### 2. Add your AgentCore MCP server to config.yaml + +```yaml title="config.yaml" showLineNumbers +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + +mcp_servers: + my_agentcore_mcp: + url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes//invocations" + transport: "http" + auth_type: "aws_sigv4" + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: "us-east-1" + aws_service_name: "bedrock-agentcore" +``` + +:::info URL encoding + +The AgentCore runtime ARN must be URL-encoded in the `url` field. For example: + +``` +arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/my-mcp-server +``` + +becomes: + +``` +arn%3Aaws%3Abedrock-agentcore%3Aus-east-1%3A123456789012%3Aruntime%2Fmy-mcp-server +``` + +::: + +### 3. Start the proxy + +```bash +litellm --config config.yaml +``` + +### 4. Use the MCP tools + +Once started, your AgentCore MCP tools are available through LiteLLM like any other MCP server: + +```bash title="List available tools" +curl http://localhost:4000/mcp-rest/tools/list \ + -H "Authorization: Bearer sk-1234" +``` + +```bash title="Call a tool" +curl http://localhost:4000/mcp-rest/tools/call \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "name": "my_agentcore_mcp_your_tool_name", + "arguments": {"key": "value"} + }' +``` + +## Config Reference + +| Field | Required | Description | +|-------|----------|-------------| +| `url` | Yes | AgentCore MCP server URL (with URL-encoded ARN) | +| `transport` | Yes | Must be `"http"` | +| `auth_type` | Yes | Must be `"aws_sigv4"` | +| `aws_access_key_id` | No | AWS access key. Supports `os.environ/VAR_NAME`. Falls back to boto3 credential chain if omitted | +| `aws_secret_access_key` | No | AWS secret key. Supports `os.environ/VAR_NAME`. Falls back to boto3 credential chain if omitted | +| `aws_region_name` | Yes | AWS region (e.g., `us-east-1`) | +| `aws_service_name` | No | AWS service name for signing. Defaults to `bedrock-agentcore` | +| `aws_session_token` | No | AWS session token for temporary credentials. Supports `os.environ/VAR_NAME` | + +## How It Works + +LiteLLM uses an `httpx.Auth` subclass (`MCPSigV4Auth`) that hooks into the HTTP request lifecycle: + +1. For every outgoing MCP request, the auth handler computes a SHA-256 hash of the request body +2. It creates a SigV4 signature using your AWS credentials, the request URL, headers, and body hash +3. The signed `Authorization` and `x-amz-date` headers are added to the request +4. AWS validates the signature and processes the MCP request + +This happens transparently — no manual token management required. + +## Using Temporary Credentials (STS) + +If you use AWS STS temporary credentials (e.g., from IAM roles or SSO), include the session token: + +```yaml title="config.yaml with STS credentials" showLineNumbers +mcp_servers: + my_agentcore_mcp: + url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes//invocations" + transport: "http" + auth_type: "aws_sigv4" + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_session_token: os.environ/AWS_SESSION_TOKEN + aws_region_name: "us-east-1" + aws_service_name: "bedrock-agentcore" +``` + +## Troubleshooting + +### 403 Forbidden from AWS + +- Verify your AWS credentials are valid and not expired +- Check that `aws_region_name` matches the region in your AgentCore URL +- Ensure `aws_service_name` is set to `bedrock-agentcore` +- If using STS credentials, confirm `aws_session_token` is set and not expired + +### Health check errors on startup + +SigV4-authenticated MCP servers skip the standard health check on proxy startup. This is expected — the proxy will still sign requests correctly when tools are invoked. + +### "botocore not found" error + +Install the `botocore` package: + +```bash +pip install botocore +``` + +`botocore` is used for SigV4 credential handling and is required when using `aws_sigv4` auth. diff --git a/docs/my-website/docs/providers/bedrock_agentcore.md b/docs/my-website/docs/providers/bedrock_agentcore.md index e3e352f7ab6..7802624fccd 100644 --- a/docs/my-website/docs/providers/bedrock_agentcore.md +++ b/docs/my-website/docs/providers/bedrock_agentcore.md @@ -13,7 +13,7 @@ Call Bedrock AgentCore in the OpenAI Request/Response format. :::info -This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers, add them as you would any other MCP server. See the [MCP documentation](https://docs.litellm.ai/docs/mcp) for details. +This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers with LiteLLM, see the [MCP AWS SigV4 Auth](https://docs.litellm.ai/docs/mcp_aws_sigv4) guide for setup instructions. ::: diff --git a/docs/my-website/docs/providers/openai.md b/docs/my-website/docs/providers/openai.md index 6817c32e9b4..bed4cd0aa5b 100644 --- a/docs/my-website/docs/providers/openai.md +++ b/docs/my-website/docs/providers/openai.md @@ -632,7 +632,22 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \ ## OpenAI Chat Completion to Responses API Bridge -Call any Responses API model from OpenAI's `/chat/completions` endpoint. +Call any Responses API model from OpenAI's `/chat/completions` endpoint. + +:::tip gpt-5.4 + reasoning_effort + function tools + +OpenAI does not support `reasoning_effort` with function tools for `gpt-5.4` in `/v1/chat/completions`. Use the responses bridge instead: + +```python +response = litellm.completion( + model="openai/responses/gpt-5.4", # routes to /v1/responses + messages=[{"role": "user", "content": "What's the weather?"}], + tools=[...], + reasoning_effort="low", +) +``` + +::: diff --git a/docs/my-website/docs/providers/openai/responses_api.md b/docs/my-website/docs/providers/openai/responses_api.md index 7799c93ccf2..0d6b9013ac8 100644 --- a/docs/my-website/docs/providers/openai/responses_api.md +++ b/docs/my-website/docs/providers/openai/responses_api.md @@ -693,6 +693,236 @@ print(final_response.output) Set `parallel_tool_calls=False` to ensure zero or one tool is called per turn. [More details](https://platform.openai.com/docs/guides/function-calling#parallel-function-calling). +## Tool Search & Namespaces + +Tool search lets models dynamically load tools at runtime instead of sending every tool definition in the prompt. Group functions into **namespaces** and mark them with `defer_loading: true` — the model only loads the schemas it actually needs, saving tokens. + +Requires `gpt-5.4` or later. See [OpenAI Tool Search docs](https://developers.openai.com/api/docs/guides/tools-tool-search) for full details. + + + + +```python showLineNumbers title="Tool Search with Namespaces" +import litellm + +# Define namespaces with deferred tools +tools = [ + {"type": "tool_search"}, # Enable tool search + { + "type": "namespace", + "name": "crm", + "description": "CRM tools for customer management", + "tools": [ + { + "type": "function", + "name": "get_customer", + "description": "Get customer details by ID", + "parameters": { + "type": "object", + "properties": { + "customer_id": {"type": "string"} + }, + "required": ["customer_id"], + }, + "defer_loading": True, + }, + { + "type": "function", + "name": "list_customers", + "description": "List customers with optional filters", + "parameters": { + "type": "object", + "properties": { + "status": {"type": "string", "enum": ["active", "inactive"]}, + }, + }, + "defer_loading": True, + }, + ], + }, + { + "type": "namespace", + "name": "billing", + "description": "Billing and invoicing tools", + "tools": [ + { + "type": "function", + "name": "get_invoice", + "description": "Get an invoice by ID", + "parameters": { + "type": "object", + "properties": { + "invoice_id": {"type": "string"} + }, + "required": ["invoice_id"], + }, + "defer_loading": True, + }, + ], + }, +] + +response = litellm.responses( + model="openai/gpt-5.4", + input="Look up invoice INV-2024-001 from the billing system", + tools=tools, +) + +# The response contains tool_search_call, tool_search_output, and function_call items +for item in response.output: + if isinstance(item, dict): + if item["type"] == "tool_search_call": + print(f"Searched namespaces: {item['arguments']['paths']}") + elif item["type"] == "tool_search_output": + print(f"Loaded {len(item['tools'])} tool(s)") + elif item["type"] == "function_call": + print(f"Called: {item.get('namespace', '')}.{item['name']}({item['arguments']})") + else: + if item.type == "function_call": + print(f"Called: {item.namespace}.{item.name}({item.arguments})") +``` + + + + +1. Set up config.yaml + +```yaml showLineNumbers title="OpenAI Proxy Configuration" +model_list: + - model_name: openai/gpt-5.4 + litellm_params: + model: openai/gpt-5.4 + api_key: os.environ/OPENAI_API_KEY +``` + +2. Start LiteLLM Proxy Server + +```bash title="Start LiteLLM Proxy Server" +litellm --config /path/to/config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +3. Test it! + +```python showLineNumbers title="Tool Search via OpenAI SDK with LiteLLM Proxy" +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:4000", + api_key="your-api-key" +) + +response = client.responses.create( + model="openai/gpt-5.4", + input="Look up invoice INV-2024-001 from the billing system", + tools=[ + {"type": "tool_search"}, + { + "type": "namespace", + "name": "billing", + "description": "Billing and invoicing tools", + "tools": [ + { + "type": "function", + "name": "get_invoice", + "description": "Get an invoice by ID", + "parameters": { + "type": "object", + "properties": {"invoice_id": {"type": "string"}}, + "required": ["invoice_id"], + }, + "defer_loading": True, + }, + ], + }, + ], +) + +print(response.output) +``` + + + + +### Tool Search via Chat Completions Bridge + +You can also use tool search through the `/v1/chat/completions` endpoint by prefixing the model with `openai/responses/`. The request is routed through the Responses API but returns a standard chat completions response. + + + + +```python showLineNumbers title="Tool Search via Chat Completions Bridge" +import litellm + +response = litellm.completion( + model="openai/responses/gpt-5.4", + messages=[{"role": "user", "content": "Look up invoice INV-2024-001"}], + tools=[ + {"type": "tool_search"}, + { + "type": "namespace", + "name": "billing", + "description": "Billing and invoicing tools", + "tools": [ + { + "type": "function", + "name": "get_invoice", + "description": "Get an invoice by ID", + "parameters": { + "type": "object", + "properties": {"invoice_id": {"type": "string"}}, + "required": ["invoice_id"], + }, + "defer_loading": True, + }, + ], + }, + ], +) + +# Standard chat completions response +for tool_call in response.choices[0].message.tool_calls: + print(f"Called: {tool_call.function.name}({tool_call.function.arguments})") +``` + + + + +```bash showLineNumbers title="Tool Search via /v1/chat/completions" +curl http://localhost:4000/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "openai/responses/gpt-5.4", + "messages": [{"role": "user", "content": "Look up invoice INV-2024-001"}], + "tools": [ + {"type": "tool_search"}, + { + "type": "namespace", + "name": "billing", + "description": "Billing and invoicing tools", + "tools": [ + { + "type": "function", + "name": "get_invoice", + "description": "Get an invoice by ID", + "parameters": { + "type": "object", + "properties": {"invoice_id": {"type": "string"}}, + "required": ["invoice_id"] + }, + "defer_loading": true + } + ] + } + ] + }' +``` + + + + ## Free-form Function Calling diff --git a/docs/my-website/docs/proxy/guardrails/team_based_guardrails.md b/docs/my-website/docs/proxy/guardrails/team_based_guardrails.md index 2d55294a711..0e610b6e445 100644 --- a/docs/my-website/docs/proxy/guardrails/team_based_guardrails.md +++ b/docs/my-website/docs/proxy/guardrails/team_based_guardrails.md @@ -1,6 +1,6 @@ import Image from '@theme/IdealImage'; -# Team-Based Guardrails +# Team Bring-Your-Own Guardrails Team-based guardrails let **developers** register a guardrail for their team via the API; an **admin** then reviews and approves or rejects it in the LiteLLM UI. Only [Generic Guardrail API](/docs/adding_provider/generic_guardrail_api) guardrails can be registered this way. diff --git a/docs/my-website/docs/reasoning_content.md b/docs/my-website/docs/reasoning_content.md index b5a5809bd4e..5dd40122c71 100644 --- a/docs/my-website/docs/reasoning_content.md +++ b/docs/my-website/docs/reasoning_content.md @@ -592,6 +592,12 @@ Expected Response +:::tip gpt-5.4: reasoning_effort + function tools + +OpenAI does not support `reasoning_effort` with function tools for `gpt-5.4` in `/v1/chat/completions`. Use `openai/responses/gpt-5.4` to route through the Responses API instead. See [Responses API Bridge](/docs/providers/openai#openai-chat-completion-to-responses-api-bridge) for details. + +::: + ## OpenAI Responses API - Auto-Summary Control When using OpenAI Responses API models (like `gpt-5`) via `/chat/completions` with `reasoning_effort`, you can control whether `summary="detailed"` is automatically added to the reasoning parameter. diff --git a/docs/my-website/docs/troubleshoot/pip_venv_upgrade.md b/docs/my-website/docs/troubleshoot/pip_venv_upgrade.md new file mode 100644 index 00000000000..6f5699e3fb0 --- /dev/null +++ b/docs/my-website/docs/troubleshoot/pip_venv_upgrade.md @@ -0,0 +1,121 @@ +# Upgrading LiteLLM Proxy (pip/venv) + +Guide for upgrading LiteLLM Proxy when installed via pip in a virtual environment. + +:::info Important +Always activate your virtual environment before running any `litellm` or `prisma` commands. All commands in this guide assume you're working inside an activated venv. +::: + +## How pip/venv Upgrades Work + +There are two pieces that need to stay in sync: + +1. **Prisma client** - Generated Python code that talks to the DB +2. **DB schema** - Tables/columns in PostgreSQL + +When you upgrade via pip, the `litellm-proxy-extras` package ships with a new `schema.prisma` and a `migrations/` directory. But unlike the Docker image, pip install does NOT automatically regenerate the Prisma client or run migrations. You have to do both manually. + +## Upgrade Workflow (pip/venv) + +### 1. Stop the proxy + +Stop your running LiteLLM proxy instance. + +### 2. (Optional) Back up your DB + +```bash +pg_dump -h -U -d -F c -f backup_$(date +%Y%m%d).dump +``` + +### 3. Upgrade the package + +```bash +pip install 'litellm[proxy]==' +``` + +### 4. Regenerate the Prisma client + +```bash +prisma generate --schema /lib/python/site-packages/litellm_proxy_extras/schema.prisma +``` + +Replace `` with your virtual environment path and `` with your Python version (e.g., `python3.11`, `python3.12`, `python3.13`). + +### 5. Apply DB migrations + +You have two options: + +**Option A: Just start the proxy** (simplest) + +The proxy automatically runs `prisma migrate deploy` on startup, which applies any new migrations. + +First, activate your virtual environment: + +```bash +source /bin/activate +``` + +Then start the proxy: + +```bash +litellm --config your_config.yaml --port 4000 +``` + +**Option B: Run manually before starting** + +Activate your virtual environment first: + +```bash +source /bin/activate +``` + +Then run the migration with the explicit schema path: + +```bash +prisma migrate deploy --schema /lib/python/site-packages/litellm_proxy_extras/schema.prisma +``` + +Replace `` with your virtual environment path and `` with your Python version (e.g., `python3.11`, `python3.12`, `python3.13`). + +### 6. Start the proxy + +If you used Option B above, now start the proxy (with venv still activated): + +```bash +litellm --config your_config.yaml --port 4000 +``` + +## How to Verify Migrations + +> **Note:** `` = `/lib/python/site-packages/litellm_proxy_extras/schema.prisma` + +### Before applying migrations: Preview what will change + +Run `pip install 'litellm[proxy]=='` first (Step 3) so the new `schema.prisma` is available. + +```bash +prisma migrate diff \ + --from-url $DATABASE_URL \ + --to-schema-datamodel \ + --script +``` + +### After applying migrations: Check status + +```bash +prisma migrate status --schema +``` + +All migrations should have a `finished_at` timestamp and no `rolled_back_at`. + +## Key Things to Know + +- **`DISABLE_SCHEMA_UPDATE=true`** env var prevents auto-migration on startup - useful if you want full manual control + +- **`prisma db push`** is the nuclear option: force-syncs the DB to match the schema, bypassing migration history. Safe when all changes are additive (new columns/tables), but always have a backup. + +- **The `schema.prisma` inside `litellm_proxy_extras` is the source of truth** - always use that one, not one from a different version or from the git repo + +## Troubleshooting + +If you encounter migration errors, see the [Prisma Migration Troubleshooting Guide](./prisma_migrations). diff --git a/docs/my-website/release_notes/v1.81.9.md b/docs/my-website/release_notes/v1.81.9.md index c7659442c4c..80be4179b46 100644 --- a/docs/my-website/release_notes/v1.81.9.md +++ b/docs/my-website/release_notes/v1.81.9.md @@ -279,7 +279,7 @@ Let's dive in. - Add HTTP support to custom code guardrails + Unified guardrails for MCP + Agent guardrail support - [PR #20619](https://github.com/BerriAI/litellm/pull/20619) - Custom Code Guardrails UI Playground - [PR #20377](https://github.com/BerriAI/litellm/pull/20377) -- **Team-Based Guardrails** +- **Team Bring-Your-Own Guardrails** - Implement team-based isolation guardrails management - [PR #20318](https://github.com/BerriAI/litellm/pull/20318) - **[OpenAI Moderations](../../docs/apply_guardrail)** diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 323d637cffb..64c8fb291be 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -614,6 +614,7 @@ const sidebars = { "mcp_usage", "mcp_openapi", "mcp_oauth", + "mcp_aws_sigv4", "mcp_public_internet", "mcp_semantic_filter", "mcp_control", @@ -1158,6 +1159,7 @@ const sidebars = { "troubleshoot/prisma_migrations", ], }, + "troubleshoot/pip_venv_upgrade", "troubleshoot/rollback", "troubleshoot", ], diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000000_add_mcp_approval_status/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000000_add_mcp_approval_status/migration.sql new file mode 100644 index 00000000000..184caef0809 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000000_add_mcp_approval_status/migration.sql @@ -0,0 +1,11 @@ +-- AlterTable: Add BYOM approval workflow fields to LiteLLM_MCPServerTable +ALTER TABLE "LiteLLM_MCPServerTable" + ADD COLUMN IF NOT EXISTS "approval_status" TEXT DEFAULT 'active', + ADD COLUMN IF NOT EXISTS "submitted_by" TEXT, + ADD COLUMN IF NOT EXISTS "submitted_at" TIMESTAMP(3), + ADD COLUMN IF NOT EXISTS "reviewed_at" TIMESTAMP(3), + ADD COLUMN IF NOT EXISTS "review_notes" TEXT; + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_MCPServerTable_approval_status_idx" + ON "LiteLLM_MCPServerTable"("approval_status"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000001_add_mcp_source_url/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000001_add_mcp_source_url/migration.sql new file mode 100644 index 00000000000..dc468b82061 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000001_add_mcp_source_url/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable: Add source_url field to LiteLLM_MCPServerTable for GitHub/docs link +ALTER TABLE "LiteLLM_MCPServerTable" + ADD COLUMN IF NOT EXISTS "source_url" TEXT; diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 849ce023109..30a1ac20d0c 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -4,7 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 -from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar, Union +from typing import Any, Awaitable, Callable, Dict, Generator, List, Optional, Tuple, TypeVar, Union import httpx from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters @@ -50,6 +50,86 @@ def to_basic_auth(auth_value: str) -> str: TSessionResult = TypeVar("TSessionResult") +class MCPSigV4Auth(httpx.Auth): + """ + httpx Auth class that signs each request with AWS SigV4. + + This is used for MCP servers that require AWS SigV4 authentication, + such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow() + for every outgoing request, enabling per-request signature computation. + """ + + requires_request_body = True + + def __init__( + self, + aws_access_key_id: Optional[str] = None, + aws_secret_access_key: Optional[str] = None, + aws_session_token: Optional[str] = None, + aws_region_name: Optional[str] = None, + aws_service_name: Optional[str] = None, + ): + try: + from botocore.credentials import Credentials + except ImportError: + raise ImportError( + "Missing botocore to use AWS SigV4 authentication. " + "Run 'pip install boto3'." + ) + + self.service_name = aws_service_name or "bedrock-agentcore" + self.region_name = aws_region_name or "us-east-1" + + # Note: os.environ/ prefixed values are already resolved by + # ProxyConfig._check_for_os_environ_vars() at config load time. + # Values arrive here as plain strings. + if aws_access_key_id and aws_secret_access_key: + self.credentials = Credentials( + access_key=aws_access_key_id, + secret_key=aws_secret_access_key, + token=aws_session_token, + ) + else: + # Fall back to default boto3 credential chain + import botocore.session + + session = botocore.session.get_session() + self.credentials = session.get_credentials() + if self.credentials is None: + raise ValueError( + "No AWS credentials found. Provide aws_access_key_id and " + "aws_secret_access_key, or configure default credentials " + "(env vars, ~/.aws/credentials, instance profile)." + ) + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + + # Build AWSRequest from the httpx Request. + # Pass all request headers so the canonical SigV4 signature covers them. + aws_request = AWSRequest( + method=request.method, + url=str(request.url), + data=request.content, + headers=dict(request.headers), + ) + + # Sign the request — SigV4Auth.add_auth() adds Authorization, + # X-Amz-Date, and X-Amz-Security-Token (if session token present). + # Host header is derived automatically from the URL. + sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name) + sigv4.add_auth(aws_request) + + # Copy SigV4 headers back to the httpx request + for header_name, header_value in aws_request.headers.items(): + request.headers[header_name] = header_value + + yield request + + class MCPClient: """ MCP Client supporting: @@ -68,6 +148,7 @@ class MCPClient: stdio_config: Optional[MCPStdioConfig] = None, extra_headers: Optional[Dict[str, str]] = None, ssl_verify: Optional[VerifyTypes] = None, + aws_auth: Optional[httpx.Auth] = None, ): self.server_url: str = server_url self.transport_type: MCPTransport = transport_type @@ -77,6 +158,7 @@ class MCPClient: self.stdio_config: Optional[MCPStdioConfig] = stdio_config self.extra_headers: Optional[Dict[str, str]] = extra_headers self.ssl_verify: Optional[VerifyTypes] = ssl_verify + self._aws_auth: Optional[httpx.Auth] = aws_auth # handle the basic auth value if provided if auth_value: self.update_auth_value(auth_value) @@ -212,8 +294,13 @@ class MCPClient: headers["Authorization"] = self._mcp_auth_value elif self.auth_type == MCPAuth.oauth2: headers["Authorization"] = f"Bearer {self._mcp_auth_value}" + elif self.auth_type == MCPAuth.token: + headers["Authorization"] = f"token {self._mcp_auth_value}" elif isinstance(self._mcp_auth_value, dict): headers.update(self._mcp_auth_value) + # Note: aws_sigv4 auth is not handled here — SigV4 requires per-request + # signing (including the body hash), so it uses httpx.Auth flow instead + # of static headers. See MCPSigV4Auth and _create_httpx_client_factory(). # update the headers with the extra headers if self.extra_headers: @@ -246,10 +333,16 @@ class MCPClient: f"MCP client using SSL configuration: {type(ssl_config).__name__}" ) + # Use SigV4 auth if configured and no explicit auth provided. + # The MCP SDK's sse_client and streamable_http_client call this + # factory without passing auth=, so self._aws_auth is used. + # For non-SigV4 clients, self._aws_auth is None — no behavior change. + effective_auth = auth if auth is not None else self._aws_auth + return httpx.AsyncClient( headers=headers, timeout=timeout, - auth=auth, + auth=effective_auth, verify=ssl_config, follow_redirects=True, ) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 5df79580d3e..67b95c7694b 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -82,8 +82,10 @@ class AnthropicCacheControlHook(CustomPromptManagement): _targetted_index: Optional[Union[int, str]] = point.get("index", None) targetted_index: Optional[int] = None if isinstance(_targetted_index, str): - if _targetted_index.isdigit(): + try: targetted_index = int(_targetted_index) + except ValueError: + pass else: targetted_index = _targetted_index diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index aed77ab2b3e..aa2a8121ee8 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -589,6 +589,16 @@ class CustomGuardrail(CustomLogger): guardrail_json_response ) + # Strip secret_fields to prevent plaintext Authorization headers from + # being persisted to spend logs, OTEL traces, or other logging backends. + # This matches the pattern used by Langfuse and Arize integrations. + if isinstance(clean_guardrail_response, dict): + clean_guardrail_response.pop("secret_fields", None) + elif isinstance(clean_guardrail_response, list): + for item in clean_guardrail_response: + if isinstance(item, dict): + item.pop("secret_fields", None) + slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index a694cec7d66..5e905da2230 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2221,6 +2221,11 @@ def sanitize_messages_for_tool_calling( Case C: Empty text content - Replace empty or whitespace-only text content with a placeholder message. + Case D: Duplicate tool_result for same tool_use (duplicate results) + - If multiple tool messages reference the same tool_call_id, keep only the last + occurrence. Anthropic requires exactly one tool_result per tool_use and rejects + with: "each tool_use must have a single result". + This function operates on OpenAI format messages before they are converted to provider-specific formats. """ @@ -2256,6 +2261,49 @@ def sanitize_messages_for_tool_calling( sanitized_messages.append(current_message) i += 1 + # Case D: Deduplicate tool results with the same tool_call_id. + # Anthropic requires exactly one tool_result per tool_use. Session history + # (e.g. from conversation resume) can contain duplicate tool_result messages + # for the same tool_call_id. Keep only the last occurrence *within each + # contiguous block of tool results following an assistant message*. This + # avoids dropping results from earlier turns if a tool_call_id is reused. + # + # NOTE: This intentionally keeps the *last* occurrence (most complete for + # session-resume duplicates), unlike _deduplicate_bedrock_content_blocks + # which keeps the *first*. The Bedrock case handles provider-side content + # block duplication where the first is authoritative; here the duplicate + # arises from history replay where the last entry is the final state. + duplicates_to_remove: Set[int] = set() + seen_in_block: Dict[str, int] = {} # tool_call_id -> index (reset per block) + for idx, msg in enumerate(sanitized_messages): + role = msg.get("role") + tcid = msg.get("tool_call_id") if role in ["tool", "function"] else None + if tcid: + if tcid in seen_in_block: + # Mark the earlier occurrence for removal (keep latest) + duplicates_to_remove.add(seen_in_block[tcid]) + verbose_logger.warning( + "sanitize_messages_for_tool_calling: dropping duplicate " + "tool_result with tool_call_id=%s. This may indicate " + "duplicate tool messages in conversation history.", + tcid, + ) + seen_in_block[tcid] = idx + elif role not in ("tool", "function"): + # Non-tool message (user, assistant, system) marks a + # conversational-turn boundary — reset tracking. + # Tool/function messages with no tool_call_id are malformed; + # they should NOT reset the block because they don't represent + # a turn boundary and would mask real within-block duplicates. + seen_in_block = {} + + if duplicates_to_remove: + sanitized_messages = [ + msg + for idx, msg in enumerate(sanitized_messages) + if idx not in duplicates_to_remove + ] + return sanitized_messages diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 04b27e87821..fd1859f7d17 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -395,6 +395,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): }, ) + # Anthropic requires input_schema.type to be "object". Normalize + # schemas from external sources (MCP servers, OpenAI callers) that + # may omit the type field or use a non-object type. + if _input_schema.get("type") != "object": + litellm.verbose_logger.debug( + "_map_tool_helper: coercing input_schema type from %r to " + "'object' for Anthropic compatibility (tool: %s)", + _input_schema.get("type"), + tool["function"].get("name"), + ) + _input_schema = dict(_input_schema) # avoid mutating caller's dict + _input_schema["type"] = "object" + if "properties" not in _input_schema: + _input_schema["properties"] = {} + _allowed_properties = set(AnthropicInputSchema.__annotations__.keys()) input_schema_filtered = { k: v for k, v in _input_schema.items() if k in _allowed_properties diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 194af4895fe..d3dd6b3d992 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19553,7 +19553,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-2025-04-14": { "cache_read_input_token_cost": 5e-07, @@ -19587,7 +19588,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, @@ -19624,7 +19626,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { "cache_read_input_token_cost": 1e-07, @@ -19658,7 +19661,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, @@ -20866,6 +20870,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21505,6 +21510,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21605,6 +21611,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21640,6 +21647,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21672,6 +21680,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true }, @@ -21707,6 +21716,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21742,6 +21752,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true }, @@ -21777,6 +21788,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21818,6 +21830,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21859,6 +21872,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21897,6 +21911,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21934,6 +21949,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -25721,7 +25737,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, @@ -25753,7 +25770,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, @@ -25786,7 +25804,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research-2025-06-26": { "cache_read_input_token_cost": 2.5e-06, @@ -25819,7 +25838,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-mini": { "cache_read_input_token_cost": 5.5e-07, @@ -25883,7 +25903,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-pro-2025-06-10": { "input_cost_per_token": 2e-05, @@ -25913,7 +25934,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini": { "cache_read_input_token_cost": 2.75e-07, @@ -25938,7 +25960,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, @@ -25957,7 +25980,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research": { "cache_read_input_token_cost": 5e-07, @@ -25990,7 +26014,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { "cache_read_input_token_cost": 5e-07, @@ -26023,7 +26048,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "oci/meta.llama-3.1-405b-instruct": { "input_cost_per_token": 1.068e-05, @@ -27639,6 +27665,92 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/qwen/qwen3.5-35b-a3b": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-35b-a3b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-27b": { + "input_cost_per_token": 3e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-27b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-122b-a10b": { + "input_cost_per_token": 4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-122b-a10b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-flash-02-23": { + "input_cost_per_token": 1e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://openrouter.ai/qwen/qwen3.5-flash-02-23", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-plus-02-15": { + "input_cost_per_token": 4e-07, + "input_cost_per_token_above_256k_tokens": 5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "output_cost_per_token_above_256k_tokens": 3e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-plus-02-15", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-397b-a17b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-397b-a17b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "openrouter/switchpoint/router": { "input_cost_per_token": 8.5e-07, "litellm_provider": "openrouter", diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 4c6735bacd3..93580b54305 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,3 +1,4 @@ +from datetime import datetime, timezone from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast from litellm._logging import verbose_proxy_logger @@ -6,6 +7,8 @@ from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, + MCPApprovalStatus, + MCPSubmissionsSummary, NewMCPServerRequest, SpecialMCPServerName, UpdateMCPServerRequest, @@ -102,12 +105,19 @@ def encrypt_credentials( async def get_all_mcp_servers( prisma_client: PrismaClient, + approval_status: Optional[str] = None, ) -> List[LiteLLM_MCPServerTable]: """ - Returns all of the mcp servers from the db + Returns mcp servers from the db, optionally filtered by approval_status. + Pass approval_status=None to return all servers regardless of approval state. """ try: - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() + where: Dict[str, Any] = {} + if approval_status is not None: + where["approval_status"] = approval_status + mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many( + where=where if where else {} + ) return [ LiteLLM_MCPServerTable(**mcp_server.model_dump()) @@ -451,3 +461,71 @@ async def delete_user_credential( await prisma_client.db.litellm_mcpusercredentials.delete( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) + + +async def approve_mcp_server( + prisma_client: PrismaClient, + server_id: str, + touched_by: str, +) -> LiteLLM_MCPServerTable: + """Set approval_status=active and record reviewed_at.""" + now = datetime.now(timezone.utc) + updated = await prisma_client.db.litellm_mcpservertable.update( + where={"server_id": server_id}, + data={ + "approval_status": MCPApprovalStatus.active, + "reviewed_at": now, + "updated_by": touched_by, + }, + ) + return LiteLLM_MCPServerTable(**updated.model_dump()) + + +async def reject_mcp_server( + prisma_client: PrismaClient, + server_id: str, + touched_by: str, + review_notes: Optional[str] = None, +) -> LiteLLM_MCPServerTable: + """Set approval_status=rejected, record reviewed_at and review_notes.""" + now = datetime.now(timezone.utc) + data: Dict[str, Any] = { + "approval_status": MCPApprovalStatus.rejected, + "reviewed_at": now, + "updated_by": touched_by, + } + if review_notes is not None: + data["review_notes"] = review_notes + updated = await prisma_client.db.litellm_mcpservertable.update( + where={"server_id": server_id}, + data=data, + ) + return LiteLLM_MCPServerTable(**updated.model_dump()) + + +async def get_mcp_submissions( + prisma_client: PrismaClient, +) -> MCPSubmissionsSummary: + """ + Returns all MCP servers that were submitted by non-admin users (submitted_at IS NOT NULL), + along with a summary count breakdown by approval_status. + Mirrors get_guardrail_submissions() from guardrail_endpoints.py. + """ + rows = await prisma_client.db.litellm_mcpservertable.find_many( + where={"submitted_at": {"not": None}}, + order={"submitted_at": "desc"}, + take=500, # safety cap; paginate if needed in a future iteration + ) + items = [LiteLLM_MCPServerTable(**r.model_dump()) for r in rows] + + pending = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review) + active = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active) + rejected = sum(1 for i in items if i.approval_status == MCPApprovalStatus.rejected) + + return MCPSubmissionsSummary( + total=len(items), + pending_review=pending, + active=active, + rejected=rejected, + items=items, + ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0b58009fcf6..8769d3952a4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -38,7 +38,7 @@ from litellm.constants import ( MCP_TOOL_LISTING_TIMEOUT, ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException -from litellm.experimental_mcp_client.client import MCPClient +from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, @@ -318,6 +318,7 @@ class MCPServerManager: # oauth specific fields client_id=server_config.get("client_id", None), client_secret=server_config.get("client_secret", None), + oauth2_flow=server_config.get("oauth2_flow", None), scopes=resolved_scopes, authorization_url=resolved_authorization_url, token_url=resolved_token_url, @@ -339,6 +340,12 @@ class MCPServerManager: available_on_public_internet=bool( server_config.get("available_on_public_internet", True) ), + # AWS SigV4 fields + aws_access_key_id=server_config.get("aws_access_key_id", None), + aws_secret_access_key=server_config.get("aws_secret_access_key", None), + aws_session_token=server_config.get("aws_session_token", None), + aws_region_name=server_config.get("aws_region_name", None), + aws_service_name=server_config.get("aws_service_name", None), ) self.config_mcp_servers[server_id] = new_server @@ -418,6 +425,8 @@ class MCPServerManager: headers["Authorization"] = f"ApiKey {server.authentication_token}" elif server.auth_type == MCPAuth.basic: headers["Authorization"] = f"Basic {server.authentication_token}" + elif server.auth_type == MCPAuth.token: + headers["Authorization"] = f"token {server.authentication_token}" # Add any static headers from server config. # @@ -588,6 +597,10 @@ class MCPServerManager: else: client_secret_value = encrypted_client_secret + # TODO: Add AWS SigV4 credential decryption here when DB-stored + # SigV4 MCP servers are supported. Requires corresponding changes + # to encrypt_credentials() in db.py and MCPCredentials TypedDict. + scopes: Optional[List[str]] = None if credentials_dict: scopes_value = credentials_dict.get("scopes") @@ -605,12 +618,17 @@ class MCPServerManager: mcp_info["description"] = mcp_server.description auth_type = cast(MCPAuthType, mcp_server.auth_type) - if mcp_server.url and auth_type == MCPAuth.oauth2: - mcp_oauth_metadata = await self._descovery_metadata( - server_url=mcp_server.url, - ) - else: - mcp_oauth_metadata = None + server_url = mcp_server.url + needs_discovery = ( + bool(server_url) + and auth_type == MCPAuth.oauth2 + and not mcp_server.authorization_url + ) + mcp_oauth_metadata = ( + await self._descovery_metadata(server_url=server_url) # type: ignore[arg-type] + if needs_discovery + else None + ) resolved_scopes = scopes or ( mcp_oauth_metadata.scopes if mcp_oauth_metadata else None @@ -632,6 +650,7 @@ class MCPServerManager: client_id=client_id_value or getattr(mcp_server, "client_id", None), client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), + oauth2_flow=getattr(mcp_server, "oauth2_flow", None), scopes=resolved_scopes, authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None), @@ -973,6 +992,18 @@ class MCPServerManager: else: # For HTTP/SSE transports server_url = server.url or "" + + # Create SigV4 auth if configured + aws_auth = None + if server.auth_type == MCPAuth.aws_sigv4: + aws_auth = MCPSigV4Auth( + aws_access_key_id=server.aws_access_key_id, + aws_secret_access_key=server.aws_secret_access_key, + aws_session_token=server.aws_session_token, + aws_region_name=server.aws_region_name, + aws_service_name=server.aws_service_name, + ) + return MCPClient( server_url=server_url, transport_type=transport, @@ -980,6 +1011,7 @@ class MCPServerManager: auth_value=auth_value, timeout=MCP_CLIENT_TIMEOUT, extra_headers=extra_headers, + aws_auth=aws_auth, ) async def _get_tools_from_server( @@ -2270,7 +2302,7 @@ class MCPServerManager: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) - db_mcp_servers = await get_all_mcp_servers(prisma_client) + db_mcp_servers = await get_all_mcp_servers(prisma_client, approval_status="active") verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in database") previous_registry = self.registry @@ -2506,9 +2538,11 @@ class MCPServerManager: if server.requires_per_user_auth: should_skip_health_check = True # Skip if auth_type is not none and authentication_token is missing + # (except aws_sigv4 which uses its own credential fields) elif ( server.auth_type and server.auth_type != MCPAuth.none + and server.auth_type != MCPAuth.aws_sigv4 and not server.authentication_token ): should_skip_health_check = True diff --git a/litellm/proxy/_experimental/out/assets/logos/figma.svg b/litellm/proxy/_experimental/out/assets/logos/figma.svg new file mode 100644 index 00000000000..2d8b70457d9 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/figma.svg @@ -0,0 +1,7 @@ + + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/gitlab.svg b/litellm/proxy/_experimental/out/assets/logos/gitlab.svg new file mode 100644 index 00000000000..18a89fa328d --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/gitlab.svg @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/gmail.svg b/litellm/proxy/_experimental/out/assets/logos/gmail.svg new file mode 100644 index 00000000000..d702890620d --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/gmail.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/google_drive.svg b/litellm/proxy/_experimental/out/assets/logos/google_drive.svg new file mode 100644 index 00000000000..7048af9915e --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/google_drive.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/hubspot.svg b/litellm/proxy/_experimental/out/assets/logos/hubspot.svg new file mode 100644 index 00000000000..b993945ac6b --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/hubspot.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/jira.svg b/litellm/proxy/_experimental/out/assets/logos/jira.svg new file mode 100644 index 00000000000..fb10ca75173 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/jira.svg @@ -0,0 +1,15 @@ + + + + + + + + + + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/linear.svg b/litellm/proxy/_experimental/out/assets/logos/linear.svg new file mode 100644 index 00000000000..83662a1f9ff --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/linear.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/notion.svg b/litellm/proxy/_experimental/out/assets/logos/notion.svg new file mode 100644 index 00000000000..170b9bb4140 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/notion.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/salesforce.svg b/litellm/proxy/_experimental/out/assets/logos/salesforce.svg new file mode 100644 index 00000000000..1a541a004f1 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/salesforce.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/sentry.svg b/litellm/proxy/_experimental/out/assets/logos/sentry.svg new file mode 100644 index 00000000000..9c3733dc43e --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/sentry.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/shopify.svg b/litellm/proxy/_experimental/out/assets/logos/shopify.svg new file mode 100644 index 00000000000..fcc7547269d --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/shopify.svg @@ -0,0 +1,4 @@ + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/slack.svg b/litellm/proxy/_experimental/out/assets/logos/slack.svg new file mode 100644 index 00000000000..801de4f70c8 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/slack.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/stripe.svg b/litellm/proxy/_experimental/out/assets/logos/stripe.svg new file mode 100644 index 00000000000..ac16a6fb170 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/stripe.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/twilio.svg b/litellm/proxy/_experimental/out/assets/logos/twilio.svg new file mode 100644 index 00000000000..3517a2824d9 --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/twilio.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_experimental/out/assets/logos/zapier.svg b/litellm/proxy/_experimental/out/assets/logos/zapier.svg new file mode 100644 index 00000000000..8428ba82a5b --- /dev/null +++ b/litellm/proxy/_experimental/out/assets/logos/zapier.svg @@ -0,0 +1,3 @@ + + + diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d797d9c7e0a..36790e9feae 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1087,6 +1087,12 @@ class SpecialMCPServerName(str, enum.Enum): all_proxy_servers = "all-proxy-mcpservers" +class MCPApprovalStatus(str, enum.Enum): + pending_review = "pending_review" + active = "active" + rejected = "rejected" + + # MCP Proxy Request Types class NewMCPServerRequest(LiteLLMPydanticObjectBase): server_id: Optional[str] = None @@ -1117,6 +1123,18 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): is_byok: bool = False byok_description: List[str] = Field(default_factory=list) byok_api_key_help_url: Optional[str] = None + source_url: Optional[str] = None + # BYOM submission fields — set by the endpoint, not by the caller. + # Any caller-provided values are silently overridden before persistence. + approval_status: Optional[str] = Field( + None, description="Server-managed: set by the endpoint; caller values are overridden." + ) + submitted_by: Optional[str] = Field( + None, description="Server-managed: set by the endpoint; caller values are overridden." + ) + submitted_at: Optional[datetime] = Field( + None, description="Server-managed: set by the endpoint; caller values are overridden." + ) @model_validator(mode="before") @classmethod @@ -1176,6 +1194,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): is_byok: bool = False byok_description: List[str] = Field(default_factory=list) byok_api_key_help_url: Optional[str] = None + source_url: Optional[str] = None @model_validator(mode="before") @classmethod @@ -1239,6 +1258,16 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): byok_description: List[str] = Field(default_factory=list) byok_api_key_help_url: Optional[str] = None has_user_credential: Optional[bool] = None + source_url: Optional[str] = None + # BYOM submission fields + approval_status: Optional[str] = Field( + default="active", + description="Approval status: 'pending_review', 'active', 'rejected'", + ) + submitted_by: Optional[str] = None + submitted_at: Optional[datetime] = None + reviewed_at: Optional[datetime] = None + review_notes: Optional[str] = None class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): @@ -1255,6 +1284,18 @@ class MCPUserCredentialResponse(LiteLLMPydanticObjectBase): has_credential: bool +class RejectMCPServerRequest(LiteLLMPydanticObjectBase): + review_notes: Optional[str] = None + + +class MCPSubmissionsSummary(LiteLLMPydanticObjectBase): + total: int + pending_review: int + active: int + rejected: int + items: List["LiteLLM_MCPServerTable"] + + ######## Skills API Types ######## @@ -2203,6 +2244,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.", ) + mcp_required_fields: Optional[List[str]] = Field( + None, + description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index ce32ebf54f8..ddeba2100c3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -36,7 +36,6 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import ( @@ -100,6 +99,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self.mock_redacted_text = mock_redacted_text self.output_parse_pii = output_parse_pii or False self.apply_to_output = apply_to_output + + # When output_parse_pii or apply_to_output is enabled, the guardrail must + # also run on post_call to unmask/mask the response. Expand the event_hook + # so should_run_guardrail returns True for both pre_call and post_call. + if (self.output_parse_pii or self.apply_to_output) and not logging_only: + current_hook = self.event_hook + if isinstance(current_hook, str) and current_hook != "post_call": + self.event_hook = [current_hook, "post_call"] + elif isinstance(current_hook, list) and "post_call" not in current_hook: + self.event_hook = current_hook + ["post_call"] self.pii_entities_config: Dict[Union[PiiEntityType, str], PiiAction] = ( pii_entities_config or {} ) @@ -475,13 +484,15 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): new_text = text if redacted_text is not None: verbose_proxy_logger.debug("redacted_text: %s", redacted_text) - for item in redacted_text["items"]: + # Process items in reverse order by start position so that + # replacing later spans first does not shift earlier coordinates. + for item in sorted( + redacted_text["items"], key=lambda x: x["start"], reverse=True + ): start = item["start"] end = item["end"] replacement = item["text"] # replacement token if item["operator"] == "replace" and output_parse_pii is True: - # check if token in dict - # if exists, add a uuid to the replacement token for swapping back to the original text in llm response output parsing if request_data is None: verbose_proxy_logger.warning( "Presidio anonymize_text called without request_data — " @@ -489,17 +500,28 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): "This may indicate a missing caller update." ) request_data = {} - if "pii_tokens" not in request_data: - request_data["pii_tokens"] = {} - pii_tokens = request_data["pii_tokens"] + # Store pii_tokens in metadata to avoid leaking to LLM providers. + # Providers like Anthropic reject unknown top-level fields. + if not request_data.get("metadata"): + request_data["metadata"] = {} + if "pii_tokens" not in request_data["metadata"]: + request_data["metadata"]["pii_tokens"] = {} + pii_tokens = request_data["metadata"]["pii_tokens"] - # Always append a UUID to ensure the replacement token is unique to this request and session. - # This prevents collisions where the LLM might hallucinate a generic token like [PHONE_NUMBER]. - replacement = f"{replacement}_{str(uuid.uuid4())[:12]}" + # Append a sequential number to make each token unique + # per request, so unmasking maps back to the correct + # original value. Format: , + # This is LLM-friendly and degrades gracefully if the + # LLM doesn't echo the token verbatim. + seq = len(pii_tokens) + 1 + if replacement.endswith(">"): + replacement = f"{replacement[:-1]}_{seq}>" + else: + replacement = f"{replacement}_{seq}" - pii_tokens[replacement] = new_text[ - start:end - ] # get text it'll replace + # Use ORIGINAL text (not new_text) since start/end + # reference the original text's coordinates. + pii_tokens[replacement] = text[start:end] new_text = new_text[:start] + replacement + new_text[end:] entity_type = item.get("entity_type", None) @@ -507,12 +529,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): masked_entity_count[entity_type] = ( masked_entity_count.get(entity_type, 0) + 1 ) - # When output_parse_pii is True, new_text contains UUID-suffixed - # tokens that match the keys in pii_tokens. Returning - # redacted_text["text"] (Presidio's original output) would send - # un-suffixed tokens to the LLM, making unmasking impossible. + # When output_parse_pii is True, new_text contains sequentially + # numbered tokens (e.g. ) that match the keys + # in pii_tokens. Returning redacted_text["text"] (Presidio's + # original output) would send un-numbered tokens to the LLM, + # making unmasking impossible. # When output_parse_pii is False, new_text == redacted_text["text"] - # because no UUID suffix is appended. + # because no suffix is appended. return new_text else: raise Exception("Invalid anonymizer response: received None") @@ -544,8 +567,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): filtered_results: List[PresidioAnalyzeResponseItem] = [] deny_list_strings = [ - getattr(x, "value", str(x)) - for x in self.presidio_entities_deny_list + getattr(x, "value", str(x)) for x in self.presidio_entities_deny_list ] for item in analyze_results: entity_type = item.get("entity_type") @@ -884,6 +906,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) if self.apply_to_output is True: + if self._is_anthropic_message_response(response): + return await self._process_anthropic_response_for_pii( + response=response, request_data=data, mode="mask" + ) return await self._mask_output_response( response=response, request_data=data ) @@ -899,6 +925,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data=data, mode="unmask", ) + elif self._is_anthropic_message_response(response): + await self._process_anthropic_response_for_pii( + response=response, request_data=data, mode="unmask" + ) return response @staticmethod @@ -927,6 +957,57 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): break return text + @staticmethod + def _is_anthropic_message_response(response: Any) -> bool: + """Check if the response is an Anthropic native message dict.""" + return ( + isinstance(response, dict) + and response.get("type") == "message" + and isinstance(response.get("content"), list) + ) + + async def _process_anthropic_response_for_pii( + self, + response: dict, + request_data: dict, + mode: Literal["mask", "unmask"], + ) -> dict: + """ + Process an Anthropic native message dict for PII masking/unmasking. + Handles content blocks with type == "text". + """ + metadata = (request_data.get("metadata") or {}) if request_data else {} + pii_tokens = metadata.get("pii_tokens", {}) + if not pii_tokens and mode == "unmask": + verbose_proxy_logger.debug( + "No pii_tokens in metadata for Anthropic response unmask" + ) + presidio_config = self.get_presidio_settings_from_request_data( + request_data or {} + ) + + content = response.get("content") + if not isinstance(content, list): + return response + + for block in content: + if not isinstance(block, dict) or block.get("type") != "text": + continue + text_value = block.get("text") + if text_value is None: + continue + if mode == "unmask": + block["text"] = self._unmask_pii_text(text_value, pii_tokens) + elif mode == "mask": + block["text"] = await self.check_pii( + text=text_value, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + + return response + async def _process_response_for_pii( self, response: ModelResponse, @@ -937,10 +1018,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Helper to recursively process a ModelResponse for PII. Handles all choices and tool calls. """ - pii_tokens = request_data.get("pii_tokens", {}) if request_data else {} + metadata = (request_data.get("metadata") or {}) if request_data else {} + pii_tokens = metadata.get("pii_tokens", {}) if not pii_tokens and mode == "unmask": verbose_proxy_logger.debug( - "No pii_tokens found in request_data — nothing to unmask" + "No pii_tokens found in request_data['metadata'] — nothing to unmask" ) presidio_config = self.get_presidio_settings_from_request_data( request_data or {} @@ -1045,7 +1127,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, response: Any, request_data: dict, - ) -> AsyncGenerator[ModelResponseStream, None]: + ) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]: """ Process streaming response chunks to unmask PII tokens when needed. """ @@ -1062,8 +1144,20 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): 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 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( @@ -1099,10 +1193,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return # --- PII unmasking path (output_parse_pii=True) --- - pii_tokens = request_data.get("pii_tokens", {}) if request_data else {} + 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 for streaming unmask path" + "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: @@ -1114,6 +1209,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async for chunk in response: 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 if not remaining_chunks: return @@ -1191,15 +1290,24 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ texts = inputs.get("texts", []) + # When input_type is "response" and pii_tokens are available, + # unmask the text instead of masking it. + metadata = (request_data.get("metadata") or {}) if request_data else {} + pii_tokens = metadata.get("pii_tokens", {}) + new_texts = [] - for text in texts: - modified_text = await self.check_pii( - text=text, - output_parse_pii=self.output_parse_pii, - presidio_config=None, - request_data=request_data or {}, - ) - new_texts.append(modified_text) + if input_type == "response" and pii_tokens: + for text in texts: + new_texts.append(self._unmask_pii_text(text, pii_tokens)) + else: + for text in texts: + modified_text = await self.check_pii( + text=text, + output_parse_pii=self.output_parse_pii, + presidio_config=None, + request_data=request_data or {}, + ) + new_texts.append(modified_text) inputs["texts"] = new_texts return inputs diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 68f997e29cb..b9dcc514d29 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -64,6 +64,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, attach_object_permission_to_dict, handle_update_object_permission_common, + validate_key_mcp_servers_against_team, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, @@ -638,6 +639,12 @@ async def _common_key_generation_helper( # noqa: PLR0915 data_json.pop("tags") + # Validate MCP servers in object_permission are within team scope + await validate_key_mcp_servers_against_team( + object_permission=data_json.get("object_permission"), + team_obj=team_table, + ) + data_json = await _set_object_permission( data_json=data_json, prisma_client=prisma_client, @@ -1947,6 +1954,27 @@ async def update_key_fn( # Set Management Endpoint Metadata Fields + # 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, + ) + non_default_values = await prepare_key_update_data( data=data, existing_key_row=existing_key_row ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f7a4cec301b..08f452859f2 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -11,14 +11,16 @@ Endpoints here: - GET `/v1/mcp/tools - lists all the tools available for a key - GET `/v1/mcp/access_groups` - lists all available MCP access groups - GET `/v1/mcp/discover` - Returns curated list of well-known MCP servers for discovery UI +- GET `/v1/mcp/openapi-registry` - Returns well-known OpenAPI APIs with OAuth 2.0 metadata """ +import functools import importlib import json import os from dataclasses import dataclass -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Literal, Optional from fastapi import ( @@ -76,11 +78,14 @@ if MCP_AVAILABLE: return _ToolNameValidationResult() from litellm.proxy._experimental.mcp_server.db import ( + approve_mcp_server, create_mcp_server, delete_mcp_server, delete_user_credential, get_all_mcp_servers_for_user, get_mcp_server, + get_mcp_submissions, + reject_mcp_server, store_user_credential, update_mcp_server, ) @@ -100,9 +105,12 @@ if MCP_AVAILABLE: LiteLLM_MCPServerTable, LitellmUserRoles, MakeMCPServersPublicRequest, + MCPApprovalStatus, + MCPSubmissionsSummary, MCPUserCredentialRequest, MCPUserCredentialResponse, NewMCPServerRequest, + RejectMCPServerRequest, SpecialMCPServerName, UpdateMCPServerRequest, UserAPIKeyAuth, @@ -155,6 +163,59 @@ if MCP_AVAILABLE: _base_validate_and_normalize_mcp_server_payload(payload) _validate_mcp_server_name_fields(payload) + _VALID_MCP_REQUIRED_FIELDS: frozenset = frozenset( + NewMCPServerRequest.model_fields + ) + + def _validate_mcp_required_fields(payload: Any) -> None: + """Validate submission payload against admin-configured mcp_required_fields.""" + from litellm.proxy.proxy_server import ( + general_settings as proxy_general_settings, + ) + + required_fields: Optional[List[str]] = proxy_general_settings.get( + "mcp_required_fields" + ) + if not required_fields: + return + + # Fail fast on unknown field names — a typo in the config would silently + # block every submission with a confusing "missing fields" error. + unknown = [f for f in required_fields if f not in _VALID_MCP_REQUIRED_FIELDS] + if unknown: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "error": f"mcp_required_fields contains unknown field names: {unknown}. " + "Check general_settings.mcp_required_fields in your proxy config." + }, + ) + + # Mirror the UI's compliance checks (MCPStandardsSettings.tsx FIELD_GROUPS): + # auth_type requires a real value — "none" is treated as absent. + _AUTH_TYPE_SENTINEL = "none" + + def _field_present(field_name: str) -> bool: + value = getattr(payload, field_name, None) + if value is None: + return False + # Treat empty string and empty list as absent (mirrors UI compliance check) + if isinstance(value, (str, list)) and not value: + return False + if field_name == "auth_type" and value == _AUTH_TYPE_SENTINEL: + return False + return True + + missing = [f for f in required_fields if not _field_present(f)] + if missing: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": f"Submission is missing required fields: {missing}. " + "Configure required fields via general_settings.mcp_required_fields." + }, + ) + def _is_public_registry_enabled() -> bool: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, @@ -554,6 +615,46 @@ if MCP_AVAILABLE: return "view_all" return "restricted" + async def _get_team_scoped_mcp_server_list( + team_id: str, + ) -> List[LiteLLM_MCPServerTable]: + """ + Return MCP servers scoped to a team: team's allowed servers + allow_all_keys servers. + Used by the Create Key UI to populate the MCP server dropdown. + """ + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.management_helpers.object_permission_utils import ( + _get_allow_all_keys_server_ids, + _get_team_allowed_mcp_servers, + ) + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + team_server_ids = await _get_team_allowed_mcp_servers(team_obj) + allow_all_server_ids = _get_allow_all_keys_server_ids() + all_allowed_ids = team_server_ids | allow_all_server_ids + + if not all_allowed_ids: + return [] + + # Collect servers from registry + servers: List[LiteLLM_MCPServerTable] = [] + for server_id in all_allowed_ids: + server = global_mcp_server_manager.get_mcp_server_by_id(server_id) + if server is not None: + mcp_server_table = global_mcp_server_manager._build_mcp_server_table( + server + ) + servers.append(mcp_server_table) + + return _redact_mcp_credentials_list(servers) + @router.get( "/server", description="Returns the mcp server list with associated teams", @@ -562,38 +663,88 @@ if MCP_AVAILABLE: ) async def fetch_all_mcp_servers( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + team_id: Optional[str] = Query( + None, + description="Filter MCP servers by team scope. When provided, returns only " + "servers the team has access to plus globally available (allow_all_keys) servers. " + "Used by the Create Key UI to show team-scoped MCP servers.", + ), ): """ Get all of the configured mcp servers for the user in the db with their associated teams ``` curl --location 'http://localhost:4000/v1/mcp/server' \ --header 'Authorization: Bearer your_api_key_here' + + # Filter by team scope (for Create Key UI) + curl --location 'http://localhost:4000/v1/mcp/server?team_id=team-123' \ + --header 'Authorization: Bearer your_api_key_here' ``` """ - user_mcp_management_mode = _get_user_mcp_management_mode() + # If team_id is provided, return team-scoped servers + allow_all_keys servers is_restricted_virtual_key = _is_restricted_virtual_key_request( user_api_key_dict ) - - if user_mcp_management_mode == "view_all" and not is_restricted_virtual_key: - servers = await global_mcp_server_manager.get_all_mcp_servers_unfiltered() - redacted_mcp_servers = _redact_mcp_credentials_list(servers) - else: - auth_contexts = await build_effective_auth_contexts(user_api_key_dict) - - aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} - for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_all_allowed_mcp_servers( - user_api_key_auth=auth_context + if team_id is not None and isinstance(team_id, str) and team_id.strip(): + # Restricted virtual keys must not use the team_id filter to + # bypass their own access limitations. + if is_restricted_virtual_key: + raise HTTPException( + status_code=403, + detail="Restricted virtual keys cannot query team-scoped MCP servers.", ) - for server in servers: - if server.server_id not in aggregated_servers: - aggregated_servers[server.server_id] = server - redacted_mcp_servers = _redact_mcp_credentials_list( - aggregated_servers.values() - ) + # Only proxy admins may query another team's MCP servers. + # Non-admins must belong to the requested team. + sanitized_team_id = team_id.strip() + is_admin = _user_has_admin_view(user_api_key_dict) + if not is_admin: + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + user_api_key_cache, + ) + + team_obj = await get_team_object( + team_id=sanitized_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + user_in_team = any( + m.user_id is not None + and m.user_id == user_api_key_dict.user_id + for m in team_obj.members_with_roles + ) + if not user_in_team: + raise HTTPException( + status_code=403, + detail="You do not have permission to view MCP servers for this team.", + ) + + redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id) + else: + user_mcp_management_mode = _get_user_mcp_management_mode() + + if user_mcp_management_mode == "view_all" and not is_restricted_virtual_key: + servers = await global_mcp_server_manager.get_all_mcp_servers_unfiltered() + redacted_mcp_servers = _redact_mcp_credentials_list(servers) + else: + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + + aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} + for auth_context in auth_contexts: + servers = await global_mcp_server_manager.get_all_allowed_mcp_servers( + user_api_key_auth=auth_context + ) + for server in servers: + if server.server_id not in aggregated_servers: + aggregated_servers[server.server_id] = server + + redacted_mcp_servers = _redact_mcp_credentials_list( + aggregated_servers.values() + ) # augment the mcp servers with public status if litellm.public_mcp_servers is not None: @@ -689,6 +840,187 @@ if MCP_AVAILABLE: for server_id, status in server_status_map.items() ] + @router.post( + "/server/register", + description="Submit a new MCP server for admin review (non-admin users). Mirrors POST /guardrails/register.", + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_MCPServerTable, + status_code=status.HTTP_201_CREATED, + ) + @management_endpoint_wrapper + async def register_mcp_server( + payload: NewMCPServerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + Allow team members to submit an MCP server for admin review. + Creates the server with approval_status=pending_review. + Requires a team-scoped API key. + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "PROXY_ADMIN users should use POST /v1/mcp/server to create servers directly instead of the submission workflow." + }, + ) + + if not user_api_key_dict.team_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "Registration requires an API key associated with a team. Use a team-scoped key." + }, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + + validate_and_normalize_mcp_server_payload(payload) + _validate_mcp_required_fields(payload) + + payload.approval_status = MCPApprovalStatus.pending_review + payload.submitted_by = user_api_key_dict.user_id + payload.submitted_at = datetime.now(timezone.utc) + + try: + new_mcp_server = await create_mcp_server( + prisma_client, + payload, + touched_by=user_api_key_dict.user_id or user_api_key_dict.team_id, + ) + except Exception as e: + verbose_proxy_logger.exception(f"Error registering mcp server: {str(e)}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": f"Error registering mcp server: {str(e)}"}, + ) + # Do NOT add to runtime registry — pending servers are not active + return _redact_mcp_credentials(new_mcp_server) + + @router.get( + "/server/submissions", + description="Returns all MCP servers submitted by non-admin users (admin review queue). Mirrors GET /guardrails/submissions.", + dependencies=[Depends(user_api_key_auth)], + response_model=MCPSubmissionsSummary, + ) + @management_endpoint_wrapper + async def get_mcp_server_submissions( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + Admin-only endpoint to view all user-submitted MCP servers pending review. + """ + if user_api_key_dict.user_role not in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Admin access required to view MCP server submissions."}, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + + return await get_mcp_submissions(prisma_client) + + @router.put( + "/server/{server_id}/approve", + description="Approve a pending MCP server submission (admin only). Mirrors PUT /guardrails/{id}/approve.", + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_MCPServerTable, + ) + @management_endpoint_wrapper + async def approve_mcp_server_submission( + server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + Admin approves a pending or previously-rejected MCP server — sets approval_status=active and loads it into the runtime registry. + """ + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Admin access required to approve MCP server submissions."}, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + + existing = await get_mcp_server(prisma_client, server_id) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP server '{server_id}' not found."}, + ) + if existing.approval_status == MCPApprovalStatus.active: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "MCP server is already active."}, + ) + + approved = await approve_mcp_server( + prisma_client, + server_id, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + ) + await global_mcp_server_manager.reload_servers_from_database() + + return _redact_mcp_credentials(approved) + + @router.put( + "/server/{server_id}/reject", + description="Reject a pending MCP server submission (admin only). Mirrors PUT /guardrails/{id}/reject.", + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_MCPServerTable, + ) + @management_endpoint_wrapper + async def reject_mcp_server_submission( + server_id: str, + payload: RejectMCPServerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + Admin rejects a pending MCP server — sets approval_status=rejected with optional review_notes. + """ + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Admin access required to reject MCP server submissions."}, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to your proxy" + ) + + existing = await get_mcp_server(prisma_client, server_id) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP server '{server_id}' not found."}, + ) + if existing.approval_status == MCPApprovalStatus.rejected: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "MCP server is already rejected."}, + ) + + was_active = existing.approval_status == MCPApprovalStatus.active + rejected = await reject_mcp_server( + prisma_client, + server_id, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + review_notes=payload.review_notes, + ) + # Only evict from the runtime registry if the server was previously active + if was_active: + await global_mcp_server_manager.reload_servers_from_database() + return _redact_mcp_credentials(rejected) + @router.get( "/server/{server_id}", description="Returns the mcp server info", @@ -829,6 +1161,13 @@ if MCP_AVAILABLE: # TODO: audit log for create + # Admin-created servers are always active — clear any submission lifecycle + # fields the caller may have provided to prevent fake entries appearing in + # the submissions queue. + payload.approval_status = MCPApprovalStatus.active + payload.submitted_by = None + payload.submitted_at = None + # Attempt to create the mcp server try: new_mcp_server = await create_mcp_server( @@ -1361,3 +1700,40 @@ if MCP_AVAILABLE: "servers": servers, "categories": categories, } + + # --- OpenAPI Registry --- + + _OPENAPI_REGISTRY_PATH = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "openapi_registry.json", + ) + + @functools.lru_cache(maxsize=1) + def _load_openapi_registry() -> Dict[str, Any]: + with open(_OPENAPI_REGISTRY_PATH, "r") as f: + data: Dict[str, Any] = json.load(f) + return data + + @router.get( + "/openapi-registry", + description="Returns well-known OpenAPI APIs with OAuth 2.0 metadata for the OpenAPI MCP picker", + ) + async def get_openapi_registry( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={ + "error": "Only proxy admins can access the OpenAPI registry. Your role={}".format( + user_api_key_dict.user_role + ) + }, + ) + try: + return _load_openapi_registry() + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to load OpenAPI registry from {_OPENAPI_REGISTRY_PATH}: {e}" + ) + return {"apis": []} diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 95b7300992c..b7714d3f866 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -21,7 +21,6 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, - compute_tag_metadata_totals, get_daily_activity, ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity @@ -554,5 +553,12 @@ async def get_tag_daily_activity( api_key=api_key, page=page, page_size=page_size, - metadata_metrics_func=compute_tag_metadata_totals, + # metadata_metrics_func=None because litellm_dailytagspend rows are + # pre-aggregated per (date, tag, model, …) and have no request_id. + # Deduplication across tags is therefore not possible at this level — + # a request tagged with N tags contributes its spend to N separate rows, + # so passing compute_tag_metadata_totals would double-count spend when + # multiple tags are present. The panel is primarily used to inspect + # individual tags, making this trade-off acceptable. + metadata_metrics_func=None, ) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 9670cdf330a..319a0b5eb79 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -4,12 +4,14 @@ organizations, teams, and keys. """ import json -from litellm._uuid import uuid -from typing import Dict, Optional, Union +from typing import Dict, List, Optional, Set, Union + +from fastapi import HTTPException, status from litellm._logging import verbose_proxy_logger -from litellm.proxy.utils import PrismaClient +from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy.utils import PrismaClient @@ -177,4 +179,178 @@ async def _set_object_permission( data_json["object_permission_id"] = created_permission.object_permission_id data_json.pop("object_permission") - return data_json \ No newline at end of file + return data_json + + +async def _resolve_team_allowed_mcp_servers( + team_object_permission: "LiteLLM_ObjectPermissionTable", +) -> Set[str]: + """ + Resolve the full set of MCP server IDs a team has access to. + + Combines: + - Direct mcp_servers list + - Servers from mcp_access_groups + - Server IDs referenced in mcp_tool_permissions keys + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + direct_servers: List[str] = team_object_permission.mcp_servers or [] + access_group_servers: List[str] = ( + await MCPRequestHandler._get_mcp_servers_from_access_groups( + team_object_permission.mcp_access_groups or [] + ) + ) + raw_tool_perms = team_object_permission.mcp_tool_permissions or {} + if isinstance(raw_tool_perms, str): + raw_tool_perms = json.loads(raw_tool_perms) + tool_perm_servers: List[str] = list(raw_tool_perms.keys()) + return set(direct_servers + access_group_servers + tool_perm_servers) + + +def _get_allow_all_keys_server_ids() -> Set[str]: + """Return the set of MCP server IDs marked with allow_all_keys=True.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return set(global_mcp_server_manager.get_allow_all_keys_server_ids()) + + +async def _get_team_allowed_mcp_servers( + team_obj: Optional["LiteLLM_TeamTableCachedObj"], +) -> Set[str]: + """ + Get the full set of MCP server IDs a team allows. + + If team has no object_permission or no MCP config, returns empty set + (meaning only allow_all_keys servers are permitted). + """ + if team_obj is None: + return set() + + team_object_permission = team_obj.object_permission + if team_object_permission is None: + return set() + + return await _resolve_team_allowed_mcp_servers(team_object_permission) + + +def _extract_requested_mcp_server_ids( + object_permission: Optional[dict], +) -> Set[str]: + """ + Extract all MCP server IDs referenced in a key's object_permission dict. + + Includes: + - mcp_servers list + - Keys from mcp_tool_permissions + """ + if not object_permission or not isinstance(object_permission, dict): + return set() + + server_ids: Set[str] = set() + mcp_servers = object_permission.get("mcp_servers") + if isinstance(mcp_servers, list): + server_ids.update(mcp_servers) + + mcp_tool_permissions = object_permission.get("mcp_tool_permissions") + if isinstance(mcp_tool_permissions, dict): + server_ids.update(mcp_tool_permissions.keys()) + + return server_ids + + +def _extract_requested_mcp_access_groups( + object_permission: Optional[dict], +) -> Set[str]: + """Extract MCP access groups from a key's object_permission dict.""" + if not object_permission or not isinstance(object_permission, dict): + return set() + + groups = object_permission.get("mcp_access_groups") + if isinstance(groups, list): + return set(groups) + return set() + + +async def validate_key_mcp_servers_against_team( + object_permission: Optional[dict], + team_obj: Optional["LiteLLM_TeamTableCachedObj"], +): + """ + Validate that MCP servers requested on a key are within the allowed scope. + + Rules: + - If key is in a team: key's mcp_servers must be a subset of + (team's allowed servers + allow_all_keys servers) + - If key is NOT in a team: key's mcp_servers must only contain + allow_all_keys servers + - If team has no MCP config: key can only use allow_all_keys servers + + Raises HTTPException(403) if validation fails. + """ + requested_servers = _extract_requested_mcp_server_ids(object_permission) + requested_access_groups = _extract_requested_mcp_access_groups(object_permission) + + # Nothing to validate + if not requested_servers and not requested_access_groups: + return + + allow_all_keys_servers = _get_allow_all_keys_server_ids() + team_allowed_servers = await _get_team_allowed_mcp_servers(team_obj) + + # Combined allowed set = team servers + allow_all_keys servers + all_allowed_servers = team_allowed_servers | allow_all_keys_servers + + # Validate requested server IDs + if requested_servers: + disallowed_servers = requested_servers - all_allowed_servers + if disallowed_servers: + if team_obj is not None: + detail = ( + f"Key requests MCP servers not allowed by team '{team_obj.team_id}': " + f"{sorted(disallowed_servers)}. " + f"Team allows: {sorted(team_allowed_servers)}. " + f"Global (allow_all_keys) servers: {sorted(allow_all_keys_servers)}." + ) + else: + detail = ( + f"Key is not in a team. Only globally available (allow_all_keys) MCP servers " + f"can be assigned: {sorted(allow_all_keys_servers)}. " + f"Disallowed servers: {sorted(disallowed_servers)}." + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": detail}, + ) + + # Validate requested access groups (must be subset of team's access groups) + if requested_access_groups: + team_access_groups: Set[str] = set() + if ( + team_obj is not None + and team_obj.object_permission is not None + and team_obj.object_permission.mcp_access_groups + ): + team_access_groups = set(team_obj.object_permission.mcp_access_groups) + + disallowed_groups = requested_access_groups - team_access_groups + if disallowed_groups: + if team_obj is not None: + detail = ( + f"Key requests MCP access groups not allowed by team '{team_obj.team_id}': " + f"{sorted(disallowed_groups)}. " + f"Team allows: {sorted(team_access_groups)}." + ) + else: + detail = ( + f"Key is not in a team. MCP access groups cannot be assigned to " + f"keys outside of a team. Disallowed groups: {sorted(disallowed_groups)}." + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": detail}, + ) \ No newline at end of file diff --git a/litellm/proxy/openapi_registry.json b/litellm/proxy/openapi_registry.json new file mode 100644 index 00000000000..d525b504a7b --- /dev/null +++ b/litellm/proxy/openapi_registry.json @@ -0,0 +1,222 @@ +{ + "apis": [ + { + "name": "github", + "title": "GitHub", + "description": "Repos, issues, PRs, and workflow automation via the GitHub REST API", + "icon_url": "https://cdn.simpleicons.org/github", + "spec_url": "https://raw.githubusercontent.com/github/rest-api-description/main/descriptions/api.github.com/api.github.com.json", + "oauth": { + "authorization_url": "https://github.com/login/oauth/authorize", + "token_url": "https://github.com/login/oauth/access_token", + "pkce": false, + "docs_url": "https://docs.github.com/en/apps/oauth-apps/building-oauth-apps/creating-an-oauth-app" + }, + "key_tools": [ + { "name": "list_repos", "description": "List repositories for a user or organization" }, + { "name": "get_file_contents", "description": "Read a file or directory from a repository" }, + { "name": "list_issues", "description": "List issues in a repository with filters" }, + { "name": "create_issue", "description": "Open a new issue in a repository" }, + { "name": "list_pull_requests", "description": "List open and merged pull requests" }, + { "name": "create_pull_request", "description": "Open a pull request between branches" }, + { "name": "search_code", "description": "Search code across all GitHub repositories" }, + { "name": "list_commits", "description": "List commits with authors and messages" } + ] + }, + { + "name": "atlassian", + "title": "Atlassian", + "description": "Jira issues, Confluence pages, and project management", + "icon_url": "https://cdn.simpleicons.org/atlassian", + "spec_url": "https://dac-static.atlassian.com/cloud/jira/platform/swagger-v3.v3.json", + "oauth": { + "authorization_url": "https://auth.atlassian.com/authorize", + "token_url": "https://auth.atlassian.com/oauth/token", + "pkce": true, + "docs_url": "https://developer.atlassian.com/cloud/jira/platform/oauth-2-3lo-apps/" + }, + "key_tools": [ + { "name": "search_issues", "description": "Search Jira issues using JQL queries" }, + { "name": "get_issue", "description": "Get full details of a Jira issue by key" }, + { "name": "create_issue", "description": "Create a new Jira issue or bug report" }, + { "name": "update_issue", "description": "Update issue fields like status, assignee, priority" }, + { "name": "add_comment", "description": "Add a comment to an issue" }, + { "name": "transition_issue", "description": "Move an issue through the workflow (e.g. In Progress → Done)" }, + { "name": "list_projects", "description": "List all Jira projects in the workspace" }, + { "name": "get_project", "description": "Get project details, board, and sprint info" } + ] + }, + { + "name": "figma", + "title": "Figma", + "description": "Design files, components, prototypes, and comments", + "icon_url": "https://cdn.simpleicons.org/figma", + "spec_url": "https://raw.githubusercontent.com/figma/rest-api-spec/main/openapi/openapi.yaml", + "oauth": { + "authorization_url": "https://www.figma.com/oauth", + "token_url": "https://www.figma.com/api/oauth/token", + "pkce": false, + "docs_url": "https://www.figma.com/developers/api#oauth2" + }, + "key_tools": [ + { "name": "get_file", "description": "Get the full node tree and structure of a Figma file" }, + { "name": "get_file_nodes", "description": "Get specific nodes by ID from a Figma file" }, + { "name": "get_image", "description": "Export nodes as PNG, SVG, or PDF" }, + { "name": "get_comments", "description": "List all comments on a Figma file" }, + { "name": "post_comment", "description": "Add a comment to a file or specific node" }, + { "name": "get_team_projects", "description": "List all projects belonging to a team" }, + { "name": "get_project_files", "description": "List all Figma files in a project" }, + { "name": "get_file_versions", "description": "Get the version history of a file" } + ] + }, + { + "name": "gmail", + "title": "Gmail", + "description": "Read, send, and manage Gmail messages and threads", + "icon_url": "https://cdn.simpleicons.org/gmail", + "spec_url": "https://raw.githubusercontent.com/APIs-guru/openapi-directory/main/APIs/googleapis.com/gmail/v1/openapi.yaml", + "oauth": { + "authorization_url": "https://accounts.google.com/o/oauth2/v2/auth", + "token_url": "https://oauth2.googleapis.com/token", + "pkce": true, + "docs_url": "https://developers.google.com/gmail/api/auth/about-auth" + }, + "key_tools": [ + { "name": "list_messages", "description": "List Gmail messages with search filters" }, + { "name": "get_message", "description": "Get the full content of a specific message" }, + { "name": "send_message", "description": "Send an email via Gmail" }, + { "name": "create_draft", "description": "Create a draft email" }, + { "name": "list_labels", "description": "List all labels in the mailbox" }, + { "name": "modify_message", "description": "Add or remove labels from a message" }, + { "name": "list_threads", "description": "List email threads" }, + { "name": "trash_message", "description": "Move a message to trash" } + ] + }, + { + "name": "stripe", + "title": "Stripe", + "description": "Payments, customers, subscriptions, and billing", + "icon_url": "https://cdn.simpleicons.org/stripe", + "spec_url": "https://raw.githubusercontent.com/stripe/openapi/master/openapi/spec3.json", + "oauth": { + "authorization_url": "https://connect.stripe.com/oauth/authorize", + "token_url": "https://connect.stripe.com/oauth/token", + "pkce": false, + "docs_url": "https://stripe.com/docs/connect/oauth-reference" + }, + "key_tools": [ + { "name": "list_customers", "description": "List and search Stripe customers" }, + { "name": "create_customer", "description": "Create a new Stripe customer record" }, + { "name": "create_payment_intent", "description": "Create a payment intent for a charge" }, + { "name": "list_subscriptions", "description": "List active and cancelled subscriptions" }, + { "name": "create_subscription", "description": "Subscribe a customer to a price plan" }, + { "name": "list_invoices", "description": "List invoices for a customer or subscription" }, + { "name": "create_refund", "description": "Refund a charge fully or partially" }, + { "name": "list_products", "description": "List products and their pricing plans" } + ] + }, + { + "name": "hubspot", + "title": "HubSpot", + "description": "CRM contacts and properties via the HubSpot Contacts API", + "icon_url": "https://cdn.simpleicons.org/hubspot", + "spec_url": "https://raw.githubusercontent.com/HubSpot/HubSpot-public-api-spec-collection/main/PublicApiSpecs/CRM/Contacts/Rollouts/424/v3/contacts.json", + "oauth": { + "authorization_url": "https://app.hubspot.com/oauth/authorize", + "token_url": "https://api.hubspot.com/oauth/v1/token", + "pkce": false, + "docs_url": "https://developers.hubspot.com/docs/api/oauth-quickstart-guide" + }, + "key_tools": [ + { "name": "search_contacts", "description": "Search CRM contacts by email, name, or custom properties" }, + { "name": "create_contact", "description": "Create a new CRM contact with properties" }, + { "name": "update_contact", "description": "Update contact properties like lifecycle stage or owner" }, + { "name": "get_contact", "description": "Get full details of a specific contact" }, + { "name": "archive_contact", "description": "Archive (soft-delete) a contact record" }, + { "name": "merge_contacts", "description": "Merge two duplicate contact records" }, + { "name": "list_contacts", "description": "List all contacts with pagination" }, + { "name": "get_contact_properties", "description": "Get available contact property definitions" } + ] + }, + { + "name": "notion", + "title": "Notion", + "description": "Pages, databases, and workspace content in Notion", + "icon_url": "https://cdn.simpleicons.org/notion", + "spec_url": "https://raw.githubusercontent.com/APIs-guru/openapi-directory/main/APIs/notion.com/1.0.0/openapi.yaml", + "oauth": { + "authorization_url": "https://api.notion.com/v1/oauth/authorize", + "token_url": "https://api.notion.com/v1/oauth/token", + "pkce": false, + "docs_url": "https://developers.notion.com/docs/authorization" + }, + "key_tools": [ + { "name": "search", "description": "Search pages and databases across the workspace" }, + { "name": "get_page", "description": "Get a page and its properties" }, + { "name": "create_page", "description": "Create a new page inside a database or as a subpage" }, + { "name": "update_page", "description": "Update page properties (title, status, dates, etc.)" }, + { "name": "query_database", "description": "Query a database with filters and sorts" }, + { "name": "create_database_item", "description": "Add a new row/item to a Notion database" }, + { "name": "append_block_children", "description": "Append content blocks (text, bullets, code) to a page" }, + { "name": "get_database", "description": "Get a database schema including all property types" } + ] + }, + { + "name": "slack", + "title": "Slack", + "description": "Channels, messages, users, and workspace management", + "icon_url": "https://raw.githubusercontent.com/simple-icons/simple-icons/develop/icons/slack.svg", + "spec_url": "https://raw.githubusercontent.com/slackapi/slack-api-specs/master/web-api/slack_web_openapi_v2.json", + "oauth": { + "authorization_url": "https://slack.com/oauth/v2/authorize", + "token_url": "https://slack.com/api/oauth.v2.access", + "pkce": false, + "docs_url": "https://api.slack.com/authentication/oauth-v2" + }, + "key_tools": [ + { "name": "chat_post_message", "description": "Send a message to a channel or DM" }, + { "name": "conversations_history", "description": "Get message history from a channel" }, + { "name": "conversations_list", "description": "List all public and private channels" }, + { "name": "conversations_replies", "description": "Get replies in a message thread" }, + { "name": "users_list", "description": "List all members of the Slack workspace" }, + { "name": "search_messages", "description": "Full-text search across all messages" }, + { "name": "files_upload", "description": "Upload a file and share it in a channel" }, + { "name": "reactions_add", "description": "Add an emoji reaction to a message" } + ] + }, + { + "name": "shopify", + "title": "Shopify", + "description": "Products, orders, customers, and store management via Shopify Admin REST API (requires your store subdomain)", + "icon_url": "https://cdn.simpleicons.org/shopify", + "spec_url": "https://raw.githubusercontent.com/Shopify/shopify-api-specs/main/admin/rest/2023-10/openapi.json", + "key_tools": [ + { "name": "list_products", "description": "List products with variants, pricing, and inventory" }, + { "name": "get_product", "description": "Get full product details including all variants" }, + { "name": "list_orders", "description": "List orders with status, customer, and line item filters" }, + { "name": "get_order", "description": "Get full order details including shipping and payment" }, + { "name": "list_customers", "description": "List customers with order history and tags" }, + { "name": "update_order", "description": "Update order notes, tags, or shipping address" }, + { "name": "create_fulfillment", "description": "Fulfill an order with tracking info" }, + { "name": "list_inventory_levels", "description": "Get stock levels across locations" } + ] + }, + { + "name": "snowflake", + "title": "Snowflake", + "description": "Data warehouse queries, database operations, and analytics via the Snowflake SQL API", + "icon_url": "https://cdn.simpleicons.org/snowflake", + "spec_url": "https://raw.githubusercontent.com/snowflakedb/snowflake-rest-api-specs/refs/heads/main/specifications/sqlapi.yaml", + "key_tools": [ + { "name": "execute_statement", "description": "Execute a SQL statement and get results" }, + { "name": "fetch_results", "description": "Fetch paginated results from a running query" }, + { "name": "cancel_statement", "description": "Cancel a running query by statement handle" }, + { "name": "list_databases", "description": "List all accessible databases in the account" }, + { "name": "list_schemas", "description": "List all schemas within a database" }, + { "name": "list_tables", "description": "List tables and views in a schema" }, + { "name": "describe_table", "description": "Get column definitions and data types for a table" }, + { "name": "list_warehouses", "description": "List virtual warehouses and their current status" } + ] + } + ] +} diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f3bc4b08037..e6bb3ee412e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -351,9 +351,6 @@ from litellm.proxy.management_endpoints.cache_settings_endpoints import ( from litellm.proxy.management_endpoints.callback_management_endpoints import ( router as callback_management_endpoints_router, ) -from litellm.proxy.management_endpoints.config_override_endpoints import ( - router as config_override_router, -) from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, admin_can_invite_user, @@ -361,6 +358,9 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.compliance_endpoints import ( router as compliance_router, ) +from litellm.proxy.management_endpoints.config_override_endpoints import ( + router as config_override_router, +) from litellm.proxy.management_endpoints.cost_tracking_settings import ( router as cost_tracking_settings_router, ) @@ -373,7 +373,9 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_update, +) from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( router as jwt_key_mapping_router, ) @@ -442,7 +444,9 @@ from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_route from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + set_files_config, +) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -541,7 +545,9 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import DeploymentTypedDict +from litellm.types.router import ( + DeploymentTypedDict, +) from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import ( RouterGeneralSettings, @@ -5788,6 +5794,8 @@ class ProxyStartupEvent: _RUNTIME_GENERAL_SETTINGS_FLAGS, ) + if prisma_client is None: + return db_record = await prisma_client.db.litellm_uisettings.find_unique( where={"id": "ui_settings"} ) @@ -11983,6 +11991,7 @@ async def get_config_list( "mcp_trusted_proxy_ranges": {"type": "List"}, "always_include_stream_usage": {"type": "Boolean"}, "forward_client_headers_to_llm_api": {"type": "Boolean"}, + "mcp_required_fields": {"type": "List"}, } return_val = [] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8d4bdffb2dd..721c3e404d2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -315,6 +315,15 @@ model LiteLLM_MCPServerTable { is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? + source_url String? + // BYOM submission lifecycle + approval_status String? @default("active") + submitted_by String? + submitted_at DateTime? + reviewed_at DateTime? + review_notes String? + + @@index([approval_status]) } // Per-user BYOK credentials for MCP servers diff --git a/litellm/router.py b/litellm/router.py index 44e9667d5c4..ecda6f4ab67 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -727,15 +727,11 @@ class Router: startup_nodes = cache_config.get("startup_nodes") if not startup_nodes: _env_cluster_nodes = get_secret("REDIS_CLUSTER_NODES") - if _env_cluster_nodes is not None and isinstance( - _env_cluster_nodes, str - ): + if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str): startup_nodes = json.loads(_env_cluster_nodes) if startup_nodes: - return RedisClusterCache( - **{**cache_config, "startup_nodes": startup_nodes} - ) + return RedisClusterCache(**{**cache_config, "startup_nodes": startup_nodes}) else: return RedisCache(**cache_config) @@ -1466,12 +1462,15 @@ class Router: silent_kwargs["metadata"]["is_silent_experiment"] = True + # Force stream=False so the response is fully consumed and callbacks fire + silent_kwargs["stream"] = False + # Pop logging objects and call IDs to ensure a fresh logging context # This prevents collisions in the Proxy's database (spend_logs) silent_kwargs.pop("litellm_call_id", None) silent_kwargs.pop("litellm_logging_obj", None) silent_kwargs.pop("standard_logging_object", None) - silent_kwargs.pop("proxy_server_request", None) + # DON'T pop proxy_server_request — it's needed for spend log metadata return silent_kwargs @@ -1494,12 +1493,31 @@ class Router: silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) - # Trigger the silent request - self.completion( - model=silent_model, - messages=cast(List[Dict[str, str]], messages), - **silent_kwargs, - ) + # Override model_group to correctly attribute metrics to the silent model + silent_kwargs["metadata"]["model_group"] = silent_model + + # Create a new event loop for this thread so that async success + # callbacks (e.g. _ProxyDBLogger) can schedule and run DB writes. + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + + async def _run_silent_completion(): + await self.acompletion( + model=silent_model, + messages=cast(List[AllMessageValues], messages), + **silent_kwargs, + ) + # Drain any fire-and-forget tasks (e.g. alerting hooks) + # scheduled via asyncio.create_task during acompletion. + pending = asyncio.all_tasks() + pending.discard(asyncio.current_task()) + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + loop.run_until_complete(_run_silent_completion()) + finally: + loop.close() except Exception as e: verbose_router_logger.error( f"Silent experiment failed for model {silent_model}: {str(e)}" @@ -1705,7 +1723,9 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - self._combine_fallback_usage(fallback_item, complete_response_object_usage) + self._combine_fallback_usage( + fallback_item, complete_response_object_usage + ) yield fallback_item else: # If fallback returns a non-streaming response, yield None @@ -1825,13 +1845,11 @@ class Router: router_self._update_kwargs_before_fallbacks( model=model_group, kwargs=initial_kwargs ) - fallback_response = ( - router_self.function_with_fallbacks( - **initial_kwargs, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) + fallback_response = router_self.function_with_fallbacks( + **initial_kwargs, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, ) if hasattr(fallback_response, "__iter__"): @@ -1841,7 +1859,9 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - router_self._combine_fallback_usage(fallback_item, complete_response_object_usage) + router_self._combine_fallback_usage( + fallback_item, complete_response_object_usage + ) yield fallback_item else: yield None @@ -1891,6 +1911,8 @@ class Router: ) silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) + # Override model_group to correctly attribute metrics to the silent model + silent_kwargs["metadata"]["model_group"] = silent_model # Trigger the silent request await self.acompletion( @@ -2753,10 +2775,9 @@ class Router: litellm_model = data.get("model", None) # litellm_agent/ prefix only strips the model name, no prompt_id needed - is_litellm_agent_model = ( - isinstance(litellm_model, str) - and litellm_model.startswith("litellm_agent/") - ) + is_litellm_agent_model = isinstance( + litellm_model, str + ) and litellm_model.startswith("litellm_agent/") prompt_id = kwargs.get("prompt_id") or prompt_management_deployment[ "litellm_params" @@ -6560,7 +6581,7 @@ class Router: tiers = complexity_router_config.get("tiers", {}) # Use MEDIUM tier as fallback default default_model = tiers.get("MEDIUM") or tiers.get("SIMPLE") - + if default_model is None: raise ValueError( "complexity_router_default_model is required for complexity-router deployments, " @@ -6793,7 +6814,9 @@ class Router: ######################################################### # Check if this is a complexity-router deployment ######################################################### - if self._is_complexity_router_deployment(litellm_params=deployment.litellm_params): + if self._is_complexity_router_deployment( + litellm_params=deployment.litellm_params + ): self.init_complexity_router_deployment(deployment=deployment) return deployment @@ -6885,9 +6908,7 @@ class Router: # zero-cost models, causing budget checks to block free models. _model_id = deployment.model_info.id if _model_id is not None: - _model_info_dict: dict = deployment.model_info.model_dump( - exclude_none=True - ) + _model_info_dict: dict = deployment.model_info.model_dump(exclude_none=True) for field in CustomPricingLiteLLMParams.model_fields.keys(): field_value = deployment.litellm_params.get(field) if field_value is not None: @@ -7179,7 +7200,10 @@ class Router: @overload def get_router_model_info( - self, deployment: Union[dict, "Deployment"], received_model_name: str, id: None = None + self, + deployment: Union[dict, "Deployment"], + received_model_name: str, + id: None = None, ) -> ModelMapInfo: pass @@ -7219,7 +7243,9 @@ class Router: ## GET BASE MODEL base_model = (deployment.get("model_info") or {}).get("base_model", None) if base_model is None: - base_model = (deployment.get("litellm_params") or {}).get("base_model", None) + base_model = (deployment.get("litellm_params") or {}).get( + "base_model", None + ) model = base_model @@ -7254,12 +7280,12 @@ class Router: if potential_models is not None: for potential_model in potential_models: try: - if (potential_model.get("model_info") or {}).get( - "id" - ) == (deployment.get("model_info") or {}).get("id"): - model = (potential_model.get("litellm_params") or {}).get( - "model" - ) + if (potential_model.get("model_info") or {}).get("id") == ( + deployment.get("model_info") or {} + ).get("id"): + model = ( + potential_model.get("litellm_params") or {} + ).get("model") break except Exception: pass @@ -8182,7 +8208,9 @@ class Router: - team_id: Optional[str] - the team id, to resolve team-specific models """ # Check if this is the no-args hot path (cacheable) - _use_cache = model_name is None and model_access_group is None and team_id is None + _use_cache = ( + model_name is None and model_access_group is None and team_id is None + ) # Return cached result for the no-args hot path if _use_cache and self._access_groups_cache is not None: diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 884dfefb42c..33e55f9bed9 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -35,6 +35,8 @@ class MCPAuth(str, enum.Enum): basic = "basic" authorization = "authorization" oauth2 = "oauth2" + aws_sigv4 = "aws_sigv4" + token = "token" # MCP Literals @@ -50,6 +52,8 @@ MCPAuthType = Optional[ MCPAuth.basic, MCPAuth.authorization, MCPAuth.oauth2, + MCPAuth.aws_sigv4, + MCPAuth.token, ] ] diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index d94795fda2e..511cfc958a2 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, ConfigDict @@ -48,6 +48,12 @@ class MCPServer(BaseModel): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + # AWS SigV4 fields + aws_access_key_id: Optional[str] = None + aws_secret_access_key: Optional[str] = None + aws_session_token: Optional[str] = None + aws_region_name: Optional[str] = None + aws_service_name: Optional[str] = None # defaults to "bedrock-agentcore" # Stdio-specific fields command: Optional[str] = None args: Optional[List[str]] = None @@ -60,12 +66,22 @@ class MCPServer(BaseModel): byok_api_key_help_url: Optional[str] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None + # OAuth2 flow type. Defaults to None (interactive / authorization_code). + # Set to "client_credentials" to enable M2M token fetching. + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None model_config = ConfigDict(arbitrary_types_allowed=True) @property def has_client_credentials(self) -> bool: - """True if this server has OAuth2 client_credentials config (client_id, client_secret, token_url).""" - return bool(self.client_id and self.client_secret and self.token_url) + """True if this server should use the OAuth2 client_credentials (M2M) flow. + + M2M flow must be opted into explicitly via ``oauth2_flow: client_credentials``. + Having client_id / client_secret / token_url present is NOT sufficient — + those fields are also used for interactive (authorization_code) OAuth, + e.g. GitHub Enterprise. Auto-detecting M2M from field presence was a + breaking regression introduced with the M2M feature. + """ + return self.oauth2_flow == "client_credentials" @property def needs_user_oauth_token(self) -> bool: diff --git a/litellm/utils.py b/litellm/utils.py index 5062bbfe152..6e07ab2975e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -780,9 +780,9 @@ def function_setup( # noqa: PLR0915 coroutine_checker = get_coroutine_checker_fn() ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = ( - kwargs.pop("callbacks", None) - ) + dynamic_callbacks: Optional[ + List[Union[str, Callable, "CustomLogger"]] + ] = kwargs.pop("callbacks", None) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -1143,6 +1143,14 @@ def function_setup( # noqa: PLR0915 litellm_params: Dict[str, Any] = {"api_base": ""} if "metadata" in kwargs: litellm_params["metadata"] = kwargs["metadata"] + if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): + litellm_params["litellm_metadata"] = kwargs["litellm_metadata"].copy() + # For endpoints like /v1/messages that use "litellm_metadata" instead + # of "metadata" (to avoid conflicting with provider API metadata fields), + # populate litellm_params["metadata"] so callbacks (e.g. Langfuse) that + # read API key info from litellm_params["metadata"] see the fields. + if not litellm_params.get("metadata"): + litellm_params["metadata"] = kwargs["litellm_metadata"].copy() logging_obj.update_environment_variables( model=model, @@ -1682,9 +1690,9 @@ def client(original_function): # noqa: PLR0915 exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops + kwargs[ + "retry_policy" + ] = reset_retry_policy() # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -1731,9 +1739,9 @@ def client(original_function): # noqa: PLR0915 exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops + kwargs[ + "retry_policy" + ] = reset_retry_policy() # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -3686,10 +3694,10 @@ def pre_process_non_default_params( if "response_format" in non_default_params: if provider_config is not None: - non_default_params["response_format"] = ( - provider_config.get_json_schema_from_pydantic_object( - response_format=non_default_params["response_format"] - ) + non_default_params[ + "response_format" + ] = provider_config.get_json_schema_from_pydantic_object( + response_format=non_default_params["response_format"] ) else: non_default_params["response_format"] = type_to_response_format_param( @@ -3818,16 +3826,16 @@ def pre_process_optional_params( True # so that main.py adds the function call to the prompt ) if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("tools") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("tools") non_default_params.pop( "tool_choice", None ) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("functions") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("functions") elif ( litellm.add_function_to_prompt ): # if user opts to add it to prompt instead @@ -7428,9 +7436,9 @@ class ModelResponseIterator: if convert_to_delta is True: _stream_response = ModelResponseStream() _stream_response.choices[0].delta.content = model_response.choices[0].message.content # type: ignore - self.model_response: Union[ModelResponse, ModelResponseStream] = ( - _stream_response - ) + self.model_response: Union[ + ModelResponse, ModelResponseStream + ] = _stream_response else: self.model_response = model_response self.is_done = False @@ -7910,7 +7918,10 @@ class ProviderConfigManager: # Simple provider mappings (no model parameter needed) LlmProviders.DEEPSEEK: (lambda: litellm.DeepSeekChatConfig(), False), LlmProviders.GROQ: (lambda: litellm.GroqChatConfig(), False), - LlmProviders.BEDROCK_MANTLE: (lambda: litellm.BedrockMantleChatConfig(), False), + LlmProviders.BEDROCK_MANTLE: ( + lambda: litellm.BedrockMantleChatConfig(), + False, + ), LlmProviders.A2A: (lambda: litellm.A2AConfig(), False), LlmProviders.BYTEZ: (lambda: litellm.BytezChatConfig(), False), LlmProviders.DATABRICKS: (lambda: litellm.DatabricksConfig(), False), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 04037c2ee81..82701c41eb6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19627,7 +19627,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-2025-04-14": { "cache_read_input_token_cost": 5e-07, @@ -19661,7 +19662,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini": { "cache_read_input_token_cost": 1e-07, @@ -19698,7 +19700,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { "cache_read_input_token_cost": 1e-07, @@ -19732,7 +19735,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, @@ -20940,6 +20944,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21579,6 +21584,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21679,6 +21685,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21714,6 +21721,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21746,6 +21754,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true }, @@ -21781,6 +21790,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21816,6 +21826,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true }, @@ -21851,6 +21862,7 @@ "supports_system_messages": false, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21892,6 +21904,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21933,6 +21946,7 @@ "supports_tool_choice": true, "supports_service_tier": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -21971,6 +21985,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -22008,6 +22023,7 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": false }, @@ -25795,7 +25811,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, @@ -25827,7 +25844,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, @@ -25860,7 +25878,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-deep-research-2025-06-26": { "cache_read_input_token_cost": 2.5e-06, @@ -25893,7 +25912,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-mini": { "cache_read_input_token_cost": 5.5e-07, @@ -25957,7 +25977,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o3-pro-2025-06-10": { "input_cost_per_token": 2e-05, @@ -25987,7 +26008,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini": { "cache_read_input_token_cost": 2.75e-07, @@ -26012,7 +26034,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, @@ -26031,7 +26054,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_service_tier": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research": { "cache_read_input_token_cost": 5e-07, @@ -26064,7 +26088,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { "cache_read_input_token_cost": 5e-07, @@ -26097,7 +26122,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "oci/meta.llama-3.1-405b-instruct": { "input_cost_per_token": 1.068e-05, @@ -27713,6 +27739,92 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/qwen/qwen3.5-35b-a3b": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-35b-a3b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-27b": { + "input_cost_per_token": 3e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-27b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-122b-a10b": { + "input_cost_per_token": 4e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-122b-a10b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-flash-02-23": { + "input_cost_per_token": 1e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://openrouter.ai/qwen/qwen3.5-flash-02-23", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-plus-02-15": { + "input_cost_per_token": 4e-07, + "input_cost_per_token_above_256k_tokens": 5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "output_cost_per_token_above_256k_tokens": 3e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-plus-02-15", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "openrouter/qwen/qwen3.5-397b-a17b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://openrouter.ai/qwen/qwen3.5-397b-a17b", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "openrouter/switchpoint/router": { "input_cost_per_token": 8.5e-07, "litellm_provider": "openrouter", diff --git a/tests/guardrails_tests/test_presidio_pii.py b/tests/guardrails_tests/test_presidio_pii.py index 0d730288e63..eda0c7bb5b5 100644 --- a/tests/guardrails_tests/test_presidio_pii.py +++ b/tests/guardrails_tests/test_presidio_pii.py @@ -1,20 +1,19 @@ import sys import os -import io, asyncio import pytest -import time from litellm import mock_completion -from unittest.mock import MagicMock, AsyncMock, patch +from unittest.mock import patch + sys.path.insert(0, os.path.abspath("../..")) import litellm -from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking, PresidioPerRequestConfig +from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + PresidioPerRequestConfig, +) from litellm.types.guardrails import PiiEntityType, PiiAction from litellm.proxy._types import UserAPIKeyAuth from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError -from litellm.types.utils import CallTypes as LitellmCallTypes - - @pytest.mark.asyncio @@ -26,42 +25,37 @@ async def test_presidio_with_entities_config(): PiiEntityType.CREDIT_CARD: PiiAction.MASK, PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Test text with different PII types test_text = "My credit card number is 4111-1111-1111-1111, my email is test@example.com, and my phone is 555-123-4567" - + # Test the analyze request configuration analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify entities were passed correctly assert "entities" in analyze_request assert set(analyze_request["entities"]) == set(pii_entities_config.keys()) - + # Test the check_pii method - this will call the actual Presidio API redacted_text = await presidio_guardrail.check_pii( - text=test_text, - output_parse_pii=True, - presidio_config=None, - request_data={} + text=test_text, output_parse_pii=True, presidio_config=None, request_data={} ) - + # Verify PII has been masked/replaced/redacted in the result assert "4111-1111-1111-1111" not in redacted_text assert "test@example.com" not in redacted_text # Since this entity is not in the config, it should not be masked assert "555-123-4567" in redacted_text - + # The specific replacements will vary based on Presidio's implementation print(f"Redacted text: {redacted_text}") @@ -73,10 +67,12 @@ async def test_presidio_apply_guardrail(): presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config={}, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + test_text = ( + "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + ) response = await presidio_guardrail.apply_guardrail( inputs={"texts": [test_text]}, request_data={}, @@ -91,6 +87,7 @@ async def test_presidio_apply_guardrail(): assert "4111-1111-1111-1111" not in modified_text assert "test@example.com" not in modified_text + @pytest.mark.asyncio async def test_presidio_with_blocked_entities(): """Test for Presidio guardrail with blocked entities - requires actual Presidio API""" @@ -100,36 +97,33 @@ async def test_presidio_with_blocked_entities(): PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Test text with blocked PII type - test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" - + test_text = ( + "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + ) + # Verify the analyze request configuration analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify entities were passed correctly assert "entities" in analyze_request assert set(analyze_request["entities"]) == set(pii_entities_config.keys()) - + # Test that BlockedPiiEntityError is raised when check_pii is called with pytest.raises(BlockedPiiEntityError) as excinfo: await presidio_guardrail.check_pii( - text=test_text, - output_parse_pii=True, - presidio_config=None, - request_data={} + text=test_text, output_parse_pii=True, presidio_config=None, request_data={} ) - + # Verify the error contains the correct entity type assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name @@ -143,37 +137,40 @@ async def test_presidio_pre_call_hook_with_blocked_entities(): PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Create a sample chat completion request with PII data data = { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com."} + { + "role": "user", + "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com.", + }, ], - "model": "gpt-3.5-turbo" + "model": "gpt-3.5-turbo", } - + # Mock objects needed for the pre-call hook user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() - + # Call the pre-call hook and expect BlockedPiiEntityError with pytest.raises(BlockedPiiEntityError) as excinfo: await presidio_guardrail.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, - call_type="completion" + call_type="completion", ) - + print(f"got error: {excinfo}") - + # Verify the error contains the correct entity type assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name @@ -188,44 +185,46 @@ async def test_presidio_pre_call_hook_with_different_call_types(call_type): PiiEntityType.CREDIT_CARD: PiiAction.MASK, PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Create a sample request with PII data data = { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567"} + { + "role": "user", + "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567", + }, ], - "model": "gpt-3.5-turbo" + "model": "gpt-3.5-turbo", } - + # Mock objects needed for the pre-call hook user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() - + # Call the pre-call hook with the specified call type modified_data = await presidio_guardrail.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=cache, - data=data, - call_type=call_type + user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type=call_type ) - + # Verify the messages have been modified to mask PII - assert modified_data["messages"][0]["content"] == "You are a helpful assistant." # System prompt should be unchanged - + assert ( + modified_data["messages"][0]["content"] == "You are a helpful assistant." + ) # System prompt should be unchanged + user_message = modified_data["messages"][1]["content"] assert "4111-1111-1111-1111" not in user_message assert "test@example.com" not in user_message # Since this entity is not in the config, it should not be masked assert "555-123-4567" in user_message - + print(f"Modified user message for call_type={call_type}: {user_message}") @@ -243,7 +242,7 @@ def test_validate_environment_missing_http(base_url): # Use patch.dict to temporarily modify environment variables only for this test env_vars = { "PRESIDIO_ANALYZER_API_BASE": f"{base_url}/analyze", - "PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize" + "PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize", } with patch.dict(os.environ, env_vars): pii_masking.validate_environment() @@ -294,8 +293,12 @@ async def test_output_parsing(): new_response = await pii_masking.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth(), data={ - "messages": [{"role": "system", "content": "You are an helpfull assistant"}], - "pii_tokens": {"": "Jane Doe", "": "034453334"}, + "messages": [ + {"role": "system", "content": "You are an helpfull assistant"} + ], + "metadata": { + "pii_tokens": {"": "Jane Doe", "": "034453334"} + }, }, response=response, ) @@ -440,24 +443,26 @@ async def test_presidio_pii_masking_logging_output_only_no_pre_api_hook(): @pytest.mark.asyncio -@patch.dict(os.environ, { - "PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002", - "PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001" -}) +@patch.dict( + os.environ, + { + "PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002", + "PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001", + }, +) async def test_presidio_pii_masking_logging_output_only_logged_response_guardrails_config(): from typing import Dict, List, Optional import litellm from litellm.proxy.guardrails.init_guardrails import initialize_guardrails from litellm.types.guardrails import ( - GuardrailItem, GuardrailItemSpec, GuardrailEventHooks, ) litellm.set_verbose = True # Environment variables are now patched via the decorator instead of setting them directly - + guardrails_config: List[Dict[str, GuardrailItemSpec]] = [ { "pii_masking": { @@ -499,60 +504,53 @@ async def test_presidio_pii_masking_logging_output_only_logged_response_guardrai async def test_presidio_language_configuration(): """Test that presidio_language parameter is properly set and used in analyze requests""" litellm._turn_on_debug() - + # Test with German language using mock testing to avoid API calls presidio_guardrail_de = _OPTIONAL_PresidioPIIMasking( pii_entities_config={}, presidio_language="de", - mock_testing=True # This bypasses the API validation + mock_testing=True, # This bypasses the API validation ) - + test_text = "Meine Telefonnummer ist +49 30 12345678" - + # Test the analyze request configuration analyze_request = presidio_guardrail_de._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify the language is set to German assert analyze_request["language"] == "de" assert analyze_request["text"] == test_text - + # Test with Spanish language presidio_guardrail_es = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - presidio_language="es", - mock_testing=True + pii_entities_config={}, presidio_language="es", mock_testing=True ) - + test_text_es = "Mi número de teléfono es +34 912 345 678" - + analyze_request_es = presidio_guardrail_es._get_presidio_analyze_request_payload( - text=test_text_es, - presidio_config=None, - request_data={} + text=test_text_es, presidio_config=None, request_data={} ) - + # Verify the language is set to Spanish assert analyze_request_es["language"] == "es" assert analyze_request_es["text"] == test_text_es - + # Test default language (English) when not specified presidio_guardrail_default = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - mock_testing=True + pii_entities_config={}, mock_testing=True ) - + test_text_en = "My phone number is +1 555-123-4567" - - analyze_request_default = presidio_guardrail_default._get_presidio_analyze_request_payload( - text=test_text_en, - presidio_config=None, - request_data={} + + analyze_request_default = ( + presidio_guardrail_default._get_presidio_analyze_request_payload( + text=test_text_en, presidio_config=None, request_data={} + ) ) - + # Verify the language defaults to English assert analyze_request_default["language"] == "en" assert analyze_request_default["text"] == test_text_en @@ -562,36 +560,30 @@ async def test_presidio_language_configuration(): async def test_presidio_language_configuration_with_per_request_override(): """Test that per-request language configuration overrides the default configured language""" litellm._turn_on_debug() - + # Set up guardrail with German as default language presidio_guardrail = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - presidio_language="de", - mock_testing=True + pii_entities_config={}, presidio_language="de", mock_testing=True ) - + test_text = "Test text with PII" - + # Test with per-request config overriding the default language presidio_config = PresidioPerRequestConfig(language="fr") - + analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=presidio_config, - request_data={} + text=test_text, presidio_config=presidio_config, request_data={} ) - + # Verify the per-request language (French) overrides the default (German) assert analyze_request["language"] == "fr" assert analyze_request["text"] == test_text - + # Test without per-request config - should use default language analyze_request_default = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify the default language (German) is used assert analyze_request_default["language"] == "de" assert analyze_request_default["text"] == test_text diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index febc7c454bd..13a09f54e68 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -11,7 +11,7 @@ sys.path.insert(0, "../../../") import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient -from litellm.types.mcp import MCPStdioConfig, MCPTransport +from litellm.types.mcp import MCPAuth, MCPStdioConfig, MCPTransport class TestMCPClient: @@ -245,6 +245,72 @@ class TestMCPClient: assert test_client.headers is not None await test_client.aclose() + def test_token_auth_header_generation(self): + """Test that token auth generates correct Authorization header""" + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.token, + auth_value="my-secret-token" + ) + + headers = client._get_auth_headers() + + assert "Authorization" in headers + assert headers["Authorization"] == "token my-secret-token" + + def test_token_auth_compatibility_with_existing_auth_types(self): + """Verify existing auth types are not affected by token auth addition""" + # Test bearer token + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.bearer_token, + auth_value="bearer-token" + ) + headers = client._get_auth_headers() + assert headers["Authorization"] == "Bearer bearer-token" + + # Test API key + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.api_key, + auth_value="api-key" + ) + headers = client._get_auth_headers() + assert headers["X-API-Key"] == "api-key" + + # Test basic auth (gets base64 encoded) + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.basic, + auth_value="user:pass" + ) + headers = client._get_auth_headers() + assert headers["Authorization"].startswith("Basic ") + + def test_token_auth_with_extra_headers(self): + """Test that token auth works alongside extra headers""" + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.token, + auth_value="my-token", + extra_headers={"X-Custom-Header": "custom-value"} + ) + + headers = client._get_auth_headers() + + assert headers["Authorization"] == "token my-token" + assert headers["X-Custom-Header"] == "custom-value" + + def test_token_auth_enum_value(self): + """Test that MCPAuth.token enum exists and has correct value""" + assert hasattr(MCPAuth, "token") + assert MCPAuth.token.value == "token" + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 4f0459d5eca..26f9a6ee941 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1003,3 +1003,73 @@ def test_gemini_cache_control_injection_list_content_detected(): cached, non_cached = separate_cached_messages(messages) assert len(cached) == 1 assert len(non_cached) == 1 +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_string_negative_index(): + """ + Test that string negative indices like "-1" are handled correctly. + + When cache_control_injection_points are stored in DB/config as JSON, indices + like -1 become the string "-1". Previously, str.isdigit() returned False for + "-1" so the cache control was silently skipped. This tests the fix. + """ + with patch.dict( + os.environ, + { + "AWS_ACCESS_KEY_ID": "fake_access_key_id", + "AWS_SECRET_ACCESS_KEY": "fake_secret_access_key", + "AWS_REGION_NAME": "us-west-2", + }, + ): + anthropic_cache_control_hook = AnthropicCacheControlHook() + litellm.callbacks = [anthropic_cache_control_hook] + + mock_response = MagicMock() + mock_response.json.return_value = { + "output": { + "message": { + "role": "assistant", + "content": "Response", + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 100, + "outputTokens": 50, + "totalTokens": 150, + }, + } + mock_response.status_code = 200 + + client = AsyncHTTPHandler() + with patch.object(client, "post", return_value=mock_response) as mock_post: + await litellm.acompletion( + model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + messages=[ + {"role": "user", "content": "First message"}, + {"role": "assistant", "content": "First response"}, + {"role": "user", "content": "Second message"}, + ], + # index is a string "-1" (as stored in DB/config JSON) + cache_control_injection_points=[ + {"location": "message", "index": "-1"}, + ], + client=client, + ) + + mock_post.assert_called_once() + request_body = json.loads(mock_post.call_args.kwargs["data"]) + + # The last user message should have cache control applied + last_message = request_body["messages"][-1] + last_message_content = last_message["content"] + assert isinstance(last_message_content, list), ( + f"Expected list content, got {type(last_message_content)}" + ) + has_cache_point = any( + isinstance(item, dict) and "cachePoint" in item + for item in last_message_content + ) + assert has_cache_point, ( + f"Expected cachePoint in last message content, got: {last_message_content}. " + "String index '-1' was not parsed correctly (str.isdigit() returns False for negative strings)." + ) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index ae4082662f9..7c60cbb52ee 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -386,6 +386,118 @@ class TestGuardrailLoggingAggregation: assert info[1]["guardrail_name"] == "test_guardrail" +class TestGuardrailSensitiveFieldStripping: + """Tests that secret_fields is stripped from guardrail responses before logging. + + Matches the pattern used by Langfuse and Arize integrations which also + pop("secret_fields") to prevent raw Authorization headers from being persisted. + """ + + def _make_guardrail(self): + from litellm.types.guardrails import GuardrailEventHooks + + return CustomGuardrail( + guardrail_name="test_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + def test_secret_fields_stripped_from_guardrail_response(self): + """Ensure secret_fields (containing raw Authorization headers) is not persisted.""" + guardrail = self._make_guardrail() + request_data = {"metadata": {}} + + guardrail_response_with_secrets = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "secret_fields": { + "raw_headers": { + "authorization": "Bearer sk-live-secret-key-12345", + "content-type": "application/json", + } + }, + "proxy_server_request": {"url": "http://localhost:4000/chat/completions"}, + } + + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=guardrail_response_with_secrets, + request_data=request_data, + guardrail_status="success", + duration=1.0, + ) + + info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(info) == 1 + logged_response = info[0]["guardrail_response"] + + # secret_fields must be stripped + assert "secret_fields" not in logged_response + + # Other fields should be preserved + assert "model" in logged_response + assert "messages" in logged_response + assert "proxy_server_request" in logged_response + + def test_string_guardrail_response_not_affected(self): + """String responses (e.g. 'allow', 'deny') should pass through unchanged.""" + guardrail = self._make_guardrail() + request_data = {"metadata": {}} + + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response="allow", + request_data=request_data, + guardrail_status="success", + duration=0.5, + ) + + info = request_data["metadata"]["standard_logging_guardrail_information"] + assert info[0]["guardrail_response"] == "allow" + + def test_no_authorization_header_in_logged_response(self): + """Verify no plaintext Authorization header ends up in the logged guardrail response.""" + import json + + guardrail = self._make_guardrail() + request_data = {"metadata": {}} + + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "model": "gpt-4", + "secret_fields": { + "raw_headers": { + "authorization": "Bearer sk-live-SHOULD-NOT-APPEAR", + } + }, + }, + request_data=request_data, + guardrail_status="success", + duration=1.0, + ) + + logged_response = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] + assert "secret_fields" not in logged_response + assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response) + + def test_secret_fields_stripped_from_list_dict_response(self): + """Ensure secret_fields is stripped from List[dict] guardrail responses too.""" + guardrail = self._make_guardrail() + request_data = {"metadata": {}} + + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=[ + {"result": "ok", "secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}}}, + {"result": "also_ok"}, + ], + request_data=request_data, + guardrail_status="success", + duration=1.0, + ) + + import json + serialized = json.dumps(request_data) + assert "secret_fields" not in serialized + assert "sk-secret" not in serialized + + class TestCustomGuardrailPassthroughSupport: """Tests for passthrough endpoint guardrail support - Issue fixes.""" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 707b5bdc777..8d68539564c 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockImageProcessor, _convert_to_bedrock_tool_call_invoke, ollama_pt, + sanitize_messages_for_tool_calling, ) @@ -1179,7 +1180,7 @@ def test_bedrock_tools_pt_does_not_handle_system_tool(): System tools (nova_grounding) should be added via web_search_options, not via the tools parameter directly. """ - + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt # Regular function tools should still work @@ -1741,3 +1742,288 @@ def test_bedrock_tool_call_invoke_multiple_normal_tools(): assert len(result) == 2 assert result[0]["toolUse"]["toolUseId"] == "call_1" assert result[1]["toolUse"]["toolUseId"] == "call_2" + + +# ======================================================================== +# Tool result deduplication tests (Case D in sanitize_messages_for_tool_calling) +# ======================================================================== + + +def test_sanitize_messages_deduplicates_tool_results(): + """ + Anthropic requires exactly one tool_result per tool_use. When conversation + history (e.g. from session resume) contains duplicate tool result messages + with the same tool_call_id, sanitize_messages_for_tool_calling should keep + only the last occurrence. + + Without this fix, Anthropic rejects with: + each tool_use must have a single result. Found multiple tool_result + blocks with id: + """ + original = litellm.modify_params + litellm.modify_params = True + try: + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "NYC"}', + }, + } + ], + }, + # First tool result (stale/duplicate) + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": "Partial result...", + }, + # Second tool result (final/complete — should be kept) + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": '{"temperature": 72, "condition": "sunny"}', + }, + ] + + result = sanitize_messages_for_tool_calling(messages) + + # Count tool messages with this ID — should be exactly 1 + tool_results = [ + m for m in result if m.get("role") == "tool" and m.get("tool_call_id") == "call_abc123" + ] + assert len(tool_results) == 1 + # Should keep the LAST occurrence (most complete) + assert tool_results[0]["content"] == '{"temperature": 72, "condition": "sunny"}' + finally: + litellm.modify_params = original + + +def test_sanitize_messages_preserves_unique_tool_results(): + """ + When each tool_call_id has exactly one tool_result, no deduplication should + occur. Messages should pass through unchanged. + """ + original = litellm.modify_params + litellm.modify_params = True + try: + messages = [ + {"role": "user", "content": "Get weather for two cities"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "NYC"}', + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "LA"}', + }, + }, + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "72F"}, + {"role": "tool", "tool_call_id": "call_2", "content": "85F"}, + ] + + result = sanitize_messages_for_tool_calling(messages) + + tool_results = [m for m in result if m.get("role") == "tool"] + assert len(tool_results) == 2 + assert tool_results[0]["tool_call_id"] == "call_1" + assert tool_results[0]["content"] == "72F" + assert tool_results[1]["tool_call_id"] == "call_2" + assert tool_results[1]["content"] == "85F" + finally: + litellm.modify_params = original + + +def test_sanitize_messages_dedup_disabled_when_modify_params_false(): + """ + When litellm.modify_params is False, messages should be returned as-is + even if they contain duplicate tool results. + """ + original = litellm.modify_params + litellm.modify_params = False + try: + messages = [ + {"role": "user", "content": "Test"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_dup", + "type": "function", + "function": {"name": "test", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_dup", "content": "first"}, + {"role": "tool", "tool_call_id": "call_dup", "content": "second"}, + ] + + result = sanitize_messages_for_tool_calling(messages) + + # Should be unchanged — no sanitization when modify_params=False + assert result == messages + finally: + litellm.modify_params = original + + +def test_sanitize_messages_dedup_scoped_per_turn_preserves_cross_turn(): + """ + When the same tool_call_id appears in two different assistant turns + (separated by a user message), both tool results must be preserved. + Deduplication should only apply within a single contiguous tool-result + block, not globally across the conversation. + + Without per-turn scoping this would incorrectly drop the first tool result, + leaving the first assistant message without its required result (which + Anthropic would reject). + """ + original = litellm.modify_params + litellm.modify_params = True + try: + messages = [ + {"role": "user", "content": "First question"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_X", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "a"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_X", "content": "result_turn_1"}, + {"role": "user", "content": "Second question"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_X", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "b"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_X", "content": "result_turn_2"}, + ] + + result = sanitize_messages_for_tool_calling(messages) + + # Both tool results must survive — one per turn + tool_results = [ + m for m in result + if m.get("role") == "tool" and m.get("tool_call_id") == "call_X" + ] + assert len(tool_results) == 2, ( + f"Expected 2 tool results (one per turn), got {len(tool_results)}. " + "Dedup may be global instead of per-turn scoped." + ) + assert tool_results[0]["content"] == "result_turn_1" + assert tool_results[1]["content"] == "result_turn_2" + finally: + litellm.modify_params = original + + +def test_sanitize_messages_combined_case_a_and_case_d(): + """ + Combined Case A + Case D: an assistant message has two tool_calls — + one with a missing result (Case A should inject a dummy) and one with + duplicate results (Case D should deduplicate to keep only the last). + + This validates that both sanitization passes compose correctly without + interfering with each other. + """ + original = litellm.modify_params + litellm.modify_params = True + try: + messages = [ + {"role": "user", "content": "Do two things"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_missing", + "type": "function", + "function": {"name": "tool_a", "arguments": "{}"}, + }, + { + "id": "call_duped", + "type": "function", + "function": {"name": "tool_b", "arguments": '{"q": "x"}'}, + }, + ], + }, + # No result for call_missing — Case A should inject a dummy + # Duplicate results for call_duped — Case D should keep last + {"role": "tool", "tool_call_id": "call_duped", "content": "stale_result"}, + {"role": "tool", "tool_call_id": "call_duped", "content": "fresh_result"}, + {"role": "user", "content": "Now summarize"}, + ] + + result = sanitize_messages_for_tool_calling(messages) + + # Collect tool results from the output + tool_results = [m for m in result if m.get("role") in ("tool", "function")] + + # Case A: call_missing should have a dummy result injected + missing_results = [ + m for m in tool_results if m.get("tool_call_id") == "call_missing" + ] + assert len(missing_results) == 1, ( + f"Expected 1 dummy result for call_missing (Case A), got {len(missing_results)}" + ) + + # Case D: call_duped should have exactly 1 result (the fresh one) + duped_results = [ + m for m in tool_results if m.get("tool_call_id") == "call_duped" + ] + assert len(duped_results) == 1, ( + f"Expected 1 result for call_duped after dedup (Case D), got {len(duped_results)}" + ) + assert duped_results[0]["content"] == "fresh_result", ( + f"Expected last-wins 'fresh_result', got '{duped_results[0]['content']}'" + ) + + # Verify tool results immediately follow the assistant message + asst_idx = next( + i for i, m in enumerate(result) if m.get("role") == "assistant" + ) + tool_msgs_after_asst = [ + m + for m in result[asst_idx + 1 :] + if m.get("role") in ("tool", "function") + ] + assert len(tool_msgs_after_asst) == 2, ( + f"Expected 2 tool results after assistant, got {len(tool_msgs_after_asst)}" + ) + # Both tool_call_ids should be present (order may vary) + tool_ids = {m["tool_call_id"] for m in tool_msgs_after_asst} + assert tool_ids == {"call_missing", "call_duped"}, ( + f"Expected tool_call_ids {{call_missing, call_duped}}, got {tool_ids}" + ) + finally: + litellm.modify_params = original diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 28624ea8b20..f4aeb27b31a 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,4 +1,3 @@ -import json import os import sys from unittest.mock import MagicMock, patch @@ -190,8 +189,13 @@ async def test_datadog_logger_not_shadowed_by_llm_obs(monkeypatch): # Regression check: we expect a distinct DataDogLogger, not the LLM Obs logger assert type(datadog_logger) is DataDogLogger - assert any(isinstance(cb, DataDogLLMObsLogger) for cb in logging_module._in_memory_loggers) - assert any(type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers) + assert any( + isinstance(cb, DataDogLLMObsLogger) + for cb in logging_module._in_memory_loggers + ) + assert any( + type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers + ) finally: logging_module._in_memory_loggers.clear() @@ -202,7 +206,9 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch): # Required env vars for Logfire integration monkeypatch.setenv("LOGFIRE_TOKEN", "test-token") - monkeypatch.setenv("LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev") # no trailing slash on purpose + monkeypatch.setenv( + "LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev" + ) # no trailing slash on purpose # Import after env vars are set (important if module-level caching exists) from litellm.integrations.opentelemetry import OpenTelemetry # logger class @@ -221,7 +227,9 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch): # Sanity: we got the right logger type and it is cached assert type(logger) is OpenTelemetry - assert any(type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers) + assert any( + type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers + ) # Core regression check: base URL env var should influence the exporter endpoint. # @@ -232,7 +240,9 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch): or getattr(logger, "config", None) or getattr(logger, "_otel_config", None) ) - assert cfg is not None, "Expected OpenTelemetry logger to keep an otel config on the instance" + assert ( + cfg is not None + ), "Expected OpenTelemetry logger to keep an otel config on the instance" endpoint = getattr(cfg, "endpoint", None) or getattr(cfg, "otlp_endpoint", None) assert endpoint is not None, "Expected otel config to expose the OTLP endpoint" @@ -297,7 +307,7 @@ async def test_logging_non_streaming_request(): mock_response="Hello, world!", ) await asyncio.sleep(1) - + # Filter calls to only count the one with the expected input message "Hey" # Bridge models may make internal calls that also log, so we filter by the actual input calls_with_expected_input = [] @@ -307,13 +317,13 @@ async def test_logging_non_streaming_request(): first_message_content = messages[0].get("content") if first_message_content == "Hey": calls_with_expected_input.append(call) - + # Assert that we have exactly one call with the expected input assert len(calls_with_expected_input) == 1, ( f"Expected 1 call with input 'Hey', but got {len(calls_with_expected_input)}. " f"Total calls: {mock_async_log_success_event.call_count}" ) - + # Use the filtered call for assertions call_args = calls_with_expected_input[0] standard_logging_object = call_args.kwargs["kwargs"][ @@ -326,14 +336,18 @@ async def test_logging_non_streaming_request(): @pytest.mark.parametrize("async_flag", ["acompletion", "aresponses"]) -def test_success_handler_skips_sync_callbacks_for_async_requests(logging_obj, async_flag): +def test_success_handler_skips_sync_callbacks_for_async_requests( + logging_obj, async_flag +): """Ensure sync success callbacks are skipped when async call type flags are set.""" from litellm.integrations.custom_logger import CustomLogger class DummyLogger(CustomLogger): pass - logging_obj.stream = False # simulate non-streaming request where sync callbacks would normally run + logging_obj.stream = ( + False # simulate non-streaming request where sync callbacks would normally run + ) logging_obj.model_call_details["litellm_params"] = {async_flag: True} logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] @@ -523,7 +537,7 @@ def test_success_handler_runs_guardrail_logging_hook_when_enabled(logging_obj): assert guardrail_call_kwargs["event_type"] == GuardrailEventHooks.logging_only guardrail.logging_hook.assert_called_once() assert logging_obj.model_call_details.get("guardrail_hook_ran") is True - + def test_get_user_agent_tags(): from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -676,21 +690,29 @@ def test_get_request_tags_does_not_mutate_original_tags(): ) # Verify the original tags list was NOT mutated - assert original_tags == ["custom-tag-1", "custom-tag-2"], ( - f"Original tags list was mutated: {original_tags}" - ) - assert metadata["tags"] == ["custom-tag-1", "custom-tag-2"], ( - f"metadata['tags'] was mutated: {metadata['tags']}" - ) + assert original_tags == [ + "custom-tag-1", + "custom-tag-2", + ], f"Original tags list was mutated: {original_tags}" + assert metadata["tags"] == [ + "custom-tag-1", + "custom-tag-2", + ], f"metadata['tags'] was mutated: {metadata['tags']}" # Verify each returned list has exactly 2 User-Agent tags (not duplicated) user_agent_count_1 = len([t for t in tags1 if t.startswith("User-Agent:")]) user_agent_count_2 = len([t for t in tags2 if t.startswith("User-Agent:")]) user_agent_count_3 = len([t for t in tags3 if t.startswith("User-Agent:")]) - assert user_agent_count_1 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_1}" - assert user_agent_count_2 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_2}" - assert user_agent_count_3 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_3}" + assert ( + user_agent_count_1 == 2 + ), f"Expected 2 User-Agent tags, got {user_agent_count_1}" + assert ( + user_agent_count_2 == 2 + ), f"Expected 2 User-Agent tags, got {user_agent_count_2}" + assert ( + user_agent_count_3 == 2 + ), f"Expected 2 User-Agent tags, got {user_agent_count_3}" # Verify all returned lists are independent (different objects) assert tags1 is not tags2 @@ -795,7 +817,6 @@ def test_get_extra_header_tags(): def test_response_cost_calculator_with_response_cost_in_hidden_params(logging_obj): from litellm import Router - from litellm.litellm_core_utils.litellm_logging import Logging router = Router( model_list=[ @@ -933,7 +954,6 @@ async def test_e2e_generate_cold_storage_object_key_successful(): with patch("litellm.cold_storage_custom_logger", return_value="s3"), patch( "litellm.integrations.s3.get_s3_object_key" ) as mock_get_s3_key: - # Mock the S3 object key generation to return a predictable result mock_get_s3_key.return_value = ( "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json" @@ -981,7 +1001,6 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path() ) as mock_get_logger, patch( "litellm.integrations.s3.get_s3_object_key" ) as mock_get_s3_key: - # Setup mocks mock_get_logger.return_value = mock_custom_logger mock_get_s3_key.return_value = ( @@ -1033,7 +1052,6 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path(): ) as mock_get_logger, patch( "litellm.integrations.s3.get_s3_object_key" ) as mock_get_s3_key: - # Setup mocks mock_get_logger.return_value = mock_custom_logger mock_get_s3_key.return_value = ( @@ -1279,9 +1297,9 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu "standard_logging_object should be set for pass-through endpoints " "even when complete_streaming_response is None" ) - assert logging_obj.model_call_details["standard_logging_object"] is not None, ( - "standard_logging_object should not be None for pass-through endpoints" - ) + assert ( + logging_obj.model_call_details["standard_logging_object"] is not None + ), "standard_logging_object should not be None for pass-through endpoints" # Verify that async_complete_streaming_response was set to prevent re-processing # This is consistent with the existing code pattern for regular streaming @@ -1289,15 +1307,15 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu "async_complete_streaming_response should be set to prevent re-processing, " "consistent with the existing code pattern" ) - assert logging_obj.model_call_details["async_complete_streaming_response"] is result, ( - "async_complete_streaming_response should be set to the result" - ) + assert ( + logging_obj.model_call_details["async_complete_streaming_response"] is result + ), "async_complete_streaming_response should be set to the result" # Verify that response_cost is set to None (cost calculation not possible for pass-through) # This is consistent with the error handling in the non-pass-through code path - assert "response_cost" in logging_obj.model_call_details, ( - "response_cost should be set for pass-through endpoints" - ) + assert ( + "response_cost" in logging_obj.model_call_details + ), "response_cost should be set for pass-through endpoints" assert logging_obj.model_call_details["response_cost"] is None, ( "response_cost should be None for pass-through endpoints since " "StandardPassThroughResponseObject doesn't have standard usage info" @@ -1356,10 +1374,14 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp # Verify first call set the values assert "standard_logging_object" in logging_obj.model_call_details assert "async_complete_streaming_response" in logging_obj.model_call_details - first_standard_logging_object = logging_obj.model_call_details["standard_logging_object"] + first_standard_logging_object = logging_obj.model_call_details[ + "standard_logging_object" + ] # Second call - should return early due to async_complete_streaming_response guard - with patch.object(logging_obj, "get_combined_callback_list", return_value=[]) as mock_callbacks: + with patch.object( + logging_obj, "get_combined_callback_list", return_value=[] + ) as mock_callbacks: await logging_obj.async_success_handler( result=result, start_time=start_time, @@ -1370,9 +1392,10 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp mock_callbacks.assert_not_called() # Verify standard_logging_object wasn't modified by second call - assert logging_obj.model_call_details["standard_logging_object"] is first_standard_logging_object, ( - "standard_logging_object should not be modified on re-processing" - ) + assert ( + logging_obj.model_call_details["standard_logging_object"] + is first_standard_logging_object + ), "standard_logging_object should not be modified on re-processing" @pytest.mark.asyncio @@ -1433,9 +1456,11 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_ "standard_logging_object should be set for streaming pass-through endpoints " "even when the response cannot be parsed into a ModelResponse" ) - assert logging_obj.model_call_details["standard_logging_object"] is not None, ( - "standard_logging_object should not be None for streaming pass-through endpoints" - ) + assert ( + logging_obj.model_call_details["standard_logging_object"] is not None + ), "standard_logging_object should not be None for streaming pass-through endpoints" + + def test_get_error_information_error_code_priority(): """ Test get_error_information prioritizes 'code' attribute over 'status_code' attribute @@ -1680,3 +1705,131 @@ async def test_async_success_handler_preserves_response_cost_for_pass_through_en slo = logging_obj.model_call_details.get("standard_logging_object") assert slo is not None assert slo["response_cost"] > 0 + + +def test_function_setup_litellm_metadata_populates_metadata(): + """ + Test that function_setup() properly handles litellm_metadata (used by /v1/messages, + /batches, /responses, /files endpoints) and populates litellm_params["metadata"] + so callbacks like Langfuse can read API key fields. + + This is the root cause of: Claude Code requests missing user_api_key_hash in Langfuse. + """ + import litellm + + test_api_key_hash = "sk-hashed-1234567890abcdef" + test_team_id = "team-test-123" + test_key_alias = "my-test-key" + + # Simulate what happens for /v1/messages: metadata is in "litellm_metadata", not "metadata" + kwargs = { + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "litellm_call_id": "test-call-id-123", + "litellm_metadata": { + "user_api_key_hash": test_api_key_hash, + "user_api_key_alias": test_key_alias, + "user_api_key_team_id": test_team_id, + "user_api_key_user_id": "user-123", + "user_api_key": test_api_key_hash, + }, + } + + logging_obj, returned_kwargs = litellm.utils.function_setup( + original_function="anthropic_messages", + rules_obj=litellm.utils.Rules(), + start_time=time.time(), + **kwargs, + ) + + # litellm_params["metadata"] must contain the API key fields + litellm_params = logging_obj.model_call_details.get("litellm_params", {}) + metadata = litellm_params.get("metadata") + assert metadata is not None, "litellm_params['metadata'] should not be None" + assert isinstance(metadata, dict), "litellm_params['metadata'] should be a dict" + assert metadata.get("user_api_key_hash") == test_api_key_hash + assert metadata.get("user_api_key_alias") == test_key_alias + assert metadata.get("user_api_key_team_id") == test_team_id + + # litellm_metadata should also be preserved + litellm_metadata = litellm_params.get("litellm_metadata") + assert litellm_metadata is not None + assert litellm_metadata.get("user_api_key_hash") == test_api_key_hash + + # metadata should be a COPY, not an alias — mutating one must not affect the other + assert ( + metadata is not litellm_metadata + ), "litellm_params['metadata'] should be a copy, not the same object" + + +def test_function_setup_metadata_takes_precedence_over_litellm_metadata(): + """ + Test that when BOTH metadata and litellm_metadata are present (e.g., user sets + Anthropic API metadata AND proxy adds litellm_metadata), metadata is used as + litellm_params["metadata"] and litellm_metadata is stored separately. + """ + import litellm + + kwargs = { + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "litellm_call_id": "test-call-id-456", + "metadata": { + "user_id": "anthropic-user-id", + }, + "litellm_metadata": { + "user_api_key_hash": "sk-hashed-xyz", + "user_api_key_team_id": "team-xyz", + }, + } + + logging_obj, _ = litellm.utils.function_setup( + original_function="anthropic_messages", + rules_obj=litellm.utils.Rules(), + start_time=time.time(), + **kwargs, + ) + + litellm_params = logging_obj.model_call_details.get("litellm_params", {}) + + # When both are present, metadata should be the explicit "metadata" dict + metadata = litellm_params.get("metadata") + assert metadata is not None + assert metadata.get("user_id") == "anthropic-user-id" + + # litellm_metadata should be preserved separately for merge_litellm_metadata() + litellm_metadata = litellm_params.get("litellm_metadata") + assert litellm_metadata is not None + assert litellm_metadata.get("user_api_key_hash") == "sk-hashed-xyz" + + +def test_function_setup_empty_metadata_falls_back_to_litellm_metadata(): + """ + Test that when metadata is explicitly set to {} (empty dict), litellm_metadata + is still used to populate litellm_params["metadata"] so API key fields are visible. + """ + import litellm + + kwargs = { + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "litellm_call_id": "test-call-id-789", + "metadata": {}, + "litellm_metadata": { + "user_api_key_hash": "sk-hashed-empty-test", + "user_api_key_team_id": "team-empty-test", + }, + } + + logging_obj, _ = litellm.utils.function_setup( + original_function="anthropic_messages", + rules_obj=litellm.utils.Rules(), + start_time=time.time(), + **kwargs, + ) + + litellm_params = logging_obj.model_call_details.get("litellm_params", {}) + metadata = litellm_params.get("metadata") + assert metadata is not None + assert metadata.get("user_api_key_hash") == "sk-hashed-empty-test" + assert metadata.get("user_api_key_team_id") == "team-empty-test" diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index b540b0d952d..6f03f630b5f 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -3175,3 +3175,128 @@ def test_map_openai_params_max_tokens_normalized_to_int(): assert "max_tokens" in result assert result["max_tokens"] == 1 + + +# ======================================================================== +# Tool schema normalization tests +# ======================================================================== + + +def test_map_tool_helper_enforces_object_type_when_missing(): + """ + Anthropic requires input_schema.type to be "object". When an OpenAI tool + has parameters without a 'type' field (common with MCP servers), LiteLLM + should inject type:"object" before forwarding to Anthropic. + + Without this fix, Anthropic rejects with: + tools.N.custom.input_schema.type: Input should be 'object' + """ + config = AnthropicConfig() + + # Tool with parameters that has properties but no 'type' field + tool = { + "type": "function", + "function": { + "name": "search_code", + "description": "Search for code patterns", + "parameters": { + "properties": { + "query": {"type": "string", "description": "Search query"} + }, + "required": ["query"], + }, + }, + } + + original_params = tool["function"]["parameters"].copy() + result, _ = config._map_tool_helper(tool) + assert result is not None + assert result["input_schema"]["type"] == "object" + assert "properties" in result["input_schema"] + assert "query" in result["input_schema"]["properties"] + # Original parameters dict must not be modified in place + assert tool["function"]["parameters"] == original_params, ( + "parameters dict was mutated; _map_tool_helper should not modify caller data" + ) + + +def test_map_tool_helper_enforces_object_type_when_wrong_type(): + """ + If a tool schema has type:"string" or type:"array" at the root level, + LiteLLM should normalize it to type:"object" for Anthropic compatibility. + """ + config = AnthropicConfig() + + tool = { + "type": "function", + "function": { + "name": "echo", + "description": "Echo input", + "parameters": { + "type": "string", + "description": "The input to echo", + }, + }, + } + + original_params = tool["function"]["parameters"].copy() + result, _ = config._map_tool_helper(tool) + assert result is not None + assert result["input_schema"]["type"] == "object" + assert result["input_schema"].get("properties") == {}, ( + "properties should be injected as {} when schema has non-object type and no properties key" + ) + # Original parameters dict must not be modified in place + assert tool["function"]["parameters"] == original_params, ( + "parameters dict was mutated; _map_tool_helper should not modify caller data" + ) + + +def test_map_tool_helper_preserves_valid_object_schema(): + """ + When a tool schema already has type:"object", it should be preserved + without modification. + """ + config = AnthropicConfig() + + tool = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + }, + "required": ["city"], + }, + }, + } + + result, _ = config._map_tool_helper(tool) + assert result is not None + assert result["input_schema"]["type"] == "object" + assert "city" in result["input_schema"]["properties"] + assert result["input_schema"]["required"] == ["city"] + + +def test_map_tool_helper_empty_parameters_get_default(): + """ + When parameters is entirely missing, the existing default should still + produce a valid {type:"object", properties:{}} schema. + """ + config = AnthropicConfig() + + tool = { + "type": "function", + "function": { + "name": "no_params_tool", + "description": "Tool with no parameters", + }, + } + + result, _ = config._map_tool_helper(tool) + assert result is not None + assert result["input_schema"]["type"] == "object" + assert result["input_schema"].get("properties") == {} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index acc76221cbb..656a9c616e8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1195,7 +1195,7 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_requires_per_user_auth_property_oauth2_with_client_creds(self): """Test that requires_per_user_auth returns False for OAuth2 with client credentials""" - # OAuth2 with client credentials + # M2M must be opted in explicitly with oauth2_flow="client_credentials" server = MCPServer( server_id="oauth-server", name="oauth-server", @@ -1205,6 +1205,7 @@ class TestMCPServerManager: client_id="client-id", client_secret="client-secret", token_url="http://oauth-server.com/token", + oauth2_flow="client_credentials", ) assert server.requires_per_user_auth is False assert server.has_client_credentials is True @@ -2393,5 +2394,94 @@ class TestMCPServerTimestamps: assert rebuilt_table.updated_at == updated +class TestHasClientCredentialsOAuth2Flow: + """ + Regression tests for the M2M auto-detection bug. + + Before the fix, has_client_credentials returned True whenever + client_id + client_secret + token_url were all set, even for + interactive OAuth setups (e.g. GitHub Enterprise). This silently + dropped user tokens and fetched M2M tokens instead. + + The fix: M2M must be opted in explicitly via oauth2_flow="client_credentials". + """ + + def _make_server(self, **kwargs) -> MCPServer: + return MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="https://github.example.com/mcp", + **kwargs, + ) + + def test_all_three_fields_set_without_oauth2_flow_is_not_m2m(self): + """ + GitHub Enterprise regression: client_id + client_secret + token_url + should NOT trigger M2M flow unless oauth2_flow is explicitly set. + """ + server = self._make_server( + client_id="gh-client-id", + client_secret="gh-client-secret", + token_url="https://github.example.com/login/oauth/access_token", + ) + assert server.has_client_credentials is False + + def test_explicit_client_credentials_flow_enables_m2m(self): + """oauth2_flow='client_credentials' opts in to M2M.""" + server = self._make_server( + client_id="svc-client-id", + client_secret="svc-client-secret", + token_url="https://idp.example.com/token", + oauth2_flow="client_credentials", + ) + assert server.has_client_credentials is True + + def test_explicit_authorization_code_flow_disables_m2m(self): + """oauth2_flow='authorization_code' always returns False.""" + server = self._make_server( + client_id="gh-client-id", + client_secret="gh-client-secret", + token_url="https://github.example.com/login/oauth/access_token", + oauth2_flow="authorization_code", + ) + assert server.has_client_credentials is False + + def test_no_fields_no_flow_is_not_m2m(self): + """No credentials configured — not M2M.""" + server = self._make_server() + assert server.has_client_credentials is False + + def test_partial_fields_without_flow_is_not_m2m(self): + """Partial credential fields without explicit flow — not M2M.""" + server = self._make_server( + client_id="only-client-id", + ) + assert server.has_client_credentials is False + + def test_needs_user_oauth_token_true_without_explicit_m2m(self): + """ + Without oauth2_flow='client_credentials', an oauth2 server with + client fields set still needs a user OAuth token (interactive flow). + """ + server = self._make_server( + client_id="gh-client-id", + client_secret="gh-client-secret", + token_url="https://github.example.com/login/oauth/access_token", + ) + assert server.needs_user_oauth_token is True + + def test_needs_user_oauth_token_false_with_explicit_m2m(self): + """With oauth2_flow='client_credentials', no per-user token needed.""" + server = self._make_server( + client_id="svc-client-id", + client_secret="svc-client-secret", + token_url="https://idp.example.com/token", + oauth2_flow="client_credentials", + ) + assert server.needs_user_oauth_token is False + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py new file mode 100644 index 00000000000..715bb8e8aee --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -0,0 +1,317 @@ +""" +Tests for AWS SigV4 authentication in MCP client. + +Tests the MCPSigV4Auth httpx.Auth subclass that enables per-request +SigV4 signing for Bedrock AgentCore MCP servers. +""" + +import pytest +from unittest.mock import patch, MagicMock + +import httpx + +from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient +from litellm.types.mcp import MCPAuth, MCPTransport + + +class TestMCPSigV4Auth: + """Unit tests for the MCPSigV4Auth class.""" + + def test_init_with_explicit_credentials(self): + """MCPSigV4Auth initializes with explicit AWS credentials.""" + auth = MCPSigV4Auth( + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_session_token="FwoGZXIvYXdzEBYaDH...", + aws_region_name="us-east-1", + aws_service_name="bedrock-agentcore", + ) + assert auth.credentials is not None + assert auth.credentials.access_key == "AKIAIOSFODNN7EXAMPLE" + assert auth.credentials.secret_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + assert auth.credentials.token == "FwoGZXIvYXdzEBYaDH..." + assert auth.region_name == "us-east-1" + assert auth.service_name == "bedrock-agentcore" + + def test_requires_request_body_flag(self): + """MCPSigV4Auth sets requires_request_body so httpx buffers the body before signing.""" + assert MCPSigV4Auth.requires_request_body is True + + def test_init_defaults(self): + """MCPSigV4Auth uses correct defaults for region and service.""" + auth = MCPSigV4Auth( + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + ) + assert auth.region_name == "us-east-1" + assert auth.service_name == "bedrock-agentcore" + + def test_init_with_resolved_env_values(self): + """MCPSigV4Auth works with pre-resolved values (os.environ/ is resolved at config load time).""" + # Values arrive already resolved by ProxyConfig._check_for_os_environ_vars(), + # so MCPSigV4Auth receives plain strings, not os.environ/ prefixed values. + auth = MCPSigV4Auth( + aws_access_key_id="RESOLVED_KEY_FROM_ENV", + aws_secret_access_key="RESOLVED_SECRET_FROM_ENV", + aws_region_name="us-west-2", + ) + assert auth.credentials.access_key == "RESOLVED_KEY_FROM_ENV" + assert auth.credentials.secret_key == "RESOLVED_SECRET_FROM_ENV" + assert auth.region_name == "us-west-2" + + def test_init_falls_back_to_boto_session(self): + """MCPSigV4Auth falls back to boto3 credential chain when no explicit creds.""" + mock_creds = MagicMock() + mock_creds.access_key = "SESSION_KEY" + mock_creds.secret_key = "SESSION_SECRET" + + mock_session = MagicMock() + mock_session.get_credentials.return_value = mock_creds + + with patch("botocore.session.get_session", return_value=mock_session): + auth = MCPSigV4Auth( + aws_region_name="eu-west-1", + aws_service_name="custom-service", + ) + assert auth.credentials == mock_creds + assert auth.region_name == "eu-west-1" + assert auth.service_name == "custom-service" + + def test_init_raises_when_no_credentials(self): + """MCPSigV4Auth raises ValueError when no credentials are available.""" + mock_session = MagicMock() + mock_session.get_credentials.return_value = None + + with patch("botocore.session.get_session", return_value=mock_session): + with pytest.raises(ValueError, match="No AWS credentials found"): + MCPSigV4Auth() + + def test_auth_flow_signs_request(self): + """MCPSigV4Auth.auth_flow adds SigV4 headers to the request.""" + auth = MCPSigV4Auth( + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name="us-east-1", + aws_service_name="bedrock-agentcore", + ) + + request = httpx.Request( + method="POST", + url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", + headers={"Content-Type": "application/json"}, + content=b'{"jsonrpc":"2.0","method":"tools/list","id":1}', + ) + + # Execute auth_flow generator + flow = auth.auth_flow(request) + signed_request = next(flow) + + # Verify SigV4 headers were added + assert "Authorization" in signed_request.headers + assert "AWS4-HMAC-SHA256" in signed_request.headers["Authorization"] + assert "x-amz-date" in signed_request.headers + assert "bedrock-agentcore" in signed_request.headers["Authorization"] + + def test_auth_flow_different_bodies_produce_different_signatures(self): + """Each request gets a unique signature based on its body.""" + auth = MCPSigV4Auth( + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name="us-east-1", + ) + + request1 = httpx.Request( + method="POST", + url="https://example.com/mcp", + headers={"Content-Type": "application/json"}, + content=b'{"jsonrpc":"2.0","method":"tools/list","id":1}', + ) + request2 = httpx.Request( + method="POST", + url="https://example.com/mcp", + headers={"Content-Type": "application/json"}, + content=b'{"jsonrpc":"2.0","method":"tools/call","id":2,"params":{"name":"search"}}', + ) + + signed1 = next(auth.auth_flow(request1)) + signed2 = next(auth.auth_flow(request2)) + + # Signatures must differ because body content differs + assert signed1.headers["Authorization"] != signed2.headers["Authorization"] + + def test_auth_flow_includes_security_token(self): + """SigV4 signing includes X-Amz-Security-Token when session token is present.""" + auth = MCPSigV4Auth( + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_session_token="SESSION_TOKEN_EXAMPLE", + aws_region_name="us-east-1", + ) + + request = httpx.Request( + method="POST", + url="https://example.com/mcp", + headers={"Content-Type": "application/json"}, + content=b'{"jsonrpc":"2.0","method":"initialize","id":0}', + ) + + signed_request = next(auth.auth_flow(request)) + assert "x-amz-security-token" in signed_request.headers + + +class TestMCPClientSigV4Integration: + """Tests for MCPClient with SigV4 auth wired through.""" + + def test_mcp_client_stores_aws_auth(self): + """MCPClient stores the aws_auth parameter.""" + mock_auth = MagicMock(spec=httpx.Auth) + client = MCPClient( + server_url="https://example.com/mcp", + transport_type=MCPTransport.http, + auth_type=MCPAuth.aws_sigv4, + aws_auth=mock_auth, + ) + assert client._aws_auth is mock_auth + + def test_mcp_client_factory_uses_aws_auth(self): + """The httpx client factory uses aws_auth when no explicit auth is passed.""" + mock_auth = MCPSigV4Auth( + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + ) + client = MCPClient( + server_url="https://example.com/mcp", + transport_type=MCPTransport.http, + aws_auth=mock_auth, + ) + + factory = client._create_httpx_client_factory() + httpx_client = factory( + headers={"Content-Type": "application/json"}, + timeout=httpx.Timeout(30.0), + ) + + # Verify the auth object was actually wired into the httpx client + assert httpx_client._auth is mock_auth + + def test_mcp_client_factory_explicit_auth_takes_precedence(self): + """When explicit auth= is passed to the factory, it takes precedence over aws_auth.""" + aws_auth = MCPSigV4Auth( + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + ) + explicit_auth = MagicMock(spec=httpx.Auth) + + client = MCPClient( + server_url="https://example.com/mcp", + transport_type=MCPTransport.http, + aws_auth=aws_auth, + ) + + factory = client._create_httpx_client_factory() + httpx_client = factory( + headers={"Content-Type": "application/json"}, + timeout=httpx.Timeout(30.0), + auth=explicit_auth, + ) + + # Explicit auth should win over aws_auth + assert httpx_client._auth is explicit_auth + + def test_mcp_client_factory_no_aws_auth(self): + """The httpx client factory works normally when no aws_auth is set.""" + client = MCPClient( + server_url="https://example.com/mcp", + transport_type=MCPTransport.http, + ) + + factory = client._create_httpx_client_factory() + httpx_client = factory( + headers={"Content-Type": "application/json"}, + timeout=httpx.Timeout(30.0), + ) + # No auth should be set when aws_auth is not configured + assert httpx_client._auth is None + + +class TestMCPServerManagerSigV4: + """Tests for MCPServerManager config loading with SigV4.""" + + @pytest.mark.asyncio + async def test_load_config_with_aws_sigv4(self): + """Config loading correctly parses aws_sigv4 auth type and AWS fields.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + config = { + "agentcore_tools": { + "url": "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", + "transport": "http", + "auth_type": "aws_sigv4", + "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "aws_region_name": "us-east-1", + "aws_service_name": "bedrock-agentcore", + } + } + + manager = MCPServerManager() + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.auth_type == MCPAuth.aws_sigv4 + assert server.aws_access_key_id == "AKIAIOSFODNN7EXAMPLE" + assert server.aws_secret_access_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + assert server.aws_region_name == "us-east-1" + assert server.aws_service_name == "bedrock-agentcore" + + @pytest.mark.asyncio + async def test_create_mcp_client_with_sigv4(self): + """_create_mcp_client creates client with SigV4 auth when auth_type is aws_sigv4.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="test-sigv4", + name="test_sigv4_server", + server_name="test_sigv4", + url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations", + transport=MCPTransport.http, + auth_type=MCPAuth.aws_sigv4, + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name="us-east-1", + ) + + manager = MCPServerManager() + client = await manager._create_mcp_client(server=server) + + assert client.auth_type == MCPAuth.aws_sigv4 + assert client._aws_auth is not None + assert isinstance(client._aws_auth, MCPSigV4Auth) + + @pytest.mark.asyncio + async def test_create_mcp_client_without_sigv4(self): + """_create_mcp_client does not create SigV4 auth for non-SigV4 servers.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="test-bearer", + name="test_bearer_server", + server_name="test_bearer", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token="test-token", + ) + + manager = MCPServerManager() + client = await manager._create_mcp_client(server=server) + + assert client._aws_auth is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 76f9c39acd0..32a8c1b1070 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -1608,3 +1608,625 @@ async def test_anonymize_text_http_error_status(): output_parse_pii=False, masked_entity_count={}, ) + + +@pytest.mark.asyncio +async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail): + """ + Regression test: pii_tokens must be stored in data['metadata']['pii_tokens'], + NOT in data['pii_tokens']. Storing at the top level leaks the field to LLM + providers like Anthropic, which reject unknown fields with + 'pii_tokens: Extra inputs are not permitted'. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + pii_entities_config={ + PiiEntityType.PERSON: PiiAction.MASK, + PiiEntityType.PHONE_NUMBER: PiiAction.MASK, + }, + ) + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + mock_cache = DualCache() + + test_data = { + "messages": [ + {"role": "user", "content": "My name is John and my phone is 555-123-4567"} + ], + "model": "claude-haiku-4-5-20251001", + "metadata": {}, + } + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + # Simulate PII masking with token storage (mimics real anonymize_text behavior) + if request_data is not None and output_parse_pii: + if "metadata" not in request_data: + request_data["metadata"] = {} + if "pii_tokens" not in request_data["metadata"]: + request_data["metadata"]["pii_tokens"] = {} + pii_tokens = request_data["metadata"]["pii_tokens"] + seq = len(pii_tokens) + 1 + token = f"" + pii_tokens[token] = "John" + text = text.replace("John", token) + return text + + guardrail.check_pii = mock_check_pii + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=test_data, + call_type="completion", + ) + + # pii_tokens must NOT be at the top level of data (would leak to providers) + assert "pii_tokens" not in result, ( + "pii_tokens must not be a top-level key in request data — " + "it would leak to LLM providers and cause 'Extra inputs are not permitted' errors" + ) + + # pii_tokens must be inside metadata (safe from provider leakage) + assert "metadata" in result + assert "pii_tokens" in result["metadata"] + assert len(result["metadata"]["pii_tokens"]) > 0 + + +@pytest.mark.asyncio +async def test_pii_tokens_in_metadata_used_for_unmasking(): + """ + Regression test: _process_response_for_pii must read pii_tokens from + data['metadata']['pii_tokens'] and correctly unmask the response. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + token_key = "" + request_data = { + "model": "claude-haiku-4-5-20251001", + "metadata": {"pii_tokens": {token_key: "John"}}, + } + + response = ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content=f"Hello {token_key}, how can I help you?", + ), + index=0, + finish_reason="stop", + ) + ] + ) + + await guardrail._process_response_for_pii( + response=response, + request_data=request_data, + mode="unmask", + ) + + assert response.choices[0].message.content == "Hello John, how can I help you?" + + +@pytest.mark.parametrize( + "initial_hook", + ["pre_call", "during_call", "pre_mcp_call"], +) +def test_event_hook_auto_expansion_for_all_string_hooks(initial_hook): + """ + Regression test: when output_parse_pii is True, the guardrail must add + 'post_call' to event_hook regardless of the initial string hook value, + not just when it's 'pre_call'. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook=initial_hook, + ) + assert isinstance(guardrail.event_hook, list) + assert initial_hook in guardrail.event_hook + assert "post_call" in guardrail.event_hook + + +def test_event_hook_no_expansion_when_already_post_call(): + """post_call alone should stay as-is — no expansion needed.""" + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + event_hook="post_call", + ) + # Should remain a string "post_call", not expanded to a list + assert guardrail.event_hook == "post_call" + + +@pytest.mark.asyncio +async def test_metadata_none_does_not_crash(): + """ + Regression test: if metadata is explicitly None in request_data, + the guardrail must not crash with TypeError on the write or read path. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + token_key = "" + # metadata explicitly None — must not crash + request_data = { + "model": "gpt-3.5-turbo", + "metadata": None, + } + + response = ModelResponse( + choices=[ + Choices( + message=Message( + role="assistant", + content=f"Hello {token_key}, how can I help you?", + ), + index=0, + finish_reason="stop", + ) + ] + ) + + # Should not raise TypeError + await guardrail._process_response_for_pii( + response=response, + request_data=request_data, + mode="unmask", + ) + + # No pii_tokens to unmask, so content stays as-is + assert ( + response.choices[0].message.content == f"Hello {token_key}, how can I help you?" + ) + + +# --------------------------------------------------------------------------- +# Tests for sequential-numbered token unmasking in _unmask_pii_text +# --------------------------------------------------------------------------- + + +def test_unmask_exact_match_with_sequential_tokens(): + """ + Normal unmasking: LLM echoes numbered tokens verbatim → original PII restored. + """ + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + pii_tokens = { + "": "John Smith", + "": "555-123-4567", + } + text = "Hello , your number is ." + result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens) + assert result == "Hello John Smith, your number is 555-123-4567." + + +def test_unmask_multiple_same_entity_type(): + """ + Two phone numbers get distinct numbered tokens and unmask correctly. + """ + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + pii_tokens = { + "": "555-111-0000", + "": "555-222-0000", + } + text = "Call or ." + result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens) + assert result == "Call 555-111-0000 or 555-222-0000." + + +def test_unmask_graceful_degradation(): + """ + If the LLM doesn't echo the token back, the numbered label stays + in the output — clean and readable, not garbage hex. + """ + from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + ) + + pii_tokens = { + "": "John", + } + # LLM paraphrased instead of echoing the token + text = "I see you provided a name." + result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens) + # No change — no garbage, just clean text + assert result == text + + +# --------------------------------------------------------------------------- +# Fix 1: Position bug — reverse sort + original text coordinates +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_anonymize_text_multiple_items_position_correctness(): + """ + Regression test: when multiple PII items exist, coordinates reference the + ORIGINAL text. Processing in reverse order prevents coordinate drift. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + # "Call John at 555-123-4567" + # "John" at [5:9], "555-123-4567" at [13:25] + anonymizer_response = { + "text": "Call at ", + "items": [ + { + "start": 5, + "end": 9, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + { + "start": 13, + "end": 25, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + ], + } + + mock_iterator = _make_mock_session_iterator(anonymizer_response) + + request_data = {"metadata": {}} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text="Call John at 555-123-4567", + analyze_results=[ + {"start": 5, "end": 9, "entity_type": "PERSON", "score": 0.9}, + {"start": 13, "end": 25, "entity_type": "PHONE_NUMBER", "score": 0.95}, + ], + output_parse_pii=True, + masked_entity_count={}, + request_data=request_data, + ) + + pii_tokens = request_data["metadata"]["pii_tokens"] + + # Verify tokens captured the correct ORIGINAL text values + person_token = [k for k in pii_tokens if "PERSON" in k][0] + phone_token = [k for k in pii_tokens if "PHONE" in k][0] + assert pii_tokens[person_token] == "John" + assert pii_tokens[phone_token] == "555-123-4567" + + # Verify both PII values are masked in the result + assert "John" not in result + assert "555-123-4567" not in result + + +# --------------------------------------------------------------------------- +# Fix 2: Anthropic native dict response handling +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_anthropic_native_response_unmasking(): + """ + Anthropic native dict responses (type='message') should be unmasked + when output_parse_pii is enabled. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + request_data = { + "model": "claude-3-haiku", + "metadata": { + "pii_tokens": { + "": "John Smith", + "": "555-123-4567", + } + }, + } + + anthropic_response = { + "type": "message", + "id": "msg_123", + "model": "claude-3-haiku", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Hello , your number is .", + } + ], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + result = await guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=mock_user_api_key, + response=anthropic_response, + ) + + assert result["content"][0]["text"] == ( + "Hello John Smith, your number is 555-123-4567." + ) + + +@pytest.mark.asyncio +async def test_anthropic_native_response_masking(): + """ + Anthropic native dict responses should be masked when + apply_to_output is enabled. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + ) + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text.replace("John Smith", "[PERSON]").replace("555-123-4567", "[PHONE]") + + guardrail.check_pii = mock_check_pii + + anthropic_response = { + "type": "message", + "id": "msg_123", + "model": "claude-3-haiku", + "role": "assistant", + "content": [{"type": "text", "text": "Hello John Smith, call 555-123-4567."}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + result = await guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=mock_user_api_key, + response=anthropic_response, + ) + + assert "[PERSON]" in result["content"][0]["text"] + assert "[PHONE]" in result["content"][0]["text"] + assert "John Smith" not in result["content"][0]["text"] + + +@pytest.mark.asyncio +async def test_anthropic_native_response_non_text_blocks_untouched(): + """ + Non-text blocks (tool_use, thinking) in Anthropic responses + should be left untouched during unmasking. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + request_data = { + "model": "claude-3-haiku", + "metadata": {"pii_tokens": {"": "John"}}, + } + + anthropic_response = { + "type": "message", + "id": "msg_123", + "content": [ + {"type": "text", "text": "Hello "}, + { + "type": "tool_use", + "id": "call_1", + "name": "search", + "input": {"q": "test"}, + }, + ], + "role": "assistant", + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + result = await guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=mock_user_api_key, + response=anthropic_response, + ) + + assert result["content"][0]["text"] == "Hello John" + assert result["content"][1]["type"] == "tool_use" + assert result["content"][1]["name"] == "search" + + +# --------------------------------------------------------------------------- +# Fix 3: Anthropic native SSE streaming — bytes passthrough +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_streaming_bytes_chunks_are_yielded_not_discarded(): + """ + Regression test: bytes chunks (Anthropic native SSE) should be yielded + through the streaming hook, not silently discarded. + """ + + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + ) + + byte_chunk = b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n' + + async def mock_stream(): + yield byte_chunk + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data={}, + ): + chunks.append(chunk) + + assert any( + isinstance(c, bytes) for c in chunks + ), "bytes chunks must not be discarded" + assert byte_chunk in chunks + + +@pytest.mark.asyncio +async def test_streaming_unmask_path_bytes_passthrough(): + """ + Bytes chunks in the unmasking path should also pass through. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + byte_chunk = b'data: {"type":"content_block_delta"}\n\n' + request_data = { + "metadata": {"pii_tokens": {"": "John"}}, + } + + async def mock_stream(): + yield byte_chunk + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data=request_data, + ): + chunks.append(chunk) + + assert len(chunks) == 1 + assert chunks[0] == byte_chunk + + +# --------------------------------------------------------------------------- +# Fix 4: apply_guardrail unmask path for input_type="response" +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_guardrail_unmask_on_response(): + """ + When input_type is 'response' and pii_tokens exist, apply_guardrail + should unmask text instead of masking it. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + mock_testing=True, + ) + + request_data = { + "model": "gpt-4o", + "metadata": { + "pii_tokens": { + "": "John Smith", + "": "555-123-4567", + } + }, + } + + inputs = { + "texts": [ + "Hello , your number is .", + ] + } + + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + + assert result["texts"][0] == "Hello John Smith, your number is 555-123-4567." + + +@pytest.mark.asyncio +async def test_apply_guardrail_masks_on_request(): + """ + When input_type is 'request', apply_guardrail should mask as before. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + mock_testing=True, + ) + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text.replace("John Smith", "") + + guardrail.check_pii = mock_check_pii + + result = await guardrail.apply_guardrail( + inputs={"texts": ["Hello John Smith"]}, + request_data={"model": "gpt-4o", "metadata": {}}, + input_type="request", + ) + + assert "" in result["texts"][0] + assert "John Smith" not in result["texts"][0] + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_bytes_only_logs_warning(): + """ + Regression test: when apply_to_output=True and the stream contains only + bytes chunks (Anthropic native SSE), output masking is skipped. + A warning must be logged so operators are aware. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + ) + + byte_chunks = [ + b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n', + b'data: {"type":"content_block_delta","delta":{"text":" world"}}\n\n', + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + + collected = [] + with patch( + "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger" + ) as mock_logger: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key, + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + # All bytes should be yielded through + assert len(collected) == len(byte_chunks) + for original, received in zip(byte_chunks, collected): + assert original == received + + # Warning must be logged about skipped masking + mock_logger.warning.assert_called_once() + warning_msg = mock_logger.warning.call_args[0][0] + assert "Output PII masking was skipped" in warning_msg diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 1e357d2f02e..54fbac1264b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -10,7 +10,6 @@ sys.path.insert( from litellm.proxy.management_endpoints.common_daily_activity import ( _is_user_agent_tag, - compute_tag_metadata_totals, get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, @@ -77,57 +76,6 @@ def test_is_user_agent_tag(): assert _is_user_agent_tag("user-agent-tag") is False # no colon -def test_compute_tag_metadata_totals(): - """Test compute_tag_metadata_totals function.""" - # Create mock records - class MockRecord: - def __init__(self, request_id, tag, spend, prompt_tokens=10, completion_tokens=5): - self.request_id = request_id - self.tag = tag - self.spend = spend - self.prompt_tokens = prompt_tokens - self.completion_tokens = completion_tokens - self.total_tokens = prompt_tokens + completion_tokens - self.cache_read_input_tokens = 0 - self.cache_creation_input_tokens = 0 - self.api_requests = 1 - self.successful_requests = 1 - self.failed_requests = 0 - - # Test deduplication by request_id (keeps max spend) - records = [ - MockRecord("req-1", "production", spend=10.0), - MockRecord("req-1", "staging", spend=20.0), # Higher spend, should be kept - MockRecord("req-2", "production", spend=15.0), - ] - result = compute_tag_metadata_totals(records) - assert result.spend == 35.0 # 20.0 + 15.0 (deduplicated req-1) - assert result.prompt_tokens == 20 # 10 + 10 (only deduplicated records) - assert result.completion_tokens == 10 # 5 + 5 (only deduplicated records) - - # Test ignoring user-agent tags - records_with_ua = [ - MockRecord("req-1", "production", spend=10.0), - MockRecord("req-1", "user-agent:chrome", spend=50.0), # Should be ignored - MockRecord("req-2", "staging", spend=15.0), - ] - result = compute_tag_metadata_totals(records_with_ua) - assert result.spend == 25.0 # 10.0 + 15.0 (user-agent ignored) - - # Test ignoring records without request_id - records_no_req_id = [ - MockRecord("req-1", "production", spend=10.0), - MockRecord(None, "staging", spend=20.0), # Should be ignored - ] - result = compute_tag_metadata_totals(records_no_req_id) - assert result.spend == 10.0 - - # Test empty records - result = compute_tag_metadata_totals([]) - assert result.spend == 0.0 - assert result.prompt_tokens == 0 - - @pytest.mark.asyncio async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): """Test that endpoint breakdown is included in aggregated daily activity.""" @@ -405,6 +353,96 @@ async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_rec assert result["old-key-hash"]["team_id"] == "latest-team" +@pytest.mark.asyncio +async def test_tag_daily_activity_metadata_totals_not_zero(): + """Test that tag daily activity returns correct metadata totals. + + Regression test: the tag endpoint previously passed metadata_metrics_func= + compute_tag_metadata_totals, which skipped every row whose request_id is + NULL. Rows in litellm_dailytagspend are pre-aggregated and always have + NULL request_id, so the totals panel showed $0. The fix is to pass + metadata_metrics_func=None so the fallback aggregation path is used instead. + """ + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + + # Create mock tag spend records (request_id is NULL for aggregated rows) + mock_record_1 = MagicMock() + mock_record_1.request_id = None # NULL in aggregated daily rows + mock_record_1.tag = "production" + mock_record_1.date = "2024-01-01" + mock_record_1.api_key = "key-1" + mock_record_1.model = "gpt-4" + mock_record_1.model_group = "gpt-4" + mock_record_1.custom_llm_provider = "openai" + mock_record_1.mcp_namespaced_tool_name = None + mock_record_1.endpoint = "/chat/completions" + mock_record_1.spend = 25.0 + mock_record_1.prompt_tokens = 500 + mock_record_1.completion_tokens = 200 + mock_record_1.cache_read_input_tokens = 0 + mock_record_1.cache_creation_input_tokens = 0 + mock_record_1.api_requests = 10 + mock_record_1.successful_requests = 9 + mock_record_1.failed_requests = 1 + + mock_record_2 = MagicMock() + mock_record_2.request_id = None + mock_record_2.tag = "staging" + mock_record_2.date = "2024-01-01" + mock_record_2.api_key = "key-2" + mock_record_2.model = "gpt-3.5-turbo" + mock_record_2.model_group = "gpt-3.5-turbo" + mock_record_2.custom_llm_provider = "openai" + mock_record_2.mcp_namespaced_tool_name = None + mock_record_2.endpoint = "/chat/completions" + mock_record_2.spend = 5.0 + mock_record_2.prompt_tokens = 300 + mock_record_2.completion_tokens = 100 + mock_record_2.cache_read_input_tokens = 0 + mock_record_2.cache_creation_input_tokens = 0 + mock_record_2.api_requests = 5 + mock_record_2.successful_requests = 5 + mock_record_2.failed_requests = 0 + + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=2) + mock_table.find_many = AsyncMock(return_value=[mock_record_1, mock_record_2]) + mock_prisma.db.litellm_dailytagspend = mock_table + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailytagspend", + entity_id_field="tag", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + page=1, + page_size=1000, + metadata_metrics_func=None, # No custom func — matches the fix + ) + + # Metadata totals must reflect actual spend, NOT be zero + assert result.metadata.total_spend == 30.0 # 25.0 + 5.0 + assert result.metadata.total_api_requests == 15 # 10 + 5 + assert result.metadata.total_successful_requests == 14 # 9 + 5 + assert result.metadata.total_failed_requests == 1 + assert result.metadata.total_tokens == 1100 # (500+200) + (300+100) + + # Verify breakdown still works + assert len(result.results) == 1 + daily = result.results[0] + assert "production" in daily.breakdown.entities + assert "staging" in daily.breakdown.entities + assert daily.breakdown.entities["production"].metrics.spend == 25.0 + assert daily.breakdown.entities["staging"].metrics.spend == 5.0 + + @pytest.mark.asyncio async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): """Test that the full aggregation pipeline should preserve metadata for deleted keys.""" 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 55366bbec2b..09dfdb81cbb 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 @@ -605,6 +605,10 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch): mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team", + AsyncMock(), + ) from litellm.proxy._types import ( GenerateKeyRequest, @@ -2346,6 +2350,9 @@ async def test_generate_key_with_object_permission(): ), patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin", + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team", + new_callable=AsyncMock, ): # Execute result = await _common_key_generation_helper( 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 e81c6264f7b..ea51965ebfc 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 @@ -797,6 +797,157 @@ class TestListMCPServers: assert result.status == "healthy" +class TestTeamScopedMCPServerAccess: + """Tests for cross-team information disclosure and restricted key bypass fixes.""" + + @pytest.mark.asyncio + async def test_non_member_cannot_query_foreign_team(self): + """Non-admin user who is NOT a member of the target team should get 403.""" + from litellm.proxy._types import Member + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="attacker_user", + ) + + # Team with a different member + mock_team_obj = MagicMock() + mock_team_obj.members_with_roles = [ + Member(user_id="legitimate_user", role="admin"), + ] + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=False, + ), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + AsyncMock(return_value=mock_team_obj), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + with pytest.raises(HTTPException) as exc_info: + await fetch_all_mcp_servers( + user_api_key_dict=mock_user_auth, team_id="foreign-team-id" + ) + assert exc_info.value.status_code == 403 + assert "permission" in str(exc_info.value.detail).lower() + + @pytest.mark.asyncio + async def test_team_member_can_query_own_team(self): + """User who IS a member of the team should be able to query it.""" + from litellm.proxy._types import Member + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="team_member", + ) + + mock_team_obj = MagicMock() + mock_team_obj.members_with_roles = [ + Member(user_id="team_member", role="user"), + ] + mock_team_obj.object_permission = MagicMock(mcp_servers=["server-1"]) + + mock_server = generate_mock_mcp_server_config_record( + server_id="server-1", name="Team Server" + ) + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) + mock_manager._build_mcp_server_table = MagicMock( + return_value=generate_mock_mcp_server_db_record(server_id="server-1") + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=False, + ), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + AsyncMock(return_value=mock_team_obj), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list", + AsyncMock( + return_value=[ + generate_mock_mcp_server_db_record(server_id="server-1") + ] + ), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers( + user_api_key_dict=mock_user_auth, team_id="my-team-id" + ) + assert len(result) == 1 + assert result[0].server_id == "server-1" + + @pytest.mark.asyncio + async def test_admin_can_query_any_team(self): + """Proxy admins should be able to query any team's MCP servers.""" + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user", + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list", + AsyncMock( + return_value=[ + generate_mock_mcp_server_db_record(server_id="server-1") + ] + ), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + # Admin should NOT need to be a team member + result = await fetch_all_mcp_servers( + user_api_key_dict=mock_user_auth, team_id="any-team-id" + ) + assert len(result) == 1 + + @pytest.mark.asyncio + async def test_restricted_virtual_key_cannot_use_team_id_filter(self): + """Restricted virtual keys must not bypass access limits via team_id.""" + mock_user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="vkey_user", + api_key="sk-restricted", + allowed_routes=["mcp_routes"], + ) + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + with pytest.raises(HTTPException) as exc_info: + await fetch_all_mcp_servers( + user_api_key_dict=mock_user_auth, team_id="some-team" + ) + assert exc_info.value.status_code == 403 + assert "Restricted virtual key" in str(exc_info.value.detail) + + class TestTemporaryMCPSessionEndpoints: def test_inherit_credentials_from_existing_server(self): payload = NewMCPServerRequest( @@ -1512,3 +1663,345 @@ class TestManagementPayloadValidation: assert len(result) == 1 assert result[0]["server_id"] == "server-1" assert result[0]["status"] == "healthy" + + +class TestMCPApprovalWorkflow: + """Tests for BYOM submission: register, list submissions, approve, reject.""" + + @pytest.mark.asyncio + async def test_register_mcp_server_requires_team_key(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + register_mcp_server, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + ) + # No team_id → should raise 400 + user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=None, + ) + with pytest.raises(HTTPException) as exc_info: + await register_mcp_server(payload=payload, user_api_key_dict=user_auth) + assert exc_info.value.status_code == 400 + assert "team" in str(exc_info.value.detail).lower() + + @pytest.mark.asyncio + async def test_register_mcp_server_sets_pending_review(self): + from litellm.proxy._types import MCPApprovalStatus + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + register_mcp_server, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + ) + user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-123", + user_id="user-abc", + ) + created_record = generate_mock_mcp_server_db_record( + alias="My Server", + url="https://example.com/mcp", + ) + created_record.approval_status = MCPApprovalStatus.pending_review + created_record.submitted_by = "user-abc" + + 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.validate_and_normalize_mcp_server_payload", + MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=created_record), + ) as mock_create, + ): + result = await register_mcp_server( + payload=payload, user_api_key_dict=user_auth + ) + + # Endpoint sets pending_review before calling create_mcp_server + call_payload = mock_create.call_args[0][1] + assert call_payload.approval_status == MCPApprovalStatus.pending_review + assert call_payload.submitted_by == "user-abc" + assert result is not None + + @pytest.mark.asyncio + async def test_get_submissions_non_admin_forbidden(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_mcp_server_submissions, + ) + + non_admin = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + ) + with pytest.raises(HTTPException) as exc_info: + await get_mcp_server_submissions(user_api_key_dict=non_admin) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_get_submissions_admin_returns_summary(self): + from litellm.proxy._types import MCPSubmissionsSummary + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_mcp_server_submissions, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + pending = generate_mock_mcp_server_db_record(alias="Pending") + pending.approval_status = "pending_review" + summary = MCPSubmissionsSummary( + total=1, pending_review=1, active=0, rejected=0, items=[pending] + ) + + 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_submissions", + AsyncMock(return_value=summary), + ), + ): + result = await get_mcp_server_submissions(user_api_key_dict=admin) + + assert result.total == 1 + assert result.pending_review == 1 + + @pytest.mark.asyncio + async def test_approve_non_pending_server_raises_400(self): + from litellm.proxy._types import MCPApprovalStatus + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + approve_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + active_server = generate_mock_mcp_server_db_record() + active_server.approval_status = MCPApprovalStatus.active + + 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=active_server), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await approve_mcp_server_submission( + server_id="server-1", user_api_key_dict=admin + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_approve_pending_server_loads_into_registry(self): + from litellm.proxy._types import MCPApprovalStatus + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + approve_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + pending_server = generate_mock_mcp_server_db_record() + pending_server.approval_status = MCPApprovalStatus.pending_review + approved_server = generate_mock_mcp_server_db_record() + approved_server.approval_status = MCPApprovalStatus.active + + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + 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=pending_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.approve_mcp_server", + AsyncMock(return_value=approved_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + result = await approve_mcp_server_submission( + server_id=pending_server.server_id, user_api_key_dict=admin + ) + + mock_manager.reload_servers_from_database.assert_awaited_once() + assert result is not None + + @pytest.mark.asyncio + async def test_reject_already_rejected_raises_400(self): + from litellm.proxy._types import MCPApprovalStatus, RejectMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + reject_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + rejected_server = generate_mock_mcp_server_db_record() + rejected_server.approval_status = MCPApprovalStatus.rejected + + 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=rejected_server), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await reject_mcp_server_submission( + server_id="server-1", + payload=RejectMCPServerRequest(review_notes="duplicate"), + user_api_key_dict=admin, + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_reject_active_server_allowed(self): + """Admin can deactivate an already-approved server via the reject endpoint.""" + from litellm.proxy._types import MCPApprovalStatus, RejectMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + reject_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + active_server = generate_mock_mcp_server_db_record() + active_server.approval_status = MCPApprovalStatus.active + now_rejected = generate_mock_mcp_server_db_record() + now_rejected.approval_status = MCPApprovalStatus.rejected + + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + 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=active_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.reject_mcp_server", + AsyncMock(return_value=now_rejected), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + result = await reject_mcp_server_submission( + server_id=active_server.server_id, + payload=RejectMCPServerRequest(review_notes="policy violation"), + user_api_key_dict=admin, + ) + assert result is not None + mock_manager.reload_servers_from_database.assert_awaited_once() + + +class TestValidateMCPRequiredFields: + """Tests for _validate_mcp_required_fields.""" + + def test_missing_required_field_raises_400(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _validate_mcp_required_fields, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + # source_url is absent + ) + with patch_proxy_general_settings({"mcp_required_fields": ["source_url"]}): + with pytest.raises(HTTPException) as exc_info: + _validate_mcp_required_fields(payload) + assert exc_info.value.status_code == 400 + assert "source_url" in str(exc_info.value.detail) + + def test_auth_type_sentinel_treated_as_absent(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _validate_mcp_required_fields, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.none, # sentinel value — treated as absent + ) + with patch_proxy_general_settings({"mcp_required_fields": ["auth_type"]}): + with pytest.raises(HTTPException) as exc_info: + _validate_mcp_required_fields(payload) + assert exc_info.value.status_code == 400 + assert "auth_type" in str(exc_info.value.detail) + + def test_all_required_fields_present_passes(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _validate_mcp_required_fields, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + source_url="https://github.com/org/repo", + auth_type=MCPAuth.bearer_token, + ) + with patch_proxy_general_settings( + {"mcp_required_fields": ["source_url", "auth_type"]} + ): + # Should not raise + _validate_mcp_required_fields(payload) + + def test_no_required_fields_configured_always_passes(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _validate_mcp_required_fields, + ) + + payload = NewMCPServerRequest( + alias="Minimal", + url="https://example.com/mcp", + transport=MCPTransport.sse, + ) + with patch_proxy_general_settings({}): + # Should not raise when no required fields are configured + _validate_mcp_required_fields(payload) + + def test_unknown_field_name_in_config_raises_500(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _validate_mcp_required_fields, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + ) + # "source_Url" is a typo — not a real field on NewMCPServerRequest + with patch_proxy_general_settings({"mcp_required_fields": ["source_Url"]}): + with pytest.raises(HTTPException) as exc_info: + _validate_mcp_required_fields(payload) + assert exc_info.value.status_code == 500 + assert "source_Url" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 07d89035dce..202b95b3199 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -3,15 +3,21 @@ import os import sys import pytest +from fastapi import HTTPException sys.path.insert( 0, os.path.abspath("../../../..") ) -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch +from litellm.proxy._types import LiteLLM_ObjectPermissionTable from litellm.proxy.management_helpers.object_permission_utils import ( + _extract_requested_mcp_access_groups, + _extract_requested_mcp_server_ids, + _resolve_team_allowed_mcp_servers, _set_object_permission, + validate_key_mcp_servers_against_team, ) @@ -82,3 +88,348 @@ async def test_set_object_permission(): assert result["user_id"] == "test_user" assert result["models"] == ["gpt-4"] + +# ---- Tests for _extract_requested_mcp_server_ids ---- + + +def test_extract_requested_mcp_server_ids_from_mcp_servers(): + obj_perm = {"mcp_servers": ["server-1", "server-2"]} + assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1", "server-2"} + + +def test_extract_requested_mcp_server_ids_from_tool_permissions(): + obj_perm = {"mcp_tool_permissions": {"server-a": ["tool1"], "server-b": ["tool2"]}} + assert _extract_requested_mcp_server_ids(obj_perm) == {"server-a", "server-b"} + + +def test_extract_requested_mcp_server_ids_combined(): + obj_perm = { + "mcp_servers": ["server-1"], + "mcp_tool_permissions": {"server-2": ["tool1"]}, + } + assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1", "server-2"} + + +def test_extract_requested_mcp_server_ids_none(): + assert _extract_requested_mcp_server_ids(None) == set() + assert _extract_requested_mcp_server_ids({}) == set() + + +# ---- Tests for _extract_requested_mcp_access_groups ---- + + +def test_extract_requested_mcp_access_groups(): + obj_perm = {"mcp_access_groups": ["group-a", "group-b"]} + assert _extract_requested_mcp_access_groups(obj_perm) == {"group-a", "group-b"} + + +def test_extract_requested_mcp_access_groups_none(): + assert _extract_requested_mcp_access_groups(None) == set() + assert _extract_requested_mcp_access_groups({}) == set() + + +# ---- Tests for validate_key_mcp_servers_against_team ---- + + +def _make_team_obj( + team_id="team-1", + mcp_servers=None, + mcp_access_groups=None, + mcp_tool_permissions=None, +): + """Create a mock team object with the given MCP permissions.""" + mock_team = MagicMock() + mock_team.team_id = team_id + + if mcp_servers is not None or mcp_access_groups is not None or mcp_tool_permissions is not None: + mock_team.object_permission = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_team.object_permission.mcp_servers = mcp_servers or [] + mock_team.object_permission.mcp_access_groups = mcp_access_groups or [] + mock_team.object_permission.mcp_tool_permissions = mcp_tool_permissions or {} + else: + mock_team.object_permission = None + + return mock_team + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_no_object_permission(mock_access_groups, mock_allow_all): + """No object_permission on key — should pass without error.""" + await validate_key_mcp_servers_against_team( + object_permission=None, + team_obj=_make_team_obj(mcp_servers=["server-1"]), + ) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_key_servers_within_team_scope(mock_access_groups, mock_allow_all): + """Key requests servers that are in the team's scope — should pass.""" + team_obj = _make_team_obj(mcp_servers=["server-1", "server-2", "server-3"]) + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-1", "server-2"]}, + team_obj=team_obj, + ) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_key_servers_outside_team_scope_raises(mock_access_groups, mock_allow_all): + """Key requests servers NOT in the team's scope — should raise 403.""" + team_obj = _make_team_obj(mcp_servers=["server-1"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-1", "server-outside"]}, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + assert "server-outside" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value={"global-server"}, +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_allow_all_keys_servers_always_allowed(mock_access_groups, mock_allow_all): + """allow_all_keys servers should be accessible even if not in team scope.""" + team_obj = _make_team_obj(mcp_servers=["server-1"]) + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-1", "global-server"]}, + team_obj=team_obj, + ) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value={"global-server"}, +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_no_team_only_allow_all_keys(mock_access_groups, mock_allow_all): + """Key without a team can only use allow_all_keys servers.""" + # This should pass — requesting a global server without a team + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["global-server"]}, + team_obj=None, + ) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value={"global-server"}, +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_no_team_non_global_server_raises(mock_access_groups, mock_allow_all): + """Key without a team requesting a non-global server — should raise 403.""" + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["private-server"]}, + team_obj=None, + ) + assert exc_info.value.status_code == 403 + assert "not in a team" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_team_no_mcp_config_blocks_all(mock_access_groups, mock_allow_all): + """Team with no object_permission — key can't use any non-global MCP servers.""" + team_obj = _make_team_obj() # No object_permission + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["some-server"]}, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_tool_permissions_validated_against_team(mock_access_groups, mock_allow_all): + """Server IDs in mcp_tool_permissions should also be validated.""" + team_obj = _make_team_obj(mcp_servers=["server-1"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={ + "mcp_tool_permissions": {"server-outside": ["tool1"]} + }, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + assert "server-outside" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_access_groups_within_team_scope(mock_access_groups, mock_allow_all): + """Key requests access groups that are in the team's scope — should pass.""" + team_obj = _make_team_obj(mcp_access_groups=["group-a", "group-b"]) + await validate_key_mcp_servers_against_team( + object_permission={"mcp_access_groups": ["group-a"]}, + team_obj=team_obj, + ) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_access_groups_outside_team_scope_raises(mock_access_groups, mock_allow_all): + """Key requests access groups NOT in the team's scope — should raise 403.""" + team_obj = _make_team_obj(mcp_access_groups=["group-a"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_access_groups": ["group-outside"]}, + team_obj=team_obj, + ) + assert exc_info.value.status_code == 403 + assert "group-outside" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_validate_access_groups_no_team_raises(mock_access_groups, mock_allow_all): + """Key without a team requesting access groups — should raise 403.""" + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_access_groups": ["group-a"]}, + team_obj=None, + ) + assert exc_info.value.status_code == 403 + assert "not in a team" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + return_value=set(), +) +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=["server-from-group"], +) +async def test_validate_team_access_groups_resolve_to_servers(mock_access_groups, mock_allow_all): + """Team access groups should resolve to server IDs and be included in allowed set.""" + team_obj = _make_team_obj(mcp_access_groups=["group-a"]) + # Key requests a server that comes from the team's access group + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-from-group"]}, + team_obj=team_obj, + ) + + +# ---- Tests for _resolve_team_allowed_mcp_servers with JSON string mcp_tool_permissions ---- + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_resolve_team_allowed_mcp_servers_string_tool_permissions(mock_access_groups): + """mcp_tool_permissions stored as a JSON string (via safe_dumps) should be deserialized correctly.""" + mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_perm.mcp_servers = ["server-1"] + mock_perm.mcp_access_groups = [] + mock_perm.mcp_tool_permissions = json.dumps({"server-2": ["tool1"]}) + + result = await _resolve_team_allowed_mcp_servers(mock_perm) + assert result == {"server-1", "server-2"} + + +@pytest.mark.asyncio +@patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], +) +async def test_resolve_team_allowed_mcp_servers_dict_tool_permissions(mock_access_groups): + """mcp_tool_permissions as a dict should work without deserialization.""" + mock_perm = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_perm.mcp_servers = [] + mock_perm.mcp_access_groups = [] + mock_perm.mcp_tool_permissions = {"server-a": ["tool1"]} + + result = await _resolve_team_allowed_mcp_servers(mock_perm) + assert result == {"server-a"} + diff --git a/tests/test_litellm/test_chat_ui_responses_session.py b/tests/test_litellm/test_chat_ui_responses_session.py new file mode 100644 index 00000000000..09ef003ebdb --- /dev/null +++ b/tests/test_litellm/test_chat_ui_responses_session.py @@ -0,0 +1,127 @@ +""" +Tests for responses API session chaining used by the chat UI. + +Verifies that: +1. previous_response_id is correctly forwarded when provided +2. Absence of previous_response_id does not break the call +3. The aresponses function signature exposes the expected parameters +""" +import inspect +import json +import os +import sys +import unittest.mock as mock + +# Use __file__ so the import path is correct regardless of the pytest working directory. +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +import httpx +import pytest + +import litellm + + +class TestResponsesSessionChaining: + """Test previous_response_id session chaining for the chat UI.""" + + def test_responses_api_signature_accepts_previous_response_id(self): + """aresponses must accept previous_response_id and onResponseId-like params.""" + sig = inspect.signature(litellm.aresponses) + assert "previous_response_id" in sig.parameters, ( + "aresponses must accept previous_response_id for multi-turn session chaining" + ) + assert "input" in sig.parameters, "aresponses must accept input" + assert "model" in sig.parameters, "aresponses must accept model" + + @pytest.mark.asyncio + async def test_previous_response_id_included_in_request_body(self): + """previous_response_id must appear in the outgoing HTTP request body.""" + captured_body: dict = {} + + async def mock_send(self_transport, request: httpx.Request, **kwargs): + try: + captured_body.update(json.loads(request.content)) + except Exception: + pass + # Return a minimal valid responses API response + response_json = { + "id": "resp_test123", + "object": "response", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_001", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + "status": "completed", + } + ], + "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, + "status": "completed", + "created_at": 1700000000, + } + return httpx.Response( + 200, + json=response_json, + request=request, + ) + + with mock.patch("httpx.AsyncClient.send", mock_send): + try: + await litellm.aresponses( + input="hello", + model="gpt-4o-mini", + previous_response_id="resp_prev_abc", + api_key="sk-test-fake", + ) + except Exception: + pass # response parsing may fail; we only care about the outgoing body + + assert captured_body.get("previous_response_id") == "resp_prev_abc", ( + f"Expected previous_response_id in request body, got: {captured_body}" + ) + + @pytest.mark.asyncio + async def test_no_previous_response_id_omitted_from_request(self): + """When previous_response_id is None, it must not appear in the request body.""" + captured_body: dict = {} + + async def mock_send(self_transport, request: httpx.Request, **kwargs): + try: + captured_body.update(json.loads(request.content)) + except Exception: + pass + response_json = { + "id": "resp_new001", + "object": "response", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_001", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + "status": "completed", + } + ], + "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, + "status": "completed", + "created_at": 1700000000, + } + return httpx.Response(200, json=response_json, request=request) + + with mock.patch("httpx.AsyncClient.send", mock_send): + try: + await litellm.aresponses( + input="hello", + model="gpt-4o-mini", + previous_response_id=None, + api_key="sk-test-fake", + ) + except Exception: + pass + + assert "previous_response_id" not in captured_body, ( + "previous_response_id must be omitted from the request body when None" + ) diff --git a/tests/test_litellm/test_router_silent_experiment.py b/tests/test_litellm/test_router_silent_experiment.py index a23ea80f7ce..67d262f83d4 100644 --- a/tests/test_litellm/test_router_silent_experiment.py +++ b/tests/test_litellm/test_router_silent_experiment.py @@ -1,4 +1,5 @@ import asyncio +import time from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -19,11 +20,20 @@ def test_get_silent_experiment_kwargs(): }, ] router = Router(model_list=model_list) - kwargs = {"metadata": {"foo": "bar"}, "litellm_call_id": "call-123"} + kwargs = { + "metadata": {"foo": "bar"}, + "litellm_call_id": "call-123", + "stream": True, + "proxy_server_request": {"body": {"model": "test"}}, + } result = router._get_silent_experiment_kwargs(**kwargs) assert result["metadata"]["is_silent_experiment"] is True assert result["metadata"]["foo"] == "bar" assert "litellm_call_id" not in result + # stream must be forced to False so callbacks fire in background + assert result["stream"] is False + # proxy_server_request must be preserved for spend log metadata + assert "proxy_server_request" in result def test_silent_experiment_completion_direct(): @@ -39,7 +49,7 @@ def test_silent_experiment_completion_direct(): ] router = Router(model_list=model_list) messages = [{"role": "user", "content": "hi"}] - with patch.object(router, "completion", return_value=None): + with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None): router._silent_experiment_completion( silent_model="gpt-3.5-turbo", messages=messages, @@ -173,12 +183,20 @@ def test_router_silent_experiment_completion(): router = Router(model_list=model_list) - # Mock litellm.completion + # Mock litellm.acompletion mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) - mock_completion = MagicMock(return_value=mock_response) + + # We need an async mock for acompletion + async def mock_acompletion(*args, **kwargs): + return mock_response + + mock_acompletion_mock = AsyncMock(side_effect=mock_acompletion) + mock_completion_mock = MagicMock(return_value=mock_response) # Patch at the litellm module level - with patch.object(litellm, "completion", mock_completion): + with patch.object(litellm, "acompletion", mock_acompletion_mock), patch.object( + litellm, "completion", mock_completion_mock + ): response = router.completion( model="primary-model", messages=[{"role": "user", "content": "hi"}], @@ -186,15 +204,13 @@ def test_router_silent_experiment_completion(): assert response.choices[0].message.content == "hello" - # The sync background call uses a thread pool. We might need to wait a bit. - import time + # The sync background call uses a thread pool. We might need to wait. + time.sleep(2.0) - time.sleep(0.5) + # Should have 1 acompletion call (the silent background call) + assert mock_acompletion_mock.call_count == 1 - # Should have 2 calls - assert mock_completion.call_count == 2 - - call_args_list = mock_completion.call_args_list + call_args_list = mock_acompletion_mock.call_args_list # Verify no silent_model in any call for call in call_args_list: @@ -212,3 +228,5 @@ def test_router_silent_experiment_completion(): ) assert silent_call is not None assert silent_call[1]["model"] == "openai/gpt-4" + # Verify model_group is set to the silent model name for correct metric attribution + assert silent_call[1]["metadata"]["model_group"] == "silent-model" diff --git a/ui/litellm-dashboard/public/assets/logos/figma.svg b/ui/litellm-dashboard/public/assets/logos/figma.svg new file mode 100644 index 00000000000..2d8b70457d9 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/figma.svg @@ -0,0 +1,7 @@ + + + + + + + diff --git a/ui/litellm-dashboard/public/assets/logos/gitlab.svg b/ui/litellm-dashboard/public/assets/logos/gitlab.svg new file mode 100644 index 00000000000..18a89fa328d --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/gitlab.svg @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/ui/litellm-dashboard/public/assets/logos/gmail.svg b/ui/litellm-dashboard/public/assets/logos/gmail.svg new file mode 100644 index 00000000000..d702890620d --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/gmail.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/google_drive.svg b/ui/litellm-dashboard/public/assets/logos/google_drive.svg new file mode 100644 index 00000000000..7048af9915e --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/google_drive.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/ui/litellm-dashboard/public/assets/logos/hubspot.svg b/ui/litellm-dashboard/public/assets/logos/hubspot.svg new file mode 100644 index 00000000000..b993945ac6b --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/hubspot.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/jira.svg b/ui/litellm-dashboard/public/assets/logos/jira.svg new file mode 100644 index 00000000000..fb10ca75173 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/jira.svg @@ -0,0 +1,15 @@ + + + + + + + + + + + + + + + diff --git a/ui/litellm-dashboard/public/assets/logos/linear.svg b/ui/litellm-dashboard/public/assets/logos/linear.svg new file mode 100644 index 00000000000..83662a1f9ff --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/linear.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/notion.svg b/ui/litellm-dashboard/public/assets/logos/notion.svg new file mode 100644 index 00000000000..170b9bb4140 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/notion.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/salesforce.svg b/ui/litellm-dashboard/public/assets/logos/salesforce.svg new file mode 100644 index 00000000000..1a541a004f1 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/salesforce.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/sentry.svg b/ui/litellm-dashboard/public/assets/logos/sentry.svg new file mode 100644 index 00000000000..9c3733dc43e --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/sentry.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/shopify.svg b/ui/litellm-dashboard/public/assets/logos/shopify.svg new file mode 100644 index 00000000000..fcc7547269d --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/shopify.svg @@ -0,0 +1,4 @@ + + + + diff --git a/ui/litellm-dashboard/public/assets/logos/slack.svg b/ui/litellm-dashboard/public/assets/logos/slack.svg new file mode 100644 index 00000000000..801de4f70c8 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/slack.svg @@ -0,0 +1,6 @@ + + + + + + diff --git a/ui/litellm-dashboard/public/assets/logos/stripe.svg b/ui/litellm-dashboard/public/assets/logos/stripe.svg new file mode 100644 index 00000000000..ac16a6fb170 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/stripe.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/twilio.svg b/ui/litellm-dashboard/public/assets/logos/twilio.svg new file mode 100644 index 00000000000..3517a2824d9 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/twilio.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/public/assets/logos/zapier.svg b/ui/litellm-dashboard/public/assets/logos/zapier.svg new file mode 100644 index 00000000000..8428ba82a5b --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/zapier.svg @@ -0,0 +1,3 @@ + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx index dbc1c4d10e2..ceb14864ad9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx @@ -20,6 +20,7 @@ import { ToolOutlined, TagsOutlined, AuditOutlined, + MessageOutlined, } from "@ant-design/icons"; // import { // all_admin_roles, @@ -47,6 +48,7 @@ interface SidebarProps { interface MenuItemCfg { key: string; + newTab?: boolean; page: string; // legacy id; we map this to a path below label: string; roles?: string[]; @@ -105,6 +107,8 @@ const routeFor = (slug: string): string => { return "guardrails"; case "policies": return "policies"; + case "chat": + return "chat"; // tools case "mcp-servers": @@ -371,19 +375,29 @@ const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelect }, [pathname, filteredMenuItems, defaultSelectedKey]); // ----- Navigation ----- - const goTo = (slug: string) => { + const goTo = (slug: string, newTab?: boolean) => { const href = toHref(slug); - router.push(href); + if (newTab) { + window.open(href, "_blank"); + } else { + router.push(href); + } }; // Wrap label in so every nav item supports right-click → "Open in new tab" // and Ctrl/Cmd+click to open in a new tab, while preserving SPA navigation for normal clicks. - const renderNavLink = (label: string, page: string): React.ReactNode => { + const renderNavLink = (label: string, page: string, newTab?: boolean): React.ReactNode => { const href = toHref(page); return ( { + if (newTab) { + e.stopPropagation(); + return; + } if (e.metaKey || e.ctrlKey || e.shiftKey || e.button === 1) { e.stopPropagation(); return; @@ -409,6 +423,8 @@ const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelect style={{ transition: "all 0.3s cubic-bezier(0.4, 0, 0.2, 1)", position: "relative", + display: "flex", + flexDirection: "column", }} > = ({ accessToken, userRole, defaultSelect borderRight: 0, backgroundColor: "transparent", fontSize: "14px", + flex: 1, + overflowY: "auto", }} items={filteredMenuItems.map((item) => ({ key: item.key, icon: item.icon, - label: renderNavLink(item.label, item.page), + label: renderNavLink(item.label, item.page, item.newTab), children: item.children?.map((child) => ({ key: child.key, icon: child.icon, - label: renderNavLink(child.label, child.page), - onClick: () => goTo(child.page), + label: renderNavLink(child.label, child.page, child.newTab), + onClick: () => goTo(child.page, child.newTab), })), - onClick: !item.children ? () => goTo(item.page) : undefined, + onClick: !item.children ? () => goTo(item.page, item.newTab) : undefined, }))} /> {isAdminRole(userRole) && !collapsed && } + + {/* Pinned "Open Chat" button at bottom */} + ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts index be910acf7e4..567e1d23013 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts @@ -38,7 +38,7 @@ describe("useMCPServerHealth", () => { vi.clearAllMocks(); }); - it("should fetch health status for given server IDs", async () => { + it("should fetch health status for all servers", async () => { const mockHealthStatuses = [ { server_id: "server-1", status: "healthy" }, { server_id: "server-2", status: "unhealthy" }, @@ -46,27 +46,6 @@ describe("useMCPServerHealth", () => { vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); - const { result } = renderHook(() => useMCPServerHealth(["server-1", "server-2"]), { - wrapper, - }); - - await waitFor(() => { - expect(result.current.isSuccess).toBe(true); - }); - - expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", ["server-1", "server-2"]); - expect(result.current.data).toEqual(mockHealthStatuses); - }); - - it("should fetch health status for all servers when no server IDs provided", async () => { - const mockHealthStatuses = [ - { server_id: "server-1", status: "healthy" }, - { server_id: "server-2", status: "healthy" }, - { server_id: "server-3", status: "unhealthy" }, - ]; - - vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); - const { result } = renderHook(() => useMCPServerHealth(), { wrapper, }); @@ -75,30 +54,15 @@ describe("useMCPServerHealth", () => { expect(result.current.isSuccess).toBe(true); }); - expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", undefined); + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123"); expect(result.current.data).toEqual(mockHealthStatuses); }); - it("should handle empty server list", async () => { - vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); - - const { result } = renderHook(() => useMCPServerHealth([]), { - wrapper, - }); - - await waitFor(() => { - expect(result.current.isSuccess).toBe(true); - }); - - expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", []); - expect(result.current.data).toEqual([]); - }); - it("should handle errors when fetching health status", async () => { const mockError = new Error("Failed to fetch health status"); vi.mocked(networking.fetchMCPServerHealth).mockRejectedValue(mockError); - const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + const { result } = renderHook(() => useMCPServerHealth(), { wrapper, }); @@ -116,7 +80,7 @@ describe("useMCPServerHealth", () => { accessToken: null, } as any); - const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + const { result } = renderHook(() => useMCPServerHealth(), { wrapper, }); @@ -124,4 +88,18 @@ describe("useMCPServerHealth", () => { expect(result.current.status).toBe("pending"); expect(networking.fetchMCPServerHealth).not.toHaveBeenCalled(); }); + + it("should use a stable query key that does not include server IDs", () => { + // Regression test: deleting a server used to pass a changing serverIds array into the + // hook, which was embedded in the query key. React Query would see a new key and fire + // a health check for every remaining server. + // + // The fix: the hook takes no serverIds parameter and uses a constant query key, so + // deleting (or adding) a server never causes an extra health check request. + // + // We verify the contract here by confirming the hook accepts no arguments. + // The stable-key behaviour is further exercised by mcp_servers.test.tsx. + const hookLength = useMCPServerHealth.length; + expect(hookLength).toBe(0); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts index 95d7f3bcee0..681bf4161ad 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts @@ -1,4 +1,5 @@ -import { useQuery } from "@tanstack/react-query"; +import { useCallback, useState } from "react"; +import { useQuery, useQueryClient } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import { fetchMCPServerHealth } from "@/components/networking"; import useAuthorized from "../useAuthorized"; @@ -10,13 +11,49 @@ interface MCPServerHealth { status: string; } -export const useMCPServerHealth = (serverIds?: string[]) => { +export const useMCPServerHealth = () => { const { accessToken } = useAuthorized(); - return useQuery({ - queryKey: [...mcpServerHealthKeys.lists(), { serverIds }], - queryFn: async () => await fetchMCPServerHealth(accessToken!, serverIds), + const queryClient = useQueryClient(); + const [recheckingServerIds, setRecheckingServerIds] = useState>(new Set()); + + const query = useQuery({ + queryKey: mcpServerHealthKeys.lists(), + queryFn: async () => await fetchMCPServerHealth(accessToken!), enabled: !!accessToken, // Refetch health status every 30 seconds to keep it up to date refetchInterval: 30000, }); + + const recheckServerHealth = useCallback(async (serverId: string) => { + if (!accessToken) return; + + setRecheckingServerIds((prev) => new Set(prev).add(serverId)); + + try { + const result: MCPServerHealth[] = await fetchMCPServerHealth(accessToken, [serverId]); + + queryClient.setQueriesData( + { queryKey: mcpServerHealthKeys.lists() }, + (oldData) => { + if (!oldData) return result; + return oldData.map((h) => { + const updated = result.find((r) => r.server_id === h.server_id); + return updated ?? h; + }); + }, + ); + } finally { + setRecheckingServerIds((prev) => { + const next = new Set(prev); + next.delete(serverId); + return next; + }); + } + }, [accessToken, queryClient]); + + return { + ...query, + recheckServerHealth, + recheckingServerIds, + }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts index 8746baae148..9210e25e1a8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts @@ -6,11 +6,11 @@ import useAuthorized from "../useAuthorized"; const mcpServersKeys = createQueryKeys("mcpServers"); -export const useMCPServers = () => { +export const useMCPServers = (teamId?: string | null) => { const { accessToken } = useAuthorized(); return useQuery({ - queryKey: mcpServersKeys.list({}), - queryFn: async () => await fetchMCPServers(accessToken!), + queryKey: mcpServersKeys.list(teamId ? { filters: { teamId } } : undefined), + queryFn: async () => await fetchMCPServers(accessToken!, teamId), enabled: !!accessToken, }); }; diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index c7c195835d0..fbfcb402766 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -1,7 +1,7 @@ import { InfoCircleOutlined, UserAddOutlined } from "@ant-design/icons"; import { useQueryClient } from "@tanstack/react-query"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import { Accordion, AccordionBody, AccordionHeader, Button as Button2, SelectItem, TextInput } from "@tremor/react"; +import { Accordion, AccordionBody, AccordionHeader, SelectItem, TextInput } from "@tremor/react"; import { Alert, Button, Form, Input, Modal, Select, Select as Select2, Space, Tooltip, Typography } from "antd"; import React, { useEffect, useMemo, useState } from "react"; import BulkCreateUsers from "./bulk_create_users_button"; @@ -229,9 +229,9 @@ export const CreateUserButton: React.FC = ({ // Original return for standalone mode return (
- setIsModalVisible(true)}> + ({ default: vi.fn(), })); +// Mock useOrganizations hook +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn().mockReturnValue({ + data: [ + { + organization_id: "org-1", + organization_alias: "Test Organization", + }, + ], + }), +})); + // Mock fetchTeams to prevent network calls vi.mock("@/app/(dashboard)/networking", async (importOriginal) => { const actual = await importOriginal(); @@ -125,6 +137,7 @@ const mockKey: KeyResponse = { user: { user_email: "user@example.com", user_id: "user-1", + user_alias: null, }, }; @@ -380,7 +393,7 @@ it("should render table headers correctly", () => { // Check that main headers are rendered (testing the header.isPlaceholder condition path) expect(screen.getByText("Key ID")).toBeInTheDocument(); expect(screen.getByText("Key Alias")).toBeInTheDocument(); - expect(screen.getByText("Team Alias")).toBeInTheDocument(); + expect(screen.getByText("Team")).toBeInTheDocument(); expect(screen.getByText("Models")).toBeInTheDocument(); expect(screen.getByText("Spend (USD)")).toBeInTheDocument(); }); @@ -463,6 +476,8 @@ it("should display 'Default Proxy Admin' for user_id when value is 'default_user const keyWithDefaultUserId = { ...mockKey, user_id: "default_user_id", + user_email: "", + user: { user_id: "default_user_id", user_email: "", user_alias: null }, }; mockUseFilterLogic.mockReturnValue({ diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index d9c590938fb..20cc1b8153c 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -1,5 +1,6 @@ "use client"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; import { @@ -25,13 +26,14 @@ import { Text, } from "@tremor/react"; import { InfoCircleOutlined, SyncOutlined } from "@ant-design/icons"; -import { Button as AntButton, Popover, Skeleton, Tooltip } from "antd"; +import { Button as AntButton, Popover, Skeleton, Tooltip, Typography } from "antd"; import React, { useEffect, useDeferredValue, useMemo, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { useFilterLogic } from "../key_team_helpers/filter_logic"; import { PaginatedKeyAliasSelect } from "../KeyAliasSelect/PaginatedKeyAliasSelect/PaginatedKeyAliasSelect"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import FilterComponent, { FilterOption } from "../molecules/filter"; +import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; import { Organization } from "../networking"; import KeyInfoView from "../templates/key_info_view"; @@ -51,6 +53,8 @@ interface VirtualKeysTableProps { */ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSort }: VirtualKeysTableProps) { + const { data: fetchedOrganizations } = useOrganizations(); + const resolvedOrganizations = fetchedOrganizations ?? organizations ?? []; const [selectedKey, setSelectedKey] = useState(null); const [sorting, setSorting] = React.useState(() => { if (currentSort) { @@ -86,6 +90,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, { sortBy: sortBy || undefined, sortOrder: sortOrder || undefined, + expand: "user", }); const [expandedAccordions, setExpandedAccordions] = useState>({}); @@ -172,11 +177,9 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo const value = info.getValue() as string; const width = info.cell.column.getSize(); return ( - - - {value ?? "-"} - - + + {value ?? "-"} + ); }, }, @@ -191,76 +194,110 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo { id: "team_alias", accessorKey: "team_id", - header: "Team Alias", + header: "Team", size: 120, enableSorting: false, - cell: ({ row, getValue }) => { - const teamId = getValue() as string; - const team = teams?.find((t) => t.team_id === teamId); - return team?.team_alias || "Unknown"; - }, - }, - { - id: "team_id", - accessorKey: "team_id", - header: "Team ID", - size: 80, - enableSorting: false, cell: (info) => { - const value = info.getValue() as string | null; + const teamId = info.getValue() as string | null; + if (!teamId) return "-"; + const team = teams?.find((t) => t.team_id === teamId); + const displayValue = team?.team_alias || teamId; const width = info.cell.column.getSize(); return ( - - - {value ?? "-"} - - + + {displayValue} + ); }, }, { - id: "organization_id", + id: "organization_alias", accessorKey: "org_id", - header: "Organization ID", + header: "Organization", size: 140, enableSorting: false, - cell: (info) => (info.getValue() ? info.renderValue() : "-"), - }, - { - id: "user_email", - accessorKey: "user", - header: "User Email", - size: 160, - enableSorting: false, cell: (info) => { - const user = info.getValue() as any; - const value = user?.user_email; + const orgId = info.getValue() as string | null; + if (!orgId) return "-"; + const org = resolvedOrganizations.find((o) => o.organization_id === orgId); + const displayValue = org?.organization_alias || orgId; const width = info.cell.column.getSize(); return ( - - - {value ?? "-"} - - + + {displayValue} + ); }, }, { - id: "user_id", - accessorKey: "user_id", - header: "User ID", - size: 70, + id: "user", + accessorKey: "user", + header: () => ( + + User + + + + + ), + size: 160, enableSorting: false, - cell: (info) => { - const userId = info.getValue() as string | null; - const displayValue = userId === "default_user_id" ? "Default Proxy Admin" : userId; - const width = info.cell.column.getSize(); + cell: ({ row }) => { + const key = row.original; + const userAlias = key.user?.user_alias ?? null; + const userEmail = key.user?.user_email ?? key.user_email ?? null; + const userId = key.user_id ?? null; + const isDefaultAdmin = userId === "default_user_id"; + const displayValue = userAlias || userEmail || userId; + const width = 160; + + const popoverContent = ( +
+ {[ + { label: "User Alias", value: userAlias }, + { label: "User Email", value: userEmail }, + { label: "User ID", value: userId }, + ].map(({ label, value }) => ( +
+ {label} + {value ? ( + + {value} + + ) : ( + - + )} +
+ ))} +
+ ); + + if (isDefaultAdmin && !userAlias && !userEmail) { + return ( + + + + + + ); + } + return ( - - - {displayValue ?? "-"} + + + {displayValue || "-"} - + ); }, }, @@ -279,18 +316,48 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo id: "created_by", accessorKey: "created_by", header: "Created By", - size: 70, + size: 160, enableSorting: false, cell: (info) => { - const value = info.getValue() as string | null; - const displayValue = value === "default_user_id" ? "Default Proxy Admin" : value; - const width = info.cell.column.getSize(); + const userId = info.getValue() as string | null; + if (!userId) return "-"; + const isDefaultAdmin = userId === "default_user_id"; + const width = 160; + + const popoverContent = ( +
+
+ User ID + + {userId} + +
+
+ ); + + if (isDefaultAdmin) { + return ( + + + + + + ); + } + return ( - - - {displayValue ?? "-"} + + + {userId} - + ); }, }, @@ -477,7 +544,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo ); }, }, - ], []); + ], [teams, resolvedOrganizations]); const filterOptions: FilterOption[] = [ { @@ -535,8 +602,6 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo }, ]; - console.log(`keys: ${JSON.stringify(keys)}`); - const table = useReactTable({ data: filteredKeys, columns: columns.filter((col) => col.id !== "expander"), @@ -548,13 +613,11 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo }, onSortingChange: (updaterOrValue) => { const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue; - console.log(`newSorting: ${JSON.stringify(newSorting)}`); setSorting(newSorting); if (newSorting && newSorting.length > 0) { const sortState = newSorting[0]; const sortBy = sortState.id; const sortOrder = sortState.desc ? "desc" : "asc"; - console.log(`sortBy: ${sortBy}, sortOrder: ${sortOrder}`); // Update filters state without triggering debouncedSearch // The useKeys hook will automatically refetch with the new sort parameters handleFilterChange( diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx index a2f964df08f..4d0eb254d0d 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx @@ -1,6 +1,6 @@ import React, { useState, useEffect } from "react"; -import { Button as TremorButton, Text } from "@tremor/react"; -import { Modal, Table, Upload, Typography } from "antd"; +import { Text } from "@tremor/react"; +import { Button, Modal, Table, Upload, Typography } from "antd"; import { UploadOutlined, DownloadOutlined, @@ -540,9 +540,9 @@ const BulkCreateUsersButton: React.FC = ({ return ( <> - setIsModalVisible(true)}> + = ({
- - Download CSV Template - +
@@ -662,14 +662,14 @@ const BulkCreateUsersButton: React.FC = ({
- } > - Remove - + Remove + {fileError ? ( @@ -694,7 +694,7 @@ const BulkCreateUsersButton: React.FC = ({

Drag and drop your CSV file here

or

- Browse files +

Only CSV files (.csv) are supported

@@ -781,21 +781,21 @@ const BulkCreateUsersButton: React.FC = ({ {!parsedData.some((user) => user.status === "success" || user.status === "failed") && (
- { setParsedData([]); setParseError(null); }} - variant="secondary" > Back - - +
)} @@ -829,40 +829,39 @@ const BulkCreateUsersButton: React.FC = ({ {!parsedData.some((user) => user.status === "success" || user.status === "failed") && (
- { setParsedData([]); setParseError(null); }} - variant="secondary" className="mr-3" > Back - - +
)} {parsedData.some((user) => user.status === "success" || user.status === "failed") && (
- { setParsedData([]); setParseError(null); }} - variant="secondary" className="mr-3" > Start New Bulk Import - - - Download User Credentials - + +
)} diff --git a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx index 10bf6b61927..2ff07c7065b 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx @@ -8,6 +8,7 @@ import remarkGfm from "remark-gfm"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; import ReasoningContent from "../playground/chat_ui/ReasoningContent"; +import MCPEventsDisplay from "../playground/chat_ui/MCPEventsDisplay"; import { ChatMessage } from "./types"; const { Panel } = Collapse; @@ -237,6 +238,8 @@ interface AssistantBubbleProps { isLastMessage: boolean; isStreaming: boolean; isTypingIndicator: boolean; + /** MCP events stored on the message — rendered inline below the response. */ + mcpEvents?: ChatMessage["mcpEvents"]; } function AssistantBubble({ @@ -244,6 +247,7 @@ function AssistantBubble({ isLastMessage, isStreaming, isTypingIndicator, + mcpEvents, }: AssistantBubbleProps) { // Ref to control ReasoningContent collapse on streaming end. // ReasoningContent manages its own expanded state; we use a key to @@ -321,6 +325,11 @@ function AssistantBubble({ + {mcpEvents && mcpEvents.length > 0 && ( +
+ +
+ )} ); } @@ -566,6 +575,7 @@ const ChatMessages: React.FC = ({ messages, isStreaming, onEditMessage }) isLastMessage={isLastMessage} isStreaming={isStreaming} isTypingIndicator={isLastMessage && isTypingIndicator} + mcpEvents={msg.mcpEvents} /> ); })} diff --git a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx index 40b41f561ce..6a67814e2c6 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx @@ -26,6 +26,8 @@ import MCPConnectPicker from "./MCPConnectPicker"; import MCPAppsPanel from "./MCPAppsPanel"; import { fetchAvailableModels } from "../playground/llm_calls/fetch_models"; import { makeOpenAIChatCompletionRequest } from "../playground/llm_calls/chat_completion"; +import { makeOpenAIResponsesRequest } from "../playground/llm_calls/responses_api"; +import type { MCPEvent } from "./types"; import { getProxyBaseUrl } from "@/components/networking"; import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig"; import { getProviderLogoAndName } from "@/components/provider_info_helpers"; @@ -135,6 +137,7 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user const [modelSearchText, setModelSearchText] = useState(""); const [selectedMCPServers, setSelectedMCPServers] = useState([]); + const [responsesSessionId, setResponsesSessionId] = useState(null); const [isStreaming, setIsStreaming] = useState(false); const [inputText, setInputText] = useState(""); const [mcpPopoverOpen, setMcpPopoverOpen] = useState(false); @@ -162,7 +165,7 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user createConversation, appendMessage, updateLastAssistantMessage, - truncateAfterMessage, + truncateFromMessage, deleteConversation, renameConversation, } = useChatHistory(activeConversationId); @@ -203,6 +206,12 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user if (staleId) router.replace(getChatUrl(uiRoot)); }, [staleId, router]); + // Reset the responses session when switching between conversations so that + // previous_response_id from conversation A is never sent for conversation B. + useEffect(() => { + setResponsesSessionId(null); + }, [activeConversationId]); + const toggleModel = useCallback((model: string) => { setSelectedModels((prev) => { let next: string[]; @@ -231,6 +240,7 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user let convId = activeConversationId; if (!convId) { convId = createConversation(model); + setResponsesSessionId(null); // new conversation starts a fresh session router.push(getChatUrl(uiRoot, convId)); } @@ -240,29 +250,56 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user setIsStreaming(true); abortControllerRef.current = new AbortController(); - const history = [ - ...(historyOverride ?? (activeConversation?.messages ?? []) - .filter((m) => m.role === "user" || m.role === "assistant") - .map((m) => ({ - role: m.role as "user" | "assistant", - content: m.content, - }))), - { role: "user" as const, content: trimmed }, - ]; + // When historyOverride is set (edit / retry), the existing server-side + // session chain covers messages that were just truncated and is no longer + // valid for the rewritten history. Eagerly clear the session so that a + // failed/aborted edit does not leave a stale session ID that contaminates + // the next regular send. + if (historyOverride) { + setResponsesSessionId(null); + } + + // On a normal continuation turn with an active session, the Responses API + // already holds the prior context server-side, so we only pass the new + // user message (sending the full history would double-count it). + // + // On the very first turn (no session yet), we send the full history. + const previousResponseId = historyOverride ? null : responsesSessionId; + + const history: Array<{ role: "user" | "assistant"; content: string }> = + historyOverride + ? [...historyOverride, { role: "user" as const, content: trimmed }] + : previousResponseId + ? [{ role: "user" as const, content: trimmed }] + : [ + // Explicitly filter to only user/assistant roles — tool messages + // lack a required tool_call_id and would cause API errors. + ...(activeConversation?.messages ?? []) + .filter((m): m is typeof m & { role: "user" | "assistant" } => + m.role === "user" || m.role === "assistant" + ) + .map((m) => ({ role: m.role, content: m.content })), + { role: "user" as const, content: trimmed }, + ]; let accumulatedContent = ""; let accumulatedReasoning = ""; + // MCP events accumulated locally so we can persist them to the message + // without relying on component state (which would cause stale closures). + const accumulatedMCPEvents: MCPEvent[] = []; + // Track clean completion so partial events are not shown on error/abort. + let streamCompletedCleanly = false; try { - await makeOpenAIChatCompletionRequest( + await makeOpenAIResponsesRequest( history, - (chunk: string) => { + (_role: string, chunk: string) => { accumulatedContent += chunk; updateLastAssistantMessage(convId!, { content: accumulatedContent }); }, model, accessToken, - undefined, + undefined, // tags abortControllerRef.current.signal, (rc: string) => { accumulatedReasoning += rc; @@ -270,7 +307,15 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user }, undefined, undefined, undefined, undefined, undefined, undefined, selectedMCPServers.length > 0 ? selectedMCPServers : undefined, + previousResponseId, + (id: string) => setResponsesSessionId(id), + (event: MCPEvent) => { + // Accumulate locally only — persisted once in finally to avoid + // one full localStorage write per MCP event during streaming. + accumulatedMCPEvents.push(event); + }, ); + streamCompletedCleanly = true; } catch (err: unknown) { if (err instanceof Error && err.name === "AbortError") { updateLastAssistantMessage(convId!, { @@ -282,12 +327,17 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user }); } } finally { + // Only persist MCP events on clean completion — partial events from an + // aborted or errored turn would show incomplete tool calls to the user. + if (accumulatedMCPEvents.length > 0 && streamCompletedCleanly) { + updateLastAssistantMessage(convId!, { mcpEvents: accumulatedMCPEvents }); + } setIsStreaming(false); abortControllerRef.current = null; } }, [activeConversationId, activeConversation, selectedModels, selectedMCPServers, accessToken, - createConversation, appendMessage, updateLastAssistantMessage, router, isStreaming], + createConversation, appendMessage, updateLastAssistantMessage, router, isStreaming, responsesSessionId], ); const handleSendComparison = useCallback( @@ -355,10 +405,10 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user const priorMessages = (idx === -1 ? msgs : msgs.slice(0, idx)) .filter((m) => m.role === "user" || m.role === "assistant") .map((m) => ({ role: m.role as "user" | "assistant", content: m.content })); - truncateAfterMessage(activeConversationId, messageId); + truncateFromMessage(activeConversationId, messageId); handleSend(newContent, priorMessages); }, - [activeConversationId, isStreaming, activeConversation, truncateAfterMessage, handleSend], + [activeConversationId, isStreaming, activeConversation, truncateFromMessage, handleSend], ); const handleSubmit = useCallback( @@ -944,10 +994,20 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user : greeting} - {isComparisonMode && ( + {isComparisonMode ? (

Send a message to see responses side-by-side

+ ) : ( +

+ Chat with 100+ LLMs + MCP tools — authenticate once, use them here.{" "} + +

)} {/* Input card */} diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 926d6697bca..625478e98a4 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -1,10 +1,10 @@ "use client"; import React, { useEffect, useState } from "react"; -import { Switch, Spin, Input, Button } from "antd"; -import { SearchOutlined, ArrowLeftOutlined, RightOutlined } from "@ant-design/icons"; +import { Spin, Input, Button, Skeleton } from "antd"; +import { SearchOutlined, ArrowLeftOutlined, RightOutlined, ToolOutlined } from "@ant-design/icons"; import { fetchMCPServers, listMCPTools } from "../networking"; -import { MCPServer } from "../mcp_tools/types"; +import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types"; import { message } from "antd"; interface Props { @@ -33,33 +33,65 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange const [activeTab, setActiveTab] = useState("all"); const [togglingOn, setTogglingOn] = useState>(new Set()); const [detailServer, setDetailServer] = useState(null); + const [detailTools, setDetailTools] = useState([]); + const [loadingTools, setLoadingTools] = useState(false); + // tool counts per server name, preloaded in background + const [toolCounts, setToolCounts] = useState>({}); + const [loadingCounts, setLoadingCounts] = useState(false); + + const nameOf = (s: MCPServer) => s.server_name ?? s.alias ?? s.server_id; useEffect(() => { let cancelled = false; setLoading(true); + + // 1. Load servers first — show the list immediately fetchMCPServers(accessToken) - .then((data) => { + .then((serverData) => { if (cancelled) return; - const list: MCPServer[] = Array.isArray(data) ? data : (data?.data ?? []); + const list: MCPServer[] = Array.isArray(serverData) ? serverData : (serverData?.data ?? []); setServers(list); + setLoading(false); + + // 2. Fetch tools per server in parallel — each resolves independently and updates counts one by one + setLoadingCounts(true); + let remaining = list.length; + if (remaining === 0) { setLoadingCounts(false); return; } + list.forEach((s) => { + listMCPTools(accessToken, s.server_id) + .then((toolsData) => { + if (cancelled) return; + const tools: MCPTool[] = Array.isArray(toolsData?.tools) ? toolsData.tools : []; + const sname = nameOf(s); + setToolCounts((prev) => ({ ...prev, [sname]: tools.length })); + }) + .catch(() => {}) + .finally(() => { + if (cancelled) return; + remaining -= 1; + if (remaining === 0) setLoadingCounts(false); + }); + }); }) .catch(() => { - if (!cancelled) setServers([]); - }) - .finally(() => { - if (!cancelled) setLoading(false); + if (!cancelled) { + setServers([]); + setLoading(false); + } }); return () => { cancelled = true; }; }, [accessToken]); - const handleToggle = async (serverName: string, checked: boolean) => { + const handleToggle = async (serverName: string, checked: boolean, serverId?: string) => { if (!checked) { onChange(selectedServers.filter((s) => s !== serverName)); return; } setTogglingOn((prev) => new Set(prev).add(serverName)); try { - const result = await listMCPTools(accessToken, serverName); + // Use UUID if available, fall back to name (for connectivity check only) + const idToFetch = serverId ?? serverName; + const result = await listMCPTools(accessToken, idToFetch); if (result?.error) { message.warning(`Could not load tools for ${serverName}`); return; @@ -76,7 +108,29 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange } }; - const nameOf = (s: MCPServer) => s.server_name ?? s.alias ?? s.server_id; + // Fetch tools for the detail view — server_id must be the UUID + useEffect(() => { + if (!detailServer) { + setDetailTools([]); + return; + } + let cancelled = false; + setLoadingTools(true); + listMCPTools(accessToken, detailServer.server_id) + .then((result) => { + if (cancelled) return; + // API returns { tools: [...], error: null } + const tools: MCPTool[] = Array.isArray(result?.tools) ? result.tools : []; + setDetailTools(tools); + }) + .catch(() => { + if (!cancelled) setDetailTools([]); + }) + .finally(() => { + if (!cancelled) setLoadingTools(false); + }); + return () => { cancelled = true; }; + }, [detailServer, accessToken]); const filtered = servers.filter((s) => { const name = nameOf(s); @@ -89,6 +143,9 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange const connectedCount = servers.filter((s) => selectedServers.includes(nameOf(s))).length; + // Total tools available across all servers (based on preloaded counts) + const totalTools = Object.values(toolCounts).reduce((sum, n) => sum + n, 0); + // ── Detail view ── if (detailServer) { const name = nameOf(detailServer); @@ -113,9 +170,25 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange {/* Avatar + name + connect */}
+ {detailServer.mcp_info?.logo_url ? ( + {`${name} { + const el = e.target as HTMLImageElement; + el.style.display = "none"; + if (el.nextElementSibling) (el.nextElementSibling as HTMLElement).style.display = "flex"; + }} + /> + ) : null}
@@ -128,7 +201,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange +
+ )} + + {/* Well-known logo grid */} +
+ {WELL_KNOWN_LOGOS.map((logo) => { + const isSelected = value === logo.url; + const hasFailed = imgErrors.has(logo.url); + if (hasFailed) return null; + return ( + + + + ); + })} +
+ + {/* Custom URL input */} + } + placeholder="Or paste a custom logo URL..." + value={value && !WELL_KNOWN_LOGOS.some((l) => l.url === value) ? value : ""} + onChange={(e) => { + const v = e.target.value.trim(); + onChange?.(v || undefined); + }} + className="rounded-lg" + size="small" + /> +
+ ); +}; + +export default MCPLogoSelector; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.tsx new file mode 100644 index 00000000000..fb38e392631 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.tsx @@ -0,0 +1,72 @@ +"use client"; + +import { MCPServer } from "./types"; + +export interface RequiredFieldDef { + key: string; + label: string; + description: string; + check: (server: MCPServer) => boolean; +} + +export interface FieldGroup { + label: string; + fields: RequiredFieldDef[]; +} + +export const FIELD_GROUPS: FieldGroup[] = [ + { + label: "Documentation", + fields: [ + { + key: "description", + label: "Description", + description: "Must have a non-empty description", + check: (s) => !!s.description?.trim(), + }, + { + key: "alias", + label: "Alias", + description: "Must have a display alias", + check: (s) => !!s.alias?.trim(), + }, + ], + }, + { + label: "Source", + fields: [ + { + key: "source_url", + label: "GitHub / Source URL", + description: "Must link to a source repository", + check: (s) => !!s.source_url?.trim(), + }, + ], + }, + { + label: "Connection", + fields: [ + { + key: "url", + label: "Server URL", + description: "Must have a URL configured", + check: (s) => !!s.url?.trim(), + }, + ], + }, + { + label: "Security", + fields: [ + { + key: "auth_type", + label: "Auth configured", + description: "Must use authentication (not 'none')", + check: (s) => !!s.auth_type && s.auth_type !== "none", + }, + ], + }, +]; + +export const MCP_REQUIRED_FIELD_DEFS: RequiredFieldDef[] = FIELD_GROUPS.flatMap((g) => g.fields); + +export const SETTINGS_KEY = "mcp_required_fields"; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPSubmissionsTab.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPSubmissionsTab.tsx new file mode 100644 index 00000000000..4ce7423f1b5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPSubmissionsTab.tsx @@ -0,0 +1,660 @@ +"use client"; + +import React, { useState, useEffect, useCallback } from "react"; +import { + SearchIcon, + CheckIcon, + XIcon, + AlertCircleIcon, + ServerIcon, + ChevronDownIcon, + ChevronUpIcon, + SettingsIcon, +} from "lucide-react"; +import { + fetchMCPSubmissions, + approveMCPServer, + rejectMCPServer, + getGeneralSettingsCall, + updateConfigFieldSetting, +} from "@/components/networking"; +import { MCPServer, MCPSubmissionsSummary } from "./types"; +import { FIELD_GROUPS, MCP_REQUIRED_FIELD_DEFS, SETTINGS_KEY } from "./MCPStandardsSettings"; +import NotificationsManager from "@/components/molecules/notifications_manager"; + +type MCPStatus = "active" | "pending_review" | "rejected"; + +const STATUS_CONFIG: Record< + MCPStatus, + { label: string; bg: string; text: string; dot: string } +> = { + active: { + label: "Active", + bg: "bg-green-50", + text: "text-green-700", + dot: "bg-green-500", + }, + pending_review: { + label: "Pending Review", + bg: "bg-yellow-50", + text: "text-yellow-700", + dot: "bg-yellow-500", + }, + rejected: { + label: "Rejected", + bg: "bg-red-50", + text: "text-red-700", + dot: "bg-red-500", + }, +}; + +function formatDate(value: string | null | undefined): string { + if (!value) return "—"; + try { + const d = new Date(value); + return isNaN(d.getTime()) ? value : d.toISOString().slice(0, 10); + } catch { + return value; + } +} + +function StatCard({ + label, + value, + color, +}: { + label: string; + value: number; + color: string; +}) { + return ( +
+
{value}
+
{label}
+
+ ); +} + +type ConfirmDialogProps = { + action: "approve" | "reject"; + serverName: string; + isCurrentlyActive?: boolean; + onConfirm: (reviewNotes?: string) => void; + onCancel: () => void; +}; + +function ConfirmDialog({ action, serverName, isCurrentlyActive, onConfirm, onCancel }: ConfirmDialogProps) { + const [reviewNotes, setReviewNotes] = useState(""); + const isApprove = action === "approve"; + const rejectBody = isCurrentlyActive + ? "This server is currently live. Rejecting it will immediately remove it from the proxy runtime." + : "This will mark the submission as rejected."; + return ( +
+
+
+ {isApprove ? ( + + ) : ( + + )} +
+

+ {isApprove ? "Approve MCP Server" : "Reject MCP Server"} +

+

+ Are you sure you want to {action}{" "} + "{serverName}"?{" "} + {isApprove + ? "This will make it active and available for use." + : rejectBody} +

+ {!isApprove && ( +