From ea80762a597bdb3327f1e4da2a8d61811484a9e1 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Tue, 3 Nov 2026 22:23:37 -0800 Subject: [PATCH 1/3] fix(mcp): read policy through the cache while a caller's policy hash is unchanged --- litellm/proxy/auth/auth_checks.py | 11 +- .../proxy/common_utils/strict_cache_refill.py | 18 + litellm/proxy/db/routing_prisma_wrapper.py | 5 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 79 +- litellm/proxy/mcp_server/catalog.py | 11 + litellm/proxy/mcp_server/policy_freshness.py | 304 +++++++ .../proxy/db/test_routing_prisma_wrapper.py | 23 + .../proxy/mcp_server/test_policy_freshness.py | 782 ++++++++++++++++++ 8 files changed, 1203 insertions(+), 30 deletions(-) create mode 100644 litellm/proxy/common_utils/strict_cache_refill.py create mode 100644 litellm/proxy/mcp_server/policy_freshness.py create mode 100644 tests/unit/proxy/mcp_server/test_policy_freshness.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9d98e2c9973..ec0a1fd15b0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -119,6 +119,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_query_params, ) from litellm.proxy.common_utils.model_listing_utils import alias_map +from litellm.proxy.common_utils.strict_cache_refill import strict_cache_refill from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, @@ -2789,7 +2790,7 @@ async def get_team_membership( return l1_membership inflight: Final[object] = _team_membership_inflight.get(_key) - if isinstance(inflight, asyncio.Task): + if isinstance(inflight, asyncio.Task) and not strict_cache_refill(): return _membership_from_shared_load(await asyncio.shield(inflight)) if prisma_client is None: @@ -3643,7 +3644,7 @@ async def _get_team_object_from_user_api_key_cache( check_db_only=use_writer, ) except Exception as e: - if use_writer: + if use_writer or strict_cache_refill(): raise verbose_proxy_logger.debug( "Failed to load object_permission for team %s with object_permission_id=%s: %s", @@ -4394,7 +4395,7 @@ async def get_key_object( check_db_only=check_db_only, ) except Exception as e: - if check_db_only: + if check_db_only or strict_cache_refill(): raise verbose_proxy_logger.debug( "Failed to load object_permission for key with object_permission_id=%s: %s", @@ -4468,7 +4469,7 @@ async def get_object_permission( ).find_unique(where={"object_permission_id": object_permission_id}) if response is None: - if check_db_only: + if check_db_only or strict_cache_refill(): raise HTTPException(status_code=403, detail="Referenced object permission does not exist") return None @@ -4482,7 +4483,7 @@ async def get_object_permission( return _perm_obj except Exception: - if check_db_only: + if check_db_only or strict_cache_refill(): raise return None diff --git a/litellm/proxy/common_utils/strict_cache_refill.py b/litellm/proxy/common_utils/strict_cache_refill.py new file mode 100644 index 00000000000..1cea40fb5f6 --- /dev/null +++ b/litellm/proxy/common_utils/strict_cache_refill.py @@ -0,0 +1,18 @@ +"""While set, a cache entry refilled from the database fails closed and is read from the writer.""" + +from contextvars import ContextVar, Token +from typing import Final + +_strict: Final[ContextVar[bool]] = ContextVar("strict_cache_refill", default=False) + + +def strict_cache_refill() -> bool: + return _strict.get() + + +def require_strict_cache_refill() -> Token[bool]: + return _strict.set(True) + + +def release_strict_cache_refill(token: Token[bool]) -> None: + _strict.reset(token) diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 50e6aa54da4..366ed976ffc 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.strict_cache_refill import strict_cache_refill from litellm.proxy.db.prisma_client import PrismaWrapper if TYPE_CHECKING: @@ -146,7 +147,7 @@ class RoutingPrismaWrapper: consult this rather than `writer`, and `__getattr__` routes through it so the two cannot drift apart. """ - return self._writer if self._reader_unavailable else self._reader + return self._writer if self._reader_unavailable or strict_cache_refill() else self._reader @property def reader_unavailable(self) -> bool: @@ -166,7 +167,7 @@ class RoutingPrismaWrapper: self._writer_unavailable = False def _should_use_reader(self) -> bool: - return not self._reader_unavailable + return not self._reader_unavailable and not strict_cache_refill() @staticmethod async def _try_connect(client: PrismaWrapper, timeout: int | timedelta | None = None) -> Exception | None: diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index be24b1630b0..4d2484a90a3 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -58,6 +58,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, ) +from litellm.proxy.common_utils.strict_cache_refill import strict_cache_refill from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, USER_NO_MCP_PERMISSION_SENTINEL, @@ -649,6 +650,13 @@ class DcrBridgeTarget: server: MCPServer +def _jwt_teams_from_claim_list() -> bool: + """True when JWT auth takes several teams from a claim list, which the policy hash cannot pin to one team.""" + from litellm.proxy.proxy_server import jwt_handler + + return hasattr(jwt_handler, "litellm_jwtauth") and bool(jwt_handler.litellm_jwtauth.team_ids_jwt_field) + + class MCPRequestHandler: """ Class to handle MCP request processing, including: @@ -1578,27 +1586,49 @@ class MCPRequestHandler: @staticmethod async def refresh_catalog_authority(auth: UserAPIKeyAuth | None) -> UserAPIKeyAuth | None: from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy.mcp_server.policy_freshness import ( + cached_policy_is_current, + policy_subject_for, + validated_policy_marks, + ) + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if auth is None: return None - refreshed: Final = auth.model_copy(update={"requires_fresh_policy": True}) - if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id: + keyless_user: Final = bool((auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id) + virtual_key: Final = bool( + not keyless_user + and auth.via_virtual_key + and auth.api_key + and auth.api_key != LITELLM_PROXY_MASTER_KEY_ALIAS + ) + subject: Final = policy_subject_for( + auth, + keyless_user=keyless_user, + virtual_key=virtual_key, + jwt_teams_from_claim_list=_jwt_teams_from_claim_list(), + ) + fresh: Final = bool(auth.requires_fresh_policy) or not await cached_policy_is_current( + subject, prisma_client=prisma_client, cache=user_api_key_cache, marks=validated_policy_marks + ) + refreshed: Final = auth.model_copy(update={"requires_fresh_policy": fresh}) + if keyless_user and auth.user_id: # A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is # sealed in the token: refresh the live user policy and keep the token's team, routes and models. - current = await MCPRequestHandler.reload_admitted_user(auth.user_id, requires_fresh_policy=True) - refreshed.object_permission = current.object_permission - refreshed.object_permission_id = current.object_permission_id - refreshed.user_role = current.user_role - refreshed.org_id = current.org_id - elif auth.via_virtual_key and auth.api_key and auth.api_key != LITELLM_PROXY_MASTER_KEY_ALIAS: - current = await MCPRequestHandler._reload_admitted_key(auth.api_key, check_db_only=True) - refreshed.object_permission = current.object_permission - refreshed.object_permission_id = current.object_permission_id - refreshed.team_id = current.team_id - refreshed.org_id = current.org_id - refreshed.project_id = current.project_id - refreshed.user_id = current.user_id - refreshed.access_group_ids = current.access_group_ids + user_policy: Final = await MCPRequestHandler.reload_admitted_user(auth.user_id, requires_fresh_policy=fresh) + refreshed.object_permission = user_policy.object_permission + refreshed.object_permission_id = user_policy.object_permission_id + refreshed.user_role = user_policy.user_role + refreshed.org_id = user_policy.org_id + elif virtual_key and auth.api_key: + key_policy: Final = await MCPRequestHandler._reload_admitted_key(auth.api_key, check_db_only=fresh) + refreshed.object_permission = key_policy.object_permission + refreshed.object_permission_id = key_policy.object_permission_id + refreshed.team_id = key_policy.team_id + refreshed.org_id = key_policy.org_id + refreshed.project_id = key_policy.project_id + refreshed.user_id = key_policy.user_id + refreshed.access_group_ids = key_policy.access_group_ids else: return refreshed if refreshed.mcp_toolset_id is not None: @@ -3136,7 +3166,7 @@ class MCPRequestHandler: unreadable="raise", ) except Exception as e: # noqa: BLE001 # a key naming groups we cannot read must deny, not widen - if user_api_key_auth.requires_fresh_policy: + if user_api_key_auth.requires_fresh_policy or strict_cache_refill(): raise UnloadableEntitlementError(f"key access groups unreadable: {e}") from e return _UNREADABLE_KEY_GROUP_TOOLS owned: Final = tuple( @@ -3182,7 +3212,7 @@ class MCPRequestHandler: unreadable="raise", ) except Exception as e: # noqa: BLE001 # team groups we cannot read grant no tools, they never widen - if user_api_key_auth.requires_fresh_policy: + if user_api_key_auth.requires_fresh_policy or strict_cache_refill(): raise UnloadableEntitlementError(f"team access groups unreadable: {e}") from e return () return MCPRequestHandler.access_group_tools_for_server(groups, server_id) @@ -3837,8 +3867,11 @@ class MCPRequestHandler: ) return list(servers) except Exception as e: - if isinstance(e, UnloadableEntitlementError) or ( - require_readable_policy or (user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy) + if ( + isinstance(e, UnloadableEntitlementError) + or require_readable_policy + or strict_cache_refill() + or (user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy) ): raise verbose_logger.warning("Failed to get allowed MCP servers for team: %s", e) @@ -4182,7 +4215,7 @@ class MCPRequestHandler: except UserNotFoundError: return None except Exception as e: # noqa: BLE001 # Legacy callers retain their existing optional user-ceiling behavior - if check_db_only: + if check_db_only or strict_cache_refill(): raise HTTPException(503, "User policy is unavailable") from e verbose_logger.warning("MCP user entitlement: link for %r unresolved, no ceiling: %s", user_id, e) return None @@ -4734,7 +4767,7 @@ class MCPRequestHandler: for server in mcp_servers: server_ids.add(server.server_id) except Exception as e: - if use_writer: + if use_writer or strict_cache_refill(): raise verbose_logger.debug("Error getting MCP servers from access groups: %s", e) return server_ids @@ -4774,7 +4807,7 @@ class MCPRequestHandler: return list(server_ids) except Exception as e: - if requires_fresh_policy: + if requires_fresh_policy or strict_cache_refill(): raise verbose_logger.warning("Failed to get MCP servers from access groups: %s", e) return [] diff --git a/litellm/proxy/mcp_server/catalog.py b/litellm/proxy/mcp_server/catalog.py index d2ff322ba5c..5332aa1ac8f 100644 --- a/litellm/proxy/mcp_server/catalog.py +++ b/litellm/proxy/mcp_server/catalog.py @@ -23,6 +23,7 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.mcp_server.policy_freshness import validated_listing from litellm.proxy.mcp_server.result_conversion import age_freshness, aggregate_freshness if TYPE_CHECKING: @@ -583,6 +584,15 @@ def catalog_operation( return decorate +def validated_listing_scope(function: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: + @wraps(function) + async def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: # kwargs-ok: preserves ParamSpec + async with validated_listing(): + return await function(*args, **kwargs) + + return wrapped + + def global_manager() -> MCPServerManager: from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager @@ -1126,6 +1136,7 @@ async def list_gateway_tools( ) +@validated_listing_scope async def list_gateway_catalog( context: OperationContext, request: CatalogListRequest, *, log_list_tools_to_spendlogs: bool = True ) -> CatalogListResult: diff --git a/litellm/proxy/mcp_server/policy_freshness.py b/litellm/proxy/mcp_server/policy_freshness.py new file mode 100644 index 00000000000..5662c48a21b --- /dev/null +++ b/litellm/proxy/mcp_server/policy_freshness.py @@ -0,0 +1,304 @@ +"""Decides, once per MCP listing, whether the caller's cached policy still matches the database. + +One query on the writer hashes every row the caller's MCP access depends on. While the hash equals the one +recorded after the last completed listing, the cache is trusted. Anything else reads policy from the database. +""" + +import json +import time +from collections import OrderedDict +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from contextlib import asynccontextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Literal + +from pydantic import BaseModel, ConfigDict, field_validator + +from litellm._logging import verbose_logger +from litellm.proxy.common_utils.strict_cache_refill import release_strict_cache_refill, require_strict_cache_refill +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + get_management_object_ttl, + object_permission_cache_key, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, + user_object_permission_id_cache_key, +) + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import PrismaClient + +_MARK_CAPACITY: Final = 4096 +_MARK_TTL_FRACTION: Final = 0.5 +_SETTLE_SECONDS: Final = 5.0 + +# Spend, activity and credential columns change on every request or are secrets, and never describe access. +_FINGERPRINT_SQL: Final = """ +WITH k AS (SELECT * FROM "LiteLLM_VerificationToken" WHERE token = $1::text), +u AS (SELECT * FROM "LiteLLM_UserTable" WHERE user_id IN (SELECT user_id FROM k) OR user_id = $2::text), +t AS (SELECT * FROM "LiteLLM_TeamTable" WHERE team_id IN (SELECT team_id FROM k WHERE team_id IS NOT NULL) + OR team_id IN (SELECT unnest(teams) FROM u) OR team_id = $3::text), +o AS (SELECT * FROM "LiteLLM_OrganizationTable" + WHERE organization_id IN (SELECT organization_id FROM k UNION SELECT organization_id FROM u + UNION SELECT organization_id FROM t) + OR organization_id = $4::text), +tm AS (SELECT * FROM "LiteLLM_TeamMembership" WHERE user_id IN (SELECT user_id FROM u) OR user_id IN (SELECT user_id FROM k)), +om AS (SELECT * FROM "LiteLLM_OrganizationMembership" WHERE user_id IN (SELECT user_id FROM u) OR user_id IN (SELECT user_id FROM k)), +p AS (SELECT * FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id IN (SELECT object_permission_id FROM k UNION SELECT object_permission_id FROM u + UNION SELECT object_permission_id FROM t UNION SELECT object_permission_id FROM o)), +ag AS (SELECT * FROM "LiteLLM_AccessGroupTable" WHERE access_group_id IN (SELECT unnest(access_group_ids) FROM k UNION SELECT unnest(access_group_ids) FROM t)), +b AS (SELECT * FROM "LiteLLM_BudgetTable" WHERE budget_id IN (SELECT budget_id FROM k UNION SELECT budget_id FROM o + UNION SELECT budget_id FROM tm UNION SELECT budget_id FROM om)) +SELECT + md5(concat_ws('#', + (SELECT COALESCE(string_agg((to_jsonb(k) - 'spend' - 'total_spend' - 'model_spend' - 'budget_reset_at' + - 'updated_at' - 'last_active' - 'last_rotation_at')::text, ',' ORDER BY token), '') FROM k), + (SELECT COALESCE(string_agg((to_jsonb(u) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at' + - 'last_breach_check_at' - 'password' - 'password_reset_required')::text, ',' ORDER BY user_id), '') FROM u), + (SELECT COALESCE(string_agg((to_jsonb(t) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY team_id), '') FROM t), + (SELECT COALESCE(string_agg((to_jsonb(o) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY organization_id), '') FROM o), + (SELECT COALESCE(string_agg((to_jsonb(tm) - 'spend' - 'total_spend' - 'updated_at')::text, ',' ORDER BY user_id, team_id), '') FROM tm), + (SELECT COALESCE(string_agg((to_jsonb(om) - 'spend' - 'updated_at')::text, ',' ORDER BY user_id, organization_id), '') FROM om), + (SELECT COALESCE(string_agg(to_jsonb(p)::text, ',' ORDER BY object_permission_id), '') FROM p), + (SELECT COALESCE(string_agg(to_jsonb(ag)::text, ',' ORDER BY access_group_id), '') FROM ag), + (SELECT COALESCE(string_agg((to_jsonb(b) - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY budget_id), '') FROM b) + )) AS digest, + EXISTS (SELECT 1 FROM p WHERE cardinality(mcp_toolsets) > 0) AS uses_toolsets, + jsonb_build_object( + 'user_id', COALESCE((SELECT user_id FROM k), $2::text), + 'project_id', (SELECT project_id FROM k), + 'team_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT team_id AS x FROM k WHERE team_id IS NOT NULL UNION SELECT unnest(teams) FROM u + UNION SELECT $3::text) q WHERE x IS NOT NULL), + 'organization_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT organization_id AS x FROM k UNION SELECT organization_id FROM u UNION SELECT organization_id FROM t + UNION SELECT $4::text) q WHERE x IS NOT NULL), + 'permission_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT object_permission_id AS x FROM k UNION SELECT object_permission_id FROM u + UNION SELECT object_permission_id FROM t UNION SELECT object_permission_id FROM o) q WHERE x IS NOT NULL), + 'access_group_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT unnest(access_group_ids) AS x FROM k UNION SELECT unnest(access_group_ids) FROM t) q WHERE x IS NOT NULL) + ) AS reachable +""" + + +@dataclass(frozen=True, slots=True) +class PolicySubject: + """Whose policy is being validated: a virtual key (its hash), or a keyless user (and the team and organization its token pins).""" + + key_hash: str | None = None + user_id: str | None = None + pinned_team_id: str | None = None + pinned_org_id: str | None = None + + @property + def mark(self) -> str: + return ( + f"key:{self.key_hash}" + if self.key_hash is not None + else f"user:{self.user_id}:{self.pinned_team_id}:{self.pinned_org_id}" + ) + + +class ReachableObjects(BaseModel): + """Every cached auth object the subject can read, including ids whose rows no longer exist.""" + + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + project_id: str | None = None + team_ids: tuple[str, ...] = () + organization_ids: tuple[str, ...] = () + permission_ids: tuple[str, ...] = () + access_group_ids: tuple[str, ...] = () + + +class PolicyFingerprint(BaseModel): + model_config = ConfigDict(frozen=True) + + digest: str + uses_toolsets: bool + reachable: ReachableObjects + + @field_validator("reachable", mode="before") + @classmethod + def _decode_reachable(cls, value: object) -> object: + return json.loads(value) if isinstance(value, str) else value + + +def policy_subject_for( + auth: "UserAPIKeyAuth", *, keyless_user: bool, virtual_key: bool, jwt_teams_from_claim_list: bool +) -> PolicySubject | None: + """The subject to validate for this caller, or None when its policy has inputs the hash cannot describe.""" + if auth.project_id or auth.agent_id or auth.end_user_id or auth.mcp_toolset_id: + return None + if keyless_user: + return PolicySubject(user_id=auth.user_id, pinned_team_id=auth.team_id) + if virtual_key: + return PolicySubject(key_hash=auth.api_key) + if auth.jwt_claims is not None and auth.user_id and not jwt_teams_from_claim_list: + return PolicySubject(user_id=auth.user_id, pinned_team_id=auth.team_id, pinned_org_id=auth.org_id) + return None + + +async def fetch_policy_fingerprint(prisma_client: "PrismaClient", subject: PolicySubject) -> PolicyFingerprint: + rows: Final[Sequence[Mapping[str, object]]] = await prisma_client.writer_db.query_raw( # pyright: ignore[reportAny] # Prisma's raw client is untyped + _FINGERPRINT_SQL, subject.key_hash, subject.user_id, subject.pinned_team_id, subject.pinned_org_id + ) + return PolicyFingerprint.model_validate(rows[0]) + + +def cached_policy_keys(subject: PolicySubject, reachable: ReachableObjects) -> tuple[str, ...]: + """The cache keys auth reads this subject's policy from; a reload refills the rows that still exist and cannot remove the rest.""" + user_id: Final = reachable.user_id or subject.user_id + return ( + *((subject.key_hash,) if subject.key_hash is not None else ()), + *((user_id, user_object_permission_id_cache_key(user_id)) if user_id is not None else ()), + *(f"team_id:{team_id}" for team_id in reachable.team_ids), + *( + team_membership_auth_cache_key(team_id=team_id, user_id=user_id) + for team_id in reachable.team_ids + if user_id is not None + ), + *( + team_membership_reservation_cache_key(team_id=team_id, user_id=user_id) + for team_id in reachable.team_ids + if user_id is not None + ), + *(f"org_id:{org_id}" for org_id in reachable.organization_ids), + *(f"org_id:{org_id}:with_budget" for org_id in reachable.organization_ids), + *(f"org_id:{org_id}:with_budget:last_known" for org_id in reachable.organization_ids), + *((f"project_id:{reachable.project_id}",) if reachable.project_id is not None else ()), + *(object_permission_cache_key(permission_id) for permission_id in reachable.permission_ids), + *(f"access_group_id:{group_id}" for group_id in reachable.access_group_ids), + ) + + +class ValidatedPolicyMarks: + """Per process, the digest each subject's policy was validated against. Trusted for ``max_age``, and only once settled: + a refill in flight or from a lagging replica can land after the first eviction, so a second one follows.""" + + def __init__(self, capacity: int = _MARK_CAPACITY, clock: Callable[[], float] = time.monotonic) -> None: + self._capacity: Final = capacity + self._clock: Final = clock + self._marks: Final[OrderedDict[str, tuple[str, float, bool]]] = OrderedDict() # mutable-ok: bounded LRU + + def matches(self, subject: PolicySubject, digest: str, max_age: float) -> bool: + mark: Final = self._marks.get(subject.mark) + if mark is None or mark[0] != digest or not mark[2] or self._clock() - mark[1] >= max_age: + return False + self._marks.move_to_end(subject.mark) + return True + + def phase(self, subject: PolicySubject, digest: str) -> Literal["other", "settling", "due", "settled"]: + """Where this digest stands: not the recorded one, still inside the settle wait, due its second eviction, or settled.""" + mark: Final = self._marks.get(subject.mark) + if mark is None or mark[0] != digest: + return "other" + if mark[2]: + return "settled" + return "due" if self._clock() - mark[1] >= _SETTLE_SECONDS else "settling" + + def record(self, subject: PolicySubject, digest: str, settled: bool = True) -> None: + self._marks[subject.mark] = (digest, self._clock(), settled) + self._marks.move_to_end(subject.mark) + while len(self._marks) > self._capacity: + self._marks.popitem(last=False) + + +validated_policy_marks: Final = ValidatedPolicyMarks() + + +@dataclass(slots=True) +class RequestFreshness: + """One listing: the digest to record once it completes. Live only inside ``validated_listing``, + because a stateful MCP session's handler task inherits its first request's context.""" + + staged: tuple[ValidatedPolicyMarks, PolicySubject, str, "PrismaClient", bool] | None = None + strict_token: Token[bool] | None = None + active: bool = True + + def trust_cache(self) -> None: + if self.strict_token is None: + self.strict_token = require_strict_cache_refill() + + async def commit(self) -> None: + """Record the staged digest only if the policy still hashes to it, so a change during the reload is not marked valid.""" + if self.staged is None: + return + marks, subject, digest, prisma_client, settled_mark = self.staged + self.staged = None + try: + confirmed: Final = await fetch_policy_fingerprint(prisma_client, subject) + except Exception as e: # noqa: BLE001 # an unconfirmed digest is simply not recorded + verbose_logger.warning("MCP policy fingerprint could not be confirmed after the listing: %s", e) + return + if confirmed.digest == digest: + marks.record(subject, digest, settled=settled_mark) + + def end(self) -> None: + self.active = False + if self.strict_token is not None: + release_strict_cache_refill(self.strict_token) + self.strict_token = None + + +_current_listing: Final[ContextVar[RequestFreshness | None]] = ContextVar("mcp_current_listing", default=None) + + +def current_listing() -> RequestFreshness | None: + listing: Final = _current_listing.get() + return listing if listing is not None and listing.active else None + + +@asynccontextmanager +async def validated_listing() -> AsyncGenerator[RequestFreshness]: + """One listing's scope. The digest it validated is recorded only when the listing completes.""" + listing: Final = RequestFreshness() + token: Final = _current_listing.set(listing) + try: + yield listing + await listing.commit() + finally: + listing.end() + _current_listing.reset(token) + + +async def cached_policy_is_current( + subject: PolicySubject | None, + *, + prisma_client: "PrismaClient | None", + cache: UserApiKeyCache, + marks: ValidatedPolicyMarks, +) -> bool: + """True only when the subject's policy hash equals the one recorded by the last completed listing. + + A changed hash evicts the caller's cached entries, and the digest is staged for the listing to record once it + completes. Every failure answers False, so the policy is read from the database as it was before this check. + """ + listing: Final = current_listing() + if subject is None or listing is None or prisma_client is None: + return False + try: + fingerprint: Final = await fetch_policy_fingerprint(prisma_client, subject) + except Exception as e: # noqa: BLE001 # an unreadable fingerprint must read policy from the database, not fail the request + verbose_logger.warning("MCP policy fingerprint unavailable, reading policy from the database: %s", e) + return False + if fingerprint.uses_toolsets: + return False + if marks.matches(subject, fingerprint.digest, get_management_object_ttl(cache) * _MARK_TTL_FRACTION): + listing.trust_cache() + return True + phase: Final = marks.phase(subject, fingerprint.digest) + if phase == "settling": + return False + if phase != "settled": + try: + await cache.async_delete_cache_keys(cached_policy_keys(subject, fingerprint.reachable)) + except Exception as e: # noqa: BLE001 # the forced reload still refills rows that exist; only deleted rows rely on this eviction + verbose_logger.warning("MCP cached policy eviction failed: %s", e) + return False + listing.staged = (marks, subject, fingerprint.digest, prisma_client, phase != "other") + return False diff --git a/tests/unit/proxy/db/test_routing_prisma_wrapper.py b/tests/unit/proxy/db/test_routing_prisma_wrapper.py index 73225acd2c8..94ad48779a3 100644 --- a/tests/unit/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/unit/proxy/db/test_routing_prisma_wrapper.py @@ -102,6 +102,29 @@ def test_per_model_reads_route_to_reader_writes_to_writer(): assert actions.delete_many is writer_inner.litellm_usertable.delete_many +def test_reads_route_to_the_writer_while_cache_refills_must_fail_closed() -> None: + from litellm.proxy.common_utils.strict_cache_refill import ( + release_strict_cache_refill, + require_strict_cache_refill, + ) + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.litellm_usertable = _model_actions_mock("writer_users") + reader_inner.litellm_usertable = _model_actions_mock("reader_users") + routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader) + + strict: Final = require_strict_cache_refill() + try: + assert routing.query_raw is writer_inner.query_raw + assert routing.litellm_usertable.find_unique is writer_inner.litellm_usertable.find_unique + finally: + release_strict_cache_refill(strict) + + assert routing.query_raw is reader_inner.query_raw + assert routing.litellm_usertable.find_unique is reader_inner.litellm_usertable.find_unique + + def test_writer_pinned_client_bypasses_reader_routing(): """Regression for #38556: read-after-write reconciles must see the writer's just-committed rows, so WriterPinnedClient must resolve reads to the writer diff --git a/tests/unit/proxy/mcp_server/test_policy_freshness.py b/tests/unit/proxy/mcp_server/test_policy_freshness.py new file mode 100644 index 00000000000..53e639aefe0 --- /dev/null +++ b/tests/unit/proxy/mcp_server/test_policy_freshness.py @@ -0,0 +1,782 @@ +import asyncio +from collections.abc import AsyncGenerator, Mapping +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException +from mcp.types import ListToolsRequest + +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _team_membership_inflight, + get_key_object, + get_object_permission, + get_team_membership, + get_team_object, + get_user_object, +) +from litellm.proxy.common_utils.strict_cache_refill import strict_cache_refill +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.mcp_server import catalog, operations +from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler, UnloadableEntitlementError +from litellm.proxy.mcp_server.policy_freshness import ( + PolicySubject, + ReachableObjects, + RequestFreshness, + ValidatedPolicyMarks, + cached_policy_is_current, + cached_policy_keys, + current_listing, + fetch_policy_fingerprint, + policy_subject_for, + validated_listing, + validated_policy_marks, +) +from litellm.proxy.mcp_server import policy_freshness as policy_freshness_module +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +KEY: Final = PolicySubject(key_hash="key-hash") +CACHE_TTL: Final = 600 +CACHED: Final = {"cached": True} + + +class Clock: + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + +def fingerprint_row( + digest: str, + *, + uses_toolsets: bool = False, + team_ids: tuple[str, ...] = ("team",), + access_group_ids: tuple[str, ...] = (), +) -> Mapping[str, object]: + return { + "digest": digest, + "uses_toolsets": uses_toolsets, + "reachable": { + "user_id": "owner", + "project_id": None, + "team_ids": list(team_ids), + "organization_ids": [], + "permission_ids": ["permission"], + "access_group_ids": list(access_group_ids), + }, + } + + +def database(*rows: Mapping[str, object]) -> MagicMock: + client: Final = MagicMock() + client.writer_db.query_raw = AsyncMock(side_effect=[[row] for row in rows]) + client.db.query_raw = AsyncMock(side_effect=AssertionError("a read replica must not serve the policy hash")) + return client + + +@pytest.fixture(autouse=True) +def settled_at_once(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(policy_freshness_module, "_SETTLE_SECONDS", 0.0) + + +@pytest.fixture +def cache() -> UserApiKeyCache: + return UserApiKeyCache(default_in_memory_ttl=CACHE_TTL) + + +async def is_current( + client: MagicMock, cache: UserApiKeyCache, marks: ValidatedPolicyMarks, subject: PolicySubject | None = KEY +) -> bool: + return await cached_policy_is_current(subject, prisma_client=client, cache=cache, marks=marks) + + +def test_marks_match_only_the_recorded_digest_and_stay_bounded() -> None: + marks: Final = ValidatedPolicyMarks(capacity=2) + first: Final = PolicySubject(key_hash="a") + second: Final = PolicySubject(key_hash="b") + third: Final = PolicySubject(key_hash="c") + assert not marks.matches(first, "d1", max_age=60) + marks.record(first, "d1") + marks.record(second, "d2") + assert marks.matches(first, "d1", max_age=60) + assert not marks.matches(first, "other", max_age=60) + marks.record(third, "d3") + assert marks.phase(second, "d2") == "other" + assert marks.matches(first, "d1", max_age=60) + assert marks.matches(third, "d3", max_age=60) + + +def test_a_mark_is_trusted_only_for_its_max_age_but_its_digest_is_still_recorded() -> None: + clock: Final = Clock() + marks: Final = ValidatedPolicyMarks(clock=clock) + marks.record(KEY, "d") + clock.now += 299 + assert marks.matches(KEY, "d", max_age=300) + clock.now += 2 + assert not marks.matches(KEY, "d", max_age=300) + assert marks.phase(KEY, "d") == "settled" + + +@pytest.mark.parametrize( + "other", + [ + PolicySubject(user_id="u", pinned_team_id="t2", pinned_org_id="o"), + PolicySubject(user_id="u", pinned_team_id="t", pinned_org_id="o2"), + PolicySubject(user_id="u2", pinned_team_id="t", pinned_org_id="o"), + ], +) +def test_a_keyless_mark_depends_on_the_user_team_and_organization(other: PolicySubject) -> None: + marks: Final = ValidatedPolicyMarks() + marks.record(PolicySubject(user_id="u", pinned_team_id="t", pinned_org_id="o"), "digest") + assert marks.matches(PolicySubject(user_id="u", pinned_team_id="t", pinned_org_id="o"), "digest", max_age=60) + assert not marks.matches(other, "digest", max_age=60) + + +def test_every_cache_family_is_evicted_including_ids_whose_rows_are_gone() -> None: + reachable: Final = ReachableObjects( + user_id="owner", + project_id="project", + team_ids=("live-team", "deleted-team"), + organization_ids=("org",), + permission_ids=("permission",), + access_group_ids=("live-group", "deleted-group"), + ) + assert set(cached_policy_keys(KEY, reachable)) == { + "key-hash", + "owner", + "user_object_permission_id:owner", + "team_id:live-team", + "team_id:deleted-team", + "live-team_owner", + "deleted-team_owner", + "team_membership:owner:live-team", + "team_membership:owner:deleted-team", + "org_id:org", + "org_id:org:with_budget", + "org_id:org:with_budget:last_known", + "project_id:project", + "object_permission_id:permission", + "access_group_id:live-group", + "access_group_id:deleted-group", + } + + +def test_a_keyless_subject_evicts_no_key_entry() -> None: + reachable: Final = ReachableObjects() + assert cached_policy_keys(PolicySubject(user_id="u"), reachable) == ("u", "user_object_permission_id:u") + + +@pytest.mark.asyncio +async def test_entries_written_by_the_real_readers_are_gone_after_a_changed_digest(cache: UserApiKeyCache) -> None: + permission_lookup: Final = AsyncMock(return_value=LiteLLM_ObjectPermissionTable(object_permission_id="permission")) + user_lookup: Final = AsyncMock(return_value=LiteLLM_UserTable(user_id="owner")) + reader: Final = MagicMock() + reader.db.litellm_objectpermissiontable.find_unique = permission_lookup + reader.db.litellm_usertable.find_unique = user_lookup + await get_object_permission("permission", reader, cache) + await get_user_object("owner", reader, cache, user_id_upsert=False) + await get_object_permission("permission", reader, cache) + await get_user_object("owner", reader, cache, user_id_upsert=False) + assert (permission_lookup.await_count, user_lookup.await_count) == (1, 1) + async with validated_listing(): + await is_current(database(fingerprint_row("v1"), fingerprint_row("v1")), cache, ValidatedPolicyMarks()) + await get_object_permission("permission", reader, cache) + await get_user_object("owner", reader, cache, user_id_upsert=False) + assert (permission_lookup.await_count, user_lookup.await_count) == (2, 2) + + +@pytest.mark.asyncio +async def test_the_fingerprint_is_read_from_the_writer_with_the_subject_bound() -> None: + client: Final = database(fingerprint_row("abc", team_ids=("t",), access_group_ids=("g",))) + subject: Final = PolicySubject(key_hash="h", user_id="u", pinned_team_id="t", pinned_org_id="o") + result: Final = await fetch_policy_fingerprint(client, subject) + assert (result.digest, result.uses_toolsets) == ("abc", False) + assert result.reachable.team_ids == ("t",) + assert result.reachable.access_group_ids == ("g",) + assert client.writer_db.query_raw.await_args.args[1:] == ("h", "u", "t", "o") + + +@pytest.mark.asyncio +async def test_the_fingerprint_accepts_a_json_text_reachable_column() -> None: + reachable_text: Final = ( + '{"user_id": "u", "project_id": null, "team_ids": ["t"], "organization_ids": [], ' + '"permission_ids": ["p"], "access_group_ids": []}' + ) + client: Final = database({"digest": "abc", "uses_toolsets": False, "reachable": reachable_text}) + assert (await fetch_policy_fingerprint(client, KEY)).reachable.permission_ids == ("p",) + + +@pytest.mark.asyncio +async def test_a_subject_that_cannot_be_fingerprinted_reads_the_database(cache: UserApiKeyCache) -> None: + client: Final = database(fingerprint_row("v1")) + async with validated_listing(): + assert not await is_current(client, cache, ValidatedPolicyMarks(), subject=None) + client.writer_db.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_check_outside_a_listing_reads_the_database(cache: UserApiKeyCache) -> None: + client: Final = database(fingerprint_row("v1")) + marks: Final = ValidatedPolicyMarks() + marks.record(KEY, "v1") + assert not await is_current(client, cache, marks) + client.writer_db.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_missing_database_reads_the_database(cache: UserApiKeyCache) -> None: + async with validated_listing(): + assert not await cached_policy_is_current(KEY, prisma_client=None, cache=cache, marks=ValidatedPolicyMarks()) + + +@pytest.mark.asyncio +async def test_a_failing_fingerprint_fails_closed_and_leaves_the_cache_alone(cache: UserApiKeyCache) -> None: + client: Final = MagicMock() + client.writer_db.query_raw = AsyncMock(side_effect=RuntimeError("db down")) + marks: Final = ValidatedPolicyMarks() + marks.record(KEY, "v1") + await cache.async_set_cache(key="team_id:team", value=CACHED) + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert await cache.async_get_cache(key="team_id:team") == CACHED + + +@pytest.mark.asyncio +async def test_a_first_listing_evicts_the_callers_entries_and_a_completed_one_validates_the_digest( + cache: UserApiKeyCache, +) -> None: + marks: Final = ValidatedPolicyMarks() + client: Final = database(*(fingerprint_row("v1"),) * 5) + await cache.async_set_cache(key="team_id:team", value=CACHED) + await cache.async_set_cache(key="team_id:someone-else", value=CACHED) + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert not strict_cache_refill() + assert await cache.async_get_cache(key="team_id:team") is None + assert await cache.async_get_cache(key="team_id:someone-else") == CACHED + async with validated_listing(): + assert not await is_current(client, cache, marks) + async with validated_listing(): + assert await is_current(client, cache, marks) + assert strict_cache_refill() + assert not strict_cache_refill() + + +@pytest.mark.asyncio +async def test_a_refill_that_lands_after_the_first_eviction_is_evicted_again_before_the_cache_is_trusted( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(policy_freshness_module, "_SETTLE_SECONDS", 5.0) + clock: Final = Clock() + marks: Final = ValidatedPolicyMarks(clock=clock) + client: Final = database(*(fingerprint_row("v1"),) * 6) + async with validated_listing(): + assert not await is_current(client, cache, marks) + await cache.async_set_cache(key="team_id:team", value=CACHED) + clock.now += 1 + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert not strict_cache_refill() + assert await cache.async_get_cache(key="team_id:team") == CACHED + clock.now += 5 + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert not strict_cache_refill() + assert await cache.async_get_cache(key="team_id:team") is None + async with validated_listing(): + assert await is_current(client, cache, marks) + + +@pytest.mark.asyncio +async def test_a_listing_that_fails_does_not_validate_the_digest(cache: UserApiKeyCache) -> None: + marks: Final = ValidatedPolicyMarks() + client: Final = database(fingerprint_row("v1"), fingerprint_row("v1"), fingerprint_row("v1")) + with pytest.raises(RuntimeError): + async with validated_listing(): + assert not await is_current(client, cache, marks) + raise RuntimeError("upstream failed") + async with validated_listing(): + assert not await is_current(client, cache, marks) + + +@pytest.mark.asyncio +async def test_a_change_made_during_the_reload_is_not_marked_valid(cache: UserApiKeyCache) -> None: + marks: Final = ValidatedPolicyMarks() + client: Final = database(fingerprint_row("v1"), fingerprint_row("v2"), fingerprint_row("v1")) + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert marks.phase(KEY, "v1") == "other" + async with validated_listing(): + assert not await is_current(client, cache, marks) + + +@pytest.mark.asyncio +async def test_a_changed_digest_evicts_again_and_stops_trusting_the_cache(cache: UserApiKeyCache) -> None: + marks: Final = ValidatedPolicyMarks() + marks.record(KEY, "v1") + client: Final = database(fingerprint_row("v2")) + await cache.async_set_cache(key="team_id:team", value=CACHED) + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert not strict_cache_refill() + assert await cache.async_get_cache(key="team_id:team") is None + + +@pytest.mark.asyncio +async def test_an_expired_validation_reads_the_database_but_keeps_the_shared_entries(cache: UserApiKeyCache) -> None: + clock: Final = Clock() + marks: Final = ValidatedPolicyMarks(clock=clock) + marks.record(KEY, "v1") + client: Final = database(*(fingerprint_row("v1"),) * 4) + await cache.async_set_cache(key="team_id:team", value=CACHED) + clock.now += 299 + async with validated_listing(): + assert await is_current(client, cache, marks) + clock.now += 2 + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert await cache.async_get_cache(key="team_id:team") == CACHED + async with validated_listing(): + assert await is_current(client, cache, marks) + + +@pytest.mark.asyncio +async def test_access_group_entries_are_evicted_with_the_rest(cache: UserApiKeyCache) -> None: + client: Final = database(fingerprint_row("v1", access_group_ids=("group",)), fingerprint_row("v1")) + await cache.async_set_cache(key="access_group_id:group", value=CACHED) + async with validated_listing(): + await is_current(client, cache, ValidatedPolicyMarks()) + assert await cache.async_get_cache(key="access_group_id:group") is None + + +@pytest.mark.asyncio +async def test_toolset_permissions_keep_the_database_read(cache: UserApiKeyCache) -> None: + marks: Final = ValidatedPolicyMarks() + marks.record(KEY, "v1") + client: Final = database(fingerprint_row("v1", uses_toolsets=True), fingerprint_row("v1", uses_toolsets=True)) + await cache.async_set_cache(key="team_id:team", value=CACHED) + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert not strict_cache_refill() + assert await cache.async_get_cache(key="team_id:team") == CACHED + + +@pytest.mark.asyncio +async def test_a_failed_eviction_fails_closed_and_validates_nothing(cache: UserApiKeyCache) -> None: + redis: Final = MagicMock() + redis.delete_cache_keys = AsyncMock(side_effect=RuntimeError("cache down")) + cache.attach_redis_cache(redis) + marks: Final = ValidatedPolicyMarks() + client: Final = database(fingerprint_row("v1"), fingerprint_row("v1")) + async with validated_listing(): + assert not await is_current(client, cache, marks) + assert marks.phase(KEY, "v1") == "other" + + +@pytest.mark.asyncio +async def test_a_task_started_inside_a_listing_sees_it_end_and_a_listing_there_opens_its_own() -> None: + gate: Final = asyncio.Event() + + async def read_after_the_listing() -> RequestFreshness | None: + await gate.wait() + return current_listing() + + async def list_after_the_listing() -> bool: + await gate.wait() + async with validated_listing() as listing: + return listing is current_listing() + + async with validated_listing(): + reader: Final = asyncio.create_task(read_after_the_listing()) + lister: Final = asyncio.create_task(list_after_the_listing()) + gate.set() + assert await reader is None + assert await lister is True + + +def caller( + *, + jwt: bool = False, + user_id: str | None = "someone", + project_id: str | None = None, + agent_id: str | None = None, + end_user_id: str | None = None, + toolset_id: str | None = None, +) -> UserAPIKeyAuth: + auth: Final = UserAPIKeyAuth( + api_key="key-hash", + user_id=user_id, + team_id="team", + org_id="org", + project_id=project_id, + agent_id=agent_id, + end_user_id=end_user_id, + jwt_claims={"sub": "someone"} if jwt else None, + ) + auth.mcp_toolset_id = toolset_id + return auth + + +def test_a_virtual_key_is_validated_by_its_hash() -> None: + subject: Final = policy_subject_for(caller(), keyless_user=False, virtual_key=True, jwt_teams_from_claim_list=False) + assert subject == PolicySubject(key_hash="key-hash") + + +def test_a_keyless_user_is_validated_by_user_and_pinned_team() -> None: + subject: Final = policy_subject_for(caller(), keyless_user=True, virtual_key=False, jwt_teams_from_claim_list=False) + assert subject == PolicySubject(user_id="someone", pinned_team_id="team") + + +def test_a_jwt_caller_is_validated_by_user_team_and_organization() -> None: + subject: Final = policy_subject_for( + caller(jwt=True), keyless_user=False, virtual_key=False, jwt_teams_from_claim_list=False + ) + assert subject == PolicySubject(user_id="someone", pinned_team_id="team", pinned_org_id="org") + + +@pytest.mark.parametrize("keyless_user", [True, False]) +@pytest.mark.parametrize( + "scoped", + [ + caller(project_id="project"), + caller(agent_id="agent"), + caller(end_user_id="customer"), + caller(toolset_id="toolset"), + ], +) +def test_callers_scoped_beyond_the_hash_keep_the_database_read(scoped: UserAPIKeyAuth, keyless_user: bool) -> None: + assert ( + policy_subject_for( + scoped, keyless_user=keyless_user, virtual_key=not keyless_user, jwt_teams_from_claim_list=False + ) + is None + ) + + +@pytest.mark.parametrize( + ("auth", "claim_list"), + [ + (caller(jwt=True), True), + (caller(jwt=True, user_id=None), False), + (caller(), False), + ], +) +def test_callers_the_hash_cannot_describe_keep_the_database_read(auth: UserAPIKeyAuth, claim_list: bool) -> None: + assert policy_subject_for(auth, keyless_user=False, virtual_key=False, jwt_teams_from_claim_list=claim_list) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing", [True, False]) +async def test_a_cache_refill_during_a_validated_listing_fails_closed(cache: UserApiKeyCache, missing: bool) -> None: + marks: Final = ValidatedPolicyMarks() + marks.record(KEY, "v1") + lookup: Final = AsyncMock(return_value=None, side_effect=None if missing else RuntimeError("unavailable")) + failing: Final = MagicMock() + failing.writer_db.litellm_objectpermissiontable.find_unique = lookup + failing.db.litellm_objectpermissiontable.find_unique = lookup + async with validated_listing(): + assert await is_current(database(fingerprint_row("v1")), cache, marks) + with pytest.raises(HTTPException if missing else RuntimeError): + await get_object_permission("referenced", failing, cache) + assert await get_object_permission("referenced", failing, cache) is None + + +@pytest.mark.asyncio +async def test_a_user_ceiling_that_cannot_be_read_denies_during_a_validated_listing( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + marks: Final = ValidatedPolicyMarks() + marks.record(KEY, "v1") + failing: Final = MagicMock() + failing.db.litellm_usertable.find_unique = AsyncMock(side_effect=RuntimeError("unavailable")) + async with validated_listing(): + assert await is_current(database(fingerprint_row("v1")), cache, marks) + with pytest.raises(HTTPException) as denied: + await MCPRequestHandler._user_object_permission_id("someone", failing) + assert denied.value.status_code == 503 + assert await MCPRequestHandler._user_object_permission_id("someone", failing) is None + + +def admitted_user(user_id: str) -> UserAPIKeyAuth: + auth: Final = UserAPIKeyAuth(user_id=user_id) + auth.mcp_admitted_user_subject = True + return auth + + +def serve_policy( + monkeypatch: pytest.MonkeyPatch, cache: UserApiKeyCache, user_id: str, *rows: Mapping[str, object] +) -> tuple[MagicMock, AsyncMock]: + user_lookup: Final = AsyncMock(return_value=LiteLLM_UserTable(user_id=user_id, user_role="internal_user")) + client: Final = MagicMock() + client.writer_db.litellm_usertable.find_unique = user_lookup + client.db.litellm_usertable.find_unique = user_lookup + client.writer_db.query_raw = AsyncMock(side_effect=[[row] for row in rows]) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + return client, user_lookup + + +async def forced_fresh(auth: UserAPIKeyAuth) -> bool | None: + async with validated_listing(): + refreshed: Final = await MCPRequestHandler.refresh_catalog_authority(auth) + return None if refreshed is None else refreshed.requires_fresh_policy + + +@pytest.mark.asyncio +async def test_a_listing_reads_the_admitted_user_from_the_cache_once_its_policy_is_validated( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + _, user_lookup = serve_policy(monkeypatch, cache, "wired-user", *(fingerprint_row("v1"),) * 6) + flags: Final = ( + await forced_fresh(admitted_user("wired-user")), + await forced_fresh(admitted_user("wired-user")), + await forced_fresh(admitted_user("wired-user")), + await forced_fresh(admitted_user("wired-user")), + ) + assert flags == (True, True, False, False) + assert user_lookup.await_count == 4 + + +@pytest.mark.asyncio +async def test_a_policy_change_between_listings_sends_the_next_one_to_the_database( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + rows: Final = (*(fingerprint_row("v1"),) * 5, *(fingerprint_row("v2"),) * 2) + serve_policy(monkeypatch, cache, "changing-user", *rows) + flags: Final = ( + await forced_fresh(admitted_user("changing-user")), + await forced_fresh(admitted_user("changing-user")), + await forced_fresh(admitted_user("changing-user")), + await forced_fresh(admitted_user("changing-user")), + ) + assert flags == (True, True, False, True) + + +@pytest.mark.asyncio +async def test_a_jwt_caller_is_trusted_to_the_cache_once_its_policy_is_validated( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + serve_policy(monkeypatch, cache, "jwt-user", *(fingerprint_row("v1"),) * 5) + + def jwt_caller() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="jwt-user", team_id="jwt-team", jwt_claims={"sub": "jwt-user"}) + + flags: Final = [await forced_fresh(jwt_caller()) for _ in range(3)] + assert flags == [True, True, False] + + +@pytest.mark.asyncio +async def test_a_project_scoped_caller_never_reaches_the_policy_hash( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + client, _ = serve_policy(monkeypatch, cache, "scoped-user", fingerprint_row("v1")) + scoped: Final = UserAPIKeyAuth(user_id="scoped-user", project_id="project", jwt_claims={"sub": "scoped-user"}) + assert await forced_fresh(scoped) is True + client.writer_db.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_jwt_caller_whose_teams_come_from_a_claim_list_keeps_the_database_read( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + proxy_server.jwt_handler, "litellm_jwtauth", LiteLLM_JWTAuth(team_ids_jwt_field="teams"), raising=False + ) + client, _ = serve_policy(monkeypatch, cache, "claim-list-user", fingerprint_row("v1")) + claim_list_caller: Final = UserAPIKeyAuth(user_id="claim-list-user", jwt_claims={"sub": "claim-list-user"}) + assert await forced_fresh(claim_list_caller) is True + client.writer_db.query_raw.assert_not_awaited() + + +@asynccontextmanager +async def trusted_listing(cache: UserApiKeyCache) -> AsyncGenerator[None]: + marks: Final = ValidatedPolicyMarks() + marks.record(KEY, "v1") + async with validated_listing(): + assert await is_current(database(fingerprint_row("v1")), cache, marks) + yield + + +def unreadable_client() -> MagicMock: + failing: Final = AsyncMock(side_effect=RuntimeError("unavailable")) + client: Final = MagicMock() + for reader in (client.db, client.writer_db): + for table in (reader.litellm_accessgrouptable, reader.litellm_teamtable, reader.litellm_mcpservertable): + table.find_unique = failing + table.find_many = failing + return client + + +def install_proxy(monkeypatch: pytest.MonkeyPatch, cache: UserApiKeyCache, client: MagicMock) -> None: + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + + +@pytest.fixture +def unreadable_database(cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch) -> MagicMock: + client: Final = unreadable_client() + install_proxy(monkeypatch, cache, client) + return client + + +@pytest.mark.asyncio +async def test_a_caller_that_already_requires_fresh_policy_stays_on_the_database_read( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + serve_policy(monkeypatch, cache, "pinned-user", *(fingerprint_row("v1"),) * 6) + flags: Final = [await forced_fresh(admitted_user("pinned-user")) for _ in range(3)] + assert flags == [True, True, False] + pinned: Final = admitted_user("pinned-user") + pinned.requires_fresh_policy = True + assert await forced_fresh(pinned) is True + + +@pytest.mark.asyncio +async def test_a_real_listing_validates_the_digest_and_later_ones_read_the_key_from_the_cache( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + grants: Final = LiteLLM_ObjectPermissionTable(object_permission_id="listing-grants", mcp_servers=["alpha"]) + get_data: Final = AsyncMock(return_value=UserAPIKeyAuth(object_permission_id="listing-grants")) + tables: Final = SimpleNamespace( + litellm_objectpermissiontable=SimpleNamespace(find_unique=AsyncMock(return_value=grants)), + query_raw=AsyncMock(side_effect=[[fingerprint_row("v1")]] * 6), + ) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables, get_data=get_data)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": []}) + monkeypatch.setitem( + operations.global_mcp_server_manager.registry, + "alpha", + MCPServer(server_id="alpha", name="alpha", transport=MCPTransport.http, spec_path="/listing.yaml"), + ) + caller: Final = UserAPIKeyAuth( + api_key="listing-key-hash", object_permission=grants, object_permission_id="listing-grants" + ) + caller.via_virtual_key = True + for _ in range(4): + await catalog.list_gateway_catalog( + operations.prepare_context(user_api_key_auth=caller), + ListToolsRequest(), + log_list_tools_to_spendlogs=False, + ) + assert get_data.await_count == 3 + assert validated_policy_marks.phase(PolicySubject(key_hash="listing-key-hash"), "v1") == "settled" + + +@pytest.mark.asyncio +async def test_unreadable_key_access_groups_deny_during_a_trusted_listing( + cache: UserApiKeyCache, unreadable_database: MagicMock +) -> None: + key: Final = UserAPIKeyAuth(api_key="group-key", access_group_ids=["group"]) + outside: Final = await MCPRequestHandler._key_access_group_tools_for_server(key, "alpha") + assert outside.owned == () and not outside.owns_server + async with trusted_listing(cache): + with pytest.raises(UnloadableEntitlementError): + await MCPRequestHandler._key_access_group_tools_for_server(key, "alpha") + + +@pytest.mark.asyncio +async def test_unreadable_team_access_groups_deny_during_a_trusted_listing( + cache: UserApiKeyCache, monkeypatch: pytest.MonkeyPatch +) -> None: + client: Final = unreadable_client() + team: Final = AsyncMock(return_value=LiteLLM_TeamTable(team_id="team", access_group_ids=["group"])) + client.db.litellm_teamtable.find_unique = team + client.writer_db.litellm_teamtable.find_unique = team + install_proxy(monkeypatch, cache, client) + member: Final = UserAPIKeyAuth(api_key="team-key", team_id="team") + assert await MCPRequestHandler._team_access_group_tools_for_server(member, "alpha") == () + async with trusted_listing(cache): + with pytest.raises(UnloadableEntitlementError): + await MCPRequestHandler._team_access_group_tools_for_server(member, "alpha") + + +@pytest.mark.asyncio +async def test_an_unreadable_team_ceiling_is_not_cached_as_no_ceiling_during_a_trusted_listing( + cache: UserApiKeyCache, +) -> None: + client: Final = unreadable_client() + team: Final = AsyncMock(return_value=LiteLLM_TeamTable(team_id="ceiling-team", object_permission_id="ceiling")) + client.db.litellm_teamtable.find_unique = team + client.writer_db.litellm_teamtable.find_unique = team + failing: Final = AsyncMock(side_effect=RuntimeError("unavailable")) + client.db.litellm_objectpermissiontable.find_unique = failing + client.writer_db.litellm_objectpermissiontable.find_unique = failing + async with trusted_listing(cache): + with pytest.raises(HTTPException): + await get_team_object(team_id="ceiling-team", prisma_client=client, user_api_key_cache=cache) + assert await cache.async_get_cache(key="team_id:ceiling-team") is None + + +@pytest.mark.asyncio +async def test_a_trusted_listing_does_not_join_a_membership_load_another_request_started( + cache: UserApiKeyCache, +) -> None: + shared: Final = LiteLLM_TeamMembership(user_id="member", team_id="shared-team") + + async def load_elsewhere() -> LiteLLM_TeamMembership: + return shared + + started_elsewhere: Final = asyncio.ensure_future(load_elsewhere()) + _team_membership_inflight["team_membership:member:shared-team"] = started_elsewhere + client: Final = MagicMock() + lookup: Final = AsyncMock(return_value=None) + client.db.litellm_teammembership.find_unique = lookup + client.writer_db.litellm_teammembership.find_unique = lookup + try: + assert await get_team_membership("member", "shared-team", client, cache) == shared + async with trusted_listing(cache): + assert await get_team_membership("member", "shared-team", client, cache) is None + finally: + _team_membership_inflight.pop("team_membership:member:shared-team", None) + + +@pytest.mark.asyncio +async def test_unreadable_access_group_servers_fail_during_a_trusted_listing( + cache: UserApiKeyCache, unreadable_database: MagicMock +) -> None: + assert await MCPRequestHandler._get_db_server_ids_for_access_groups(unreadable_database, ["group"]) == set() + assert await MCPRequestHandler.get_mcp_servers_from_access_groups(["group"]) == [] + async with trusted_listing(cache): + with pytest.raises(RuntimeError): + await MCPRequestHandler._get_db_server_ids_for_access_groups(unreadable_database, ["group"]) + with pytest.raises(RuntimeError): + await MCPRequestHandler.get_mcp_servers_from_access_groups(["group"]) + + +@pytest.mark.asyncio +async def test_an_unreadable_team_fails_during_a_trusted_listing( + cache: UserApiKeyCache, unreadable_database: MagicMock +) -> None: + member: Final = UserAPIKeyAuth(api_key="team-key", team_id="team") + assert await MCPRequestHandler._allowed_mcp_servers_for_single_team("team", member) == [] + async with trusted_listing(cache): + with pytest.raises(HTTPException) as unreadable: + await MCPRequestHandler._allowed_mcp_servers_for_single_team("team", member) + assert unreadable.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_a_key_whose_permission_cannot_be_read_fails_during_a_trusted_listing(cache: UserApiKeyCache) -> None: + client: Final = MagicMock() + client.get_data = AsyncMock(return_value=UserAPIKeyAuth(object_permission_id="permission")) + client.db.litellm_objectpermissiontable.find_unique = AsyncMock(side_effect=RuntimeError("unavailable")) + assert (await get_key_object("key-outside", client, cache)).object_permission_id == "permission" + async with trusted_listing(cache): + with pytest.raises(RuntimeError): + await get_key_object("key-inside", client, cache) From 38ea301ea12753debd556859e133fd4cfaac4ae6 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 4 Nov 2026 20:29:23 +0000 Subject: [PATCH 2/3] test(mcp): keep the failing listing out of the pytest.raises block --- tests/unit/proxy/mcp_server/test_policy_freshness.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit/proxy/mcp_server/test_policy_freshness.py b/tests/unit/proxy/mcp_server/test_policy_freshness.py index 53e639aefe0..e642faaddc9 100644 --- a/tests/unit/proxy/mcp_server/test_policy_freshness.py +++ b/tests/unit/proxy/mcp_server/test_policy_freshness.py @@ -305,10 +305,14 @@ async def test_a_refill_that_lands_after_the_first_eviction_is_evicted_again_bef async def test_a_listing_that_fails_does_not_validate_the_digest(cache: UserApiKeyCache) -> None: marks: Final = ValidatedPolicyMarks() client: Final = database(fingerprint_row("v1"), fingerprint_row("v1"), fingerprint_row("v1")) - with pytest.raises(RuntimeError): + + async def failing_listing() -> None: async with validated_listing(): assert not await is_current(client, cache, marks) raise RuntimeError("upstream failed") + + with pytest.raises(RuntimeError): + await failing_listing() async with validated_listing(): assert not await is_current(client, cache, marks) From 426b2b36b8e611d3e17d1af040b6b788c57244cf Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 4 Nov 2026 20:46:32 +0000 Subject: [PATCH 3/3] refactor(mcp): move the policy fingerprint query into MCPPolicyRepository --- litellm/models/mcp_policy.py | 48 ++++++++ litellm/proxy/mcp_server/policy_freshness.py | 111 +----------------- litellm/repositories/mcp_policy_repository.py | 67 +++++++++++ .../proxy/mcp_server/test_policy_freshness.py | 9 +- 4 files changed, 124 insertions(+), 111 deletions(-) create mode 100644 litellm/models/mcp_policy.py create mode 100644 litellm/repositories/mcp_policy_repository.py diff --git a/litellm/models/mcp_policy.py b/litellm/models/mcp_policy.py new file mode 100644 index 00000000000..a39b86399c5 --- /dev/null +++ b/litellm/models/mcp_policy.py @@ -0,0 +1,48 @@ +import json +from dataclasses import dataclass + +from pydantic import BaseModel, ConfigDict, field_validator + + +@dataclass(frozen=True, slots=True) +class PolicySubject: + """Whose policy is being validated: a virtual key (its hash), or a keyless user (and the team and organization its token pins).""" + + key_hash: str | None = None + user_id: str | None = None + pinned_team_id: str | None = None + pinned_org_id: str | None = None + + @property + def mark(self) -> str: + return ( + f"key:{self.key_hash}" + if self.key_hash is not None + else f"user:{self.user_id}:{self.pinned_team_id}:{self.pinned_org_id}" + ) + + +class ReachableObjects(BaseModel): + """Every cached auth object the subject can read, including ids whose rows no longer exist.""" + + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + project_id: str | None = None + team_ids: tuple[str, ...] = () + organization_ids: tuple[str, ...] = () + permission_ids: tuple[str, ...] = () + access_group_ids: tuple[str, ...] = () + + +class PolicyFingerprint(BaseModel): + model_config = ConfigDict(frozen=True) + + digest: str + uses_toolsets: bool + reachable: ReachableObjects + + @field_validator("reachable", mode="before") + @classmethod + def _decode_reachable(cls, value: object) -> object: + return json.loads(value) if isinstance(value, str) else value diff --git a/litellm/proxy/mcp_server/policy_freshness.py b/litellm/proxy/mcp_server/policy_freshness.py index 5662c48a21b..8af93552f30 100644 --- a/litellm/proxy/mcp_server/policy_freshness.py +++ b/litellm/proxy/mcp_server/policy_freshness.py @@ -4,18 +4,16 @@ One query on the writer hashes every row the caller's MCP access depends on. Whi recorded after the last completed listing, the cache is trusted. Anything else reads policy from the database. """ -import json import time from collections import OrderedDict -from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Callable from contextlib import asynccontextmanager from contextvars import ContextVar, Token from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Literal -from pydantic import BaseModel, ConfigDict, field_validator - from litellm._logging import verbose_logger +from litellm.models.mcp_policy import PolicySubject, ReachableObjects from litellm.proxy.common_utils.strict_cache_refill import release_strict_cache_refill, require_strict_cache_refill from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, @@ -25,6 +23,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_reservation_cache_key, user_object_permission_id_cache_key, ) +from litellm.repositories.mcp_policy_repository import MCPPolicyRepository if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth @@ -34,99 +33,6 @@ _MARK_CAPACITY: Final = 4096 _MARK_TTL_FRACTION: Final = 0.5 _SETTLE_SECONDS: Final = 5.0 -# Spend, activity and credential columns change on every request or are secrets, and never describe access. -_FINGERPRINT_SQL: Final = """ -WITH k AS (SELECT * FROM "LiteLLM_VerificationToken" WHERE token = $1::text), -u AS (SELECT * FROM "LiteLLM_UserTable" WHERE user_id IN (SELECT user_id FROM k) OR user_id = $2::text), -t AS (SELECT * FROM "LiteLLM_TeamTable" WHERE team_id IN (SELECT team_id FROM k WHERE team_id IS NOT NULL) - OR team_id IN (SELECT unnest(teams) FROM u) OR team_id = $3::text), -o AS (SELECT * FROM "LiteLLM_OrganizationTable" - WHERE organization_id IN (SELECT organization_id FROM k UNION SELECT organization_id FROM u - UNION SELECT organization_id FROM t) - OR organization_id = $4::text), -tm AS (SELECT * FROM "LiteLLM_TeamMembership" WHERE user_id IN (SELECT user_id FROM u) OR user_id IN (SELECT user_id FROM k)), -om AS (SELECT * FROM "LiteLLM_OrganizationMembership" WHERE user_id IN (SELECT user_id FROM u) OR user_id IN (SELECT user_id FROM k)), -p AS (SELECT * FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id IN (SELECT object_permission_id FROM k UNION SELECT object_permission_id FROM u - UNION SELECT object_permission_id FROM t UNION SELECT object_permission_id FROM o)), -ag AS (SELECT * FROM "LiteLLM_AccessGroupTable" WHERE access_group_id IN (SELECT unnest(access_group_ids) FROM k UNION SELECT unnest(access_group_ids) FROM t)), -b AS (SELECT * FROM "LiteLLM_BudgetTable" WHERE budget_id IN (SELECT budget_id FROM k UNION SELECT budget_id FROM o - UNION SELECT budget_id FROM tm UNION SELECT budget_id FROM om)) -SELECT - md5(concat_ws('#', - (SELECT COALESCE(string_agg((to_jsonb(k) - 'spend' - 'total_spend' - 'model_spend' - 'budget_reset_at' - - 'updated_at' - 'last_active' - 'last_rotation_at')::text, ',' ORDER BY token), '') FROM k), - (SELECT COALESCE(string_agg((to_jsonb(u) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at' - - 'last_breach_check_at' - 'password' - 'password_reset_required')::text, ',' ORDER BY user_id), '') FROM u), - (SELECT COALESCE(string_agg((to_jsonb(t) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY team_id), '') FROM t), - (SELECT COALESCE(string_agg((to_jsonb(o) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY organization_id), '') FROM o), - (SELECT COALESCE(string_agg((to_jsonb(tm) - 'spend' - 'total_spend' - 'updated_at')::text, ',' ORDER BY user_id, team_id), '') FROM tm), - (SELECT COALESCE(string_agg((to_jsonb(om) - 'spend' - 'updated_at')::text, ',' ORDER BY user_id, organization_id), '') FROM om), - (SELECT COALESCE(string_agg(to_jsonb(p)::text, ',' ORDER BY object_permission_id), '') FROM p), - (SELECT COALESCE(string_agg(to_jsonb(ag)::text, ',' ORDER BY access_group_id), '') FROM ag), - (SELECT COALESCE(string_agg((to_jsonb(b) - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY budget_id), '') FROM b) - )) AS digest, - EXISTS (SELECT 1 FROM p WHERE cardinality(mcp_toolsets) > 0) AS uses_toolsets, - jsonb_build_object( - 'user_id', COALESCE((SELECT user_id FROM k), $2::text), - 'project_id', (SELECT project_id FROM k), - 'team_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( - SELECT team_id AS x FROM k WHERE team_id IS NOT NULL UNION SELECT unnest(teams) FROM u - UNION SELECT $3::text) q WHERE x IS NOT NULL), - 'organization_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( - SELECT organization_id AS x FROM k UNION SELECT organization_id FROM u UNION SELECT organization_id FROM t - UNION SELECT $4::text) q WHERE x IS NOT NULL), - 'permission_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( - SELECT object_permission_id AS x FROM k UNION SELECT object_permission_id FROM u - UNION SELECT object_permission_id FROM t UNION SELECT object_permission_id FROM o) q WHERE x IS NOT NULL), - 'access_group_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( - SELECT unnest(access_group_ids) AS x FROM k UNION SELECT unnest(access_group_ids) FROM t) q WHERE x IS NOT NULL) - ) AS reachable -""" - - -@dataclass(frozen=True, slots=True) -class PolicySubject: - """Whose policy is being validated: a virtual key (its hash), or a keyless user (and the team and organization its token pins).""" - - key_hash: str | None = None - user_id: str | None = None - pinned_team_id: str | None = None - pinned_org_id: str | None = None - - @property - def mark(self) -> str: - return ( - f"key:{self.key_hash}" - if self.key_hash is not None - else f"user:{self.user_id}:{self.pinned_team_id}:{self.pinned_org_id}" - ) - - -class ReachableObjects(BaseModel): - """Every cached auth object the subject can read, including ids whose rows no longer exist.""" - - model_config = ConfigDict(frozen=True) - - user_id: str | None = None - project_id: str | None = None - team_ids: tuple[str, ...] = () - organization_ids: tuple[str, ...] = () - permission_ids: tuple[str, ...] = () - access_group_ids: tuple[str, ...] = () - - -class PolicyFingerprint(BaseModel): - model_config = ConfigDict(frozen=True) - - digest: str - uses_toolsets: bool - reachable: ReachableObjects - - @field_validator("reachable", mode="before") - @classmethod - def _decode_reachable(cls, value: object) -> object: - return json.loads(value) if isinstance(value, str) else value - def policy_subject_for( auth: "UserAPIKeyAuth", *, keyless_user: bool, virtual_key: bool, jwt_teams_from_claim_list: bool @@ -143,13 +49,6 @@ def policy_subject_for( return None -async def fetch_policy_fingerprint(prisma_client: "PrismaClient", subject: PolicySubject) -> PolicyFingerprint: - rows: Final[Sequence[Mapping[str, object]]] = await prisma_client.writer_db.query_raw( # pyright: ignore[reportAny] # Prisma's raw client is untyped - _FINGERPRINT_SQL, subject.key_hash, subject.user_id, subject.pinned_team_id, subject.pinned_org_id - ) - return PolicyFingerprint.model_validate(rows[0]) - - def cached_policy_keys(subject: PolicySubject, reachable: ReachableObjects) -> tuple[str, ...]: """The cache keys auth reads this subject's policy from; a reload refills the rows that still exist and cannot remove the rest.""" user_id: Final = reachable.user_id or subject.user_id @@ -231,7 +130,7 @@ class RequestFreshness: marks, subject, digest, prisma_client, settled_mark = self.staged self.staged = None try: - confirmed: Final = await fetch_policy_fingerprint(prisma_client, subject) + confirmed: Final = await MCPPolicyRepository(prisma_client).get_fingerprint(subject) except Exception as e: # noqa: BLE001 # an unconfirmed digest is simply not recorded verbose_logger.warning("MCP policy fingerprint could not be confirmed after the listing: %s", e) return @@ -282,7 +181,7 @@ async def cached_policy_is_current( if subject is None or listing is None or prisma_client is None: return False try: - fingerprint: Final = await fetch_policy_fingerprint(prisma_client, subject) + fingerprint: Final = await MCPPolicyRepository(prisma_client).get_fingerprint(subject) except Exception as e: # noqa: BLE001 # an unreadable fingerprint must read policy from the database, not fail the request verbose_logger.warning("MCP policy fingerprint unavailable, reading policy from the database: %s", e) return False diff --git a/litellm/repositories/mcp_policy_repository.py b/litellm/repositories/mcp_policy_repository.py new file mode 100644 index 00000000000..c0c1d7e07cb --- /dev/null +++ b/litellm/repositories/mcp_policy_repository.py @@ -0,0 +1,67 @@ +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.models.mcp_policy import PolicyFingerprint, PolicySubject + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +# Spend, activity and credential columns change on every request or are secrets, and never describe access. +_FINGERPRINT_SQL: Final = """ +WITH k AS (SELECT * FROM "LiteLLM_VerificationToken" WHERE token = $1::text), +u AS (SELECT * FROM "LiteLLM_UserTable" WHERE user_id IN (SELECT user_id FROM k) OR user_id = $2::text), +t AS (SELECT * FROM "LiteLLM_TeamTable" WHERE team_id IN (SELECT team_id FROM k WHERE team_id IS NOT NULL) + OR team_id IN (SELECT unnest(teams) FROM u) OR team_id = $3::text), +o AS (SELECT * FROM "LiteLLM_OrganizationTable" + WHERE organization_id IN (SELECT organization_id FROM k UNION SELECT organization_id FROM u + UNION SELECT organization_id FROM t) + OR organization_id = $4::text), +tm AS (SELECT * FROM "LiteLLM_TeamMembership" WHERE user_id IN (SELECT user_id FROM u) OR user_id IN (SELECT user_id FROM k)), +om AS (SELECT * FROM "LiteLLM_OrganizationMembership" WHERE user_id IN (SELECT user_id FROM u) OR user_id IN (SELECT user_id FROM k)), +p AS (SELECT * FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id IN (SELECT object_permission_id FROM k UNION SELECT object_permission_id FROM u + UNION SELECT object_permission_id FROM t UNION SELECT object_permission_id FROM o)), +ag AS (SELECT * FROM "LiteLLM_AccessGroupTable" WHERE access_group_id IN (SELECT unnest(access_group_ids) FROM k UNION SELECT unnest(access_group_ids) FROM t)), +b AS (SELECT * FROM "LiteLLM_BudgetTable" WHERE budget_id IN (SELECT budget_id FROM k UNION SELECT budget_id FROM o + UNION SELECT budget_id FROM tm UNION SELECT budget_id FROM om)) +SELECT + md5(concat_ws('#', + (SELECT COALESCE(string_agg((to_jsonb(k) - 'spend' - 'total_spend' - 'model_spend' - 'budget_reset_at' + - 'updated_at' - 'last_active' - 'last_rotation_at')::text, ',' ORDER BY token), '') FROM k), + (SELECT COALESCE(string_agg((to_jsonb(u) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at' + - 'last_breach_check_at' - 'password' - 'password_reset_required')::text, ',' ORDER BY user_id), '') FROM u), + (SELECT COALESCE(string_agg((to_jsonb(t) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY team_id), '') FROM t), + (SELECT COALESCE(string_agg((to_jsonb(o) - 'spend' - 'model_spend' - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY organization_id), '') FROM o), + (SELECT COALESCE(string_agg((to_jsonb(tm) - 'spend' - 'total_spend' - 'updated_at')::text, ',' ORDER BY user_id, team_id), '') FROM tm), + (SELECT COALESCE(string_agg((to_jsonb(om) - 'spend' - 'updated_at')::text, ',' ORDER BY user_id, organization_id), '') FROM om), + (SELECT COALESCE(string_agg(to_jsonb(p)::text, ',' ORDER BY object_permission_id), '') FROM p), + (SELECT COALESCE(string_agg(to_jsonb(ag)::text, ',' ORDER BY access_group_id), '') FROM ag), + (SELECT COALESCE(string_agg((to_jsonb(b) - 'budget_reset_at' - 'updated_at')::text, ',' ORDER BY budget_id), '') FROM b) + )) AS digest, + EXISTS (SELECT 1 FROM p WHERE cardinality(mcp_toolsets) > 0) AS uses_toolsets, + jsonb_build_object( + 'user_id', COALESCE((SELECT user_id FROM k), $2::text), + 'project_id', (SELECT project_id FROM k), + 'team_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT team_id AS x FROM k WHERE team_id IS NOT NULL UNION SELECT unnest(teams) FROM u + UNION SELECT $3::text) q WHERE x IS NOT NULL), + 'organization_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT organization_id AS x FROM k UNION SELECT organization_id FROM u UNION SELECT organization_id FROM t + UNION SELECT $4::text) q WHERE x IS NOT NULL), + 'permission_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT object_permission_id AS x FROM k UNION SELECT object_permission_id FROM u + UNION SELECT object_permission_id FROM t UNION SELECT object_permission_id FROM o) q WHERE x IS NOT NULL), + 'access_group_ids', (SELECT COALESCE(jsonb_agg(x), '[]'::jsonb) FROM ( + SELECT unnest(access_group_ids) AS x FROM k UNION SELECT unnest(access_group_ids) FROM t) q WHERE x IS NOT NULL) + ) AS reachable +""" + + +class MCPPolicyRepository: + def __init__(self, prisma_client: "PrismaClient") -> None: + self._prisma_client: Final = prisma_client + + async def get_fingerprint(self, subject: PolicySubject) -> PolicyFingerprint: + rows: Final[Sequence[Mapping[str, object]]] = await self._prisma_client.writer_db.query_raw( # pyright: ignore[reportAny] # Prisma's raw client is untyped + _FINGERPRINT_SQL, subject.key_hash, subject.user_id, subject.pinned_team_id, subject.pinned_org_id + ) + return PolicyFingerprint.model_validate(rows[0]) diff --git a/tests/unit/proxy/mcp_server/test_policy_freshness.py b/tests/unit/proxy/mcp_server/test_policy_freshness.py index e642faaddc9..f8706aeb08d 100644 --- a/tests/unit/proxy/mcp_server/test_policy_freshness.py +++ b/tests/unit/proxy/mcp_server/test_policy_freshness.py @@ -30,20 +30,19 @@ from litellm.proxy.common_utils.strict_cache_refill import strict_cache_refill from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.mcp_server import catalog, operations from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler, UnloadableEntitlementError +from litellm.models.mcp_policy import PolicySubject, ReachableObjects from litellm.proxy.mcp_server.policy_freshness import ( - PolicySubject, - ReachableObjects, RequestFreshness, ValidatedPolicyMarks, cached_policy_is_current, cached_policy_keys, current_listing, - fetch_policy_fingerprint, policy_subject_for, validated_listing, validated_policy_marks, ) from litellm.proxy.mcp_server import policy_freshness as policy_freshness_module +from litellm.repositories.mcp_policy_repository import MCPPolicyRepository from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -203,7 +202,7 @@ async def test_entries_written_by_the_real_readers_are_gone_after_a_changed_dige async def test_the_fingerprint_is_read_from_the_writer_with_the_subject_bound() -> None: client: Final = database(fingerprint_row("abc", team_ids=("t",), access_group_ids=("g",))) subject: Final = PolicySubject(key_hash="h", user_id="u", pinned_team_id="t", pinned_org_id="o") - result: Final = await fetch_policy_fingerprint(client, subject) + result: Final = await MCPPolicyRepository(client).get_fingerprint(subject) assert (result.digest, result.uses_toolsets) == ("abc", False) assert result.reachable.team_ids == ("t",) assert result.reachable.access_group_ids == ("g",) @@ -217,7 +216,7 @@ async def test_the_fingerprint_accepts_a_json_text_reachable_column() -> None: '"permission_ids": ["p"], "access_group_ids": []}' ) client: Final = database({"digest": "abc", "uses_toolsets": False, "reachable": reachable_text}) - assert (await fetch_policy_fingerprint(client, KEY)).reachable.permission_ids == ("p",) + assert (await MCPPolicyRepository(client).get_fingerprint(KEY)).reachable.permission_ids == ("p",) @pytest.mark.asyncio