diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2a385c4c42a..3f2196a2e15 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -93,6 +93,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, attach_object_permission_to_dict, handle_update_object_permission_common, + invalidate_cached_object_permissions, validate_key_mcp_servers_against_team, validate_key_search_tools_against_team, validate_key_vector_stores_against_team, @@ -2339,6 +2340,13 @@ async def _process_single_key_update( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) + await invalidate_cached_object_permissions( + object_permission_ids=( + existing_key_row.object_permission_id, + non_default_values.get("object_permission_id"), + ), + user_api_key_cache=user_api_key_cache, + ) # Trigger async hook asyncio.create_task( @@ -2820,6 +2828,13 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) + await invalidate_cached_object_permissions( + object_permission_ids=( + existing_key_row.object_permission_id, + non_default_values.get("object_permission_id"), + ), + user_api_key_cache=user_api_key_cache, + ) if data.spend is not None: from litellm.proxy.proxy_server import spend_counter_cache diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 33b84545915..a0e9ac00aa9 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -4,7 +4,7 @@ organizations, teams, and keys. """ import json -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Optional @@ -14,6 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ObjectPermissionDict, SpecialMCPServerName, SpecialMCPServerNames +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key from litellm.proxy.utils import PrismaClient from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import MCPServerRepository @@ -171,6 +172,26 @@ async def handle_update_object_permission_common( return created_object_permission_row.object_permission_id +async def invalidate_cached_object_permissions( + object_permission_ids: Iterable[object], + user_api_key_cache: UserApiKeyCache, +) -> None: + """Drop permission rows an entitlement change makes stale. + + ``get_object_permission`` caches a row under its own id and an upsert keeps that id, so clearing only + the entity's cache entry re-attaches the old grants until the management-object TTL expires. Pass both + the outgoing and incoming ids, since a change can also mint a new row. Non-string ids, which untyped + update payloads can carry, are ignored. + """ + for object_permission_id in dict.fromkeys(pid for pid in object_permission_ids if isinstance(pid, str)): + try: + await user_api_key_cache.async_delete_cache(key=object_permission_cache_key(object_permission_id)) + except Exception as e: # noqa: BLE001 # a cache we cannot clear still expires; never fail the write + verbose_proxy_logger.warning( + "Failed to invalidate cached object permission %r: %s", object_permission_id, e + ) + + async def _set_object_permission( data_json: dict, prisma_client: PrismaClient | None, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 939607dd139..dfab243f8cd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15600,3 +15600,89 @@ async def test_unblock_key_stamps_settings_updated_at(monkeypatch): sent = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs["data"] assert sent["blocked"] is False assert before <= sent["settings_updated_at"] <= after + +@pytest.mark.asyncio +async def test_key_update_invalidates_cached_object_permission(monkeypatch): + """Regression: /key/update must drop the cached permission row, not just the cached key. + + The permission row is cached under its own id and the upsert keeps that id, so a key read + after the update re-attached the OLD grants until the management-object TTL expired, which + served revoked MCP tools and withheld newly granted ones. + """ + from litellm.proxy._types import LiteLLM_ObjectPermissionBase + from litellm.proxy.auth.auth_checks import get_object_permission + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + permission_id = "objperm-lit5479" + grants = {"old": ["tool_a"], "new": ["tool_a", "tool_b"]} + + def _row(tools): + row = MagicMock() + row.dict.return_value = { + "object_permission_id": permission_id, + "mcp_tool_permissions": {"server-1": tools}, + } + return row + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( + side_effect=lambda **kwargs: _row(grants["old"]) + ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( + return_value=MagicMock(object_permission_id=permission_id) + ) + existing_key_row = LiteLLM_VerificationToken( + token="hashed-sk-lit5479", + user_id="user-123", + object_permission_id=permission_id, + ) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=existing_key_row + ) + updated_key = MagicMock() + updated_key.model_dump.return_value = {"user_id": "user-123"} + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_key}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + user_api_key_cache = UserApiKeyCache() + assert ( + await get_object_permission( + object_permission_id=permission_id, + prisma_client=mock_prisma_client, + user_api_key_cache=user_api_key_cache, + ) + ).mcp_tool_permissions == {"server-1": grants["old"]} + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" + ): + await _process_single_key_update( + update_key_request=UpdateKeyRequest( + key="sk-lit5479", + object_permission=LiteLLM_ObjectPermissionBase( + mcp_tool_permissions={"server-1": grants["new"]} + ), + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ), + litellm_changed_by=None, + prisma_client=mock_prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=AsyncMock(), + llm_router=None, + existing_key_row=existing_key_row, + ) + + mock_prisma_client.db.litellm_objectpermissiontable.find_unique.side_effect = ( + lambda **kwargs: _row(grants["new"]) + ) + reread = await get_object_permission( + object_permission_id=permission_id, + prisma_client=mock_prisma_client, + user_api_key_cache=user_api_key_cache, + ) + assert reread is not None + assert reread.mcp_tool_permissions == {"server-1": grants["new"]}