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/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b40eb451979..b21bb8737af 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, @@ -2800,7 +2801,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: @@ -3657,7 +3658,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", @@ -4408,7 +4409,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", @@ -4482,7 +4483,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 @@ -4496,7 +4497,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 c71e2206f45..2fb62560d98 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: @@ -149,7 +150,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: @@ -169,7 +170,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..8af93552f30 --- /dev/null +++ b/litellm/proxy/mcp_server/policy_freshness.py @@ -0,0 +1,203 @@ +"""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 time +from collections import OrderedDict +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 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, + get_management_object_ttl, + object_permission_cache_key, + team_membership_auth_cache_key, + 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 + from litellm.proxy.utils import PrismaClient + +_MARK_CAPACITY: Final = 4096 +_MARK_TTL_FRACTION: Final = 0.5 +_SETTLE_SECONDS: Final = 5.0 + + +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 + + +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 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 + 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 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 + 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/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/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..f8706aeb08d --- /dev/null +++ b/tests/unit/proxy/mcp_server/test_policy_freshness.py @@ -0,0 +1,785 @@ +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.models.mcp_policy import PolicySubject, ReachableObjects +from litellm.proxy.mcp_server.policy_freshness import ( + RequestFreshness, + ValidatedPolicyMarks, + cached_policy_is_current, + cached_policy_keys, + current_listing, + 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 + +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 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",) + 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 MCPPolicyRepository(client).get_fingerprint(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")) + + 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) + + +@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)