fix(proxy): broadcast user entitlement cache eviction to every worker

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

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -1451,11 +1451,7 @@ async def _invalidate_cached_user_entitlement(user_id: str | None, object_permis
         *(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 ()),
     )
-    for key in keys:
-        try:
-            await user_api_key_cache.async_delete_cache(key=key)
-        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 entitlement key %r: %s", key, e)
+    await evict_and_broadcast(cache_keys=keys, user_api_key_cache=user_api_key_cache)
 
 
 async def _update_single_user_helper(

tests/integration/authorization/test_rag_query_vector_store_allowlist.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import os
 import time
 import uuid
 from collections.abc import Iterator, Mapping
@@ -13,16 +14,18 @@ import jwt
 import pytest
 import yaml
 from cryptography.hazmat.primitives.asymmetric import rsa
-from integration._support.client import Gateway, Scenario, gateway_from_environment, object_value
+from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value
 from integration._support.process import owned_proxy
 from integration._support.wire import Reply, Request, wire_server
 from integration.authorization._guardrail_opt_out import upstream_observations
 from pydantic import JsonValue
+from redis import Redis
 
 CONFIG_STORE_ID: Final = "vs_integration_config_store"
 PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
 REMOVE_OPENAI_API_BASE: Final = ("OPENAI_API_BASE",)
 JWT_KEY_ID: Final = "integration-vector-store-jwt-key"
+AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
 JsonObject: TypeAlias = dict[str, JsonValue]
 
 
@@ -106,10 +109,19 @@ def no_registry_gateways(tmp_path_factory: pytest.TempPathFactory) -> Iterator[t
 
 
 class StrictGateway:
-    def __init__(self, gateway: Gateway, upstream: Gateway, signing_key: rsa.RSAPrivateKey) -> None:
+    def __init__(
+        self,
+        gateway: Gateway,
+        upstream: Gateway,
+        signing_key: rsa.RSAPrivateKey,
+        config: Path,
+        environment: Mapping[str, str],
+    ) -> None:
         self.gateway: Final = gateway
         self.upstream: Final = upstream
         self._signing_key: Final = signing_key
+        self.config: Final = config
+        self.environment: Final = environment
 
     def jwt(self, subject: str, groups: tuple[str, ...] = ()) -> str:
         claims: Final[JsonObject] = {
@@ -154,14 +166,16 @@ def strict_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[StrictG
 
     with gateway_from_environment() as upstream_gateway, wire_server(respond) as jwks:
         directory: Final = tmp_path_factory.mktemp("rag_query_deny_by_default")
+        config: Final = _strict_config(directory)
+        environment: Final = MappingProxyType({**_openai_environment(upstream_gateway), "JWT_PUBLIC_KEY_URL": jwks.url})
         with owned_proxy(
             upstream_gateway,
             directory,
-            {**_openai_environment(upstream_gateway), "JWT_PUBLIC_KEY_URL": jwks.url},
-            config=_strict_config(directory),
+            environment,
+            config=config,
             remove_environment=REMOVE_OPENAI_API_BASE,
         ) as gateway:
-            yield StrictGateway(gateway, upstream_gateway, signing_key)
+            yield StrictGateway(gateway, upstream_gateway, signing_key, config, environment)
 
 
 StrictCase: TypeAlias = Literal[
@@ -185,9 +199,7 @@ def _strict_denied_request(
         standalone_key: Final = scenario.key(models=models)
         return strict.gateway.request(
             "POST", "/v1/rag/query", _rag_query_body(model, marker, store_id), key=standalone_key
-        ), (
-            "key_vector_store_access_denied"
-        )
+        ), ("key_vector_store_access_denied")
     if case in ("team_key_empty_key_grants", "team_key_empty_team_grants"):
         key_grants_store: Final = case == "team_key_empty_team_grants"
         team: Final = scenario.team(models=models, object_permission=empty if key_grants_store else granted)
@@ -197,9 +209,7 @@ def _strict_denied_request(
         error_type: Final = "team_vector_store_access_denied" if key_grants_store else "key_vector_store_access_denied"
         return strict.gateway.request(
             "POST", "/v1/rag/query", _rag_query_body(model, marker, store_id), key=team_key
-        ), (
-            error_type
-        )
+        ), (error_type)
     if case == "multi_store_one_ungranted":
         partial_key: Final = scenario.key(models=models, object_permission=_permission_for_stores(store_id))
         body: Final[JsonObject] = {
@@ -539,3 +549,38 @@ def test_deny_by_default_without_registry_checks_search_route_and_file_search_to
             assert denied_observations == ()
             assert search_granted.status_code == 200, search_granted.text
             assert len(granted_searches) == 1, granted_searches
+
+
+def _auth_cache_subscribers() -> int:
+    with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
+        return int(cache.pubsub_numsub(AUTH_CACHE_INVALIDATION_CHANNEL)[0][1])
+
+
+def test_revoked_user_grant_stops_working_on_another_proxy(strict_gateway: StrictGateway, tmp_path: Path) -> None:
+    subscribers_before_peer: Final = _auth_cache_subscribers()
+    with (
+        owned_proxy(
+            strict_gateway.upstream,
+            tmp_path,
+            strict_gateway.environment,
+            config=strict_gateway.config,
+            remove_environment=REMOVE_OPENAI_API_BASE,
+        ) as peer,
+        strict_gateway.gateway.scenario() as scenario,
+    ):
+        eventually(_auth_cache_subscribers, lambda count: count > subscribers_before_peer)
+        model: Final = scenario.model()
+        store_id: Final = f"vs_unregistered_{uuid.uuid4().hex}"
+        user: Final = scenario.user(user_role="internal_user", object_permission=_permission_for_stores(store_id))
+        token: Final = strict_gateway.jwt(user)
+
+        def peer_status() -> int:
+            marker: Final = f"lit6035 revoked user grant {uuid.uuid4().hex}"
+            return peer.request(
+                "POST", "/v1/rag/query", _rag_query_body(model, marker, store_id), key=token
+            ).status_code
+
+        assert eventually(peer_status, lambda status: status == 200, seconds=30, return_last_on_timeout=True) == 200
+        strict_gateway.gateway.post("/user/update", {"user_id": user, "object_permission": _permission_for_stores()})
+
+        assert eventually(peer_status, lambda status: status == 401, seconds=10, return_last_on_timeout=True) == 401

tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -4003,6 +4003,34 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker):
     }
 
 
+@pytest.mark.asyncio
+async def test_user_update_broadcasts_the_entitlement_invalidation_to_other_workers(mocker: MockerFixture):
+    from litellm.proxy.management_endpoints.internal_user_endpoints import (
+        _update_single_user_helper,
+    )
+
+    _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
+    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,
+    )
+
+    await _update_single_user_helper(
+        user_request=UpdateUserRequest(user_id="target-user", object_permission={"vector_stores": []}),
+        user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
+    )
+
+    broadcast_keys: Final = {call.kwargs["cache_key"] for call in broadcast.await_args_list}
+    assert broadcast_keys == {
+        "object_permission_id:perm-new",
+        "user_object_permission_id:target-user",
+        "target-user",
+    }
+
+
 @pytest.mark.asyncio
 async def test_admin_can_clear_a_users_mcp_entitlement(mocker):
     """An explicit empty object_permission means "no object permission", so it must unlink.