Merge branch 'main' into litellm_oss_staging_03_04_2026
2
.github/workflows/codeql.yml
vendored
|
|
@ -34,8 +34,6 @@ jobs:
|
|||
build-mode: none
|
||||
- language: python
|
||||
build-mode: none
|
||||
- language: ruby
|
||||
build-mode: none
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
|
|
@ -35,6 +35,14 @@ deploymentLabels: {}
|
|||
podAnnotations: {}
|
||||
podLabels: {}
|
||||
|
||||
# -- Deployment strategy configuration
|
||||
# Example:
|
||||
# type: RollingUpdate
|
||||
# rollingUpdate:
|
||||
# maxUnavailable: 0
|
||||
# maxSurge: 1
|
||||
strategy: {}
|
||||
|
||||
terminationGracePeriodSeconds: 90
|
||||
topologySpreadConstraints:
|
||||
[]
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
144
docs/my-website/docs/mcp_aws_sigv4.md
Normal 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.
|
||||
|
|
@ -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.
|
||||
|
||||
:::
|
||||
|
||||
|
|
|
|||
|
|
@ -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">
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
121
docs/my-website/docs/troubleshoot/pip_venv_upgrade.md
Normal 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).
|
||||
|
|
@ -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)**
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
],
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
@ -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;
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
7
litellm/proxy/_experimental/out/assets/logos/figma.svg
Normal 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 |
8
litellm/proxy/_experimental/out/assets/logos/gitlab.svg
Normal 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 |
3
litellm/proxy/_experimental/out/assets/logos/gmail.svg
Normal 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 |
|
|
@ -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 |
3
litellm/proxy/_experimental/out/assets/logos/hubspot.svg
Normal 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 |
15
litellm/proxy/_experimental/out/assets/logos/jira.svg
Normal 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 |
3
litellm/proxy/_experimental/out/assets/logos/linear.svg
Normal 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 |
3
litellm/proxy/_experimental/out/assets/logos/notion.svg
Normal 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 |
|
|
@ -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 |
3
litellm/proxy/_experimental/out/assets/logos/sentry.svg
Normal 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 |
4
litellm/proxy/_experimental/out/assets/logos/shopify.svg
Normal 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 |
6
litellm/proxy/_experimental/out/assets/logos/slack.svg
Normal 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 |
3
litellm/proxy/_experimental/out/assets/logos/stripe.svg
Normal 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 |
3
litellm/proxy/_experimental/out/assets/logos/twilio.svg
Normal 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 |
3
litellm/proxy/_experimental/out/assets/logos/zapier.svg
Normal 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 |
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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": []}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
)
|
||||
222
litellm/proxy/openapi_registry.json
Normal 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" }
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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__])
|
||||
|
|
|
|||
|
|
@ -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)."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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") == {}
|
||||
|
|
|
|||
|
|
@ -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__])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
||||
|
|
|
|||
127
tests/test_litellm/test_chat_ui_responses_session.py
Normal 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"
|
||||
)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
7
ui/litellm-dashboard/public/assets/logos/figma.svg
Normal 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 |
8
ui/litellm-dashboard/public/assets/logos/gitlab.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/gmail.svg
Normal 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 |
|
|
@ -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 |
3
ui/litellm-dashboard/public/assets/logos/hubspot.svg
Normal 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 |
15
ui/litellm-dashboard/public/assets/logos/jira.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/linear.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/notion.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/salesforce.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/sentry.svg
Normal 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 |
4
ui/litellm-dashboard/public/assets/logos/shopify.svg
Normal 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 |
6
ui/litellm-dashboard/public/assets/logos/slack.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/stripe.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/twilio.svg
Normal 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 |
3
ui/litellm-dashboard/public/assets/logos/zapier.svg
Normal 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 |
|
|
@ -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>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
/>
|
||||
);
|
||||
})}
|
||||
|
|
|
|||
|
|
@ -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 */}
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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={{
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -99,6 +99,7 @@ export interface KeyResponse {
|
|||
user?: {
|
||||
user_id: string;
|
||||
user_email: string;
|
||||
user_alias: string | null;
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||