refactor(redis): migrate user MCP permission links to typed store

This commit is contained in:
Devin AI 2026-10-30 20:21:18 +00:00
parent 06affb432f
commit e03e42fedd
7 changed files with 484 additions and 105 deletions

View file

@ -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:

View file

@ -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(

View file

@ -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

View file

@ -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")

View file

@ -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:

View file

@ -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

View file

@ -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"