diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index e64fd755ef9..f86920ebd92 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -3,16 +3,28 @@ from __future__ import annotations import asyncio import re from collections.abc import Sequence +from contextlib import suppress +from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_batch import active_request_redis_batch from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.redis.stores.mcp_user_permission_links import ( + USER_NO_MCP_PERMISSION_SENTINEL as _USER_NO_PERMISSION, +) +from litellm.redis.stores.mcp_user_permission_links import ( + McpUserPermissionLinkStore, + TransportMcpUserPermissionLinkStore, + user_permission_link_key, +) +from litellm.redis.transport.contract import RedisFailure if TYPE_CHECKING: from opentelemetry.trace import Span @@ -291,21 +303,64 @@ class UserApiKeyCache(DualCache): await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs) -#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row, -#: so a human without an entitlement costs no DB read per request. Lives beside the key builder -#: because it is part of the same cache protocol: a reader that knows the key must know this value. -USER_NO_MCP_PERMISSION_SENTINEL: Final = "__user_no_mcp_permission__" +_PERMISSION_LINK_VALUE: Final = TypeAdapter(object) +USER_NO_MCP_PERMISSION_SENTINEL: Final = _USER_NO_PERMISSION def user_object_permission_id_cache_key(user_id: str) -> str: - """Cache key for the ``user_id -> object_permission_id`` link. + return user_permission_link_key(user_id) - Lives here rather than next to either user because two modules own the two halves: the MCP auth - resolver writes it on read, and ``/user/update`` deletes it after changing the link. A key format - duplicated across those two drifts silently, and the failure is an entitlement change that never - takes effect. - """ - return f"user_object_permission_id:{user_id}" + +async def get_user_permission_link(cache: UserApiKeyCache, user_id: str) -> str | None: + key: Final = user_permission_link_key(user_id) + local: Final = _PERMISSION_LINK_VALUE.validate_python(await cache.async_get_cache(key=key, local_only=True)) + if local is not None: + return local if isinstance(local, str) else None + if cache.redis_cache is None: + return None + batch: Final = active_request_redis_batch(cache.redis_cache) + if batch is not None and batch.read_as_missing(key): + return None + with suppress(Exception): + store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore(cache.redis_cache.transport()) + result: Final = await store.get(user_id) + if isinstance(result, RedisFailure): + return None + if result is not None: + await cache.async_set_cache(key=key, value=result, local_only=True) + return result + return None + + +async def fill_user_permission_link( + cache: UserApiKeyCache, user_id: str, permission_id: str | None, *, ttl: float +) -> None: + await cache.async_set_cache( + key=user_permission_link_key(user_id), + value=permission_id or USER_NO_MCP_PERMISSION_SENTINEL, + ttl=ttl, + local_only=True, + ) + if cache.redis_cache is not None: + with suppress(Exception): + store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore( + cache.redis_cache.transport() + ) + _ = await store.fill(user_id, permission_id, ttl=timedelta(seconds=int(ttl))) + + +async def invalidate_user_permission_link(cache: UserApiKeyCache, user_id: str) -> None: + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + + key: Final = user_permission_link_key(user_id) + cache.in_memory_cache.delete_cache(key) + if cache.redis_cache is not None: + with suppress(Exception): + store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore( + cache.redis_cache.transport() + ) + _ = await store.invalidate(user_id) + await publish_auth_cache_invalidation(cache_key=key) def object_permission_cache_key(object_permission_id: str) -> str: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 54a2d4062d2..34b6e1480de 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -50,8 +50,8 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_ from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, + invalidate_user_permission_link, object_permission_cache_key, - user_object_permission_id_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks @@ -1481,9 +1481,11 @@ async def _invalidate_cached_user_entitlement(user_id: str | None, object_permis keys: Final = ( *(object_permission_cache_key(permission_id) for permission_id in dict.fromkeys(object_permission_ids)), - *((user_object_permission_id_cache_key(user_id), user_id) if user_id is not None else ()), ) await evict_and_broadcast(cache_keys=keys, user_api_key_cache=user_api_key_cache) + if user_id is not None: + await invalidate_user_permission_link(user_api_key_cache, user_id) + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) async def _update_single_user_helper( 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 db640b6d689..67a3f2fb590 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 @@ -61,8 +61,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, USER_NO_MCP_PERMISSION_SENTINEL, + fill_user_permission_link, get_management_object_ttl, - user_object_permission_id_cache_key, + get_user_permission_link, ) from litellm.proxy.mcp_server.catalog import global_manager from litellm.proxy.mcp_server.oauth_utils import ( @@ -4149,9 +4150,8 @@ class MCPRequestHandler: from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import user_api_key_cache - cache_key: Final = user_object_permission_id_cache_key(user_id) try: - cached: Final[object] = None if check_db_only else await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final = None if check_db_only else await get_user_permission_link(user_api_key_cache, user_id) if cached == USER_NO_MCP_PERMISSION_SENTINEL: return None if isinstance(cached, str) and cached: @@ -4165,9 +4165,10 @@ class MCPRequestHandler: ) linked: Final[object] = getattr(user_row, "object_permission_id", None) if user_row is not None else None object_permission_id: Final = linked if isinstance(linked, str) and linked else None - await user_api_key_cache.async_set_cache( - key=cache_key, - value=object_permission_id or USER_NO_MCP_PERMISSION_SENTINEL, + await fill_user_permission_link( + user_api_key_cache, + user_id, + object_permission_id, ttl=get_management_object_ttl(user_api_key_cache), ) return object_permission_id diff --git a/litellm/redis/stores/mcp_user_permission_links.py b/litellm/redis/stores/mcp_user_permission_links.py new file mode 100644 index 00000000000..565325b6130 --- /dev/null +++ b/litellm/redis/stores/mcp_user_permission_links.py @@ -0,0 +1,64 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +USER_NO_MCP_PERMISSION_SENTINEL: Final = "__user_no_mcp_permission__" +_VALUE: Final = TypeAdapter(object) + + +def user_permission_link_key(user_id: str) -> str: + return f"user_object_permission_id:{user_id}" + + +class McpUserPermissionLinkStore(Protocol): + async def get(self, user_id: str) -> str | None | RedisFailure: ... + + async def fill(self, user_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: ... + + async def invalidate(self, user_id: str) -> int | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpUserPermissionLinkStore: + transport: RedisTransport + + async def get(self, user_id: str) -> str | None | RedisFailure: + reply: Final = await self.transport.get(user_permission_link_key(user_id)) + if reply is None or isinstance(reply, RedisFailure): + return reply + try: + decoded: Final = reply.decode("utf-8") + try: + value: Final = _VALUE.validate_json(decoded) + except ValidationError: + literal: Final = _VALUE.validate_python(ast.literal_eval(decoded)) + return _permission_id(literal) + return _permission_id(value) + except (UnicodeDecodeError, ValueError, SyntaxError): + return RedisFailure(kind="decode", message="Invalid user permission link") + + async def fill(self, user_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: + seconds: Final = int(ttl.total_seconds()) + key: Final = user_permission_link_key(user_id) + value: Final = json.dumps(permission_id or USER_NO_MCP_PERMISSION_SENTINEL) + expiry: Final = timedelta(seconds=seconds) + batch: Final = self.transport.request_batch() + if batch is not None: + _ = batch.set_seconds(key, value, ttl=expiry) + return None + return await self.transport.set_seconds(key, value, ttl=expiry) + + async def invalidate(self, user_id: str) -> int | RedisFailure: + return await self.transport.delete(user_permission_link_key(user_id)) + + +def _permission_id(value: object) -> str | RedisFailure: + if isinstance(value, str): + return value + return RedisFailure(kind="decode", message="User permission link is not a string") diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 39586a1a204..299a0bf29a1 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -9,6 +9,7 @@ from types import SimpleNamespace from typing import TYPE_CHECKING, Final, cast from unittest.mock import AsyncMock, MagicMock +import fakeredis import httpx import pytest from fastapi import HTTPException @@ -41,6 +42,7 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ( from litellm.proxy.proxy_server import app from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains +from tests._support.fake_redis import fake_redis_cache from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( CascadingJWTMappingTable, JWTMappingRow, @@ -870,7 +872,6 @@ async def test_new_user_license_over_limit(mocker): """ Test that /user/new endpoint raises an error when license is over the user limit """ - from fastapi import HTTPException from litellm.proxy._types import NewUserRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.internal_user_endpoints import new_user @@ -1389,9 +1390,7 @@ async def test_user_info_url_encoding_plus_character(mocker): # Mock get_data to return user when called with user_id, empty list for keys async def mock_get_data(*args, **kwargs): - if kwargs.get("table_name") == "key": - return [] - elif kwargs.get("table_name") == "team": + if kwargs.get("table_name") == "key" or kwargs.get("table_name") == "team": return [] elif kwargs.get("user_id") is not None: return mock_user @@ -1524,7 +1523,7 @@ async def test_new_user_default_teams_flow(mocker): - Teams ARE sent to _add_user_to_team """ import litellm - from litellm.proxy._types import NewUserRequest, NewUserRequestTeam, UserAPIKeyAuth + from litellm.proxy._types import NewUserRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.internal_user_endpoints import new_user # Mock the prisma client @@ -1672,9 +1671,7 @@ async def test_new_user_clears_recent_missing_user_lookup(mocker: MockerFixture) ) mocker.patch( "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", - new=mocker.AsyncMock( - return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None} - ), + new=mocker.AsyncMock(return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None}), ) try: @@ -1884,8 +1881,6 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): assert user_email_clause.get("mode") == "insensitive" assert user_email_clause.get("equals") == "newuser@example.com" - return None # No existing user found - mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate # Should not raise any exception when no duplicate is found @@ -1938,7 +1933,6 @@ async def test_check_duplicate_user_id(mocker): user_id_clause = where_clause.get("user_id", {}) assert user_id_clause.get("equals") == "new-user-id" assert "mode" not in user_id_clause - return None mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate @@ -2301,7 +2295,11 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo "user_alias": "Spruce", } - replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0} + replacement: Final = { + "model-spruce": {"budget_limit": 0, "time_period": "1d"}, + "model-birch": 5.0, + "model-cedar": 0, + } request: Final = UpdateUserRequest( user_id="user-spruce", model_max_budget=replacement, @@ -2343,10 +2341,14 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2382,10 +2384,14 @@ async def test_user_status_update_refreshes_cached_user( prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2419,10 +2425,14 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke prisma_client: Final = mocker.MagicMock() prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2454,7 +2464,9 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke @pytest.mark.asyncio @pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"]) -async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None: +async def test_user_max_budget_update_evicts_cached_user_on_every_worker( + mocker: MockerFixture, all_users: bool +) -> None: from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update @@ -2467,10 +2479,14 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -2875,24 +2891,12 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) # Mock all delete_many calls - mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock( - return_value=[] - ) - mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock( - return_value=1 - ) - mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1) + mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) _serve_locked_user_delete(mock_prisma_client, mock_user_row) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -2985,9 +2989,15 @@ async def test_delete_user_evicts_jwt_key_mapping_cache_of_its_keys(mocker): mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) _serve_locked_user_delete(mock_prisma_client, user_row) - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: substitute the database dependency - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache - mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj", None + ) # test-quality-ok: delete_user reads it off proxy_server at call time await delete_user( data=DeleteUserRequest(user_ids=["jwt-user"]), @@ -3160,9 +3170,7 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker: MockerFixture return None mock_prisma_client: Final = _PrismaClient( - db=_Database( - litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique)) - ) + db=_Database(litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique))) ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -3556,7 +3564,7 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): """ from fastapi import Request - from litellm.proxy._types import LiteLLM_TeamTable, UserInfoV2Response + from litellm.proxy._types import UserInfoV2Response from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2 mock_prisma_client = mocker.MagicMock() @@ -4419,8 +4427,7 @@ async def test_user_update_persists_mcp_entitlement_and_links_it(mocker): ) mock_prisma_client = _object_permission_mocks(mocker) - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4456,8 +4463,10 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker): ) _object_permission_mocks(mocker) - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + keys: Final = ("object_permission_id:perm-new", "user_object_permission_id:target-user", "target-user") + for key in keys: + await cache.async_set_cache(key, "stale", ttl=60) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4468,12 +4477,8 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker): user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) - deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} - assert deleted == { - "object_permission_id:perm-new", - "user_object_permission_id:target-user", - "target-user", - } + for key in keys: + assert await cache.async_get_cache(key) is None @pytest.mark.asyncio @@ -4483,9 +4488,10 @@ async def test_user_update_broadcasts_the_entitlement_invalidation_to_other_work ) _object_permission_mocks(mocker) - cache: Final = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: substitute the cache dependency + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: substitute the cache dependency broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -4520,8 +4526,10 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker): ) mock_prisma_client = _object_permission_mocks(mocker, "perm-existing") - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + keys: Final = ("object_permission_id:perm-existing", "user_object_permission_id:target-user", "target-user") + for key in keys: + await cache.async_set_cache(key, "stale", ttl=60) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4534,12 +4542,8 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker): assert "object_permission" not in written mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_called() - deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} - assert deleted == { - "object_permission_id:perm-existing", - "user_object_permission_id:target-user", - "target-user", - } + for key in keys: + assert await cache.async_get_cache(key) is None @pytest.mark.asyncio @@ -4554,8 +4558,15 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock ) _object_permission_mocks(mocker, "perm-existing") - cache = mocker.MagicMock() - cache.async_delete_cache = mocker.AsyncMock() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements")) + keys: Final = ( + "object_permission_id:perm-existing", + "object_permission_id:perm-new", + "user_object_permission_id:target-user", + "target-user", + ) + for key in keys: + await cache.async_set_cache(key, "stale", ttl=60) mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) await _update_single_user_helper( @@ -4566,13 +4577,8 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) - deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} - assert deleted == { - "object_permission_id:perm-existing", - "object_permission_id:perm-new", - "user_object_permission_id:target-user", - "target-user", - } + for key in keys: + assert await cache.async_get_cache(key) is None @pytest.mark.asyncio @@ -4755,7 +4761,7 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget, async def get_data(self, *args, **kwargs): return None - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server monkeypatch.setattr(proxy_server, "prisma_client", _FakePrisma(), raising=False) # model_max_budget is an enterprise feature; without this the call is rejected @@ -5032,12 +5038,18 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) _serve_locked_user_delete(prisma_client, deleted) - mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() for row in (deleted, survivor): await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable) - mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache - mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + mocker.patch( + "litellm.proxy.proxy_server.user_api_key_cache", cache + ) # test-quality-ok: exercise a real isolated cache + mocker.patch( + "litellm.proxy.proxy_server.proxy_logging_obj", None + ) # test-quality-ok: delete_user reads it off proxy_server at call time broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", new_callable=mocker.AsyncMock, @@ -5143,7 +5155,9 @@ async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_ assert raised.value.code == "503" assert raised.value.type == ProxyErrorTypes.no_db_connection assert isinstance(raised.value.__cause__, httpx.ConnectError) - outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()] + outage_logs: Final = [ + r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage() + ] assert outage_logs == ["Database unavailable during user search: ConnectError"] @@ -5301,7 +5315,9 @@ async def test_user_update_without_password_revokes_nothing(_admin_prisma, mocke @pytest.mark.asyncio -async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(_admin_prisma: MagicMock) -> None: +async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error( + _admin_prisma: MagicMock, +) -> None: from litellm.proxy.management_endpoints.internal_user_endpoints import user_update with pytest.raises(ProxyException) as exc_info: diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index ee3ef3f5323..230a3c84790 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -8,6 +8,7 @@ from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import pytest from fastapi import HTTPException from fastapi.testclient import TestClient @@ -26,6 +27,7 @@ from litellm.proxy._types import ( SpecialMCPServerNames, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, UnloadableEntitlementError, @@ -33,6 +35,48 @@ from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import ( is_mcp_admitted_user_subject, ) from litellm.types.agents import AgentCaller +from tests._support.fake_redis import fake_redis_cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize("permission_id", ["policy", None]) +async def test_user_permission_link_fills_from_cached_user_row_and_survives_redis_outage( + monkeypatch: pytest.MonkeyPatch, permission_id: str | None +) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="auth") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=41) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + await cache.async_set_cache( + "link-user", LiteLLM_UserTable(user_id="link-user", object_permission_id=permission_id), local_only=True + ) + prisma: Final = MagicMock() + assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) == permission_id + assert await redis.async_get_cache("user_object_permission_id:link-user") == ( + permission_id or "__user_no_mcp_permission__" + ) + cache.in_memory_cache.delete_cache("user_object_permission_id:link-user") + server.connected = False + assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) == permission_id + assert await cache.async_get_cache("user_object_permission_id:link-user", local_only=True) == ( + permission_id or "__user_no_mcp_permission__" + ) + prisma.db.litellm_usertable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_fresh_user_policy_does_not_use_redis_sentinel(monkeypatch: pytest.MonkeyPatch) -> None: + redis: Final = fake_redis_cache(fakeredis.FakeServer(), namespace="auth") + cache: Final = UserApiKeyCache(redis_cache=redis) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + await redis.async_set_cache("user_object_permission_id:link-user", "__user_no_mcp_permission__", ttl=60) + prisma: Final = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) is None # pyright: ignore[reportPrivateUsage] # resolver internals under test + prisma.db.litellm_usertable.find_unique.assert_not_called() + with pytest.raises(HTTPException) as error: + await MCPRequestHandler._user_object_permission_id("link-user", prisma, check_db_only=True) # pyright: ignore[reportPrivateUsage] # resolver internals under test + assert error.value.status_code == 503 @pytest.mark.parametrize( @@ -11309,11 +11353,11 @@ async def test_fresh_mcp_user_permission_link_ignores_cached_and_replica_grants( database.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=current) database.db.litellm_usertable.find_unique = AsyncMock(return_value=cached) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) - assert await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) == "current" + assert await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) == "current" # pyright: ignore[reportPrivateUsage] # resolver internals under test database.db.litellm_usertable.find_unique.assert_not_awaited() database.writer_db.litellm_usertable.find_unique.side_effect = failure with pytest.raises(HTTPException) as denied: - await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) + await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) # pyright: ignore[reportPrivateUsage] # resolver internals under test assert denied.value.status_code == 503 @@ -12451,7 +12495,6 @@ async def test_personal_key_display_name_billing_uses_original_tool_grant( async def test_personal_key_display_name_billing_uses_allowed_servers( monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth ) -> None: - from litellm.proxy import proxy_server from litellm.proxy.auth import auth_checks from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager from litellm.types.mcp import MCPTransport diff --git a/tests/unit/redis/stores/test_mcp_user_permission_links.py b/tests/unit/redis/stores/test_mcp_user_permission_links.py new file mode 100644 index 00000000000..c669da07863 --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_user_permission_links.py @@ -0,0 +1,198 @@ +from datetime import timedelta +from typing import Final + +import fakeredis +import pytest +import pytest_asyncio +from pytest import MonkeyPatch +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + fill_user_permission_link, + get_user_permission_link, + invalidate_user_permission_link, +) +from litellm.redis.stores.mcp_user_permission_links import TransportMcpUserPermissionLinkStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + + +@pytest_asyncio.fixture(loop_scope="function") +async def link_cache(monkeypatch: MonkeyPatch) -> tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache]: + monkeypatch.setattr("time.time", lambda: 1000.0) + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="links") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=37.9) + return server, redis, cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", ['policy é "quoted"', "__user_no_mcp_permission__", ""]) +async def test_old_writer_new_reader( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", value, ttl=37.9) + assert await get_user_permission_link(cache, "u") == value + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) == value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "value,wire", + [(None, b'"__user_no_mcp_permission__"'), ('policy é "quoted"', b'"policy \\u00e9 \\"quoted\\""')], +) +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_new_writer_old_reader_and_whole_second_expiry( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str | None, wire: bytes, ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await fill_user_permission_link(cache, "u", value, ttl=ttl) + assert await raw.get("links:user_object_permission_id:u") == wire + assert await raw.pttl("links:user_object_permission_id:u") == 37000 + assert await redis.async_get_cache("user_object_permission_id:u") == (value or "__user_no_mcp_permission__") + assert await raw.keys("*") == [b"links:user_object_permission_id:u"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wire", [b"broken", b"42", b"null", b"{}", b"[]", b"\xff"]) +async def test_bad_payload_is_decode_failure_and_caller_cache_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], wire: bytes +) -> None: + server, redis, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:user_object_permission_id:u", wire) + result: Final = await TransportMcpUserPermissionLinkStore(redis.transport()).get("u") + assert isinstance(result, RedisFailure) and result.kind == "decode" + assert await get_user_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_literal_legacy_string_and_local_non_string( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, _, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:user_object_permission_id:u", b"'legacy'") + assert await get_user_permission_link(cache, "u") == "legacy" + await cache.async_set_cache("user_object_permission_id:u", {"invalid": True}, local_only=True) + assert await get_user_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_outage_preserves_local_fill_and_invalidation( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, redis, cache = link_cache + server.connected = False + assert await get_user_permission_link(cache, "u") is None + await fill_user_permission_link(cache, "u", "policy", ttl=30) + assert await get_user_permission_link(cache, "u") == "policy" + await invalidate_user_permission_link(cache, "u") + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) is None + store: Final = TransportMcpUserPermissionLinkStore(redis.transport()) + for _ in range(10): + assert isinstance(await store.get("u"), RedisFailure) + result: Final = await store.get("u") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_request_fill_stays_queued_and_clears_prior_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + with request_redis_batch_scope(): + batch: Final = active_request_redis_batch(redis) + assert batch is not None + assert await batch.mget(("user_object_permission_id:u",)) == {"user_object_permission_id:u": None} + await redis.async_set_cache("user_object_permission_id:u", "external", ttl=60) + assert await get_user_permission_link(cache, "u") is None + await fill_user_permission_link(cache, "u", "policy", ttl=ttl) + assert await get_user_permission_link(cache, "u") == "policy" + assert await raw.get("links:user_object_permission_id:u") == b'"external"' + await batch.flush() + assert await raw.get("links:user_object_permission_id:u") == b'"policy"' + assert await raw.pttl("links:user_object_permission_id:u") == 37000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.9, 0.999999999, 0, -1]) +async def test_rejected_expiry_keeps_existing_redis_value_and_new_local_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", "old", ttl=60) + await fill_user_permission_link(cache, "u", "new", ttl=ttl) + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) == ("new" if ttl >= 0 else None) + assert await redis.async_get_cache("user_object_permission_id:u") == "old" + result: Final = await TransportMcpUserPermissionLinkStore(redis.transport()).fill( + "u", "new", ttl=timedelta(seconds=int(ttl)) + ) + assert isinstance(result, RedisFailure) and result.kind == "server" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.999999999, 0, -1]) +async def test_rejected_batched_expiry_keeps_existing_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", "old", ttl=60) + with request_redis_batch_scope() as scope: + await fill_user_permission_link(cache, "u", "new", ttl=ttl) + await scope.flush_all() + assert await redis.async_get_cache("user_object_permission_id:u") == "old" + + +@pytest.mark.asyncio +async def test_unclassified_client_failure_keeps_cache_optional( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + _, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def unavailable_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RuntimeError("client callback failed") + + monkeypatch.setattr(connection, "read_response", unavailable_reply) + await client.connection_pool.release(connection) + assert await cache.async_get_cache("user_object_permission_id:u") is None + assert await get_user_permission_link(cache, "u") is None + await fill_user_permission_link(cache, "u", "policy", ttl=30) + assert await get_user_permission_link(cache, "u") == "policy" + await invalidate_user_permission_link(cache, "u") + assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) is None + + +@pytest.mark.asyncio +async def test_lost_set_reply_does_not_roll_back_applied_write( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + server, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def lost_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RedisConnectionError("reply lost after write") + + monkeypatch.setattr(connection, "read_response", lost_reply) + await client.connection_pool.release(connection) + await fill_user_permission_link(cache, "u", "applied", ttl=37.9) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get("links:user_object_permission_id:u") == b'"applied"' + assert await raw.pttl("links:user_object_permission_id:u") == 37000 + assert await get_user_permission_link(cache, "u") == "applied"