diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 1f693526d1f..66fac8d76ee 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -19,6 +19,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import delete_cached_project_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field from litellm.proxy.management_helpers.utils import ( @@ -514,6 +515,7 @@ async def update_project( litellm_proxy_admin_name, premium_user, prisma_client, + user_api_key_cache, ) try: @@ -672,6 +674,11 @@ async def update_project( include={"litellm_budget_table": True, "object_permission": True}, ) + await delete_cached_project_object( + project_id=data.project_id, + user_api_key_cache=user_api_key_cache, + ) + return updated_project except Exception as e: verbose_proxy_logger.exception( @@ -710,7 +717,7 @@ async def delete_project( }' ``` """ - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache try: if not premium_user: @@ -773,6 +780,11 @@ async def delete_project( prisma_models.LiteLLM_ProjectTable | None ) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id}) + await delete_cached_project_object( + project_id=project_id, + user_api_key_cache=user_api_key_cache, + ) + deleted_projects.append(deleted_project) return deleted_projects diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3fba464dd23..01034d0cf58 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4262,6 +4262,10 @@ async def _project_soft_budget_check( ) +def _project_cache_key(project_id: str) -> str: + return f"project_id:{project_id}" + + async def get_project_object( project_id: str, prisma_client: PrismaClient | None, @@ -4279,7 +4283,7 @@ async def get_project_object( return None # Check cache first - cache_key: Final = f"project_id:{project_id}" + cache_key: Final = _project_cache_key(project_id) deserialized_project: Final = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_ProjectTableCachedObj, @@ -4310,6 +4314,32 @@ async def get_project_object( return project_obj +async def delete_cached_project_object( + project_id: str, + user_api_key_cache: UserApiKeyCache, +) -> None: + """ + Every endpoint that mutates litellm_projecttable must call this: get_project_object + serves auth cache-first with no freshness check, so without invalidation a stale + project (e.g. a pre-update empty model allowlist) keeps being enforced until the + TTL expires (LIT-3803). Best-effort on both steps: the DB write has already + committed, so a cache backend error must not fail the endpoint; the stale entry + then expires via TTL. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + + cache_key: Final = _project_cache_key(project_id) + try: + await user_api_key_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # best-effort eviction: any cache backend error must not fail the mutation + verbose_proxy_logger.warning( + "Failed to evict cached project entry %s; a stale project may be served until its TTL expires: %s", + cache_key, + e, + ) + await publish_auth_cache_invalidation(cache_key=cache_key) + + async def _organization_max_budget_check( valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py new file mode 100644 index 00000000000..7fc8da42a3d --- /dev/null +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -0,0 +1,153 @@ +import asyncio +import json +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.config_sync_pubsub import ( + _ConfigSyncPubSub, + _pubsub_capable_client, + coordination_redis_cache, +) + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + +AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" +_POLL_TIMEOUT_SECONDS: Final = 1.0 +_BACKOFF_INITIAL_SECONDS: Final = 5.0 +_BACKOFF_MAX_SECONDS: Final = 60.0 + + +def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: + if redis_cache.namespace is None: + return AUTH_CACHE_INVALIDATION_CHANNEL + return f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}" + + +@dataclass(frozen=True, slots=True) +class _CacheInvalidationMessage: + cache_key: str + + +def _cache_invalidation_message_json(cache_key: str) -> str: + return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key))) + + +def _cache_key_from_message_data(data: object) -> str | None: + if isinstance(data, bytes): + data = data.decode("utf-8", errors="replace") + if not isinstance(data, str): + return None + try: + parsed: Final = json.loads(data) + except json.JSONDecodeError: + return None + if not isinstance(parsed, dict): + return None + cache_key: Final = parsed.get("cache_key") + return cache_key if isinstance(cache_key, str) else None + + +async def publish_auth_cache_invalidation(cache_key: str) -> None: + """ + Best-effort broadcast so every worker drops its local in-memory copy of a + mutated management object; without this, only the handling worker and Redis + are evicted and other workers keep serving the stale object until its TTL. + """ + redis_cache: Final = coordination_redis_cache() + if redis_cache is None: + return + try: + client: Final = _pubsub_capable_client(redis_cache) + if client is None: + verbose_proxy_logger.debug( + "auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support", + cache_key, + ) + return + await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key)) + except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors + verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) + + +class AuthCacheInvalidationSubscriber: + __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") + + def __init__( + self, + redis_cache: "RedisCache", + user_api_key_cache: "UserApiKeyCache", + ) -> None: + self._redis_cache = redis_cache + self._user_api_key_cache = user_api_key_cache + self._task: asyncio.Task[None] | None = None + + def start(self) -> None: + if self._task is not None: + return + self._task = asyncio.create_task(self._run()) + + async def stop(self) -> None: + task: Final = self._task + if task is None: + return + self._task = None + _ = task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def _run(self) -> None: + backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects + while True: + try: + client = _pubsub_capable_client(self._redis_cache) # rebind-ok: re-resolved on every reconnect + if client is None: + verbose_proxy_logger.warning( + "auth cache invalidation subscriber disabled: cluster redis client has no pub/sub support; " + "cross-worker eviction falls back to the local cache TTL" + ) + return + pubsub = client.pubsub() # rebind-ok: fresh pubsub per reconnect + try: + await pubsub.subscribe(auth_cache_invalidation_channel(self._redis_cache)) + backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: reset after successful subscribe + await self._consume(pubsub) + finally: + await self._close_pubsub(pubsub) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect + verbose_proxy_logger.warning( + "auth cache invalidation subscriber redis error: %s; reconnecting in %.0fs", + e, + backoff_seconds, + ) + await asyncio.sleep(backoff_seconds) + backoff_seconds = min(backoff_seconds * 2, _BACKOFF_MAX_SECONDS) # rebind-ok: backoff accumulator + + async def _consume(self, pubsub: _ConfigSyncPubSub) -> None: + while True: + message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) + if message is None: + continue + self._apply_message(message) + + def _apply_message(self, message: object) -> None: + data: Final = message.get("data") if isinstance(message, dict) else None + cache_key: Final = _cache_key_from_message_data(data) + if cache_key is None: + return + in_memory_cache: Final = self._user_api_key_cache.in_memory_cache + if in_memory_cache is not None: + in_memory_cache.delete_cache(cache_key) + + @staticmethod + async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: + try: + await pubsub.aclose() + except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + verbose_proxy_logger.debug("auth cache invalidation pubsub close failed: %s", e) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 539b68c1aee..46b03f4cc5c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -297,6 +297,9 @@ def generate_feedback_box(): _should_return_raw_model_name, create_response, ) +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AuthCacheInvalidationSubscriber, +) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers @@ -1189,6 +1192,8 @@ async def _run_pw_migration(): await proxy_config.stop_config_sync_subscriber() + await proxy_config.stop_auth_cache_invalidation_subscriber() + await proxy_shutdown_event() @@ -3900,6 +3905,7 @@ def __init__(self) -> None: self._last_hashicorp_vault_config: dict[str, Any] | None = None self.worker_registry: list[WorkerRegistryEntry] = [] self.config_sync_subscriber: ConfigSyncSubscriber | None = None + self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None from litellm.litellm_core_utils.get_model_cost_map import ( get_model_cost_map_loaded_at, ) @@ -6370,6 +6376,30 @@ async def stop_config_sync_subscriber(self) -> None: except Exception as e: verbose_proxy_logger.error("Error stopping config sync subscriber: %s", e) + def start_auth_cache_invalidation_subscriber( + self, + redis_cache: RedisCache | None, + user_api_key_cache: UserApiKeyCache, + ) -> None: + if redis_cache is None or self.auth_cache_invalidation_subscriber is not None: + return + subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=redis_cache, + user_api_key_cache=user_api_key_cache, + ) + self.auth_cache_invalidation_subscriber = subscriber + subscriber.start() + + async def stop_auth_cache_invalidation_subscriber(self) -> None: + subscriber: Final = self.auth_cache_invalidation_subscriber + if subscriber is None: + return + self.auth_cache_invalidation_subscriber = None + try: + await subscriber.stop() + except Exception as e: # noqa: BLE001 # best-effort: a failing stop must not break proxy shutdown + verbose_proxy_logger.error("Error stopping auth cache invalidation subscriber: %s", e) + async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): """ Use this to read non-llm objects from the db and initialize them @@ -8296,6 +8326,11 @@ async def initialize_scheduled_background_jobs( misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + proxy_config.start_auth_cache_invalidation_subscriber( + redis_cache=redis_usage_cache, + user_api_key_cache=user_api_key_cache, + ) + if store_model_in_db is True: ### GET STORED CREDENTIALS ### scheduler.add_job( diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c55b66b402b..c29b4c68bb0 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -864,3 +864,178 @@ def test_litellm_project_table_has_timestamp_fields(): fields = LiteLLM_ProjectTable.model_fields assert "created_at" in fields, "LiteLLM_ProjectTable must have created_at field" assert "updated_at" in fields, "LiteLLM_ProjectTable must have updated_at field" + + +@pytest.mark.asyncio +async def test_update_project_invalidates_cached_project_object(monkeypatch): + """ + LIT-3803 regression: auth reads projects cache-first with no freshness check, + so /project/update must evict the cached project. Before the fix, a project + cached with models=[] kept bypassing the new allowlist until the TTL expired. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.auth.auth_checks import get_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + cache = UserApiKeyCache() + + stale_row = MagicMock() + stale_row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": []} + + mock_prisma = MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=stale_row) + + seeded = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert seeded is not None and seeded.models == [] + + updated_models = ["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"] + existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None) + updated_row = MagicMock() + updated_row.model_dump = lambda: { + "project_id": project_id, + "team_id": None, + "models": updated_models, + } + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + + await update_project( + data=UpdateProjectRequest(project_id=project_id, models=updated_models), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=updated_row) + refreshed = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert refreshed is not None + assert refreshed.models == updated_models + + +@pytest.mark.asyncio +async def test_delete_project_invalidates_cached_project_object(monkeypatch): + """ + LIT-3803 regression: /project/delete must evict the cached project so auth + stops enforcing (or trusting) a project that no longer exists. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.auth.auth_checks import get_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + cache = UserApiKeyCache() + + row = MagicMock() + row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": ["gpt-5.5"]} + + mock_prisma = MagicMock() + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=row) + + seeded = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert seeded is not None + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_projecttable.delete = AsyncMock(return_value=row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + + await delete_project( + data=DeleteProjectRequest(project_ids=[project_id]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=None) + assert ( + await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + is None + ) + + +@pytest.mark.asyncio +async def test_update_project_succeeds_when_cache_eviction_fails(monkeypatch): + """ + The DB write has already committed when eviction runs, so a cache backend + error must not turn a successful update into a 500; the stale entry is + bounded by the TTL instead. + """ + from unittest.mock import AsyncMock, MagicMock + + project_id = f"project-{uuid.uuid4()}" + failing_cache = MagicMock() + failing_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down")) + + existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None) + updated_row = MagicMock() + + mock_prisma = MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", failing_cache) + + response = await update_project( + data=UpdateProjectRequest(project_id=project_id, models=["gpt-oss-120b"]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + assert response is updated_row + failing_cache.async_delete_cache.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_project_eviction_publishes_cross_worker_invalidation(monkeypatch): + """ + Single-worker eviction only fixes the handling worker; the broadcast is what + lets every other worker drop its in-memory copy instead of serving the stale + project until the TTL expires. + """ + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import delete_cached_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new=AsyncMock(), + ) as mock_publish: + await delete_cached_project_object( + project_id=project_id, user_api_key_cache=UserApiKeyCache() + ) + + mock_publish.assert_awaited_once_with(cache_key=f"project_id:{project_id}") diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index a5211ba83e7..757991b8ff3 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5464,6 +5464,63 @@ async def test_get_project_object_db_fetch_returns_cached_obj(): assert result.project_alias == "proj" +@pytest.mark.asyncio +async def test_project_allowlist_enforced_when_key_models_empty(): + """ + LIT-3803: a project-bound key with models=[] has no key-level restriction, + but the project allowlist must still 403 team models outside it. + """ + from litellm.proxy._types import ( + LiteLLM_ProjectTableCachedObj, + ProxyErrorTypes, + ProxyException, + ) + from litellm.proxy.auth.auth_checks import _run_project_checks, can_key_call_model + + valid_token = UserAPIKeyAuth( + api_key="hashed-key", + project_id="p-1", + team_id="t-1", + models=[], + ) + project = LiteLLM_ProjectTableCachedObj( + project_id="p-1", + team_id="t-1", + models=["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"], + ) + + assert ( + await can_key_call_model( + model="gemini-2.5-flash", + llm_model_list=None, + valid_token=valid_token, + llm_router=None, + ) + is True + ) + + await _run_project_checks( + project_object=project, + _model="gemini-2.5-flash-image", + llm_router=None, + skip_budget_checks=True, + valid_token=valid_token, + proxy_logging_obj=MagicMock(), + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_project_checks( + project_object=project, + _model="gemini-2.5-flash", + llm_router=None, + skip_budget_checks=True, + valid_token=valid_token, + proxy_logging_obj=MagicMock(), + ) + assert exc_info.value.type == ProxyErrorTypes.project_model_access_denied + assert exc_info.value.code == "403" + + def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py new file mode 100644 index 00000000000..468e8aabae8 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -0,0 +1,161 @@ +import asyncio +import json +from typing import Iterable, List, Optional, Tuple +from unittest.mock import patch + +import pytest +from redis.asyncio import Redis + +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AUTH_CACHE_INVALIDATION_CHANNEL, + AuthCacheInvalidationSubscriber, + publish_auth_cache_invalidation, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +class _RecordingRedisClient(Redis): + def __init__(self) -> None: + self.published: List[Tuple[str, str]] = [] + + async def publish(self, channel: str, message: str) -> int: + self.published.append((channel, message)) + return 1 + + +class _FailingPublishRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + raise ConnectionError("redis down") + + +class _QueuePubSub: + def __init__(self, initial_messages: Iterable[object] = ()) -> None: + self.queue: "asyncio.Queue[object]" = asyncio.Queue() + for message in initial_messages: + self.queue.put_nowait(message) + self.subscribed_channels: List[str] = [] + self.closed = False + + async def subscribe(self, *channels: str) -> None: + self.subscribed_channels.extend(channels) + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]: + try: + return await asyncio.wait_for(self.queue.get(), timeout) + except asyncio.TimeoutError: + return None + + async def aclose(self) -> None: + self.closed = True + + +class _ScriptedPubSubRedisClient(Redis): + def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None: + self._scripted_pubsubs = iter(pubsubs) + + def pubsub(self) -> _QueuePubSub: + return next(self._scripted_pubsubs) + + +class _FakeRedisCache: + def __init__(self, client: object, namespace: Optional[str] = None) -> None: + self._client = client + self.namespace = namespace + + def init_async_client(self) -> object: + return self._client + + +def _invalidation_message(cache_key: str) -> dict: + return {"type": "message", "data": json.dumps({"cache_key": cache_key}).encode()} + + +@pytest.mark.asyncio +async def test_publish_sends_cache_key_json_on_channel() -> None: + client = _RecordingRedisClient() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + assert client.published == [(AUTH_CACHE_INVALIDATION_CHANNEL, json.dumps({"cache_key": "project_id:p-1"}))] + + +@pytest.mark.asyncio +async def test_publish_uses_namespaced_channel() -> None: + client = _RecordingRedisClient() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client, namespace="ns1"), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + assert client.published[0][0] == f"ns1:{AUTH_CACHE_INVALIDATION_CHANNEL}" + + +@pytest.mark.asyncio +async def test_publish_noops_without_coordination_redis() -> None: + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=None, + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + +@pytest.mark.asyncio +async def test_publish_swallows_redis_errors() -> None: + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=_FailingPublishRedisClient()), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + +@pytest.mark.asyncio +async def test_subscriber_deletes_local_cache_entry_on_message() -> None: + """ + The cross-worker half of LIT-3803: a worker that did not handle the project + mutation must drop its in-memory copy when the invalidation broadcast lands, + instead of serving the stale object until the TTL expires. + """ + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + for _ in range(200): + if cache.in_memory_cache.get_cache("project_id:p-1") is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert cache.in_memory_cache.get_cache("project_id:p-1") is None + assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] + + +@pytest.mark.asyncio +async def test_subscriber_ignores_malformed_messages() -> None: + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), + user_api_key_cache=cache, + ) + subscriber._apply_message({"type": "message", "data": b"not json"}) + subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()}) + subscriber._apply_message("raw string") + subscriber._apply_message(None) + + assert cache.in_memory_cache.get_cache("project_id:p-1") is not None