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 4090666e16f..baac6b0cbd3 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 @@ -5,7 +5,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timezone -from types import MappingProxyType +from types import EllipsisType, MappingProxyType from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import HTTPException @@ -464,6 +464,7 @@ def _admission_failure_fallback( class MCPServerAccess: server_ids: tuple[str, ...] scope: Literal["unscoped", "scoped", "unresolved"] = "unscoped" + end_user_denies_mcp: bool = False @dataclass(frozen=True, slots=True) @@ -1830,7 +1831,21 @@ class MCPRequestHandler: Returns: List[str]: List of allowed MCP servers by server id """ - from litellm.proxy.proxy_server import general_settings + from litellm.proxy.proxy_server import general_settings, prisma_client + + try: + end_user_permission: Final = ( + await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if user_api_key_auth and user_api_key_auth.end_user_id and prisma_client is not None + else None + ) + except UnloadableEntitlementError as error: + verbose_logger.warning("Denying MCP access, customer entitlement unreadable: %s", error) + return MCPServerAccess(server_ids=(), scope="scoped", end_user_denies_mcp=True) + if end_user_permission is not None and SpecialMCPServerNames.no_mcp_servers.value in ( + end_user_permission.mcp_servers or [] + ): + return MCPServerAccess(server_ids=(), scope="scoped", end_user_denies_mcp=True) if managed_agent_policy(user_api_key_auth) is not None: from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_servers @@ -1911,7 +1926,7 @@ class MCPRequestHandler: ######################################################### if user_api_key_auth and user_api_key_auth.end_user_id: allowed_mcp_servers_for_end_user: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( - user_api_key_auth + user_api_key_auth, object_permission=end_user_permission ) # If end_user has explicit MCP server permissions, apply intersection @@ -3428,6 +3443,8 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_end_user( user_api_key_auth: UserAPIKeyAuth | None = None, + *, + object_permission: LiteLLM_ObjectPermissionTable | None | EllipsisType = ..., ) -> list[str]: """ Get allowed MCP servers for an end user. @@ -3445,9 +3462,19 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return [] - object_permission = await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) - if object_permission is None: + permission: Final = ( + await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if object_permission is ... + else object_permission + ) + if permission is None: return [] + if SpecialMCPServerNames.no_mcp_servers.value in (permission.mcp_servers or []): + return [SpecialMCPServerNames.no_mcp_servers.value] + + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) try: # Permission entries may be server_ids OR names/aliases — expand to ids. @@ -3455,25 +3482,30 @@ class MCPRequestHandler: global_mcp_server_manager, ) - direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permission.mcp_servers or []) + direct_mcp_servers: Final = global_mcp_server_manager.expand_permission_list( + permission.mcp_servers or [] + ) # Get MCP servers from access groups access_group_servers: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups( - object_permission.mcp_access_groups or [], + permission.mcp_access_groups or [], requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) # servers referenced in tool permissions should also be accessible tool_perm_servers: Final = list( - global_mcp_server_manager.expand_tool_permissions(object_permission.mcp_tool_permissions).keys() + global_mcp_server_manager.expand_tool_permissions(permission.mcp_tool_permissions).keys() ) - # Combine all lists - all_servers: Final = direct_mcp_servers + access_group_servers + tool_perm_servers - return list(MCPRequestHandler._server_grants_with_group_scope(set(all_servers), object_permission)) + all_servers: Final = frozenset( + (*direct_mcp_servers, *access_group_servers, *tool_perm_servers, *toolset_grants) + ) + return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, permission)) except Exception as e: - verbose_logger.warning("Failed to get allowed MCP servers for end_user: %s", e) - return [] + raise UnloadableEntitlementError( + f"end user {user_api_key_auth.end_user_id!r} has an object permission " + "whose MCP servers could not be read" + ) from e @staticmethod async def _get_user_object_permission( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b724a61cd94..b83e5a1b496 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3691,6 +3691,14 @@ class MCPServerManager: ) try: + customer_access: Final = ( + access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) + if user_api_key_auth and user_api_key_auth.end_user_id + else None + ) + if customer_access is not None and customer_access.end_user_denies_mcp: + return [] + # Admin view with no explicit object permission and no entitlement ceiling resolves the # whole registry, for keys AND admitted session subjects alike (one predicate owns the # question). Seeded into the union rather than returned early so the session resource @@ -3699,7 +3707,7 @@ class MCPServerManager: resolved_access: Final = ( MCPServerAccess(server_ids=()) if admin_unscoped - else access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) + else customer_access or access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) ) resolved_server_ids: Final = ( set(self.get_registry().keys()) if admin_unscoped else set(resolved_access.server_ids) diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index 12f707cf6f1..7fd51bd6e8c 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -11,15 +11,17 @@ Each case asserts the feature actually happened, not just a 200. Coverage matrix - Vertex (gemini-2.5-flash): explicit context caching via ``cache_control`` with a 5-minute ttl. litellm builds the Vertex cache before the generate call, so a never-seen prefix must come back cached on its very first call - (Gemini's implicit caching cannot hit a cold prefix), the cached count must - cover the marked block, and the spend row must be billed below the uncached - price of the cached tokens. Vertex's cache create is nondeterministic: + (Gemini's implicit caching cannot hit a cold prefix). That call bills the + cache write as ``cache_creation_input_tokens`` on top of Vertex's own + prompt count, the cached count must cover the marked block, the spend row + must carry the write at the input rate and the read below it, and an + identical second call must read the cache without building a new one. + Vertex's cache create is nondeterministic: identical bodies come back 200 or with the minimum-token 400 ("The cached content is of 1 tokens"), in failure bursts of 45 seconds and more, so up to eight never-seen prefixes are tried with a pause after each rejection. The - billing check prices the cached tokens rather than prompt_tokens, which - Vertex reports inclusive of the cached prefix on some calls and exclusive - of it on others. + checks hold whether Vertex reports its prompt count inclusive of the cached + prefix or exclusive of it, which it has done both ways. - Anthropic (claude-haiku-4-5, direct): the same ``cache_control`` prefix over the OpenAI-compatible route; the second call must report cache-read tokens > 0. - OpenAI (gpt-5.6): automatic prompt caching needs no request marker, so the @@ -42,6 +44,8 @@ from __future__ import annotations import time from collections.abc import Callable +from dataclasses import dataclass +from functools import partial from typing import Final import pytest @@ -60,6 +64,7 @@ from models import ( ChatTool, ChatToolFunction, ChatToolResultTurn, + CustomPricing, LiteLLMParamsBody, PromptCacheOptions, RichMessage, @@ -203,19 +208,26 @@ def _assert_openai_writes_then_reads_cache(client: PassthroughClient, key: str, assert read_tokens > 0, f"{body.model}: the identical second call read no cache ({second.usage})" -def _cold_cache_call(send: Callable[[str], Result[ChatResponse]]) -> ChatResponse | None: - result: Final = send(_cacheable_prefix()) +@dataclass(frozen=True, slots=True) +class EngagedColdCall: + prefix: str + completion: ChatResponse + + +def _cold_cache_call(send: Callable[[str], Result[ChatResponse]]) -> EngagedColdCall | None: + prefix: Final = _cacheable_prefix() + result: Final = send(prefix) match result: case UnknownApiError(status_code=400, body=body) if VERTEX_CACHE_REJECTION_MARKER in body: return None case _: - return unwrap(result) + return EngagedColdCall(prefix=prefix, completion=unwrap(result)) -def _first_engaged_cold_call(send: Callable[[str], Result[ChatResponse]]) -> ChatResponse | None: +def _first_engaged_cold_call(send: Callable[[str], Result[ChatResponse]]) -> EngagedColdCall | None: for attempt in range(1, VERTEX_COLD_CALL_ATTEMPTS + 1): candidate = _cold_cache_call(send) - if candidate is not None and _cached_read_tokens(candidate.usage) >= VERTEX_MINIMUM_CACHED_TOKENS: + if candidate is not None and _cached_read_tokens(candidate.completion.usage) >= VERTEX_MINIMUM_CACHED_TOKENS: return candidate if attempt < VERTEX_COLD_CALL_ATTEMPTS: print( @@ -227,32 +239,61 @@ def _first_engaged_cold_call(send: Callable[[str], Result[ChatResponse]]) -> Cha return None -def _first_cold_call_reads_cache(model: str, send: Callable[[str], Result[ChatResponse]]) -> ChatResponse: - completion: Final = _first_engaged_cold_call(send) - assert completion is not None, ( +def _cache_creation_tokens(usage: Usage | None) -> int: + if usage is None: + return 0 + return usage.cache_creation_input_tokens or 0 + + +def _first_cold_call_creates_and_reads_cache( + model: str, send: Callable[[str], Result[ChatResponse]] +) -> EngagedColdCall: + engaged: Final = _first_engaged_cold_call(send) + assert engaged is not None, ( f"{model}: {VERTEX_COLD_CALL_ATTEMPTS} never-seen prompts marked with cache_control, spread over " f"{VERTEX_COLD_CALL_PAUSE_SECONDS * (VERTEX_COLD_CALL_ATTEMPTS - 1):.0f}s, were each either rejected by " f"Vertex's minimum-token check or served with fewer than {VERTEX_MINIMUM_CACHED_TOKENS} cached tokens on " "their first call; explicit context caching did not engage" ) + completion: Final = engaged.completion assert completion.choices, f"{model}: cached call returned no choices: {completion}" usage: Final = completion.usage - cached: Final = _cached_read_tokens(usage) assert usage and usage.prompt_tokens, f"{model}: cached completion carried no prompt_tokens: {usage}" - assert cached >= CACHED_SHARE_OF_PROMPT * usage.prompt_tokens, ( - f"{model}: only {cached} of {usage.prompt_tokens} prompt tokens were served from the cache; the " - "cache_control block was not cached whole" + created: Final = _cache_creation_tokens(usage) + assert created > 0, ( + f"{model}: the call that built the Vertex cache reported no cache_creation_input_tokens ({usage}); " + "the cache write went unbilled" ) - return completion + generate_prompt_tokens: Final = usage.prompt_tokens - created + assert generate_prompt_tokens > 0, ( + f"{model}: prompt_tokens {usage.prompt_tokens} does not cover the {created} cache-creation tokens plus " + "the generate call's own prompt" + ) + cached: Final = _cached_read_tokens(usage) + assert cached >= CACHED_SHARE_OF_PROMPT * generate_prompt_tokens, ( + f"{model}: only {cached} of the generate call's {generate_prompt_tokens} prompt tokens were served from " + "the cache; the cache_control block was not cached whole" + ) + return engaged -def _input_rate(client: PassthroughClient, model: str) -> float: +def _assert_second_call_reuses_cache(model: str, send: Callable[[str], Result[ChatResponse]], prefix: str) -> None: + usage: Final = unwrap(send(prefix)).usage + assert _cached_read_tokens(usage) > 0, f"{model}: the identical second call read no cache ({usage})" + assert _cache_creation_tokens(usage) == 0, ( + f"{model}: the identical second call built and billed a new cache ({usage}) instead of reusing the first one" + ) + + +def _pricing(client: PassthroughClient, model: str) -> CustomPricing: entry: Final = next((row for row in client.proxy.model_info() if row.model_name == model), None) assert entry and entry.model_info.input_cost_per_token, f"/model/info resolved no input rate for {model}" - return entry.model_info.input_cost_per_token + return entry.model_info -def _assert_billed_below_uncached_prompt(client: PassthroughClient, model: str, completion: ChatResponse) -> None: +def _assert_billed_for_creation_and_discounted_read( + client: PassthroughClient, model: str, completion: ChatResponse +) -> None: assert completion.id, f"{model}: cached completion carried no id to find its spend row by" usage: Final = completion.usage assert usage and usage.prompt_tokens, f"{model}: cached completion carried no prompt_tokens: {usage}" @@ -262,11 +303,30 @@ def _assert_billed_below_uncached_prompt(client: PassthroughClient, model: str, assert row.prompt_tokens == usage.prompt_tokens, ( f"{model}: spend row prompt_tokens {row.prompt_tokens} != response prompt_tokens {usage.prompt_tokens}" ) + assert row.spend is not None, f"{model}: spend row carried no spend: {row}" + pricing: Final = _pricing(client, model) + input_rate: Final = pricing.input_cost_per_token or 0 + created: Final = _cache_creation_tokens(usage) cached: Final = _cached_read_tokens(usage) - uncached_read_cost: Final = cached * _input_rate(client, model) - assert row.spend is not None and row.spend < uncached_read_cost, ( - f"{model}: spend {row.spend} is not below the uncached price of the {cached} tokens read from the cache " - f"({uncached_read_cost}); cache-read pricing was not applied" + # Google bills cache creation at the standard input price (read 2026-10-16): + # https://docs.cloud.google.com/vertex-ai/generative-ai/docs/context-cache/context-cache-overview + creation_cost: Final = created * input_rate + assert row.spend >= creation_cost, ( + f"{model}: spend {row.spend} is below the {creation_cost} the {created} cache-creation tokens cost at the " + "input rate; the cache write was not billed" + ) + read_rate: Final = pricing.cache_read_input_token_cost + assert read_rate is not None and read_rate < input_rate, ( + f"/model/info advertises no discounted cache-read rate for {model}: {pricing}" + ) + undiscounted_cost: Final = pricing.token_cost( + created + max(usage.prompt_tokens - created, cached), usage.completion_tokens or 0 + ) + half_the_read_discount: Final = cached * (input_rate - read_rate) / 2 + assert row.spend < undiscounted_cost - half_the_read_discount, ( + f"{model}: spend {row.spend} saved less than half the cache-read discount on its {cached} cached tokens " + f"against {undiscounted_cost}, the price of creating the cache and reading it uncached; cache-read pricing " + "was not applied" ) @@ -326,10 +386,10 @@ class TestCacheControl: ) resources.defer(lambda: client.proxy.delete_model(model_id)) key = resources.key() - completion: Final = _first_cold_call_reads_cache( - model, lambda prefix: _cache_chat(client, key, model, prefix, ttl=VERTEX_CACHE_TTL) - ) - _assert_billed_below_uncached_prompt(client, model, completion) + send: Final = partial(_cache_chat, client, key, model, ttl=VERTEX_CACHE_TTL) + engaged: Final = _first_cold_call_creates_and_reads_cache(model, send) + _assert_billed_for_creation_and_discounted_read(client, model, engaged.completion) + _assert_second_call_reuses_cache(model, send, engaged.prefix) @pytest.mark.covers( "llm.chat_completions.anthropic.prompt_cache_5m.nonstream.works", diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c4414003ca5..324f4a4f4e9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -13,7 +13,7 @@ from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers -from litellm.caching.caching import DualCache +from litellm.caching.dual_cache import DualCache from litellm.constants import MCP_ALL_TOOLS_WILDCARD from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( @@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( _agent_capped_servers, is_mcp_admitted_user_subject, ) +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, @@ -32,6 +33,8 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.types.agents import AgentCaller +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.mcp_server.mcp_toolset import MCPToolset @pytest.mark.parametrize( @@ -10656,26 +10659,32 @@ class TestEntitlementFaultSemantics: assert set(allowed) == {"srv1"} -async def _cache_with_end_user(end_user_id, *, mcp_tool_permissions=None, object_permission_id=None): +async def _cache_with_end_user( + end_user_id: str, + *, + mcp_tool_permissions: dict[str, list[str]] | None = None, + object_permission_id: str | None = None, + object_permission_override: LiteLLM_ObjectPermissionTable | None = None, +) -> DualCache: """A real DualCache already holding the end user row, so ``get_end_user_object`` answers from cache and no ``litellm.`` internal has to be patched. ``object_permission_id`` without a permission body models a row that NAMES an entitlement the DB then fails to serve.""" - from litellm.caching.dual_cache import DualCache from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key - cache = DualCache() + permission: Final = object_permission_override or ( + LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_tool_permissions=mcp_tool_permissions) + if mcp_tool_permissions + else None + ) + cache: Final = DualCache() await cache.async_set_cache( key=end_user_cache_key(end_user_id), value=LiteLLM_EndUserTable( user_id=end_user_id, blocked=False, - object_permission_id=object_permission_id or ("op-eu" if mcp_tool_permissions else None), - object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id="op-eu", mcp_tool_permissions=mcp_tool_permissions - ) - if mcp_tool_permissions - else None, + object_permission_id=object_permission_id or (permission.object_permission_id if permission else None), + object_permission=permission, ), ) return cache @@ -10686,6 +10695,207 @@ class TestEndUserToolCeiling: """The end user (customer) level narrows the TOOLS axis exactly as it narrows the servers axis, so `object_permission.mcp_tool_permissions` on `/customer/new` is enforced, not just stored.""" + @pytest.mark.parametrize("require_defined", (False, True)) + @pytest.mark.parametrize( + "key_servers,expected_servers", ((("server-a", "server-b"), ("server-a",)), (("server-b",), ())) + ) + async def test_end_user_toolsets_cap_shared_key_servers( + self, require_defined: bool, key_servers: tuple[str, ...], expected_servers: tuple[str, ...] + ) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"server-a": ["read"]}) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", mcp_servers=list(key_servers) + ), + ) + with ( + _entitlement_fault_globals(user_api_key_cache=cache), + patch( + "litellm.proxy.proxy_server.general_settings", {"require_end_user_mcp_access_defined": require_defined} + ), + ): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + allowed_tools: Final = await MCPRequestHandler.get_allowed_tools_for_server("server-a", auth) + read_allowed: Final = await MCPRequestHandler.is_tool_allowed_for_server("read", "server-a", auth) + write_allowed: Final = await MCPRequestHandler.is_tool_allowed_for_server("write", "server-a", auth) + assert set(allowed_servers) == set(expected_servers) + assert allowed_tools == ["read"] + assert read_allowed is True + assert write_allowed is False + + async def test_end_user_toolset_servers_union_with_existing_grants(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="op-eu", + mcp_servers=["direct"], + mcp_access_groups=["readers"], + mcp_tool_permissions={"inline": ["read"]}, + mcp_toolsets=["toolset-a"], + ) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"toolset": ["read"]}) + manager: Final = MCPServerManager() + manager.config_mcp_servers = { + "group": MCPServer(server_id="group", name="group", transport="http", access_groups=["readers"]), + } + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + mcp_servers=["direct", "group", "inline", "toolset", "other"], + ), + ) + with ( + _entitlement_fault_globals(user_api_key_cache=cache), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", manager), + ): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(allowed_servers) == {"direct", "group", "inline", "toolset"} + + @pytest.mark.parametrize( + "toolset_rows", + ((), (MCPToolset(toolset_id="toolset-a", toolset_name="Empty", tools=[]),)), + ids=("missing-toolset", "empty-toolset"), + ) + async def test_end_user_unresolvable_toolsets_deny_shared_key_servers( + self, toolset_rows: tuple[MCPToolset, ...] + ) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + prisma: Final = MagicMock() + prisma.configure_mock(**{"db.litellm_mcptoolsettable.find_many": AsyncMock(return_value=list(toolset_rows))}) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["server-a"]), + ) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + with pytest.raises(UnloadableEntitlementError): + await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(auth) + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == [] + + async def test_end_user_toolset_ceiling_reads_fresh_policy(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"old-server": ["read"]}) + prisma: Final = MagicMock() + prisma.configure_mock( + **{ + "writer_db.litellm_mcptoolsettable.find_many": AsyncMock( + return_value=[ + MCPToolset( + toolset_id="toolset-a", + toolset_name="Updated", + tools=[ + {"server_id": "new-server", "tool_name": "read"}, + ], + ), + ] + ), + } + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + mcp_servers=["old-server", "new-server"], + ), + ).model_copy(update={"requires_fresh_policy": True}) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(auth) + assert allowed_servers == ["new-server"] + + async def test_end_user_toolset_writer_fault_does_not_drop_ceiling(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"old-server": ["read"]}) + prisma: Final = MagicMock() + prisma.configure_mock( + **{ + "writer_db.litellm_mcptoolsettable.find_many": AsyncMock( + side_effect=RuntimeError("writer unavailable") + ), + } + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", mcp_servers=["other-server"] + ), + ).model_copy(update={"requires_fresh_policy": True}) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == [] + + async def test_end_user_access_group_writer_fault_does_not_drop_toolset_ceiling(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="op-eu", mcp_toolsets=["toolset-a"], mcp_access_groups=["group-x"] + ) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + prisma: Final = MagicMock() + prisma.configure_mock( + **{ + "writer_db.litellm_mcptoolsettable.find_many": AsyncMock( + return_value=[ + MCPToolset( + toolset_id="toolset-a", + toolset_name="Read only", + tools=[{"server_id": "server-a", "tool_name": "read"}], + ), + ] + ), + "writer_db.litellm_mcpservertable.find_many": AsyncMock(side_effect=RuntimeError("writer unavailable")), + } + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", mcp_servers=["server-a", "server-b"] + ), + ).model_copy(update={"requires_fresh_policy": True}) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert "server-b" not in allowed_servers + + async def test_end_user_mcp_opt_out_overrides_toolset_servers(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="op-eu", + mcp_servers=[SpecialMCPServerNames.no_mcp_servers.value], + mcp_toolsets=["toolset-a"], + ) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"server-a": ["read"]}) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["server-a"]), + ) + with _entitlement_fault_globals(user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == [] + + async def test_end_user_empty_permission_keeps_inherited_servers(self) -> None: + cache: Final = await _cache_with_end_user( + "eu-1", + object_permission_override=LiteLLM_ObjectPermissionTable(object_permission_id="op-eu"), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["server-a"]), + ) + with _entitlement_fault_globals(user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == ["server-a"] + async def test_end_user_tool_permissions_intersect_key_tools(self): auth = _key_auth_reaching("srv1", tools=["tool_a", "tool_b"], end_user_id="eu-1") cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv1": ["tool_a"]}) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 40901d1dd96..9a0129a06b4 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -20292,3 +20292,114 @@ async def test_legacy_upstream_elicitation_rejects_modern_downstream_without_con finally: active_mcp_request_ctx_var.reset(request_token) legacy_server.active_mcp_session_var.reset(session_token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["manager", "mcp", "rest"]) +@pytest.mark.parametrize("admin", [False, True]) +@pytest.mark.parametrize("permission_state", ["denied", "stale-grants", "unreadable"]) +async def test_customer_denial_blocks_operator_open_servers( + monkeypatch: pytest.MonkeyPatch, + surface: Literal["manager", "mcp", "rest"], + admin: bool, + permission_state: Literal["denied", "stale-grants", "unreadable"], +) -> None: + from starlette.requests import Request + + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager, operations, rest_endpoints + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="operator-open", name="operator_open", transport=MCPTransport.http, allow_all_keys=True + ) + manager.config_mcp_servers = {server.server_id: server} + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="customer-denial", + mcp_servers=["no-mcp-servers"], + mcp_tool_permissions={server.server_id: ["delete"]} if permission_state == "stale-grants" else {}, + mcp_toolsets=["stale-toolset"] if permission_state == "stale-grants" else [], + mcp_access_groups=["stale-group"] if permission_state == "stale-grants" else [], + ) + cache: Final = DualCache() + await cache.async_set_cache( + key=end_user_cache_key("denied-customer"), + value=LiteLLM_EndUserTable( + user_id="denied-customer", + blocked=False, + object_permission_id=permission.object_permission_id, + object_permission=None if permission_state == "unreadable" else permission, + ), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="denied-customer", + user_role=LitellmUserRoles.PROXY_ADMIN if admin else None, + object_permission=None if admin else LiteLLM_ObjectPermissionTable( + object_permission_id="key-grant", mcp_servers=[server.server_id] + ), + ).model_copy(update={"via_virtual_key": True}) + _catalog_database(monkeypatch, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + + if surface == "manager": + assert await manager.get_allowed_mcp_servers(auth) == [] + elif surface == "mcp": + assert await operations._get_allowed_mcp_servers(auth, None) == [] + else: + request: Final = Request({"type": "http", "method": "POST", "path": "/mcp-rest/tools/call", "headers": []}) + with pytest.raises(HTTPException) as denied: + await rest_endpoints._resolve_allowed_mcp_servers_with_ip_filter(request, auth, server.server_id) + assert denied.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("customer", ["absent", "unrestricted", "empty-group", "limited"]) +@pytest.mark.parametrize("admin", [False, True]) +async def test_customer_without_explicit_deny_keeps_operator_open_servers( + monkeypatch: pytest.MonkeyPatch, + customer: Literal["absent", "unrestricted", "empty-group", "limited"], + admin: bool, +) -> None: + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="operator-open", name="operator_open", transport=MCPTransport.http, allow_all_keys=True + ) + manager.config_mcp_servers = {server.server_id: server} + cache: Final = DualCache() + await cache.async_set_cache( + key=end_user_cache_key("unrestricted-customer"), + value=LiteLLM_EndUserTable( + user_id="unrestricted-customer", + blocked=False, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="ordinary-customer", + mcp_servers=["another-server"] if customer == "limited" else [], + mcp_access_groups=["empty-group"] if customer == "empty-group" else [], + ), + ), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id=None if customer == "absent" else "unrestricted-customer", + user_role=LitellmUserRoles.PROXY_ADMIN if admin else None, + ).model_copy(update={"via_virtual_key": True}) + _catalog_database(monkeypatch, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + + assert await manager.get_allowed_mcp_servers(auth) == [server.server_id]