Merge branch 'main' into litellm_oss_staging_03_04_2026

This commit is contained in:
Sameer Kankute 2026-03-11 18:31:20 +05:30 committed by GitHub
commit 3dab62023c
127 changed files with 8522 additions and 1026 deletions

View file

@ -34,8 +34,6 @@ jobs:
build-mode: none
- language: python
build-mode: none
- language: ruby
build-mode: none
steps:
- name: Checkout repository

View file

@ -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 }}

View file

@ -35,6 +35,14 @@ deploymentLabels: {}
podAnnotations: {}
podLabels: {}
# -- Deployment strategy configuration
# Example:
# type: RollingUpdate
# rollingUpdate:
# maxUnavailable: 0
# maxSurge: 1
strategy: {}
terminationGracePeriodSeconds: 90
topologySpreadConstraints:
[]

View file

@ -217,6 +217,7 @@ mcp_servers:
| `bearer_token` | `Authorization: Bearer <auth_value>` |
| `basic` | `Authorization: Basic <auth_value>` |
| `authorization` | `Authorization: <auth_value>` |
| `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/<url-encoded-ARN>/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"

View file

@ -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/<url-encoded-ARN>/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/<url-encoded-ARN>/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.

View file

@ -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.
:::

View file

@ -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",
)
```
:::
<Tabs>
<TabItem value="sdk" label="SDK">

View file

@ -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.
<Tabs>
<TabItem value="sdk" label="LiteLLM Python SDK">
```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})")
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
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)
```
</TabItem>
</Tabs>
### 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.
<Tabs>
<TabItem value="sdk" label="LiteLLM Python SDK">
```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})")
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
```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
}
]
}
]
}'
```
</TabItem>
</Tabs>
## Free-form Function Calling
<Tabs>

View file

@ -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.

View file

@ -592,6 +592,12 @@ Expected Response
</TabItem>
</Tabs>
:::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.

View file

@ -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 <host> -U <user> -d <db> -F c -f backup_$(date +%Y%m%d).dump
```
### 3. Upgrade the package
```bash
pip install 'litellm[proxy]==<version>'
```
### 4. Regenerate the Prisma client
```bash
prisma generate --schema <venv>/lib/python<version>/site-packages/litellm_proxy_extras/schema.prisma
```
Replace `<venv>` with your virtual environment path and `<version>` 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 <venv>/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 <venv>/bin/activate
```
Then run the migration with the explicit schema path:
```bash
prisma migrate deploy --schema <venv>/lib/python<version>/site-packages/litellm_proxy_extras/schema.prisma
```
Replace `<venv>` with your virtual environment path and `<version>` 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:** `<schema-path>` = `<venv>/lib/python<version>/site-packages/litellm_proxy_extras/schema.prisma`
### Before applying migrations: Preview what will change
Run `pip install 'litellm[proxy]==<version>'` first (Step 3) so the new `schema.prisma` is available.
```bash
prisma migrate diff \
--from-url $DATABASE_URL \
--to-schema-datamodel <schema-path> \
--script
```
### After applying migrations: Check status
```bash
prisma migrate status --schema <schema-path>
```
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).

View file

@ -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)**

View file

@ -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",
],

View file

@ -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");

View file

@ -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;

View file

@ -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,
)

View file

@ -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

View file

@ -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,

View file

@ -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

View file

@ -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

View file

@ -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",

View file

@ -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,
)

View file

@ -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

View file

@ -0,0 +1,7 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M8 24c2.208 0 4-1.792 4-4v-4H8c-2.208 0-4 1.792-4 4s1.792 4 4 4z" fill="#0ACF83"/>
<path d="M4 12c0-2.208 1.792-4 4-4h4v8H8c-2.208 0-4-1.792-4-4z" fill="#A259FF"/>
<path d="M4 4c0-2.208 1.792-4 4-4h4v8H8C5.792 8 4 6.208 4 4z" fill="#F24E1E"/>
<path d="M12 0h4c2.208 0 4 1.792 4 4s-1.792 4-4 4h-4V0z" fill="#FF7262"/>
<path d="M20 12c0 2.208-1.792 4-4 4s-4-1.792-4-4 1.792-4 4-4 4 1.792 4 4z" fill="#1ABCFE"/>
</svg>

After

Width:  |  Height:  |  Size: 496 B

View file

@ -0,0 +1,8 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M23.955 13.587l-1.342-4.135-2.664-8.189a.455.455 0 0 0-.867 0L16.418 9.45H7.582L4.918 1.263a.455.455 0 0 0-.867 0L1.386 9.452.044 13.587a.924.924 0 0 0 .331 1.03L12 23.054l11.625-8.436a.92.92 0 0 0 .33-1.03z" fill="#E24329"/>
<path d="M12 23.054L16.418 9.45H7.582L12 23.054z" fill="#FC6D26"/>
<path d="M12 23.054L7.582 9.452H1.386L12 23.054z" fill="#FCA326"/>
<path d="M1.386 9.452L.044 13.587a.924.924 0 0 0 .331 1.03L12 23.054 1.386 9.452z" fill="#E24329"/>
<path d="M12 23.054l4.418-13.602h5.036L12 23.054z" fill="#FCA326"/>
<path d="M22.614 9.452l1.341 4.135a.924.924 0 0 1-.33 1.03L12 23.054l10.614-13.602z" fill="#E24329"/>
</svg>

After

Width:  |  Height:  |  Size: 719 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M24 5.457v13.909c0 .904-.732 1.636-1.636 1.636h-3.819V11.73L12 16.64l-6.545-4.91v9.273H1.636A1.636 1.636 0 0 1 0 19.366V5.457c0-2.023 2.309-3.178 3.927-1.964L5.455 4.64 12 9.548l6.545-4.91 1.528-1.145C21.69 2.28 24 3.434 24 5.457z" fill="#EA4335"/>
</svg>

After

Width:  |  Height:  |  Size: 328 B

View file

@ -0,0 +1,6 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M7.71 14.29L2 22h7.65l5.71-7.71H7.71z" fill="#0066DA"/>
<path d="M22 22l-5.71-7.71H8.65L14.35 22H22z" fill="#00AC47"/>
<path d="M8.16 2L2.45 14.29h7.65L15.81 2H8.16z" fill="#FFBA00"/>
<path d="M15.84 2l-5.71 12.29h7.65L23.49 2H15.84z" fill="#EA4335"/>
</svg>

After

Width:  |  Height:  |  Size: 337 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M18.164 7.93V5.084a2.198 2.198 0 0 0 1.267-1.984v-.066A2.2 2.2 0 0 0 17.237.84h-.066a2.2 2.2 0 0 0-2.194 2.194v.066c0 .844.48 1.574 1.18 1.94V7.93a6.152 6.152 0 0 0-2.866 1.388L5.962 3.72a2.385 2.385 0 0 0 .07-.557A2.37 2.37 0 0 0 3.662.793a2.37 2.37 0 0 0-2.37 2.37 2.37 2.37 0 0 0 2.37 2.37c.432 0 .836-.12 1.183-.325l7.47 5.5A6.175 6.175 0 0 0 11.19 14.1a6.2 6.2 0 0 0 1.097 3.504l-2.12 2.12a1.786 1.786 0 0 0-.52-.082A1.803 1.803 0 1 0 11.45 21.445l2.172-2.172a6.175 6.175 0 0 0 3.572 1.135 6.2 6.2 0 1 0 .97-12.478zM17.204 17.2a3.004 3.004 0 1 1 0-6.008 3.004 3.004 0 0 1 0 6.008z" fill="#FF7A59"/>
</svg>

After

Width:  |  Height:  |  Size: 683 B

View file

@ -0,0 +1,15 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M11.571 11.513H0a5.218 5.218 0 0 0 5.232 5.215h2.13v2.057A5.215 5.215 0 0 0 12.575 24V12.518a1.005 1.005 0 0 0-1.005-1.005z" fill="#2684FF"/>
<path d="M6.262 6.259H17.793a5.218 5.218 0 0 0-5.232-5.214H6.26A5.218 5.218 0 0 0 1.03 6.259v6.268a1.005 1.005 0 0 0 1.005 1.005h9.527V8.318a2.06 2.06 0 0 0-2.06-2.059H6.262z" fill="url(#jiraGrad1)"/>
<path d="M17.53 6.259a5.218 5.218 0 0 1 5.232 5.214v6.268a1.005 1.005 0 0 1-1.005 1.005H12.23v-5.214a2.06 2.06 0 0 1 2.06-2.059h3.24V6.259z" fill="url(#jiraGrad2)"/>
<defs>
<linearGradient id="jiraGrad1" x1="6" y1="1" x2="1" y2="13">
<stop stop-color="#0052CC"/>
<stop offset="1" stop-color="#2684FF"/>
</linearGradient>
<linearGradient id="jiraGrad2" x1="18" y1="6" x2="23" y2="18">
<stop stop-color="#0052CC"/>
<stop offset="1" stop-color="#2684FF"/>
</linearGradient>
</defs>
</svg>

After

Width:  |  Height:  |  Size: 949 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M2.886 10.449a.64.64 0 0 1 .078.848L1.337 13.56a.32.32 0 0 1-.526-.036A11.044 11.044 0 0 1 .103 11.53a.32.32 0 0 1 .238-.37l2.186-.535a.64.64 0 0 1 .36-.176zm-.803 3.753a.32.32 0 0 0-.492.097 11.127 11.127 0 0 0-.747 2.148.32.32 0 0 0 .108.331l1.638 1.319a.32.32 0 0 0 .476-.07l1.18-1.87a.64.64 0 0 0-.033-.733L2.083 14.2zm.79 5.275a.32.32 0 0 0-.42.02l-.367.354a.32.32 0 0 0-.022.43 11.18 11.18 0 0 0 2.854 2.67.32.32 0 0 0 .42-.044l.404-.431a.32.32 0 0 0-.006-.445L2.873 19.477zM7.55 22.553a.32.32 0 0 0 .014.458l.28.232a.32.32 0 0 0 .427-.023 11.2 11.2 0 0 0 3.084-4.807.32.32 0 0 0-.19-.394l-.672-.25a.32.32 0 0 0-.405.18A9.614 9.614 0 0 1 7.55 22.553zM23.988 12c0 6.627-5.373 12-12 12-.742 0-1.47-.067-2.176-.197a.32.32 0 0 1-.218-.494L22.006 8.523a.32.32 0 0 1 .552.108c.283.882.43 1.818.43 2.79V12zm-.836-4.488a.32.32 0 0 0-.56-.05L10.163 22.68a.32.32 0 0 0 .096.472c.83.43 1.73.74 2.679.914a.32.32 0 0 0 .333-.14L23.195 8.204a.32.32 0 0 0-.043-.412v-.28zM20.684 5.216a.32.32 0 0 0 .484.013l.257-.269a.32.32 0 0 0 .017-.427A11.955 11.955 0 0 0 12 .007C5.373.007 0 5.38 0 12.007c0 .414.021.824.063 1.229a.32.32 0 0 0 .547.197l18.13-20.07a.32.32 0 0 1 .453-.014l1.49 1.368v.5z" fill="#5E6AD2"/>
</svg>

After

Width:  |  Height:  |  Size: 1.2 KiB

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M4.459 4.208c.746.606 1.026.56 2.428.466l13.215-.793c.28 0 .047-.28-.046-.326L18.57 2.33c-.42-.326-.98-.7-2.055-.607L3.62 2.79c-.466.046-.56.28-.374.466l1.213.952zm.793 3.08v13.904c0 .747.373 1.027 1.214.98l14.523-.84c.84-.046.933-.56.933-1.167V6.354c0-.606-.233-.933-.746-.886l-15.177.886c-.56.047-.747.327-.747.934zm14.337.745c.093.42 0 .84-.42.888l-.7.14v10.264c-.608.327-1.168.514-1.635.514-.747 0-.933-.234-1.494-.934l-4.577-7.186v6.952l1.448.327s0 .84-1.168.84l-3.222.186c-.093-.186 0-.653.327-.746l.84-.233V9.854L7.822 9.76c-.094-.42.14-1.026.793-1.073l3.456-.233 4.764 7.279v-6.44l-1.215-.14c-.093-.513.28-.886.747-.933l3.222-.187zM2.1 1.424L15.856.466c1.68-.14 2.1.093 2.8.606l3.876 2.753c.467.326.607.42.607.793v16.844c0 1.026-.374 1.633-1.68 1.726l-15.457.933c-.98.047-1.448-.093-1.962-.747l-3.13-4.06c-.56-.747-.793-1.306-.793-1.96V2.917c0-.84.374-1.54 1.782-1.493z" fill="#000000"/>
</svg>

After

Width:  |  Height:  |  Size: 975 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M10.006 5.15a4.678 4.678 0 0 1 3.358-1.426 4.7 4.7 0 0 1 4.273 2.775 5.476 5.476 0 0 1 2.163-.444C22.108 6.055 24 7.954 24 10.288a4.258 4.258 0 0 1-3.8 4.225h-.2a3.756 3.756 0 0 1-3.468 2.308 3.726 3.726 0 0 1-1.76-.44 4.418 4.418 0 0 1-3.89 2.32 4.418 4.418 0 0 1-3.798-2.15 3.678 3.678 0 0 1-.844.098A3.68 3.68 0 0 1 2.56 12.97c0-.77.237-1.484.642-2.074A4.448 4.448 0 0 1 0 7.2a4.448 4.448 0 0 1 4.448-4.448c1.11 0 2.13.41 2.91 1.086A4.672 4.672 0 0 1 10.006 5.15z" fill="#00A1E0"/>
</svg>

After

Width:  |  Height:  |  Size: 564 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M13.91 2.505c-.873-1.553-3.066-1.553-3.94 0L7.092 7.67a10.783 10.783 0 0 1 4.376 3.033l1.72-3.027a4.23 4.23 0 0 0-.207-.257L13.91 2.505zM3.6 20.3h2.685a8.782 8.782 0 0 0-.094-5.725l-2.59 4.613c-.25.445.109 1.003.613 1.003L3.6 20.3zm16.595.108c.5 0 .862-.557.613-1.003L14.075 7.5l-1.72 3.027A8.782 8.782 0 0 1 17.35 20.3h2.845v.108z" fill="#362D59"/>
</svg>

After

Width:  |  Height:  |  Size: 429 B

View file

@ -0,0 +1,4 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M15.337 23.979l7.216-1.561s-2.604-17.613-2.625-17.73c-.018-.116-.114-.2-.2-.2s-1.848-.138-1.848-.138-1.225-1.197-1.363-1.338c-.04-.04-.085-.06-.132-.074l-.793 18.867 -.255.174zm-3.07-16.938s-.718-.378-1.594-.378c-1.29 0-1.353.809-1.353 1.013 0 1.112 2.9 1.538 2.9 4.145 0 2.051-1.3 3.37-3.053 3.37-2.105 0-3.18-1.31-3.18-1.31l.564-1.865s1.105.949 2.036.949c.607 0 .856-.479.856-.829 0-1.453-2.38-1.519-2.38-3.906 0-2.008 1.441-3.953 4.351-3.953 1.12 0 1.674.321 1.674.321l-.82 2.443z" fill="#95BF47"/>
<path d="M14.998 6.268c-.082-.025-.18-.048-.252-.048-.028 0-.062.002-.092.006l-.76 18.07.663-.143 3.377-23.182c-.138.14-1.363 1.337-1.363 1.337s-.95-.093-1.573-.04z" fill="#5E8E3E"/>
</svg>

After

Width:  |  Height:  |  Size: 766 B

View file

@ -0,0 +1,6 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none">
<path d="M5.042 15.165a2.528 2.528 0 0 1-2.52 2.523A2.528 2.528 0 0 1 0 15.165a2.527 2.527 0 0 1 2.522-2.52h2.52v2.52zm1.271 0a2.527 2.527 0 0 1 2.521-2.52 2.527 2.527 0 0 1 2.521 2.52v6.313A2.528 2.528 0 0 1 8.834 24a2.528 2.528 0 0 1-2.521-2.522v-6.313z" fill="#E01E5A"/>
<path d="M8.834 5.042a2.528 2.528 0 0 1-2.521-2.52A2.528 2.528 0 0 1 8.834 0a2.528 2.528 0 0 1 2.521 2.522v2.52H8.834zm0 1.271a2.528 2.528 0 0 1 2.521 2.521 2.528 2.528 0 0 1-2.521 2.521H2.522A2.528 2.528 0 0 1 0 8.834a2.528 2.528 0 0 1 2.522-2.521h6.312z" fill="#36C5F0"/>
<path d="M18.956 8.834a2.528 2.528 0 0 1 2.522-2.521A2.528 2.528 0 0 1 24 8.834a2.528 2.528 0 0 1-2.522 2.521h-2.522V8.834zm-1.27 0a2.528 2.528 0 0 1-2.522 2.521 2.527 2.527 0 0 1-2.521-2.521V2.522A2.527 2.527 0 0 1 15.165 0a2.528 2.528 0 0 1 2.522 2.522v6.312z" fill="#2EB67D"/>
<path d="M15.165 18.956a2.528 2.528 0 0 1 2.522 2.522A2.528 2.528 0 0 1 15.165 24a2.527 2.527 0 0 1-2.521-2.522v-2.522h2.521zm0-1.27a2.527 2.527 0 0 1-2.521-2.522 2.527 2.527 0 0 1 2.521-2.521h6.313A2.528 2.528 0 0 1 24 15.165a2.528 2.528 0 0 1-2.522 2.521h-6.313z" fill="#ECB22E"/>
</svg>

After

Width:  |  Height:  |  Size: 1.2 KiB

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M13.976 9.15c-2.172-.806-3.356-1.426-3.356-2.409 0-.831.683-1.305 1.901-1.305 2.227 0 4.515.858 6.09 1.631l.89-5.494C18.252.975 15.697 0 12.165 0 9.667 0 7.589.654 6.104 1.872 4.56 3.147 3.757 4.992 3.757 7.218c0 4.039 2.467 5.76 6.476 7.219 2.585.92 3.445 1.574 3.445 2.583 0 .98-.84 1.545-2.354 1.545-1.875 0-4.965-.921-6.99-2.109l-.9 5.555C5.175 22.99 8.385 24 11.714 24c2.641 0 4.843-.624 6.328-1.813 1.664-1.305 2.525-3.236 2.525-5.732 0-4.128-2.524-5.851-6.591-7.305z" fill="#635BFF"/>
</svg>

After

Width:  |  Height:  |  Size: 571 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M12 0C5.381 0 0 5.381 0 12s5.381 12 12 12 12-5.381 12-12S18.619 0 12 0zm0 20.4c-4.636 0-8.4-3.764-8.4-8.4S7.364 3.6 12 3.6s8.4 3.764 8.4 8.4-3.764 8.4-8.4 8.4zm3.6-10.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8zm0 4.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8zm-4.8-4.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8zm0 4.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8z" fill="#F22F46"/>
</svg>

After

Width:  |  Height:  |  Size: 587 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M15.535 8.465l-1.263 1.264a4.502 4.502 0 0 0-4.544 0L8.465 8.465a6.51 6.51 0 0 1 2.534-1.796V4.5h2v2.17a6.51 6.51 0 0 1 2.536 1.795zM12 10.5a1.5 1.5 0 1 0 0 3 1.5 1.5 0 0 0 0-3zm3.535 5.035l1.263-1.264a6.51 6.51 0 0 1-1.796 2.536V19h-2v-2.17a6.51 6.51 0 0 1-2.534-1.795l1.263-1.264a4.502 4.502 0 0 0 4.544 0h-.74zM19 11h-2.17a6.51 6.51 0 0 1-1.795 2.534l1.264 1.263a4.502 4.502 0 0 0 0-4.544l-1.264 1.263A6.51 6.51 0 0 1 16.831 11H19v2zm-12.17 2H5v-2h2.17a6.51 6.51 0 0 1 1.795-2.534L7.7 7.203a4.502 4.502 0 0 0 0 4.544l1.264-1.263A6.51 6.51 0 0 1 7.17 13H6.83z" fill="#FF4A00"/>
</svg>

After

Width:  |  Height:  |  Size: 659 B

View file

@ -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):

View file

@ -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: <PHONE_NUMBER_1>, <PHONE_NUMBER_2>
# 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. <PHONE_NUMBER_1>) 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

View file

@ -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
)

View file

@ -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": []}

View file

@ -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,
)

View file

@ -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
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},
)

View file

@ -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" }
]
}
]
}

View file

@ -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 = []

View file

@ -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

View file

@ -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:

View file

@ -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,
]
]

View file

@ -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:

View file

@ -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),

View file

@ -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",

View file

@ -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": {"<PERSON>": "Jane Doe", "<PHONE_NUMBER>": "034453334"},
"messages": [
{"role": "system", "content": "You are an helpfull assistant"}
],
"metadata": {
"pii_tokens": {"<PERSON>": "Jane Doe", "<PHONE_NUMBER>": "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

View file

@ -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__])

View file

@ -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)."
)

View file

@ -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."""

View file

@ -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: <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

View file

@ -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"

View file

@ -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") == {}

View file

@ -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__])

View file

@ -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

View file

@ -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"<PERSON_{seq}>"
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 = "<PERSON_1>"
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 = "<PERSON_1>"
# 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 = {
"<PERSON_1>": "John Smith",
"<PHONE_NUMBER_1>": "555-123-4567",
}
text = "Hello <PERSON_1>, your number is <PHONE_NUMBER_1>."
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 = {
"<PHONE_NUMBER_1>": "555-111-0000",
"<PHONE_NUMBER_2>": "555-222-0000",
}
text = "Call <PHONE_NUMBER_1> or <PHONE_NUMBER_2>."
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 = {
"<PERSON_1>": "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 <PERSON> at <PHONE_NUMBER>",
"items": [
{
"start": 5,
"end": 9,
"entity_type": "PERSON",
"text": "<PERSON>",
"operator": "replace",
},
{
"start": 13,
"end": 25,
"entity_type": "PHONE_NUMBER",
"text": "<PHONE_NUMBER>",
"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": {
"<PERSON_1>": "John Smith",
"<PHONE_NUMBER_1>": "555-123-4567",
}
},
}
anthropic_response = {
"type": "message",
"id": "msg_123",
"model": "claude-3-haiku",
"role": "assistant",
"content": [
{
"type": "text",
"text": "Hello <PERSON_1>, your number is <PHONE_NUMBER_1>.",
}
],
"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": {"<PERSON_1>": "John"}},
}
anthropic_response = {
"type": "message",
"id": "msg_123",
"content": [
{"type": "text", "text": "Hello <PERSON_1>"},
{
"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": {"<PERSON_1>": "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": {
"<PERSON_1>": "John Smith",
"<PHONE_NUMBER_1>": "555-123-4567",
}
},
}
inputs = {
"texts": [
"Hello <PERSON_1>, your number is <PHONE_NUMBER_1>.",
]
}
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", "<PERSON>")
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 "<PERSON>" 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

View file

@ -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."""

View file

@ -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(

View file

@ -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)

View file

@ -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"}

View file

@ -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"
)

View file

@ -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"

View file

@ -0,0 +1,7 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M8 24c2.208 0 4-1.792 4-4v-4H8c-2.208 0-4 1.792-4 4s1.792 4 4 4z" fill="#0ACF83"/>
<path d="M4 12c0-2.208 1.792-4 4-4h4v8H8c-2.208 0-4-1.792-4-4z" fill="#A259FF"/>
<path d="M4 4c0-2.208 1.792-4 4-4h4v8H8C5.792 8 4 6.208 4 4z" fill="#F24E1E"/>
<path d="M12 0h4c2.208 0 4 1.792 4 4s-1.792 4-4 4h-4V0z" fill="#FF7262"/>
<path d="M20 12c0 2.208-1.792 4-4 4s-4-1.792-4-4 1.792-4 4-4 4 1.792 4 4z" fill="#1ABCFE"/>
</svg>

After

Width:  |  Height:  |  Size: 496 B

View file

@ -0,0 +1,8 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M23.955 13.587l-1.342-4.135-2.664-8.189a.455.455 0 0 0-.867 0L16.418 9.45H7.582L4.918 1.263a.455.455 0 0 0-.867 0L1.386 9.452.044 13.587a.924.924 0 0 0 .331 1.03L12 23.054l11.625-8.436a.92.92 0 0 0 .33-1.03z" fill="#E24329"/>
<path d="M12 23.054L16.418 9.45H7.582L12 23.054z" fill="#FC6D26"/>
<path d="M12 23.054L7.582 9.452H1.386L12 23.054z" fill="#FCA326"/>
<path d="M1.386 9.452L.044 13.587a.924.924 0 0 0 .331 1.03L12 23.054 1.386 9.452z" fill="#E24329"/>
<path d="M12 23.054l4.418-13.602h5.036L12 23.054z" fill="#FCA326"/>
<path d="M22.614 9.452l1.341 4.135a.924.924 0 0 1-.33 1.03L12 23.054l10.614-13.602z" fill="#E24329"/>
</svg>

After

Width:  |  Height:  |  Size: 719 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M24 5.457v13.909c0 .904-.732 1.636-1.636 1.636h-3.819V11.73L12 16.64l-6.545-4.91v9.273H1.636A1.636 1.636 0 0 1 0 19.366V5.457c0-2.023 2.309-3.178 3.927-1.964L5.455 4.64 12 9.548l6.545-4.91 1.528-1.145C21.69 2.28 24 3.434 24 5.457z" fill="#EA4335"/>
</svg>

After

Width:  |  Height:  |  Size: 328 B

View file

@ -0,0 +1,6 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M7.71 14.29L2 22h7.65l5.71-7.71H7.71z" fill="#0066DA"/>
<path d="M22 22l-5.71-7.71H8.65L14.35 22H22z" fill="#00AC47"/>
<path d="M8.16 2L2.45 14.29h7.65L15.81 2H8.16z" fill="#FFBA00"/>
<path d="M15.84 2l-5.71 12.29h7.65L23.49 2H15.84z" fill="#EA4335"/>
</svg>

After

Width:  |  Height:  |  Size: 337 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M18.164 7.93V5.084a2.198 2.198 0 0 0 1.267-1.984v-.066A2.2 2.2 0 0 0 17.237.84h-.066a2.2 2.2 0 0 0-2.194 2.194v.066c0 .844.48 1.574 1.18 1.94V7.93a6.152 6.152 0 0 0-2.866 1.388L5.962 3.72a2.385 2.385 0 0 0 .07-.557A2.37 2.37 0 0 0 3.662.793a2.37 2.37 0 0 0-2.37 2.37 2.37 2.37 0 0 0 2.37 2.37c.432 0 .836-.12 1.183-.325l7.47 5.5A6.175 6.175 0 0 0 11.19 14.1a6.2 6.2 0 0 0 1.097 3.504l-2.12 2.12a1.786 1.786 0 0 0-.52-.082A1.803 1.803 0 1 0 11.45 21.445l2.172-2.172a6.175 6.175 0 0 0 3.572 1.135 6.2 6.2 0 1 0 .97-12.478zM17.204 17.2a3.004 3.004 0 1 1 0-6.008 3.004 3.004 0 0 1 0 6.008z" fill="#FF7A59"/>
</svg>

After

Width:  |  Height:  |  Size: 683 B

View file

@ -0,0 +1,15 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M11.571 11.513H0a5.218 5.218 0 0 0 5.232 5.215h2.13v2.057A5.215 5.215 0 0 0 12.575 24V12.518a1.005 1.005 0 0 0-1.005-1.005z" fill="#2684FF"/>
<path d="M6.262 6.259H17.793a5.218 5.218 0 0 0-5.232-5.214H6.26A5.218 5.218 0 0 0 1.03 6.259v6.268a1.005 1.005 0 0 0 1.005 1.005h9.527V8.318a2.06 2.06 0 0 0-2.06-2.059H6.262z" fill="url(#jiraGrad1)"/>
<path d="M17.53 6.259a5.218 5.218 0 0 1 5.232 5.214v6.268a1.005 1.005 0 0 1-1.005 1.005H12.23v-5.214a2.06 2.06 0 0 1 2.06-2.059h3.24V6.259z" fill="url(#jiraGrad2)"/>
<defs>
<linearGradient id="jiraGrad1" x1="6" y1="1" x2="1" y2="13">
<stop stop-color="#0052CC"/>
<stop offset="1" stop-color="#2684FF"/>
</linearGradient>
<linearGradient id="jiraGrad2" x1="18" y1="6" x2="23" y2="18">
<stop stop-color="#0052CC"/>
<stop offset="1" stop-color="#2684FF"/>
</linearGradient>
</defs>
</svg>

After

Width:  |  Height:  |  Size: 949 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M2.886 10.449a.64.64 0 0 1 .078.848L1.337 13.56a.32.32 0 0 1-.526-.036A11.044 11.044 0 0 1 .103 11.53a.32.32 0 0 1 .238-.37l2.186-.535a.64.64 0 0 1 .36-.176zm-.803 3.753a.32.32 0 0 0-.492.097 11.127 11.127 0 0 0-.747 2.148.32.32 0 0 0 .108.331l1.638 1.319a.32.32 0 0 0 .476-.07l1.18-1.87a.64.64 0 0 0-.033-.733L2.083 14.2zm.79 5.275a.32.32 0 0 0-.42.02l-.367.354a.32.32 0 0 0-.022.43 11.18 11.18 0 0 0 2.854 2.67.32.32 0 0 0 .42-.044l.404-.431a.32.32 0 0 0-.006-.445L2.873 19.477zM7.55 22.553a.32.32 0 0 0 .014.458l.28.232a.32.32 0 0 0 .427-.023 11.2 11.2 0 0 0 3.084-4.807.32.32 0 0 0-.19-.394l-.672-.25a.32.32 0 0 0-.405.18A9.614 9.614 0 0 1 7.55 22.553zM23.988 12c0 6.627-5.373 12-12 12-.742 0-1.47-.067-2.176-.197a.32.32 0 0 1-.218-.494L22.006 8.523a.32.32 0 0 1 .552.108c.283.882.43 1.818.43 2.79V12zm-.836-4.488a.32.32 0 0 0-.56-.05L10.163 22.68a.32.32 0 0 0 .096.472c.83.43 1.73.74 2.679.914a.32.32 0 0 0 .333-.14L23.195 8.204a.32.32 0 0 0-.043-.412v-.28zM20.684 5.216a.32.32 0 0 0 .484.013l.257-.269a.32.32 0 0 0 .017-.427A11.955 11.955 0 0 0 12 .007C5.373.007 0 5.38 0 12.007c0 .414.021.824.063 1.229a.32.32 0 0 0 .547.197l18.13-20.07a.32.32 0 0 1 .453-.014l1.49 1.368v.5z" fill="#5E6AD2"/>
</svg>

After

Width:  |  Height:  |  Size: 1.2 KiB

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M4.459 4.208c.746.606 1.026.56 2.428.466l13.215-.793c.28 0 .047-.28-.046-.326L18.57 2.33c-.42-.326-.98-.7-2.055-.607L3.62 2.79c-.466.046-.56.28-.374.466l1.213.952zm.793 3.08v13.904c0 .747.373 1.027 1.214.98l14.523-.84c.84-.046.933-.56.933-1.167V6.354c0-.606-.233-.933-.746-.886l-15.177.886c-.56.047-.747.327-.747.934zm14.337.745c.093.42 0 .84-.42.888l-.7.14v10.264c-.608.327-1.168.514-1.635.514-.747 0-.933-.234-1.494-.934l-4.577-7.186v6.952l1.448.327s0 .84-1.168.84l-3.222.186c-.093-.186 0-.653.327-.746l.84-.233V9.854L7.822 9.76c-.094-.42.14-1.026.793-1.073l3.456-.233 4.764 7.279v-6.44l-1.215-.14c-.093-.513.28-.886.747-.933l3.222-.187zM2.1 1.424L15.856.466c1.68-.14 2.1.093 2.8.606l3.876 2.753c.467.326.607.42.607.793v16.844c0 1.026-.374 1.633-1.68 1.726l-15.457.933c-.98.047-1.448-.093-1.962-.747l-3.13-4.06c-.56-.747-.793-1.306-.793-1.96V2.917c0-.84.374-1.54 1.782-1.493z" fill="#000000"/>
</svg>

After

Width:  |  Height:  |  Size: 975 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M10.006 5.15a4.678 4.678 0 0 1 3.358-1.426 4.7 4.7 0 0 1 4.273 2.775 5.476 5.476 0 0 1 2.163-.444C22.108 6.055 24 7.954 24 10.288a4.258 4.258 0 0 1-3.8 4.225h-.2a3.756 3.756 0 0 1-3.468 2.308 3.726 3.726 0 0 1-1.76-.44 4.418 4.418 0 0 1-3.89 2.32 4.418 4.418 0 0 1-3.798-2.15 3.678 3.678 0 0 1-.844.098A3.68 3.68 0 0 1 2.56 12.97c0-.77.237-1.484.642-2.074A4.448 4.448 0 0 1 0 7.2a4.448 4.448 0 0 1 4.448-4.448c1.11 0 2.13.41 2.91 1.086A4.672 4.672 0 0 1 10.006 5.15z" fill="#00A1E0"/>
</svg>

After

Width:  |  Height:  |  Size: 564 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M13.91 2.505c-.873-1.553-3.066-1.553-3.94 0L7.092 7.67a10.783 10.783 0 0 1 4.376 3.033l1.72-3.027a4.23 4.23 0 0 0-.207-.257L13.91 2.505zM3.6 20.3h2.685a8.782 8.782 0 0 0-.094-5.725l-2.59 4.613c-.25.445.109 1.003.613 1.003L3.6 20.3zm16.595.108c.5 0 .862-.557.613-1.003L14.075 7.5l-1.72 3.027A8.782 8.782 0 0 1 17.35 20.3h2.845v.108z" fill="#362D59"/>
</svg>

After

Width:  |  Height:  |  Size: 429 B

View file

@ -0,0 +1,4 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M15.337 23.979l7.216-1.561s-2.604-17.613-2.625-17.73c-.018-.116-.114-.2-.2-.2s-1.848-.138-1.848-.138-1.225-1.197-1.363-1.338c-.04-.04-.085-.06-.132-.074l-.793 18.867 -.255.174zm-3.07-16.938s-.718-.378-1.594-.378c-1.29 0-1.353.809-1.353 1.013 0 1.112 2.9 1.538 2.9 4.145 0 2.051-1.3 3.37-3.053 3.37-2.105 0-3.18-1.31-3.18-1.31l.564-1.865s1.105.949 2.036.949c.607 0 .856-.479.856-.829 0-1.453-2.38-1.519-2.38-3.906 0-2.008 1.441-3.953 4.351-3.953 1.12 0 1.674.321 1.674.321l-.82 2.443z" fill="#95BF47"/>
<path d="M14.998 6.268c-.082-.025-.18-.048-.252-.048-.028 0-.062.002-.092.006l-.76 18.07.663-.143 3.377-23.182c-.138.14-1.363 1.337-1.363 1.337s-.95-.093-1.573-.04z" fill="#5E8E3E"/>
</svg>

After

Width:  |  Height:  |  Size: 766 B

View file

@ -0,0 +1,6 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none">
<path d="M5.042 15.165a2.528 2.528 0 0 1-2.52 2.523A2.528 2.528 0 0 1 0 15.165a2.527 2.527 0 0 1 2.522-2.52h2.52v2.52zm1.271 0a2.527 2.527 0 0 1 2.521-2.52 2.527 2.527 0 0 1 2.521 2.52v6.313A2.528 2.528 0 0 1 8.834 24a2.528 2.528 0 0 1-2.521-2.522v-6.313z" fill="#E01E5A"/>
<path d="M8.834 5.042a2.528 2.528 0 0 1-2.521-2.52A2.528 2.528 0 0 1 8.834 0a2.528 2.528 0 0 1 2.521 2.522v2.52H8.834zm0 1.271a2.528 2.528 0 0 1 2.521 2.521 2.528 2.528 0 0 1-2.521 2.521H2.522A2.528 2.528 0 0 1 0 8.834a2.528 2.528 0 0 1 2.522-2.521h6.312z" fill="#36C5F0"/>
<path d="M18.956 8.834a2.528 2.528 0 0 1 2.522-2.521A2.528 2.528 0 0 1 24 8.834a2.528 2.528 0 0 1-2.522 2.521h-2.522V8.834zm-1.27 0a2.528 2.528 0 0 1-2.522 2.521 2.527 2.527 0 0 1-2.521-2.521V2.522A2.527 2.527 0 0 1 15.165 0a2.528 2.528 0 0 1 2.522 2.522v6.312z" fill="#2EB67D"/>
<path d="M15.165 18.956a2.528 2.528 0 0 1 2.522 2.522A2.528 2.528 0 0 1 15.165 24a2.527 2.527 0 0 1-2.521-2.522v-2.522h2.521zm0-1.27a2.527 2.527 0 0 1-2.521-2.522 2.527 2.527 0 0 1 2.521-2.521h6.313A2.528 2.528 0 0 1 24 15.165a2.528 2.528 0 0 1-2.522 2.521h-6.313z" fill="#ECB22E"/>
</svg>

After

Width:  |  Height:  |  Size: 1.2 KiB

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M13.976 9.15c-2.172-.806-3.356-1.426-3.356-2.409 0-.831.683-1.305 1.901-1.305 2.227 0 4.515.858 6.09 1.631l.89-5.494C18.252.975 15.697 0 12.165 0 9.667 0 7.589.654 6.104 1.872 4.56 3.147 3.757 4.992 3.757 7.218c0 4.039 2.467 5.76 6.476 7.219 2.585.92 3.445 1.574 3.445 2.583 0 .98-.84 1.545-2.354 1.545-1.875 0-4.965-.921-6.99-2.109l-.9 5.555C5.175 22.99 8.385 24 11.714 24c2.641 0 4.843-.624 6.328-1.813 1.664-1.305 2.525-3.236 2.525-5.732 0-4.128-2.524-5.851-6.591-7.305z" fill="#635BFF"/>
</svg>

After

Width:  |  Height:  |  Size: 571 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M12 0C5.381 0 0 5.381 0 12s5.381 12 12 12 12-5.381 12-12S18.619 0 12 0zm0 20.4c-4.636 0-8.4-3.764-8.4-8.4S7.364 3.6 12 3.6s8.4 3.764 8.4 8.4-3.764 8.4-8.4 8.4zm3.6-10.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8zm0 4.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8zm-4.8-4.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8zm0 4.8c0 .993-.807 1.8-1.8 1.8s-1.8-.807-1.8-1.8.807-1.8 1.8-1.8 1.8.807 1.8 1.8z" fill="#F22F46"/>
</svg>

After

Width:  |  Height:  |  Size: 587 B

View file

@ -0,0 +1,3 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
<path d="M15.535 8.465l-1.263 1.264a4.502 4.502 0 0 0-4.544 0L8.465 8.465a6.51 6.51 0 0 1 2.534-1.796V4.5h2v2.17a6.51 6.51 0 0 1 2.536 1.795zM12 10.5a1.5 1.5 0 1 0 0 3 1.5 1.5 0 0 0 0-3zm3.535 5.035l1.263-1.264a6.51 6.51 0 0 1-1.796 2.536V19h-2v-2.17a6.51 6.51 0 0 1-2.534-1.795l1.263-1.264a4.502 4.502 0 0 0 4.544 0h-.74zM19 11h-2.17a6.51 6.51 0 0 1-1.795 2.534l1.264 1.263a4.502 4.502 0 0 0 0-4.544l-1.264 1.263A6.51 6.51 0 0 1 16.831 11H19v2zm-12.17 2H5v-2h2.17a6.51 6.51 0 0 1 1.795-2.534L7.7 7.203a4.502 4.502 0 0 0 0 4.544l1.264-1.263A6.51 6.51 0 0 1 7.17 13H6.83z" fill="#FF4A00"/>
</svg>

After

Width:  |  Height:  |  Size: 659 B

View file

@ -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<SidebarProps> = ({ 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 <a> 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 (
<a
href={href}
target={newTab ? "_blank" : undefined}
rel={newTab ? "noopener noreferrer" : undefined}
onClick={(e) => {
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<SidebarProps> = ({ accessToken, userRole, defaultSelect
style={{
transition: "all 0.3s cubic-bezier(0.4, 0, 0.2, 1)",
position: "relative",
display: "flex",
flexDirection: "column",
}}
>
<ConfigProvider
@ -431,22 +447,60 @@ const Sidebar2: React.FC<SidebarProps> = ({ 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,
}))}
/>
</ConfigProvider>
{isAdminRole(userRole) && !collapsed && <UsageIndicator accessToken={accessToken} width={220} />}
{/* Pinned "Open Chat" button at bottom */}
<div style={{
padding: collapsed ? "10px 8px" : "10px 12px",
borderTop: "1px solid #f0f0f0",
flexShrink: 0,
}}>
<a
href={toHref("chat")}
target="_blank"
rel="noopener noreferrer"
style={{
display: "flex",
alignItems: "center",
justifyContent: collapsed ? "center" : "flex-start",
gap: 8,
padding: collapsed ? "8px 0" : "8px 10px",
borderRadius: 8,
background: "#1677ff",
color: "#fff",
textDecoration: "none",
fontSize: 13,
fontWeight: 600,
transition: "background 0.15s",
}}
onMouseEnter={(e) => {
(e.currentTarget as HTMLAnchorElement).style.background = "#0958d9";
}}
onMouseLeave={(e) => {
(e.currentTarget as HTMLAnchorElement).style.background = "#1677ff";
}}
>
<MessageOutlined style={{ fontSize: 16, flexShrink: 0 }} />
{!collapsed && <span>Open Chat</span>}
</a>
</div>
</Sider>
</Layout>
);

View file

@ -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);
});
});

View file

@ -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<MCPServerHealth[]>({
queryKey: [...mcpServerHealthKeys.lists(), { serverIds }],
queryFn: async () => await fetchMCPServerHealth(accessToken!, serverIds),
const queryClient = useQueryClient();
const [recheckingServerIds, setRecheckingServerIds] = useState<Set<string>>(new Set());
const query = useQuery<MCPServerHealth[]>({
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<MCPServerHealth[]>(
{ 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,
};
};

View file

@ -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<MCPServer[]>({
queryKey: mcpServersKeys.list({}),
queryFn: async () => await fetchMCPServers(accessToken!),
queryKey: mcpServersKeys.list(teamId ? { filters: { teamId } } : undefined),
queryFn: async () => await fetchMCPServers(accessToken!, teamId),
enabled: !!accessToken,
});
};

View file

@ -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<CreateuserProps> = ({
// Original return for standalone mode
return (
<div className="flex gap-2">
<Button2 className="mb-0" onClick={() => setIsModalVisible(true)}>
<Button type="primary" className="mb-0" onClick={() => setIsModalVisible(true)}>
+ Invite User
</Button2>
</Button>
<BulkCreateUsers accessToken={accessToken} teams={teams} possibleUIRoles={possibleUIRoles} />
<Modal
title="Invite User"

View file

@ -57,6 +57,18 @@ vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({
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<typeof import("@/app/(dashboard)/networking")>();
@ -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({

View file

@ -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<KeyResponse | null>(null);
const [sorting, setSorting] = React.useState<SortingState>(() => {
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<Record<string, boolean>>({});
@ -172,11 +177,9 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
const value = info.getValue() as string;
const width = info.cell.column.getSize();
return (
<Tooltip title={value}>
<span className={`font-mono text-xs truncate block`} style={{ maxWidth: width, overflow: "hidden" }}>
{value ?? "-"}
</span>
</Tooltip>
<span className="font-mono text-xs truncate block" style={{ maxWidth: width, overflow: "hidden" }}>
{value ?? "-"}
</span>
);
},
},
@ -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 (
<Tooltip title={value}>
<span className={`font-mono text-xs truncate block`} style={{ maxWidth: width, overflow: "hidden" }}>
{value ?? "-"}
</span>
</Tooltip>
<span className="font-mono text-xs truncate block" style={{ maxWidth: width, overflow: "hidden" }}>
{displayValue}
</span>
);
},
},
{
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 (
<Tooltip title={value}>
<span className={`font-mono text-xs truncate block`} style={{ maxWidth: width, overflow: "hidden" }}>
{value ?? "-"}
</span>
</Tooltip>
<span className="font-mono text-xs truncate block" style={{ maxWidth: width, overflow: "hidden" }}>
{displayValue}
</span>
);
},
},
{
id: "user_id",
accessorKey: "user_id",
header: "User ID",
size: 70,
id: "user",
accessorKey: "user",
header: () => (
<span className="flex items-center gap-1">
User
<Popover
content="Displays the first available value: User Alias, User Email, or User ID."
trigger="hover"
>
<InfoCircleOutlined className="text-gray-400 text-xs cursor-help" />
</Popover>
</span>
),
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 = (
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
{[
{ label: "User Alias", value: userAlias },
{ label: "User Email", value: userEmail },
{ label: "User ID", value: userId },
].map(({ label, value }) => (
<div key={label} className="flex flex-col min-w-0">
<span className="text-gray-400">{label}</span>
{value ? (
<Typography.Text
className="font-mono text-xs"
ellipsis={{ tooltip: value }}
copyable
>
{value}
</Typography.Text>
) : (
<span className="font-mono">-</span>
)}
</div>
))}
</div>
);
if (isDefaultAdmin && !userAlias && !userEmail) {
return (
<Popover content={popoverContent} trigger="hover" placement="bottomLeft">
<span className="cursor-default">
<DefaultProxyAdminTag userId={userId} />
</span>
</Popover>
);
}
return (
<Tooltip title={displayValue}>
<span className={`font-mono text-xs truncate block`} style={{ maxWidth: width, overflow: "hidden" }}>
{displayValue ?? "-"}
<Popover content={popoverContent} trigger="hover" placement="bottomLeft">
<span
className="font-mono text-xs truncate block cursor-default"
style={{ maxWidth: width, overflow: "hidden" }}
>
{displayValue || "-"}
</span>
</Tooltip>
</Popover>
);
},
},
@ -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 = (
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
<div className="flex flex-col min-w-0">
<span className="text-gray-400">User ID</span>
<Typography.Text
className="font-mono text-xs"
ellipsis={{ tooltip: userId }}
copyable
>
{userId}
</Typography.Text>
</div>
</div>
);
if (isDefaultAdmin) {
return (
<Popover content={popoverContent} trigger="hover" placement="bottomLeft">
<span className="cursor-default">
<DefaultProxyAdminTag userId={userId} />
</span>
</Popover>
);
}
return (
<Tooltip title={displayValue}>
<span className={`font-mono text-xs truncate block`} style={{ maxWidth: width, overflow: "hidden" }}>
{displayValue ?? "-"}
<Popover content={popoverContent} trigger="hover" placement="bottomLeft">
<span
className="font-mono text-xs truncate block cursor-default"
style={{ maxWidth: width, overflow: "hidden" }}
>
{userId}
</span>
</Tooltip>
</Popover>
);
},
},
@ -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(

View file

@ -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<BulkCreateUsersProps> = ({
return (
<>
<TremorButton className="mb-0" onClick={() => setIsModalVisible(true)}>
<Button type="primary" className="mb-0" onClick={() => setIsModalVisible(true)}>
+ Bulk Invite Users
</TremorButton>
</Button>
<Modal
title="Bulk Invite Users"
@ -629,9 +629,9 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
</div>
</div>
<TremorButton onClick={downloadTemplate} size="lg" className="w-full md:w-auto">
<DownloadOutlined className="mr-2" /> Download CSV Template
</TremorButton>
<Button type="primary" size="large" className="w-full md:w-auto" icon={<DownloadOutlined />}>
Download CSV Template
</Button>
</div>
<div className="flex items-center mb-4">
@ -662,14 +662,14 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
</Typography.Text>
</div>
</div>
<TremorButton
size="xs"
variant="secondary"
<Button
size="small"
onClick={removeSelectedFile}
className="flex items-center"
icon={<DeleteOutlined />}
>
<DeleteOutlined className="mr-1" /> Remove
</TremorButton>
Remove
</Button>
</div>
{fileError ? (
@ -694,7 +694,7 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
<UploadOutlined className="text-3xl text-gray-400 mb-2" />
<p className="mb-1">Drag and drop your CSV file here</p>
<p className="text-sm text-gray-500 mb-3">or</p>
<TremorButton size="sm">Browse files</TremorButton>
<Button size="small">Browse files</Button>
<p className="text-xs text-gray-500 mt-4">Only CSV files (.csv) are supported</p>
</div>
</Upload>
@ -781,21 +781,21 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
{!parsedData.some((user) => user.status === "success" || user.status === "failed") && (
<div className="flex space-x-3">
<TremorButton
<Button
onClick={() => {
setParsedData([]);
setParseError(null);
}}
variant="secondary"
>
Back
</TremorButton>
<TremorButton
</Button>
<Button
type="primary"
onClick={handleBulkCreate}
disabled={parsedData.filter((d) => d.isValid).length === 0 || isProcessing}
>
{isProcessing ? "Creating..." : `Create ${parsedData.filter((d) => d.isValid).length} Users`}
</TremorButton>
</Button>
</div>
)}
</div>
@ -829,40 +829,39 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
{!parsedData.some((user) => user.status === "success" || user.status === "failed") && (
<div className="flex justify-end mt-4">
<TremorButton
<Button
onClick={() => {
setParsedData([]);
setParseError(null);
}}
variant="secondary"
className="mr-3"
>
Back
</TremorButton>
<TremorButton
</Button>
<Button
type="primary"
onClick={handleBulkCreate}
disabled={parsedData.filter((d) => d.isValid).length === 0 || isProcessing}
>
{isProcessing ? "Creating..." : `Create ${parsedData.filter((d) => d.isValid).length} Users`}
</TremorButton>
</Button>
</div>
)}
{parsedData.some((user) => user.status === "success" || user.status === "failed") && (
<div className="flex justify-end mt-4">
<TremorButton
<Button
onClick={() => {
setParsedData([]);
setParseError(null);
}}
variant="secondary"
className="mr-3"
>
Start New Bulk Import
</TremorButton>
<TremorButton onClick={downloadResults} variant="primary" className="flex items-center">
<DownloadOutlined className="mr-2" /> Download User Credentials
</TremorButton>
</Button>
<Button type="primary" onClick={downloadResults} icon={<DownloadOutlined />}>
Download User Credentials
</Button>
</div>
)}
</div>

View file

@ -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({
</div>
<CopyButton text={mainContent} />
{mcpEvents && mcpEvents.length > 0 && (
<div style={{ marginTop: 8, maxWidth: "100%" }}>
<MCPEventsDisplay events={mcpEvents} />
</div>
)}
</div>
);
}
@ -566,6 +575,7 @@ const ChatMessages: React.FC<Props> = ({ messages, isStreaming, onEditMessage })
isLastMessage={isLastMessage}
isStreaming={isStreaming}
isTypingIndicator={isLastMessage && isTypingIndicator}
mcpEvents={msg.mcpEvents}
/>
);
})}

View file

@ -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<ChatPageProps> = ({ accessToken, userRole, userId, user
const [modelSearchText, setModelSearchText] = useState("");
const [selectedMCPServers, setSelectedMCPServers] = useState<string[]>([]);
const [responsesSessionId, setResponsesSessionId] = useState<string | null>(null);
const [isStreaming, setIsStreaming] = useState(false);
const [inputText, setInputText] = useState("");
const [mcpPopoverOpen, setMcpPopoverOpen] = useState(false);
@ -162,7 +165,7 @@ const ChatPage: React.FC<ChatPageProps> = ({ accessToken, userRole, userId, user
createConversation,
appendMessage,
updateLastAssistantMessage,
truncateAfterMessage,
truncateFromMessage,
deleteConversation,
renameConversation,
} = useChatHistory(activeConversationId);
@ -203,6 +206,12 @@ const ChatPage: React.FC<ChatPageProps> = ({ 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<ChatPageProps> = ({ 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<ChatPageProps> = ({ 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<ChatPageProps> = ({ 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<ChatPageProps> = ({ 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<ChatPageProps> = ({ 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<ChatPageProps> = ({ accessToken, userRole, userId, user
: greeting}
</h1>
{isComparisonMode && (
{isComparisonMode ? (
<p style={{ margin: "-16px 0 24px", fontSize: 14, color: "#6b7280", textAlign: "center" }}>
Send a message to see responses side-by-side
</p>
) : (
<p style={{ margin: "-16px 0 28px", fontSize: 14, color: "#6b7280", textAlign: "center", maxWidth: 520, lineHeight: 1.6 }}>
Chat with 100+ LLMs + MCP tools authenticate once, use them here.{" "}
<button
onClick={() => setSidebarView("apps")}
style={{ background: "none", border: "none", cursor: "pointer", color: "#1677ff", fontSize: 14, padding: 0, fontWeight: 500 }}
>
Open Apps
</button>
</p>
)}
{/* Input card */}

View file

@ -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<Props> = ({ accessToken, selectedServers, onChange
const [activeTab, setActiveTab] = useState<TabKey>("all");
const [togglingOn, setTogglingOn] = useState<Set<string>>(new Set());
const [detailServer, setDetailServer] = useState<MCPServer | null>(null);
const [detailTools, setDetailTools] = useState<MCPTool[]>([]);
const [loadingTools, setLoadingTools] = useState(false);
// tool counts per server name, preloaded in background
const [toolCounts, setToolCounts] = useState<Record<string, number>>({});
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<Props> = ({ 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<Props> = ({ 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<Props> = ({ accessToken, selectedServers, onChange
{/* Avatar + name + connect */}
<div style={{ display: "flex", alignItems: "flex-start", gap: 20, marginBottom: 28 }}>
{detailServer.mcp_info?.logo_url ? (
<img
src={detailServer.mcp_info.logo_url}
alt={`${name} logo`}
style={{
width: 64, height: 64, borderRadius: 16,
objectFit: "contain", flexShrink: 0,
background: "#f9fafb",
}}
onError={(e) => {
const el = e.target as HTMLImageElement;
el.style.display = "none";
if (el.nextElementSibling) (el.nextElementSibling as HTMLElement).style.display = "flex";
}}
/>
) : null}
<div style={{
width: 64, height: 64, borderRadius: 16,
background: color, display: "flex",
background: color, display: detailServer.mcp_info?.logo_url ? "none" : "flex",
alignItems: "center", justifyContent: "center",
color: "#fff", fontWeight: 700, fontSize: 28, flexShrink: 0,
}}>
@ -128,7 +201,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
<Button
type={isConnected ? "default" : "primary"}
loading={isTogglingOn}
onClick={() => handleToggle(name, !isConnected)}
onClick={() => handleToggle(name, !isConnected, detailServer.server_id)}
style={{ borderRadius: 8, fontWeight: 600, height: 38, minWidth: 110 }}
>
{isConnected ? "Disconnect" : "Connect"}
@ -137,10 +210,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
{/* Info table */}
<h3 style={{ margin: "0 0 12px", fontSize: 15, fontWeight: 600, color: "#111827" }}>Information</h3>
<div style={{ border: "1px solid #e5e7eb", borderRadius: 10, overflow: "hidden" }}>
<div style={{ border: "1px solid #e5e7eb", borderRadius: 10, overflow: "hidden", marginBottom: 28 }}>
{[
["Server ID", detailServer.server_id],
["Transport", (detailServer as MCPServer & { mcp_info?: { server_url?: string } }).mcp_info?.server_url ? "HTTP" : "stdio"],
["Transport", handleTransport(detailServer.transport, detailServer.spec_path)],
["Status", isConnected ? "Connected" : "Not connected"],
].filter(([, v]) => v).map(([label, value], i, arr) => (
<div key={label} style={{
@ -154,6 +227,43 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
</div>
))}
</div>
{/* Tools section */}
<div style={{ display: "flex", alignItems: "center", gap: 8, marginBottom: 12 }}>
<h3 style={{ margin: 0, fontSize: 15, fontWeight: 600, color: "#111827" }}>Available Tools</h3>
{!loadingTools && (
<span style={{
fontSize: 11, fontWeight: 600, color: "#6b7280",
background: "#f3f4f6", borderRadius: 4, padding: "1px 6px",
}}>{detailTools.length}</span>
)}
</div>
{loadingTools ? (
<div style={{ display: "flex", justifyContent: "center", padding: "24px 0" }}>
<Spin size="small" />
</div>
) : detailTools.length === 0 ? (
<div style={{ color: "#9ca3af", fontSize: 13, padding: "8px 0" }}>
No tools available
</div>
) : (
<div style={{ display: "flex", flexDirection: "column", gap: 8 }}>
{detailTools.map((tool) => (
<div key={tool.name} style={{
border: "1px solid #e5e7eb", borderRadius: 8,
padding: "10px 14px", background: "#fafafa",
}}>
<div style={{ display: "flex", alignItems: "center", gap: 8, marginBottom: tool.description ? 4 : 0 }}>
<ToolOutlined style={{ fontSize: 13, color: "#6b7280" }} />
<span style={{ fontSize: 13, fontWeight: 600, color: "#111827", fontFamily: "monospace" }}>{tool.name}</span>
</div>
{tool.description && (
<p style={{ margin: 0, fontSize: 12, color: "#6b7280", paddingLeft: 21 }}>{tool.description}</p>
)}
</div>
))}
</div>
)}
</div>
);
}
@ -165,7 +275,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
{/* Header row */}
<div style={{ display: "flex", alignItems: "center", justifyContent: "space-between", marginBottom: 20, gap: 16, flexWrap: "wrap" }}>
<div>
<div style={{ display: "flex", alignItems: "center", gap: 8, marginBottom: 2 }}>
<div style={{ display: "flex", alignItems: "center", gap: 8, marginBottom: 4 }}>
<h2 style={{ margin: 0, fontSize: 18, fontWeight: 600, color: "#111827" }}>MCP Servers</h2>
<span style={{
fontSize: 10, fontWeight: 600, color: "#1677ff",
@ -173,9 +283,22 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
letterSpacing: "0.05em", textTransform: "uppercase",
}}>Beta</span>
</div>
<p style={{ margin: 0, fontSize: 13, color: "#6b7280" }}>
Connect tools to your chat.
</p>
<div style={{ display: "flex", alignItems: "center", gap: 12 }}>
<p style={{ margin: 0, fontSize: 13, color: "#6b7280" }}>
Browse tools, authenticate once, use in chat no setup needed.
</p>
{loadingCounts ? (
<span style={{ display: "flex", alignItems: "center", gap: 5, fontSize: 12, color: "#9ca3af" }}>
<Spin size="small" style={{ transform: "scale(0.7)" }} />
Loading tools...
</span>
) : totalTools > 0 ? (
<span style={{ display: "flex", alignItems: "center", gap: 4, fontSize: 12, color: "#6b7280" }}>
<ToolOutlined style={{ fontSize: 11 }} />
{totalTools} tool{totalTools !== 1 ? "s" : ""} available
</span>
) : null}
</div>
</div>
<Input
prefix={<SearchOutlined style={{ color: "#9ca3af", fontSize: 13 }} />}
@ -227,6 +350,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
const isConnected = selectedServers.includes(name);
const color = getAvatarColor(name);
const isLeftCol = idx % 2 === 0;
const count = toolCounts[name];
return (
<div
@ -243,9 +367,26 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
onMouseEnter={(e) => { (e.currentTarget as HTMLDivElement).style.background = "#fafafa"; }}
onMouseLeave={(e) => { (e.currentTarget as HTMLDivElement).style.background = "#fff"; }}
>
{server.mcp_info?.logo_url ? (
<img
src={server.mcp_info.logo_url}
alt={`${name} logo`}
style={{
width: 38, height: 38, borderRadius: 10,
objectFit: "contain", flexShrink: 0,
background: "#f9fafb",
}}
onError={(e) => {
const el = e.target as HTMLImageElement;
el.style.display = "none";
if (el.nextElementSibling) (el.nextElementSibling as HTMLElement).style.display = "flex";
}}
/>
) : null}
<div style={{
width: 38, height: 38, borderRadius: 10, background: color,
display: "flex", alignItems: "center", justifyContent: "center",
display: server.mcp_info?.logo_url ? "none" : "flex",
alignItems: "center", justifyContent: "center",
color: "#fff", fontWeight: 700, fontSize: 16, flexShrink: 0,
}}>
{name.charAt(0).toUpperCase()}
@ -254,13 +395,33 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
<div style={{ fontSize: 14, fontWeight: 500, color: "#111827", overflow: "hidden", textOverflow: "ellipsis", whiteSpace: "nowrap" }}>
{name}
</div>
<div style={{ fontSize: 12, color: "#9ca3af", marginTop: 1, overflow: "hidden", textOverflow: "ellipsis", whiteSpace: "nowrap" }}>
{server.description ?? "MCP server"}
<div style={{ fontSize: 12, color: "#9ca3af", marginTop: 1, display: "flex", alignItems: "center", gap: 6 }}>
<span style={{ overflow: "hidden", textOverflow: "ellipsis", whiteSpace: "nowrap" }}>
{server.description ?? "MCP server"}
</span>
{count !== undefined ? (
count > 0 ? (
<span style={{ flexShrink: 0, display: "flex", alignItems: "center", gap: 3, color: "#9ca3af" }}>
· <ToolOutlined style={{ fontSize: 10 }} /> {count}
</span>
) : null
) : loadingCounts ? (
<Skeleton.Input active size="small" style={{ width: 28, height: 12, minWidth: 28, flexShrink: 0 }} />
) : null}
</div>
</div>
{isConnected && (
<span style={{ width: 7, height: 7, borderRadius: "50%", background: "#1677ff", flexShrink: 0 }} />
)}
{server.auth_type === AUTH_TYPE.OAUTH2 && (
<span style={{
fontSize: 10, fontWeight: 600, color: "#7c3aed",
background: "#f3e8ff", borderRadius: 4, padding: "1px 5px",
letterSpacing: "0.03em", flexShrink: 0, whiteSpace: "nowrap",
}}>
OAuth2
</span>
)}
<RightOutlined style={{ fontSize: 11, color: "#d1d5db", flexShrink: 0 }} />
</div>
);

View file

@ -112,6 +112,18 @@ const MCPConnectPicker: React.FC<Props> = ({ accessToken, selectedServers, onCha
gap: 12,
}}
>
{server.mcp_info?.logo_url && (
<img
src={server.mcp_info.logo_url}
alt={`${name} logo`}
style={{
width: 24, height: 24, borderRadius: 6,
objectFit: "contain", flexShrink: 0,
marginTop: 1,
}}
onError={(e) => { (e.target as HTMLImageElement).style.display = "none"; }}
/>
)}
<div style={{ flex: 1, minWidth: 0 }}>
<div
style={{

View file

@ -1,8 +1,12 @@
export type { MCPEvent } from "../mcp_tools/types";
export interface ChatMessage {
id: string;
role: "user" | "assistant" | "tool";
content: string;
reasoningContent?: string;
/** MCP tool events that occurred during this assistant turn, in order. */
mcpEvents?: MCPEvent[];
toolName?: string;
toolArgs?: Record<string, unknown>;
toolResult?: string;

View file

@ -51,8 +51,9 @@ export function useChatHistory(activeConversationId: string | null): {
staleId: boolean;
createConversation: (model: string) => string;
appendMessage: (conversationId: string, message: Omit<ChatMessage, "id" | "timestamp">) => void;
updateLastAssistantMessage: (conversationId: string, updates: Partial<Pick<ChatMessage, "content" | "reasoningContent">>) => void;
truncateAfterMessage: (conversationId: string, messageId: string) => void;
updateLastAssistantMessage: (conversationId: string, updates: Partial<Pick<ChatMessage, "content" | "reasoningContent" | "mcpEvents">>) => void;
/** Remove the message with `messageId` and all subsequent messages from the conversation. */
truncateFromMessage: (conversationId: string, messageId: string) => void;
deleteConversation: (id: string) => void;
renameConversation: (id: string, newTitle: string) => void;
setActiveConversationId: (id: string | null) => void;
@ -148,7 +149,7 @@ export function useChatHistory(activeConversationId: string | null): {
const updateLastAssistantMessage = useCallback(
(
conversationId: string,
updates: Partial<Pick<ChatMessage, "content" | "reasoningContent">>,
updates: Partial<Pick<ChatMessage, "content" | "reasoningContent" | "mcpEvents">>,
) => {
setConversations((prev) => {
const updated = prev.map((conv) => {
@ -168,7 +169,7 @@ export function useChatHistory(activeConversationId: string | null): {
[],
);
const truncateAfterMessage = useCallback(
const truncateFromMessage = useCallback(
(conversationId: string, messageId: string) => {
setConversations((prev) => {
const updated = prev.map((conv) => {
@ -222,7 +223,7 @@ export function useChatHistory(activeConversationId: string | null): {
createConversation,
appendMessage,
updateLastAssistantMessage,
truncateAfterMessage,
truncateFromMessage,
deleteConversation,
renameConversation,
setActiveConversationId,

View file

@ -140,7 +140,7 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
<Tab>Guardrail Garden</Tab>
<Tab>Guardrails</Tab>
<Tab disabled={!accessToken || guardrailsList.length === 0}>Test Playground</Tab>
<Tab>Team Guardrails</Tab>
<Tab>Submitted Guardrails</Tab>
</TabList>
<TabPanels>

View file

@ -99,6 +99,7 @@ export interface KeyResponse {
user?: {
user_id: string;
user_email: string;
user_alias: string | null;
};
}

View file

@ -13,6 +13,7 @@ interface MCPServerSelectorProps {
accessToken: string;
placeholder?: string;
disabled?: boolean;
teamId?: string | null;
}
const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
@ -22,8 +23,9 @@ const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
accessToken,
placeholder = "Select MCP servers",
disabled = false,
teamId,
}) => {
const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers();
const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(teamId);
const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups();
const loading = serversLoading || groupsLoading;

Some files were not shown because too many files have changed in this diff Show more