From 6825f92806515f41319709496aca7bf569a6bdba Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 18:22:49 -0400 Subject: [PATCH 1/7] fix(proxy): enforce customer toolset server ceilings --- .../mcp_server/auth/user_api_key_auth_mcp.py | 21 +- .../auth/test_user_api_key_auth_mcp.py | 198 +++++++++++++++++- 2 files changed, 205 insertions(+), 14 deletions(-) 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 18b7c279606..934fa3fe565 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 @@ -3189,9 +3189,17 @@ 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) + object_permission: Final = await MCPRequestHandler._get_end_user_object_permission( + user_api_key_auth, prisma_client + ) if object_permission is None: return [] + if SpecialMCPServerNames.no_mcp_servers.value in (object_permission.mcp_servers or []): + return [SpecialMCPServerNames.no_mcp_servers.value] + + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + object_permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) try: # Permission entries may be server_ids OR names/aliases — expand to ids. @@ -3199,7 +3207,9 @@ 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( + object_permission.mcp_servers or [] + ) # Get MCP servers from access groups access_group_servers: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups( @@ -3212,9 +3222,10 @@ class MCPRequestHandler: global_mcp_server_manager.expand_tool_permissions(object_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, object_permission)) except Exception as e: verbose_logger.warning("Failed to get allowed MCP servers for end_user: %s", e) return [] 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 078fff0266b..9010b8584a6 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 @@ -10,12 +10,14 @@ from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers +from litellm.caching.dual_cache import DualCache from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, UnloadableEntitlementError, _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, ProxyException, @@ -24,6 +26,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("headers,expected", [ @@ -10535,26 +10539,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 @@ -10565,6 +10575,176 @@ 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_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"]}) From b1e0be1f997599d4c2ac32108dc21dbcc78266ff Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 18:53:55 -0400 Subject: [PATCH 2/7] fix(proxy): deny MCP access when a customer's server ceiling cannot be read --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +-- .../auth/test_user_api_key_auth_mcp.py | 31 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) 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 934fa3fe565..f36a12b58ef 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 @@ -3227,8 +3227,9 @@ class MCPRequestHandler: ) return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, object_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/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 9010b8584a6..88d5f0cd387 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 @@ -10714,6 +10714,37 @@ class TestEndUserToolCeiling: 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", From 13a9e88608206e5449e359042c90fbeee0d95e8e Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 19:37:05 -0400 Subject: [PATCH 3/7] ci(test-unit): run the MCP auth unit tests in the proxy-auth shard --- .github/workflows/test-unit.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 23a340cd436..2e2de270745 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -457,6 +457,7 @@ jobs: - shard: proxy-auth test-path: |- tests/unit/proxy/auth + tests/unit/proxy/_experimental/mcp_server/auth --ignore=tests/unit/proxy/auth/test_auth_checks.py --ignore=tests/unit/proxy/auth/test_user_api_key_auth.py --ignore=tests/unit/proxy/auth/test_default_end_user_budget_simple.py From 1e87e5c8c796b628cb87a2db27a6d70774e611a5 Mon Sep 17 00:00:00 2001 From: lewis Date: Wed, 14 Oct 2026 19:01:29 -0400 Subject: [PATCH 4/7] revert: remove duplicate MCP auth CI coverage This reverts commit 13a9e88608. The MCP auth directory already runs in the mcp-server shard. --- .github/workflows/test-unit.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index da50f589014..f521b2c8edd 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -463,7 +463,6 @@ jobs: - shard: proxy-auth test-path: |- tests/unit/proxy/auth - tests/unit/proxy/_experimental/mcp_server/auth --ignore=tests/unit/proxy/auth/test_auth_checks.py --ignore=tests/unit/proxy/auth/test_user_api_key_auth.py --ignore=tests/unit/proxy/auth/test_default_end_user_budget_simple.py From 6116a2384010e320d9f6bf26a04b02e4a499ef73 Mon Sep 17 00:00:00 2001 From: lewis Date: Wed, 14 Oct 2026 19:05:26 -0400 Subject: [PATCH 5/7] style(proxy): wrap customer MCP ceiling error --- .../_experimental/mcp_server/auth/user_api_key_auth_mcp.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 7de0ef11ea3..ad9dce91f5e 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 @@ -3434,7 +3434,8 @@ class MCPRequestHandler: return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, object_permission)) except Exception as e: raise UnloadableEntitlementError( - f"end user {user_api_key_auth.end_user_id!r} has an object permission whose MCP servers could not be read" + 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 From 81a3802a9216b4e79951687fd2f6e44be7e90b29 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Oct 2026 18:51:38 -0700 Subject: [PATCH 6/7] test(e2e): assert the Vertex cache-creation billing contract in the prompt caching test (#47027) The cache-creating call now reports cache_creation_input_tokens on top of Vertex's own prompt count and bills the write at the input rate, as Google does. The old assertions assumed neither, so the test went red on every run. Check the cache write, the discounted read, and that a second identical call reuses the cache instead. --- .../e2e/llm_translation/test_cache_control.py | 120 +++++++++++++----- 1 file changed, 90 insertions(+), 30 deletions(-) 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", From 0f52da9232aa0604bb0114782d676b319f021ff2 Mon Sep 17 00:00:00 2001 From: lewis Date: Sat, 17 Oct 2026 21:11:21 -0400 Subject: [PATCH 7/7] fix(proxy): enforce customer MCP denial before open grants --- .../mcp_server/auth/user_api_key_auth_mcp.py | 43 +++++-- .../mcp_server/mcp_server_manager.py | 10 +- .../mcp_server/test_mcp_server_manager.py | 111 ++++++++++++++++++ 3 files changed, 151 insertions(+), 13 deletions(-) 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 235580c18ad..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,16 +3462,18 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return [] - object_permission: Final = await MCPRequestHandler._get_end_user_object_permission( - user_api_key_auth, prisma_client + permission: Final = ( + await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if object_permission is ... + else object_permission ) - if object_permission is None: + if permission is None: return [] - if SpecialMCPServerNames.no_mcp_servers.value in (object_permission.mcp_servers or []): + 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( - object_permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy ) try: @@ -3464,24 +3483,24 @@ class MCPRequestHandler: ) direct_mcp_servers: Final = global_mcp_server_manager.expand_permission_list( - object_permission.mcp_servers or [] + 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() ) 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, object_permission)) + return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, permission)) except Exception as e: raise UnloadableEntitlementError( f"end user {user_api_key_auth.end_user_id!r} has an 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/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]