fix(key_management): invalidate cached object permissions on key update

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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,

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"]}