From 29b743ca2f3befe565e79887b41c70696e177e1b Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 17 Feb 2026 22:49:47 -0800 Subject: [PATCH] =?UTF-8?q?Revert=20"End=20users=20-=20Allow=20giving=20en?= =?UTF-8?q?d=20users=20access=20to=20specific=20mcp=20servers=20=20(#?= =?UTF-8?q?=E2=80=A6"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 1f521be0f2dde533d5781cca50f042168565d59d. --- docs/my-website/docs/mcp.md | 62 -- docs/my-website/docs/proxy/customers.md | 198 +----- docs/my-website/docs/proxy/request_headers.md | 2 - .../migration.sql | 6 - .../litellm_proxy_extras/schema.prisma | 3 - .../mcp_server/auth/user_api_key_auth_mcp.py | 172 +---- litellm/proxy/_types.py | 103 +-- litellm/proxy/auth/auth_checks.py | 90 ++- litellm/proxy/litellm_pre_call_utils.py | 67 +- .../customer_endpoints.py | 182 +---- litellm/proxy/schema.prisma | 3 - schema.prisma | 3 - .../auth/test_user_api_key_auth_mcp.py | 670 +----------------- 13 files changed, 206 insertions(+), 1355 deletions(-) delete mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 8be507c6a6a..84d10c25931 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -808,68 +808,6 @@ 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. diff --git a/docs/my-website/docs/proxy/customers.md b/docs/my-website/docs/proxy/customers.md index 8eb8aea6fb4..1101884c36b 100644 --- a/docs/my-website/docs/proxy/customers.md +++ b/docs/my-website/docs/proxy/customers.md @@ -2,7 +2,7 @@ import Image from '@theme/IdealImage'; import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; -# Customers / End-Users +# Customers / End-User Budgets Track spend, set budgets for your customers. @@ -10,11 +10,9 @@ Track spend, set budgets for your customers. ### 1. Make LLM API call w/ Customer ID -You can pass the customer ID in two ways: +Make a /chat/completions call, pass 'user' - First call Works -**Option 1: In the request body** (using the `user` field) - -```bash showLineNumbers title="Make request with customer ID in body" +```bash showLineNumbers title="Make request with customer ID" curl -X POST 'http://0.0.0.0:4000/chat/completions' \ --header 'Content-Type: application/json' \ --header 'Authorization: Bearer sk-1234' \ # 👈 YOUR PROXY KEY @@ -30,30 +28,6 @@ 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. @@ -149,171 +123,7 @@ Expected Response -## 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 +## Setting Customer Budgets Set customer budgets (e.g. monthly budgets, tpm/rpm limits) on LiteLLM Proxy diff --git a/docs/my-website/docs/proxy/request_headers.md b/docs/my-website/docs/proxy/request_headers.md index b26d21823f2..090c201f884 100644 --- a/docs/my-website/docs/proxy/request_headers.md +++ b/docs/my-website/docs/proxy/request_headers.md @@ -20,8 +20,6 @@ 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. diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql deleted file mode 100644 index 5c5dc6fd6f1..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql +++ /dev/null @@ -1,6 +0,0 @@ --- 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; - diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8bd46672ae6..441c2cdf70d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -233,7 +233,6 @@ model LiteLLM_ObjectPermissionTable { verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] - end_users LiteLLM_EndUserTable[] } // Holds the MCP server configuration @@ -404,9 +403,7 @@ 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) } diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 0d1a11e760c..ed4fb133478 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -6,8 +6,12 @@ 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 @@ -332,28 +336,15 @@ 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. - - 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 + Get list of allowed MCP servers for the given user/key based on 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: - # Get allowed servers from key and team + allowed_mcp_servers: List[str] = [] allowed_mcp_servers_for_key = ( await MCPRequestHandler._get_allowed_mcp_servers_for_key( user_api_key_auth @@ -366,9 +357,8 @@ class MCPRequestHandler: ) ######################################################### - # Calculate key/team allowed servers using inheritance and intersection logic + # If team has mcp_servers, handle 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 @@ -381,52 +371,6 @@ 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)}") @@ -458,9 +402,11 @@ 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}" @@ -595,11 +541,12 @@ 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, @@ -667,64 +614,6 @@ 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] @@ -771,8 +660,9 @@ 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( @@ -828,9 +718,11 @@ 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 [] @@ -866,9 +758,11 @@ 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 [] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e1cf5a1e8ba..c327dd130a5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,40 +1,60 @@ 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 @@ -1389,13 +1409,12 @@ 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 - ) - object_permission: Optional[LiteLLM_ObjectPermissionBase] = 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 @model_validator(mode="before") @classmethod @@ -1417,13 +1436,12 @@ 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 - ) - object_permission: Optional[LiteLLM_ObjectPermissionBase] = 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 class DeleteCustomerRequest(LiteLLMPydanticObjectBase): @@ -2321,8 +2339,7 @@ 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, @@ -2354,8 +2371,7 @@ 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, @@ -2515,8 +2531,6 @@ 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 @@ -2626,8 +2640,7 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): @model_validator(mode="after") def mask_api_keys(self): - from litellm.litellm_core_utils.sensitive_data_masker import \ - SensitiveDataMasker + from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker masker = SensitiveDataMasker(sensitive_patterns={"key"}) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5a8a9f4e4df..3eb6f28ddfc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,8 +11,7 @@ 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 @@ -21,27 +20,41 @@ 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 @@ -353,8 +366,7 @@ 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: @@ -780,7 +792,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, "object_permission": True}, + include={"litellm_budget_table": True}, ) if response is None: @@ -1800,8 +1812,9 @@ 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") @@ -1847,8 +1860,9 @@ 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") @@ -1887,8 +1901,9 @@ 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" @@ -2135,11 +2150,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 - 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 + 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, + ) prisma_client = prisma_client or _prisma_client user_api_key_cache = user_api_key_cache or _user_api_key_cache @@ -2921,8 +2936,7 @@ 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 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b1d4c8e507b..4d77af513a8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -10,10 +10,15 @@ 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( @@ -23,9 +28,12 @@ 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 @@ -387,24 +395,6 @@ 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 @@ -650,7 +640,8 @@ 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 = {} @@ -881,15 +872,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915 general_settings, user_api_key_dict, _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) + # Parse user info from headers user = LiteLLMProxyRequestSetup.get_user_from_headers(_headers, general_settings) if user is not None: if user_api_key_dict.end_user_id is None: @@ -1547,9 +1530,12 @@ 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) @@ -1684,8 +1670,9 @@ 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 diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 43eb1ea51ec..9ff0fe6e590 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -19,13 +19,11 @@ 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.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity router = APIRouter() @@ -109,8 +107,9 @@ 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, @@ -165,43 +164,6 @@ 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"], @@ -238,16 +200,6 @@ 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 @@ -262,22 +214,9 @@ 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" + "default_model": "azure/gpt-3.5-turbo-eu" <- all calls from this user, use this model? }' - # 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 ``` @@ -294,8 +233,11 @@ 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( @@ -347,34 +289,13 @@ 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, "object_permission": True}, + include={"litellm_budget_table": True}, ) - # 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 + return end_user_record except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.management_endpoints.customer_endpoints.new_end_user(): Exception occured - {}".format( @@ -430,7 +351,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, "object_permission": True} + where={"user_id": end_user_id}, include={"litellm_budget_table": True} ) if user_info is None: @@ -440,15 +361,7 @@ async def end_user_info( code=404, param="end_user_id", ) - - # 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 + return user_info.model_dump(exclude_none=True) except Exception as e: verbose_proxy_logger.exception( @@ -488,16 +401,6 @@ 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: ``` @@ -509,24 +412,11 @@ async def update_end_user( "budget_id": "paid_tier" }' - # 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 + 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() @@ -577,14 +467,6 @@ 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: @@ -616,12 +498,11 @@ 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, "object_permission": True} # type: ignore + where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True} # type: ignore ) if response is None: raise ValueError( @@ -630,15 +511,7 @@ async def update_end_user( verbose_proxy_logger.debug( f"received response from updating prisma client. response={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 + return response else: raise ValueError(f"user_id is required, passed user_id = {data.user_id}") @@ -790,17 +663,12 @@ async def list_end_user( ) response = await prisma_client.db.litellm_endusertable.find_many( - include={"litellm_budget_table": True, "object_permission": True} + include={"litellm_budget_table": True} ) returned_response: List[LiteLLM_EndUserTable] = [] for item in response: - 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)) + returned_response.append(LiteLLM_EndUserTable(**item.model_dump())) return returned_response except Exception as e: @@ -838,7 +706,9 @@ 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( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8bd46672ae6..441c2cdf70d 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -233,7 +233,6 @@ model LiteLLM_ObjectPermissionTable { verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] - end_users LiteLLM_EndUserTable[] } // Holds the MCP server configuration @@ -404,9 +403,7 @@ 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) } diff --git a/schema.prisma b/schema.prisma index 8bd46672ae6..441c2cdf70d 100644 --- a/schema.prisma +++ b/schema.prisma @@ -233,7 +233,6 @@ model LiteLLM_ObjectPermissionTable { verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] - end_users LiteLLM_EndUserTable[] } // Holds the MCP server configuration @@ -404,9 +403,7 @@ 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) } diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index da5b06f92f5..c2dbc94f721 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -15,8 +15,9 @@ 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 @@ -1175,8 +1176,7 @@ 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,8 +1277,7 @@ 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( @@ -1341,8 +1340,7 @@ 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( @@ -1597,659 +1595,3 @@ 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)