ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

End users - Allow giving end users access to specific mcp servers (#21411)

Browse Source 
* feat(schema.prisma): add object permissions for end users

allows controlling if end user can call specific mcp servers

* feat: cleanup for customer_endpoints support of object permission id

* fix: cleanup str

* feat(customers/): enforce end user can only call allowed mcps - if configured

* docs: document customer/end user object permission usage

* feat: address greptile comments
This commit is contained in:
Krish Dholakia 2026-02-17 22:45:49 -08:00 committed by GitHub
parent 371cabfebd
commit 1f521be0f2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 1355 additions and 206 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
docs/my-website/docs/mcp.md
View File

@ -808,6 +808,68 @@ If your stdio MCP server needs per-request credentials, you can map HTTP headers
In this example, when a client makes a request with the `X-GITHUB_PERSONAL_ACCESS_TOKEN` header, the proxy forwards that value into the stdio process as the `GITHUB_PERSONAL_ACCESS_TOKEN` environment variable.
## Control MCP Access for End Users
Control which MCP servers end users of your AI application can access (e.g. users of an internal chat UI). Pass the customer ID in the `x-litellm-end-user` header to:
- Enforce object permissions (limit which MCP servers they can access)
- Apply customer-specific budgets
- Track spend per customer
**FastMCP Client Example:**
```python title="Track customer spend with x-litellm-end-user" showLineNumbers
from fastmcp import Client
import asyncio
# MCP client configuration with customer tracking
config = {
"mcpServers": {
"github": {
"url": "http://localhost:4000/github_mcp/mcp",
"headers": {
"x-litellm-api-key": "Bearer sk-1234",
"x-litellm-end-user": "customer_123", # 👈 CUSTOMER ID
"Authorization": "Bearer gho_token"
}
}
}
}
client = Client(config)
async def main():
async with client:
# All MCP calls will be tracked under customer_123
tools = await client.list_tools()
result = await client.call_tool(tools[0].name, {})
print(f"Tool result: {result}")
asyncio.run(main())
```
**Cursor IDE Example:**
```json title="Cursor config with customer tracking" showLineNumbers
{
"mcpServers": {
"GitHub": {
"url": "http://localhost:4000/github_mcp/mcp",
"headers": {
"x-litellm-api-key": "Bearer $LITELLM_API_KEY",
"x-litellm-end-user": "customer_123"
}
}
}
}
```
**What happens:**
- Customer-specific object permissions are enforced (only allowed MCP servers are accessible)
- Customer budgets are applied
- All tool calls are tracked under `customer_123`
[Learn more about customer management →](./proxy/customers)
## Using your MCP with client side credentials
Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP.

198
docs/my-website/docs/proxy/customers.md
View File

@ -2,7 +2,7 @@ import Image from '@theme/IdealImage';
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Customers / End-User Budgets
# Customers / End-Users
Track spend, set budgets for your customers.
@ -10,9 +10,11 @@ Track spend, set budgets for your customers.
### 1. Make LLM API call w/ Customer ID
Make a /chat/completions call, pass 'user' - First call Works
You can pass the customer ID in two ways:
```bash showLineNumbers title="Make request with customer ID"
**Option 1: In the request body** (using the `user` field)
```bash showLineNumbers title="Make request with customer ID in body"
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--header 'Authorization: Bearer sk-1234' \ # 👈 YOUR PROXY KEY
@ -28,6 +30,30 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
}'
```
**Option 2: In the request headers** (using `x-litellm-end-user`)
```bash showLineNumbers title="Make request with customer ID in header"
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--header 'Authorization: Bearer sk-1234' \
--header 'x-litellm-end-user: ishaan3' \ # 👈 CUSTOMER ID IN HEADER
--data ' {
"model": "azure-gpt-3.5",
"messages": [
{
"role": "user",
"content": "what time is it"
}
]
}'
```
**Use `x-litellm-end-user` to control permissions for end users of your AI application** (e.g. users of an internal chat UI):
- Apply customer-specific object permissions (limit which MCP servers they can access)
- Enforce customer budgets
- Works with all endpoints (chat/completions, embeddings, MCP, etc.)
- No need to modify request body
The customer_id will be upserted into the DB with the new spend.
If the customer_id already exists, spend will be incremented.
@ -123,7 +149,171 @@ Expected Response
</Tabs>
## Setting Customer Budgets
## Setting Customer Object Permissions
Control which resources (MCP servers, vector stores, agents) a customer can access.
### What are Object Permissions?
Object permissions allow you to restrict customer access to specific:
- **MCP Servers**: Limit which MCP servers the customer can call
- **MCP Access Groups**: Assign customers to predefined groups of MCP servers
- **MCP Tool Permissions**: Granular control over which tools within an MCP server the customer can use
- **Vector Stores**: Control which vector stores the customer can query
- **Agents**: Restrict which agents the customer can interact with
- **Agent Access Groups**: Assign customers to predefined groups of agents
### Creating a Customer with Object Permissions
```bash showLineNumbers title="Create customer with object permissions"
curl -L -X POST 'http://localhost:4000/customer/new' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"user_id": "user_1",
"object_permission": {
"mcp_servers": ["server_1", "server_2"],
"mcp_access_groups": ["public_group"],
"mcp_tool_permissions": {
"server_1": ["tool_a", "tool_b"]
},
"vector_stores": ["vector_store_1"],
"agents": ["agent_1"],
"agent_access_groups": ["basic_agents"]
}
}'
```
**Parameters:**
- `mcp_servers` (Optional[List[str]]): List of allowed MCP server IDs
- `mcp_access_groups` (Optional[List[str]]): List of MCP access group names
- `mcp_tool_permissions` (Optional[Dict[str, List[str]]]): Map of server ID to allowed tool names
- `vector_stores` (Optional[List[str]]): List of allowed vector store IDs
- `agents` (Optional[List[str]]): List of allowed agent IDs
- `agent_access_groups` (Optional[List[str]]): List of agent access group names
**Note:** If `object_permission` is `null` or `{}`, the customer has no object-level restrictions.
### Updating Customer Object Permissions
You can update object permissions for existing customers:
```bash showLineNumbers title="Update customer object permissions"
curl -L -X POST 'http://localhost:4000/customer/update' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"user_id": "user_1",
"object_permission": {
"mcp_servers": ["server_3"],
"vector_stores": ["vector_store_2", "vector_store_3"]
}
}'
```
### Viewing Customer Object Permissions
When you query customer info, object permissions are included in the response:
```bash showLineNumbers title="Get customer info with object permissions"
curl -X GET 'http://0.0.0.0:4000/customer/info?end_user_id=user_1' \
-H 'Authorization: Bearer sk-1234'
```
**Response:**
```json showLineNumbers title="Response with object permissions"
{
"user_id": "user_1",
"blocked": false,
"alias": "John Doe",
"spend": 0.0,
"object_permission": {
"object_permission_id": "perm_abc123",
"mcp_servers": ["server_1", "server_2"],
"mcp_access_groups": ["public_group"],
"mcp_tool_permissions": {
"server_1": ["tool_a", "tool_b"]
},
"vector_stores": ["vector_store_1"],
"agents": ["agent_1"],
"agent_access_groups": ["basic_agents"]
},
"litellm_budget_table": null
}
```
### Use Cases
**1. Tiered Access Control**
Create different permission tiers for your customers:
```bash showLineNumbers title="Free tier customer"
# Free tier - limited access
curl -L -X POST 'http://localhost:4000/customer/new' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"user_id": "free_user",
"budget_id": "free_tier",
"object_permission": {
"mcp_access_groups": ["public_group"],
"agent_access_groups": ["basic_agents"]
}
}'
```
```bash showLineNumbers title="Premium tier customer"
# Premium tier - full access
curl -L -X POST 'http://localhost:4000/customer/new' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"user_id": "premium_user",
"budget_id": "premium_tier",
"object_permission": {
"mcp_servers": ["server_1", "server_2", "server_3"],
"vector_stores": ["vector_store_1", "vector_store_2"],
"agents": ["agent_1", "agent_2", "agent_3"]
}
}'
```
**2. Department-Specific Access**
Restrict customers to resources relevant to their department:
```bash showLineNumbers title="Sales team customer"
curl -L -X POST 'http://localhost:4000/customer/new' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"user_id": "sales_user",
"object_permission": {
"mcp_servers": ["crm_server", "email_server"],
"agents": ["sales_assistant"],
"vector_stores": ["sales_knowledge_base"]
}
}'
```
**3. Tool-Level Restrictions**
Grant access to specific tools within an MCP server:
```bash showLineNumbers title="Limited tool access"
curl -L -X POST 'http://localhost:4000/customer/new' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"user_id": "restricted_user",
"object_permission": {
"mcp_servers": ["database_server"],
"mcp_tool_permissions": {
"database_server": ["read_only_query", "get_table_schema"]
}
}
}'
```
## Setting Customer Budgets
Set customer budgets (e.g. monthly budgets, tpm/rpm limits) on LiteLLM Proxy

2
docs/my-website/docs/proxy/request_headers.md
View File

@ -20,6 +20,8 @@ By default, LiteLLM does not forward client headers to LLM provider APIs. Howeve
`x-litellm-spend-logs-metadata`: Optional[str]: JSON string containing custom metadata to include in spend logs. Example: `{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}`. [Learn More](../proxy/enterprise#tracking-spend-with-custom-metadata)
`x-litellm-end-user`: Optional[str]: The customer/end-user ID to track spend and apply budgets/permissions. Alternative to passing `user` in the request body. Works with all endpoints including MCP. [Learn More](./customers)
## Anthropic Headers
`anthropic-version` Optional[str]: The version of the Anthropic API to use.

6
litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql Normal file
View File

@ -0,0 +1,6 @@
-- AlterTable
ALTER TABLE "LiteLLM_EndUserTable" ADD COLUMN "object_permission_id" TEXT;
-- AddForeignKey
ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;

3
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View File

@ -233,6 +233,7 @@ model LiteLLM_ObjectPermissionTable {
verification_tokens LiteLLM_VerificationToken[]
organizations LiteLLM_OrganizationTable[]
users LiteLLM_UserTable[]
end_users LiteLLM_EndUserTable[]
}
// Holds the MCP server configuration
@ -403,7 +404,9 @@ model LiteLLM_EndUserTable {
allowed_model_region String? // require all user requests to use models in this specific region
default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model.
budget_id String?
object_permission_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
blocked Boolean @default(false)
}

172
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View File

@ -6,12 +6,8 @@ from starlette.requests import Request
from starlette.types import Scope
from litellm._logging import verbose_logger
from litellm.proxy._types import (
LiteLLM_TeamTable,
ProxyException,
SpecialHeaders,
UserAPIKeyAuth,
)
from litellm.proxy._types import (LiteLLM_TeamTable, ProxyException,
SpecialHeaders, UserAPIKeyAuth)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -336,15 +332,28 @@ class MCPRequestHandler:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
) -> List[str]:
"""
Get list of allowed MCP servers for the given user/key based on permissions
Get list of allowed MCP servers for the given user/key based on permissions.
Permission hierarchy:
1. Calculate key/team allowed servers:
- If team has permissions: key inherits from team (or intersects if key has own permissions)
- If team has no permissions: use key permissions
2. Apply end_user permissions if end_user_id is set:
- If require_end_user_mcp_access_defined=True and end_user has no permissions: block all access
- If end_user has permissions:
* If key/team has no restrictions (empty): use end_user permissions
* If key/team has restrictions: intersect key/team AND end_user
- If end_user has no permissions and flag is False: use key/team permissions
Returns:
List[str]: List of allowed MCP servers by server id
"""
from typing import List
from litellm.proxy.proxy_server import general_settings
try:
allowed_mcp_servers: List[str] = []
# Get allowed servers from key and team
allowed_mcp_servers_for_key = (
await MCPRequestHandler._get_allowed_mcp_servers_for_key(
user_api_key_auth
@ -357,8 +366,9 @@ class MCPRequestHandler:
)
#########################################################
# If team has mcp_servers, handle inheritance and intersection logic
# Calculate key/team allowed servers using inheritance and intersection logic
#########################################################
allowed_mcp_servers: List[str] = []
if len(allowed_mcp_servers_for_team) > 0:
if len(allowed_mcp_servers_for_key) > 0:
# Key has its own MCP permissions - use intersection with team permissions
@ -371,6 +381,52 @@ class MCPRequestHandler:
else:
allowed_mcp_servers = allowed_mcp_servers_for_key
#########################################################
# Check end_user permissions if end_user_id is set
#########################################################
if user_api_key_auth and user_api_key_auth.end_user_id:
allowed_mcp_servers_for_end_user = (
await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(
user_api_key_auth
)
)
# Check if require_end_user_mcp_access flag is enabled
require_end_user_mcp_access = general_settings.get(
"require_end_user_mcp_access_defined", False
)
# If the flag is enabled and end_user has no permissions, block all access
if require_end_user_mcp_access and len(allowed_mcp_servers_for_end_user) == 0:
verbose_logger.debug(
f"require_end_user_mcp_access_defined=True and end_user {user_api_key_auth.end_user_id} has no MCP permissions - blocking MCP access"
)
return []
# If end_user has explicit MCP server permissions, apply intersection logic
if len(allowed_mcp_servers_for_end_user) > 0:
verbose_logger.debug(
f"End user {user_api_key_auth.end_user_id} has explicit MCP permissions: {allowed_mcp_servers_for_end_user}"
)
# If key/team has no restrictions (empty list), use end_user permissions
if len(allowed_mcp_servers) == 0:
verbose_logger.debug(
"No key/team MCP restrictions - using end_user permissions"
)
allowed_mcp_servers = allowed_mcp_servers_for_end_user
else:
# Apply intersection: key/team AND end_user
# This ensures key/team restrictions are ALWAYS respected
filtered_servers = []
for _mcp_server in allowed_mcp_servers:
if _mcp_server in allowed_mcp_servers_for_end_user:
filtered_servers.append(_mcp_server)
allowed_mcp_servers = filtered_servers
verbose_logger.debug(
f"Applied end_user intersection with key/team. Final allowed servers: {allowed_mcp_servers}"
)
return list(set(allowed_mcp_servers))
except Exception as e:
verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}")
@ -402,11 +458,9 @@ class MCPRequestHandler:
get_team_object() in litellm/proxy/auth/auth_checks.py
"""
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.proxy_server import (prisma_client,
proxy_logging_obj,
user_api_key_cache)
verbose_logger.debug(
f"MCP team permission lookup: team_id={user_api_key_auth.team_id if user_api_key_auth else None}"
@ -541,12 +595,11 @@ class MCPRequestHandler:
user_api_key_auth
)
if key_object_permission is None and user_api_key_auth and user_api_key_auth.object_permission_id:
from litellm.proxy.auth.auth_checks import get_object_permission
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.auth.auth_checks import \
get_object_permission
from litellm.proxy.proxy_server import (prisma_client,
proxy_logging_obj,
user_api_key_cache)
if prisma_client is not None:
key_object_permission = await get_object_permission(
object_permission_id=user_api_key_auth.object_permission_id,
@ -614,6 +667,64 @@ class MCPRequestHandler:
)
return []
@staticmethod
async def _get_allowed_mcp_servers_for_end_user(
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
) -> List[str]:
"""
Get allowed MCP servers for an end user.
Returns the MCP servers from the end_user's object_permission.
"""
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.proxy.proxy_server import (prisma_client,
proxy_logging_obj,
user_api_key_cache)
if not user_api_key_auth or not user_api_key_auth.end_user_id:
return []
if prisma_client is None:
verbose_logger.debug("prisma_client is None")
return []
try:
# Use optimized get_end_user_object function with caching
end_user_obj = await get_end_user_object(
end_user_id=user_api_key_auth.end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route="/mcp",
)
if end_user_obj is None or end_user_obj.object_permission is None:
return []
# Get direct MCP servers
direct_mcp_servers = end_user_obj.object_permission.mcp_servers or []
# Get MCP servers from access groups
access_group_servers = (
await MCPRequestHandler._get_mcp_servers_from_access_groups(
end_user_obj.object_permission.mcp_access_groups or []
)
)
# Combine both lists
all_servers = direct_mcp_servers + access_group_servers
return list(set(all_servers))
except Exception as e:
verbose_logger.warning(
f"Failed to get allowed MCP servers for end_user: {str(e)}"
)
return []
@staticmethod
def _get_config_server_ids_for_access_groups(
config_mcp_servers, access_groups: List[str]
@ -660,9 +771,8 @@ class MCPRequestHandler:
try:
# Import here to avoid circular import
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import \
global_mcp_server_manager
# Use the new helper for config-loaded servers
server_ids = MCPRequestHandler._get_config_server_ids_for_access_groups(
@ -718,11 +828,9 @@ class MCPRequestHandler:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
) -> List[str]:
from litellm.proxy.auth.auth_checks import get_object_permission
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.proxy_server import (prisma_client,
proxy_logging_obj,
user_api_key_cache)
if user_api_key_auth is None:
return []
@ -758,11 +866,9 @@ class MCPRequestHandler:
Get MCP access groups for the team
"""
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.proxy_server import (prisma_client,
proxy_logging_obj,
user_api_key_cache)
if user_api_key_auth is None:
return []

103
litellm/proxy/_types.py
View File

@ -1,60 +1,40 @@
import enum
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal,
Optional, Union)
import httpx
from pydantic import (
BaseModel,
ConfigDict,
Field,
Json,
field_validator,
model_validator,
)
from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator,
model_validator)
from typing_extensions import Required, TypedDict
from litellm._uuid import uuid
from litellm.types.integrations.slack_alerting import AlertType
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIFileObject,
ResponsesAPIResponse,
)
from litellm.types.mcp import (
MCPAuth,
MCPAuthType,
MCPCredentials,
MCPTransport,
MCPTransportType,
)
from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject,
ResponsesAPIResponse)
from litellm.types.mcp import (MCPAuth, MCPAuthType, MCPCredentials,
MCPTransport, MCPTransportType)
from litellm.types.mcp_server.mcp_server_manager import MCPInfo
from litellm.types.router import RouterErrors, UpdateRouterConfig
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (
CallTypes,
CostBreakdown,
EmbeddingResponse,
GenericBudgetConfigType,
ImageResponse,
LiteLLMBatch,
LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase,
ModelResponse,
ProviderField,
StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse,
)
from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse,
GenericBudgetConfigType, ImageResponse,
LiteLLMBatch, LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase, ModelResponse,
ProviderField, StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse)
from litellm.types.videos.main import VideoObject
from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type
from .types_utils.utils import (get_instance_fn,
validate_custom_validate_return_type)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -1409,12 +1389,13 @@ class NewCustomerRequest(BudgetNewRequest):
blocked: bool = False # allow/disallow requests for this end-user
budget_id: Optional[str] = None # give either a budget_id or max_budget
spend: Optional[float] = None
allowed_model_region: Optional[
AllowedModelRegion
] = None # require all user requests to use models in this specific region
default_model: Optional[
str
] = None # if no equivalent model in allowed region - default all requests to this model
allowed_model_region: Optional[AllowedModelRegion] = (
None # require all user requests to use models in this specific region
)
default_model: Optional[str] = (
None # if no equivalent model in allowed region - default all requests to this model
)
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
@model_validator(mode="before")
@classmethod
@ -1436,12 +1417,13 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase):
blocked: bool = False # allow/disallow requests for this end-user
max_budget: Optional[float] = None
budget_id: Optional[str] = None # give either a budget_id or max_budget
allowed_model_region: Optional[
AllowedModelRegion
] = None # require all user requests to use models in this specific region
default_model: Optional[
str
] = None # if no equivalent model in allowed region - default all requests to this model
allowed_model_region: Optional[AllowedModelRegion] = (
None # require all user requests to use models in this specific region
)
default_model: Optional[str] = (
None # if no equivalent model in allowed region - default all requests to this model
)
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
class DeleteCustomerRequest(LiteLLMPydanticObjectBase):
@ -2339,7 +2321,8 @@ class UserAPIKeyAuth(
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
from litellm.constants import \
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
@ -2371,7 +2354,8 @@ class UserAPIKeyAuth(
This is used to track actions performed by automated system jobs.
"""
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
from litellm.constants import \
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
@ -2531,6 +2515,8 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase):
allowed_model_region: Optional[AllowedModelRegion] = None
default_model: Optional[str] = None
litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
object_permission_id: Optional[str] = None
object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
@model_validator(mode="before")
@classmethod
@ -2640,7 +2626,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):
@model_validator(mode="after")
def mask_api_keys(self):
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.litellm_core_utils.sensitive_data_masker import \
SensitiveDataMasker
masker = SensitiveDataMasker(sensitive_patterns={"key"})

90
litellm/proxy/auth/auth_checks.py
View File

@ -11,7 +11,8 @@ Run checks for:
import asyncio
import re
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union,
cast)
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
@ -20,41 +21,27 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.dual_cache import LimitedSizeOrderedDict
from litellm.constants import (
CLI_JWT_EXPIRATION_HOURS,
CLI_JWT_TOKEN_NAME,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
)
from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE)
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.proxy._types import (
RBAC_ROLES,
CallInfo,
LiteLLM_AccessGroupTable,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_JWTAuth,
LiteLLM_ObjectPermissionTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_OrganizationTable,
LiteLLM_TagTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LiteLLMRoutes,
LitellmUserRoles,
NewTeamRequest,
ProxyErrorTypes,
ProxyException,
RoleBasedPermissions,
SpecialModelNames,
UserAPIKeyAuth,
)
from litellm.proxy._types import (RBAC_ROLES, CallInfo,
LiteLLM_AccessGroupTable,
LiteLLM_BudgetTable, LiteLLM_EndUserTable,
Litellm_EntityType, LiteLLM_JWTAuth,
LiteLLM_ObjectPermissionTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_OrganizationTable, LiteLLM_TagTable,
LiteLLM_TeamMembership, LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable, LiteLLMRoutes,
LitellmUserRoles, NewTeamRequest,
ProxyErrorTypes, ProxyException,
RoleBasedPermissions, SpecialModelNames,
UserAPIKeyAuth)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
@ -366,7 +353,8 @@ async def common_checks(
_request_metadata: dict = request_body.get("metadata", {}) or {}
if _request_metadata.get("guardrails"):
# check if team allowed to modify guardrails
from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails
from litellm.proxy.guardrails.guardrail_helpers import \
can_modify_guardrails
can_modify: bool = can_modify_guardrails(team_object)
if can_modify is False:
@ -792,7 +780,7 @@ async def get_end_user_object(
try:
response = await prisma_client.db.litellm_endusertable.find_unique(
where={"user_id": end_user_id},
include={"litellm_budget_table": True},
include={"litellm_budget_table": True, "object_permission": True},
)
if response is None:
@ -1812,9 +1800,8 @@ class ExperimentalUIJWTToken:
def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str:
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
encrypt_value_helper
if user_info.user_role is None:
raise Exception("User role is required for experimental UI login")
@ -1860,9 +1847,8 @@ class ExperimentalUIJWTToken:
"""
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
encrypt_value_helper
if user_info.user_role is None:
raise Exception("User role is required for CLI JWT login")
@ -1901,9 +1887,8 @@ class ExperimentalUIJWTToken:
import json
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
decrypt_value_helper
decrypted_token = decrypt_value_helper(
hashed_token, key="ui_hash_key", exception_type="debug"
@ -2150,11 +2135,11 @@ async def _get_resources_from_access_groups(
# Lazy import to avoid circular imports
if prisma_client is None or user_api_key_cache is None:
from litellm.proxy.proxy_server import (
prisma_client as _prisma_client,
proxy_logging_obj as _proxy_logging_obj,
user_api_key_cache as _user_api_key_cache,
)
from litellm.proxy.proxy_server import prisma_client as _prisma_client
from litellm.proxy.proxy_server import \
proxy_logging_obj as _proxy_logging_obj
from litellm.proxy.proxy_server import \
user_api_key_cache as _user_api_key_cache
prisma_client = prisma_client or _prisma_client
user_api_key_cache = user_api_key_cache or _user_api_key_cache
@ -2936,7 +2921,8 @@ async def _tag_max_budget_check(
BudgetExceededError if any tag is over its max budget.
Triggers a budget alert if any tag is over its max budget.
"""
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.common_utils.http_parsing_utils import \
get_tags_from_request_body
if prisma_client is None:
return

67
litellm/proxy/litellm_pre_call_utils.py
View File

@ -10,15 +10,10 @@ import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles,
SpecialHeaders,
TeamCallbackMetadata,
UserAPIKeyAuth,
)
from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles, SpecialHeaders,
TeamCallbackMetadata, UserAPIKeyAuth)
# Cache special headers as a frozenset for O(1) lookup performance
_SPECIAL_HEADERS_CACHE = frozenset(
@ -28,12 +23,9 @@ from litellm.proxy.auth.route_checks import RouteChecks
from litellm.router import Router
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
LlmProviders,
ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls,
)
from litellm.types.utils import (LlmProviders, ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls)
service_logger_obj = ServiceLogging() # used for tracking latency on OTEL
@ -395,6 +387,24 @@ class LiteLLMProxyRequestSetup:
return user
@staticmethod
def get_end_user_from_headers(headers: dict) -> Optional[str]:
"""
Get the end user ID from the x-litellm-end-user header.
This header allows you to track customer/end-user spend and apply customer-specific
budgets and permissions without modifying the request body.
Returns:
Optional[str]: The end user ID if found in headers, None otherwise
"""
end_user = LiteLLMProxyRequestSetup._get_case_insensitive_header(
headers, "x-litellm-end-user"
)
if end_user is not None:
verbose_logger.info(f'found end_user "{end_user}" in x-litellm-end-user header')
return end_user
@staticmethod
def get_openai_org_id_from_headers(
headers: dict, general_settings: Optional[Dict] = None
@ -640,8 +650,7 @@ class LiteLLMProxyRequestSetup:
return data
from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
)
LiteLLM_ManagementEndpoint_MetadataFields_Premium)
# ignore any special fields
added_metadata = {}
@ -872,7 +881,15 @@ async def add_litellm_data_to_request( # noqa: PLR0915
general_settings, user_api_key_dict, _headers
)
# Parse user info from headers
# Parse end user ID from x-litellm-end-user header (takes precedence)
end_user_from_header = LiteLLMProxyRequestSetup.get_end_user_from_headers(_headers)
if end_user_from_header is not None:
if user_api_key_dict.end_user_id is None:
user_api_key_dict.end_user_id = end_user_from_header
if "user" not in data:
data["user"] = end_user_from_header
# Parse user info from headers (fallback to general_settings.user_header_name)
user = LiteLLMProxyRequestSetup.get_user_from_headers(_headers, general_settings)
if user is not None:
if user_api_key_dict.end_user_id is None:
@ -1530,12 +1547,9 @@ def _match_and_track_policies(
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.callback_utils import (
add_policy_sources_to_metadata,
add_policy_to_applied_policies_header,
)
from litellm.proxy.policy_engine.attachment_registry import (
get_attachment_registry,
)
add_policy_sources_to_metadata, add_policy_to_applied_policies_header)
from litellm.proxy.policy_engine.attachment_registry import \
get_attachment_registry
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
# Get matching policies via attachments (with match reasons for attribution)
@ -1670,9 +1684,8 @@ def add_guardrails_from_policy_engine(
user_api_key_dict: The user's API key authentication info
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.http_parsing_utils import (
get_tags_from_request_body,
)
from litellm.proxy.common_utils.http_parsing_utils import \
get_tags_from_request_body
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import PolicyMatchContext

182
litellm/proxy/management_endpoints/customer_endpoints.py
View File

@ -19,11 +19,13 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_daily_activity import \
get_daily_activity
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission, handle_update_object_permission_common)
from litellm.proxy.utils import handle_exception_on_proxy
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
from litellm.types.proxy.management_endpoints.common_daily_activity import \
SpendAnalyticsPaginatedResponse
router = APIRouter()
@ -107,9 +109,8 @@ async def unblock_user(data: BlockUsers):
```
"""
try:
from enterprise.enterprise_hooks.blocked_user_list import (
_ENTERPRISE_BlockedUserList,
)
from enterprise.enterprise_hooks.blocked_user_list import \
_ENTERPRISE_BlockedUserList
except ImportError:
raise HTTPException(
status_code=400,
@ -164,6 +165,43 @@ def new_budget_request(data: NewCustomerRequest) -> Optional[BudgetNewRequest]:
return None
async def _handle_customer_object_permission_update(
non_default_values: dict,
end_user_table_data_typed: Optional[LiteLLM_EndUserTable],
update_end_user_table_data: dict,
prisma_client,
) -> None:
"""
Handle object permission updates for customer endpoints.
Updates the update_end_user_table_data dict in place with the new object_permission_id
and removes the object_permission field to prevent it from being sent to the database.
Args:
non_default_values: Dictionary containing the update values including object_permission
end_user_table_data_typed: Existing end user table data
update_end_user_table_data: Dictionary to update with new object_permission_id
prisma_client: Prisma database client
"""
if "object_permission" in non_default_values:
existing_object_permission_id = (
end_user_table_data_typed.object_permission_id
if end_user_table_data_typed is not None
else None
)
object_permission_id = await handle_update_object_permission_common(
data_json=non_default_values,
existing_object_permission_id=existing_object_permission_id,
prisma_client=prisma_client,
)
if object_permission_id is not None:
update_end_user_table_data["object_permission_id"] = object_permission_id
# Remove object_permission from update_end_user_table_data to prevent DB write attempts
# object_permission is a read-only relationship field, not a writable column
update_end_user_table_data.pop("object_permission", None)
@router.post(
"/end_user/new",
tags=["Customer Management"],
@ -200,6 +238,16 @@ async def new_end_user(
- soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests.
- spend: Optional[float] - Specify initial spend for a given customer.
- budget_reset_at: Optional[str] - Specify the date and time when the budget should be reset.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - Customer-specific object permissions to control access to resources.
Supported fields:
* mcp_servers: List[str] - List of allowed MCP server IDs
* mcp_access_groups: List[str] - List of MCP access group names
* mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names (e.g., {"server_1": ["tool_a", "tool_b"]})
* vector_stores: List[str] - List of allowed vector store IDs
* agents: List[str] - List of allowed agent IDs
* agent_access_groups: List[str] - List of agent access group names
Example: {"mcp_servers": ["server_1", "server_2"], "vector_stores": ["vector_store_1"], "agents": ["agent_1"]}
IF null or {} then no object-level restrictions apply.
- Allow specifying allowed regions
@ -214,9 +262,22 @@ async def new_end_user(
"user_id" : "ishaan-jaff-3",
"allowed_region": "eu",
"budget_id": "free_tier",
"default_model": "azure/gpt-3.5-turbo-eu" <- all calls from this user, use this model?
"default_model": "azure/gpt-3.5-turbo-eu"
}'
# With object permissions
curl -L -X POST 'http://localhost:4000/customer/new' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"user_id": "user_1",
"object_permission": {
"mcp_servers": ["server_1"],
"mcp_access_groups": ["public_group"],
"vector_stores": ["vector_store_1"]
}
}'
# return end-user object
```
@ -233,11 +294,8 @@ async def new_end_user(
- end-user object
- currently allowed models
"""
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
llm_router,
prisma_client,
)
from litellm.proxy.proxy_server import (litellm_proxy_admin_name,
llm_router, prisma_client)
if prisma_client is None:
raise HTTPException(
@ -289,13 +347,34 @@ async def new_end_user(
if k not in BudgetNewRequest.model_fields.keys():
new_end_user_obj[k] = v
## Handle Object Permission - MCP Servers, Vector Stores etc.
new_end_user_obj = await _set_object_permission(
data_json=new_end_user_obj,
prisma_client=prisma_client,
)
# Ensure object_permission is not in the data being sent to create
# It should have been converted to object_permission_id by _set_object_permission
if "object_permission" in new_end_user_obj:
verbose_proxy_logger.warning(
f"object_permission still in new_end_user_obj after _set_object_permission: {new_end_user_obj.get('object_permission')}"
)
new_end_user_obj.pop("object_permission", None)
## WRITE TO DB ##
end_user_record = await prisma_client.db.litellm_endusertable.create(
data=new_end_user_obj, # type: ignore
include={"litellm_budget_table": True},
include={"litellm_budget_table": True, "object_permission": True},
)
return end_user_record
# Convert to dict and clean up recursive fields
response_dict = end_user_record.model_dump()
if response_dict.get("object_permission"):
# Remove reverse relations from object_permission
for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]:
response_dict["object_permission"].pop(field, None)
return response_dict
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.management_endpoints.customer_endpoints.new_end_user(): Exception occured - {}".format(
@ -351,7 +430,7 @@ async def end_user_info(
)
user_info = await prisma_client.db.litellm_endusertable.find_first(
where={"user_id": end_user_id}, include={"litellm_budget_table": True}
where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}
)
if user_info is None:
@ -361,7 +440,15 @@ async def end_user_info(
code=404,
param="end_user_id",
)
return user_info.model_dump(exclude_none=True)
# Convert to dict and clean up recursive fields
response_dict = user_info.model_dump(exclude_none=True)
if response_dict.get("object_permission"):
# Remove reverse relations from object_permission
for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]:
response_dict["object_permission"].pop(field, None)
return response_dict
except Exception as e:
verbose_proxy_logger.exception(
@ -401,6 +488,16 @@ async def update_end_user(
- default_model: Optional[str] = (
None # if no equivalent model in allowed region - default all requests to this model
)
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - Customer-specific object permissions to control access to resources.
Supported fields:
* mcp_servers: List[str] - List of allowed MCP server IDs
* mcp_access_groups: List[str] - List of MCP access group names
* mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names
* vector_stores: List[str] - List of allowed vector store IDs
* agents: List[str] - List of allowed agent IDs
* agent_access_groups: List[str] - List of agent access group names
Example: {"mcp_servers": ["server_1"], "vector_stores": ["vector_store_1"]}
IF null or {} then no object-level restrictions apply.
Example curl:
```
@ -412,11 +509,24 @@ async def update_end_user(
"budget_id": "paid_tier"
}'
See below for all params
# Updating object permissions
curl -L -X POST 'http://localhost:4000/customer/update' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"user_id": "user_1",
"object_permission": {
"mcp_servers": ["server_3"],
"vector_stores": ["vector_store_2", "vector_store_3"]
}
}'
See below for all params
```
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
from litellm.proxy.proxy_server import (litellm_proxy_admin_name,
prisma_client)
try:
data_json: dict = data.json()
@ -467,6 +577,14 @@ async def update_end_user(
elif k in LiteLLM_EndUserTable.model_fields.keys():
update_end_user_table_data[k] = v
## Handle object permission updates (MCP servers, vector stores, etc.)
await _handle_customer_object_permission_update(
non_default_values=non_default_values,
end_user_table_data_typed=end_user_table_data_typed,
update_end_user_table_data=update_end_user_table_data,
prisma_client=prisma_client,
)
## Check if we need to create a new budget (only if budget fields are provided, not just budget_id) ##
if budget_table_data:
if end_user_budget_table is None:
@ -498,11 +616,12 @@ async def update_end_user(
## Update user table, with update params + new budget id (if set) ##
verbose_proxy_logger.debug("/customer/update: Received data = %s", data)
if data.user_id is not None and len(data.user_id) > 0:
update_end_user_table_data["user_id"] = data.user_id # type: ignore
verbose_proxy_logger.debug("In update customer, user_id condition block.")
response = await prisma_client.db.litellm_endusertable.update(
where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True} # type: ignore
where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True, "object_permission": True} # type: ignore
)
if response is None:
raise ValueError(
@ -511,7 +630,15 @@ async def update_end_user(
verbose_proxy_logger.debug(
f"received response from updating prisma client. response={response}"
)
return response
# Convert to dict and clean up recursive fields
response_dict = response.model_dump()
if response_dict.get("object_permission"):
# Remove reverse relations from object_permission
for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]:
response_dict["object_permission"].pop(field, None)
return response_dict
else:
raise ValueError(f"user_id is required, passed user_id = {data.user_id}")
@ -663,12 +790,17 @@ async def list_end_user(
)
response = await prisma_client.db.litellm_endusertable.find_many(
include={"litellm_budget_table": True}
include={"litellm_budget_table": True, "object_permission": True}
)
returned_response: List[LiteLLM_EndUserTable] = []
for item in response:
returned_response.append(LiteLLM_EndUserTable(**item.model_dump()))
item_dict = item.model_dump()
# Remove reverse relations from object_permission
if item_dict.get("object_permission"):
for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]:
item_dict["object_permission"].pop(field, None)
returned_response.append(LiteLLM_EndUserTable(**item_dict))
return returned_response
except Exception as e:
@ -706,9 +838,7 @@ async def get_customer_daily_activity(
"""
Get daily activity for specific organizations or all accessible organizations.
"""
from litellm.proxy.proxy_server import (
prisma_client,
)
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(

3
litellm/proxy/schema.prisma
View File

@ -233,6 +233,7 @@ model LiteLLM_ObjectPermissionTable {
verification_tokens LiteLLM_VerificationToken[]
organizations LiteLLM_OrganizationTable[]
users LiteLLM_UserTable[]
end_users LiteLLM_EndUserTable[]
}
// Holds the MCP server configuration
@ -403,7 +404,9 @@ model LiteLLM_EndUserTable {
allowed_model_region String? // require all user requests to use models in this specific region
default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model.
budget_id String?
object_permission_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
blocked Boolean @default(false)
}

3
schema.prisma
View File

@ -233,6 +233,7 @@ model LiteLLM_ObjectPermissionTable {
verification_tokens LiteLLM_VerificationToken[]
organizations LiteLLM_OrganizationTable[]
users LiteLLM_UserTable[]
end_users LiteLLM_EndUserTable[]
}
// Holds the MCP server configuration
@ -403,7 +404,9 @@ model LiteLLM_EndUserTable {
allowed_model_region String? // require all user requests to use models in this specific region
default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model.
budget_id String?
object_permission_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
blocked Boolean @default(false)
}

670
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View File

@ -15,9 +15,8 @@ sys.path.insert(
from starlette.datastructures import Headers
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import \
MCPRequestHandler
from litellm.proxy._types import SpecialHeaders, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -1176,7 +1175,8 @@ class TestMCPAccessGroupsE2E:
@pytest.mark.asyncio
def test_mcp_path_based_server_segregation(monkeypatch):
# Import the MCP server FastAPI app and context getter
from litellm.proxy._experimental.mcp_server.server import app, get_auth_context
from litellm.proxy._experimental.mcp_server.server import (
app, get_auth_context)
captured_mcp_servers = {}
@ -1277,7 +1277,8 @@ async def test_get_team_object_permission_with_already_loaded_permission():
Test that _get_team_object_permission returns the already loaded object_permission
from the team object without making an additional DB call.
"""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable
from litellm.proxy._types import (LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable)
# Create mock object permission
mock_object_permission = LiteLLM_ObjectPermissionTable(
@ -1340,7 +1341,8 @@ async def test_get_team_object_permission_with_core_auth_auto_loading():
the team object returned by get_team_object() should already have object_permission loaded
when an object_permission_id exists.
"""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable
from litellm.proxy._types import (LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable)
# Create mock object permission
mock_object_permission = LiteLLM_ObjectPermissionTable(
@ -1595,3 +1597,659 @@ async def test_get_allowed_mcp_servers_for_key_prefers_in_memory_permission():
assert set(result) == {"direct-server", "group-server"}
mock_get_perm.assert_not_called()
mock_access_groups.assert_called_once_with(["grp-alpha"])
@pytest.mark.asyncio
class TestEndUserMCPPermissions:
"""Test suite for end_user MCP permission functionality"""
@pytest.mark.parametrize(
"key_servers,team_servers,end_user_servers,require_flag,expected_result,scenario",
[
# Test case 1: End user with permissions, key/team with permissions - intersection
(
["server1", "server2", "server3"],
["server1", "server2", "server3", "server4"],
["server2", "server3", "server5"],
False,
["server2", "server3"],
"end_user_intersects_with_key_team",
),
# Test case 2: End user with permissions, key/team with no permissions (empty) - use end_user permissions
(
[],
[],
["server1", "server2"],
False,
["server1", "server2"],
"end_user_with_empty_key_team",
),
# Test case 3: End user with no permissions, require_flag=True - fall back to key/team restrictions
(
["server1", "server2"],
["server1", "server2"],
[],
True,
["server1", "server2"],
"end_user_no_perms_with_require_flag_true_fallback_to_key_team",
),
# Test case 4: End user with no permissions, require_flag=False - use key/team permissions
(
["server1", "server2"],
["server1", "server2"],
[],
False,
["server1", "server2"],
"end_user_no_perms_with_require_flag_false",
),
# Test case 5: End user with permissions, require_flag=True - still intersect with key/team
(
["server1", "server2"],
["server1", "server2", "server3"],
["server2", "server3", "server4"],
True,
["server2"],
"end_user_with_require_flag_true_still_intersects",
),
# Test case 6: End user has permissions but no overlap with key/team
(
["server1", "server2"],
["server1", "server2"],
["server3", "server4"],
False,
[],
"end_user_no_overlap_with_key_team",
),
# Test case 7: Key has servers, team empty, end_user has overlap with key
(
["server1", "server2", "server3"],
[],
["server2", "server3", "server4"],
False,
["server2", "server3"],
"key_only_intersects_with_end_user",
),
# Test case 8: All have different servers
(
["server1"],
["server1", "server2"],
["server3"],
False,
[],
"all_different_servers",
),
# Test case 9: require_flag=True, end_user has permissions but key/team empty
(
[],
[],
["server1", "server2"],
True,
["server1", "server2"],
"require_flag_true_with_empty_key_team_uses_end_user",
),
],
)
async def test_get_allowed_mcp_servers_with_end_user(
self,
key_servers,
team_servers,
end_user_servers,
require_flag,
expected_result,
scenario,
):
"""Test get_allowed_mcp_servers with end_user permissions"""
# Create a mock user with end_user_id
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id="test-end-user",
)
# Mock the helper methods
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
) as mock_team_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_end_user"
) as mock_end_user_servers:
with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings:
# Set up return values
mock_key_servers.return_value = key_servers
mock_team_servers.return_value = team_servers
mock_end_user_servers.return_value = end_user_servers
mock_general_settings.get.return_value = require_flag
# Call the method
result = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=mock_user_auth
)
# Assert the result (order-independent comparison)
assert sorted(result) == sorted(expected_result)
# Verify helper methods were called
mock_key_servers.assert_called_once_with(mock_user_auth)
mock_team_servers.assert_called_once_with(mock_user_auth)
mock_end_user_servers.assert_called_once_with(mock_user_auth)
@pytest.mark.parametrize(
"key_servers,team_servers,expected_result,scenario",
[
# Test case 1: No end_user_id - should use key/team logic only
(
["server1", "server2"],
["server1", "server2", "server3"],
["server1", "server2"],
"no_end_user_id_uses_key_team",
),
# Test case 2: No end_user_id with empty key
(
[],
["server1", "server2"],
["server1", "server2"],
"no_end_user_id_inherits_from_team",
),
],
)
async def test_get_allowed_mcp_servers_without_end_user_id(
self,
key_servers,
team_servers,
expected_result,
scenario,
):
"""Test get_allowed_mcp_servers when end_user_id is not set"""
# Create a mock user WITHOUT end_user_id
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id=None,
)
# Mock the helper methods
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
) as mock_team_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_end_user"
) as mock_end_user_servers:
# Set up return values
mock_key_servers.return_value = key_servers
mock_team_servers.return_value = team_servers
# Call the method
result = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=mock_user_auth
)
# Assert the result (order-independent comparison)
assert sorted(result) == sorted(expected_result)
# Verify helper methods were called correctly
mock_key_servers.assert_called_once_with(mock_user_auth)
mock_team_servers.assert_called_once_with(mock_user_auth)
# end_user method should NOT be called since no end_user_id
mock_end_user_servers.assert_not_called()
async def test_get_allowed_mcp_servers_for_end_user_with_valid_user(self):
"""Test _get_allowed_mcp_servers_for_end_user with valid end_user"""
from litellm.proxy._types import (LiteLLM_EndUserTable,
LiteLLM_ObjectPermissionTable)
# Create mock object permission
mock_object_permission = LiteLLM_ObjectPermissionTable(
object_permission_id="perm-123",
mcp_servers=["server1", "server2"],
mcp_access_groups=["group1"],
)
# Create mock end_user object
mock_end_user = LiteLLM_EndUserTable(
user_id="test-end-user",
object_permission=mock_object_permission,
object_permission_id="perm-123",
)
# Create mock user auth
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
end_user_id="test-end-user",
)
# Mock the dependencies
mock_prisma = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user:
with patch.object(
MCPRequestHandler, "_get_mcp_servers_from_access_groups"
) as mock_get_access_groups:
# Configure mocks
mock_get_end_user.return_value = mock_end_user
mock_get_access_groups.return_value = ["group-server1", "group-server2"]
# Call the method
result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(
mock_user_auth
)
# Assert the result contains both direct and access group servers
assert set(result) == {"server1", "server2", "group-server1", "group-server2"}
# Verify methods were called
mock_get_end_user.assert_called_once()
mock_get_access_groups.assert_called_once_with(["group1"])
async def test_get_allowed_mcp_servers_for_end_user_with_no_permission(self):
"""Test _get_allowed_mcp_servers_for_end_user when end_user has no object_permission"""
from litellm.proxy._types import LiteLLM_EndUserTable
# Create mock end_user object without object_permission
mock_end_user = LiteLLM_EndUserTable(
user_id="test-end-user",
object_permission=None,
)
# Create mock user auth
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
end_user_id="test-end-user",
)
# Mock the dependencies
mock_prisma = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user:
# Configure mock
mock_get_end_user.return_value = mock_end_user
# Call the method
result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(
mock_user_auth
)
# Assert empty list is returned
assert result == []
# Verify method was called
mock_get_end_user.assert_called_once()
async def test_get_allowed_mcp_servers_for_end_user_without_end_user_id(self):
"""Test _get_allowed_mcp_servers_for_end_user when no end_user_id is set"""
# Create mock user auth without end_user_id
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
end_user_id=None,
)
# Call the method
result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(
mock_user_auth
)
# Assert empty list is returned
assert result == []
async def test_get_allowed_mcp_servers_for_end_user_without_prisma_client(self):
"""Test _get_allowed_mcp_servers_for_end_user when prisma_client is None"""
# Create mock user auth
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
end_user_id="test-end-user",
)
# Mock prisma_client as None
with patch("litellm.proxy.proxy_server.prisma_client", None):
# Call the method
result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(
mock_user_auth
)
# Assert empty list is returned
assert result == []
@pytest.mark.asyncio
class TestEndUserMCPPermissionsFallbackBehavior:
"""Test suite for end_user MCP permission fallback behavior.
Tests the updated behavior where when end_user has NO explicit MCP permissions defined,
the system falls back to key/team restrictions instead of blocking all access.
"""
@pytest.mark.parametrize(
"key_servers,team_servers,end_user_servers,require_flag,expected_result,scenario",
[
# Test case 1: End user with NO permissions, require_flag=True - should fall back to key/team
(
["server1", "server2"],
["server1", "server2", "server3"],
[],
True,
["server1", "server2"],
"no_end_user_perms_require_flag_true_fallback_to_key_team_intersection",
),
# Test case 2: End user with NO permissions, require_flag=False - should fall back to key/team
(
["server1", "server2"],
["server1", "server2", "server3"],
[],
False,
["server1", "server2"],
"no_end_user_perms_require_flag_false_fallback_to_key_team",
),
# Test case 3: End user with NO permissions, only key has permissions - should use key permissions
(
["server1", "server2", "server3"],
[],
[],
True,
["server1", "server2", "server3"],
"no_end_user_perms_only_key_has_permissions",
),
# Test case 4: End user with NO permissions, only team has permissions - should inherit from team
(
[],
["team_server1", "team_server2"],
[],
True,
["team_server1", "team_server2"],
"no_end_user_perms_only_team_has_permissions",
),
# Test case 5: End user with NO permissions, both key and team empty - should return empty
(
[],
[],
[],
True,
[],
"no_permissions_anywhere_returns_empty",
),
# Test case 6: End user with NO permissions, key and team have different servers
(
["server1", "server2"],
["server3", "server4"],
[],
True,
[],
"no_end_user_perms_key_team_no_overlap_returns_empty",
),
# Test case 7: End user HAS permissions, require_flag=True - should still intersect
(
["server1", "server2", "server3"],
["server1", "server2", "server3"],
["server2", "server3", "server4"],
True,
["server2", "server3"],
"end_user_has_perms_require_flag_true_still_intersects",
),
# Test case 8: End user HAS permissions but empty key/team - should use end_user permissions
(
[],
[],
["server1", "server2"],
True,
["server1", "server2"],
"end_user_has_perms_empty_key_team_uses_end_user",
),
# Test case 9: End user NO permissions, complex key/team intersection
(
["server1", "server2", "server3", "server4"],
["server2", "server3", "server5"],
[],
True,
["server2", "server3"],
"no_end_user_perms_complex_key_team_intersection",
),
# Test case 10: End user HAS permissions with no overlap - should return empty
(
["server1", "server2"],
["server1", "server2"],
["server5", "server6"],
True,
[],
"end_user_has_perms_no_overlap_returns_empty",
),
],
)
async def test_end_user_permission_fallback_to_key_team(
self,
key_servers,
team_servers,
end_user_servers,
require_flag,
expected_result,
scenario,
):
"""Test that end_user permissions fall back to key/team when not defined"""
# Create a mock user with end_user_id
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id="test-end-user",
)
# Mock the helper methods
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
) as mock_team_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_end_user"
) as mock_end_user_servers:
with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings:
# Set up return values
mock_key_servers.return_value = key_servers
mock_team_servers.return_value = team_servers
mock_end_user_servers.return_value = end_user_servers
mock_general_settings.get.return_value = require_flag
# Call the method
result = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=mock_user_auth
)
# Assert the result (order-independent comparison)
assert sorted(result) == sorted(expected_result), (
f"Test scenario '{scenario}' failed: "
f"Expected {sorted(expected_result)}, got {sorted(result)}"
)
# Verify helper methods were called
mock_key_servers.assert_called_once_with(mock_user_auth)
mock_team_servers.assert_called_once_with(mock_user_auth)
mock_end_user_servers.assert_called_once_with(mock_user_auth)
async def test_end_user_no_permissions_with_require_flag_true_detailed(self):
"""
Detailed test: When require_end_user_mcp_access_defined=True and end_user has NO permissions,
should fall back to key/team restrictions (not block all access).
This is the main behavior change being tested.
"""
from litellm.proxy._types import LiteLLM_EndUserTable
# Create mock end_user object without object_permission
mock_end_user = LiteLLM_EndUserTable(
user_id="test-end-user",
object_permission=None,
)
# Create mock user auth
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id="test-end-user",
)
# Mock the dependencies
mock_prisma = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
) as mock_team_servers:
with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings:
# Configure mocks
mock_get_end_user.return_value = mock_end_user
mock_key_servers.return_value = ["server1", "server2"]
mock_team_servers.return_value = ["server1", "server2", "server3"]
mock_general_settings.get.return_value = True # require_flag=True
# Call the method
result = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=mock_user_auth
)
# CRITICAL: Should fall back to key/team intersection, not block all
assert sorted(result) == ["server1", "server2"], (
"When end_user has NO permissions and require_flag=True, "
"should fall back to key/team restrictions, not block all access"
)
# Verify methods were called
mock_get_end_user.assert_called_once()
mock_key_servers.assert_called_once_with(mock_user_auth)
mock_team_servers.assert_called_once_with(mock_user_auth)
async def test_end_user_no_permissions_without_require_flag(self):
"""
Test: When require_end_user_mcp_access_defined=False and end_user has NO permissions,
should use key/team restrictions (default behavior).
"""
from litellm.proxy._types import LiteLLM_EndUserTable
# Create mock end_user object without object_permission
mock_end_user = LiteLLM_EndUserTable(
user_id="test-end-user",
object_permission=None,
)
# Create mock user auth
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id="test-end-user",
)
# Mock the dependencies
mock_prisma = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
) as mock_team_servers:
with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings:
# Configure mocks
mock_get_end_user.return_value = mock_end_user
mock_key_servers.return_value = ["server1", "server2", "server3"]
mock_team_servers.return_value = []
mock_general_settings.get.return_value = False # require_flag=False
# Call the method
result = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=mock_user_auth
)
# Should use key permissions when team is empty
assert sorted(result) == ["server1", "server2", "server3"]
# Verify methods were called
mock_get_end_user.assert_called_once()
mock_key_servers.assert_called_once_with(mock_user_auth)
mock_team_servers.assert_called_once_with(mock_user_auth)
async def test_end_user_with_permissions_still_enforces_intersection(self):
"""
Test: When end_user HAS explicit permissions, should still enforce intersection
with key/team restrictions (behavior unchanged).
"""
from litellm.proxy._types import (LiteLLM_EndUserTable,
LiteLLM_ObjectPermissionTable)
# Create mock object permission for end_user
mock_object_permission = LiteLLM_ObjectPermissionTable(
object_permission_id="perm-123",
mcp_servers=["server2", "server3", "server4"],
mcp_access_groups=[],
)
# Create mock end_user object WITH object_permission
mock_end_user = LiteLLM_EndUserTable(
user_id="test-end-user",
object_permission=mock_object_permission,
object_permission_id="perm-123",
)
# Create mock user auth
mock_user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
end_user_id="test-end-user",
)
# Mock the dependencies
mock_prisma = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
) as mock_team_servers:
with patch.object(
MCPRequestHandler, "_get_mcp_servers_from_access_groups"
) as mock_access_groups:
with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings:
# Configure mocks
mock_get_end_user.return_value = mock_end_user
mock_key_servers.return_value = ["server1", "server2", "server3"]
mock_team_servers.return_value = ["server1", "server2", "server3"]
mock_access_groups.return_value = [] # No access group servers
mock_general_settings.get.return_value = True
# Call the method
result = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=mock_user_auth
)
# Should return intersection of key/team AND end_user
# key/team intersection = ["server1", "server2", "server3"]
# end_user = ["server2", "server3", "server4"]
# final intersection = ["server2", "server3"]
assert sorted(result) == ["server2", "server3"]
# Verify methods were called
mock_get_end_user.assert_called_once()
mock_key_servers.assert_called_once_with(mock_user_auth)
mock_team_servers.assert_called_once_with(mock_user_auth)