From d67a4a6142c9277dc25387dc482addd98f0f6f85 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:45:30 -0700 Subject: [PATCH] fix(proxy): stop serving stale team model allowlist after /team/update Mirror of #34266 onto stable/1.93.x (v1.93.0) for an ad hoc release. Applied from PR head 098cb5dd9712897a6d2da370ceda3c7c280d3a8d (upstream PR still open against litellm_internal_staging). --- litellm/proxy/auth/auth_checks.py | 29 ++- litellm/proxy/auth/user_api_key_auth.py | 11 - .../proxy/auth/test_auth_checks.py | 160 ++++++++++++- .../proxy/auth/test_user_api_key_auth.py | 211 ++++++++++-------- .../test_team_endpoints.py | 1 + 5 files changed, 299 insertions(+), 113 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 230b9b70ff0..93c9deb08b8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1744,9 +1744,22 @@ async def _cache_team_object( ## CACHE REFRESH TIME! team_table.last_refreshed_at = time.time() + key = "team_id:{}".format(team_id) + + if proxy_logging_obj is not None: + try: + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the write + verbose_proxy_logger.warning( + "Failed to invalidate internal usage cache entry %s; " + "a stale team object may be served until its TTL expires: %s", + key, + e, + ) + # team_id is the table primary key — guaranteed unique, safe to write. await _cache_management_object( - key="team_id:{}".format(team_id), + key=key, value=team_table, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, @@ -1768,9 +1781,17 @@ async def _cache_team_object( # the cache from a verified single row. if team_table.team_alias: alias_key = "team_alias:{}".format(team_table.team_alias) - user_api_key_cache.delete_cache(key=alias_key) - if proxy_logging_obj is not None: - await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=alias_key) + try: + user_api_key_cache.delete_cache(key=alias_key) + if proxy_logging_obj is not None: + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=alias_key) + except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the mutation + verbose_proxy_logger.warning( + "Failed to invalidate cached team alias entry %s; " + "a stale team object may be served until its TTL expires: %s", + alias_key, + e, + ) async def _cache_key_object( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2613510bd0c..45909f8e9e1 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1936,17 +1936,6 @@ async def _user_api_key_auth_builder( else: valid_token.team_object_permission = None - # Cache under the canonical "team_id:{id}" key so get_team_object and - # _update_team_cache serve this write from the L2 cache. The guard keeps a - # non-team (personal) key, whose team_id is None, from reaching the cache - # layer, which Redis rejects with a NoneType key error. - if valid_token.team_id is not None and _team_obj is not None: - await user_api_key_cache.async_set_cache( - key=f"team_id:{valid_token.team_id}", - value=_team_obj, - model_type=LiteLLM_TeamTableCachedObj, - ) - # Fetch project object if key belongs to a project _project_obj = None if valid_token.project_id is not None: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d12ff20ee5b..f337bd15138 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import ( vector_store_access_check, ) from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -3944,6 +3945,11 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): len(teams)==1 before populating the cache. 3. When team_alias is None, NO alias-key operation happens (no delete of an empty-keyed entry, no spurious write). + 4. DELETES the team_id-keyed entry from the internal usage cache + BEFORE the fresh write (LIT-4391). `_get_team_object_from_cache` + consults the internal usage cache first, so a leftover copy there + (backfilled from a Redis shared with `user_api_key_cache`) would + keep serving the pre-update team allowlist. """ from unittest.mock import AsyncMock, MagicMock @@ -3990,9 +3996,14 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): # (2) team_alias-keyed entry is deleted in BOTH the in-memory cache # and the Redis dual cache (mirrors _delete_cache_key_object pattern). cache.delete_cache.assert_called_once_with(key="team_alias:H-Capacity") - logging_obj.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with( - key="team_alias:H-Capacity" - ) + + # (4) internal usage cache: team_id entry deleted BEFORE the fresh + # write, alias entry deleted as before. + internal_deleted_keys = [ + (c.kwargs.get("key") or c.args[0]) + for c in logging_obj.internal_usage_cache.dual_cache.async_delete_cache.await_args_list + ] + assert internal_deleted_keys == ["team_id:team-1234", "team_alias:H-Capacity"] # ===== team_alias is None: no alias-key operation ===== aliasless = LiteLLM_TeamTableCachedObj(**{**base_team_row, "team_alias": None}) @@ -4010,7 +4021,9 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): ) cache2.delete_cache.assert_not_called() - logging_obj2.internal_usage_cache.dual_cache.async_delete_cache.assert_not_awaited() + logging_obj2.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with( + key="team_id:team-no-alias" + ) written_keys_aliasless = [ (c.kwargs.get("key") or c.args[0]) for c in cache2.async_set_cache.await_args_list @@ -4018,6 +4031,145 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): assert written_keys_aliasless == ["team_id:team-no-alias"] +class _SharedFakeRedis(RedisCache): + """Dict-backed stand-in for the single Redis that both + ``user_api_key_cache`` (enable_redis_auth_cache) and + ``proxy_logging_obj.internal_usage_cache.dual_cache`` share in the + LIT-4391 deployment topology. Only the methods DualCache calls are + implemented; ``super().__init__`` is skipped intentionally.""" + + def __init__(self): + self._store: dict = {} + + async def async_set_cache(self, key, value, **kwargs): + self._store[key] = json.dumps(value) + + async def async_get_cache(self, key, **kwargs): + raw = self._store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key): + self._store.pop(key, None) + + +@pytest.mark.asyncio +async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): + """ + Regression test for LIT-4391: keys with models=["all-team-models"] kept + getting 403 team_model_access_denied for models added via /team/update. + + `_get_team_object_from_cache` consults the internal usage cache BEFORE + `user_api_key_cache`. When both share one Redis (enable_redis_auth_cache), + any team read backfills the internal cache's in-memory tier with the team + object. `_cache_team_object` (the /team/update refresh) only wrote + `user_api_key_cache`, so that backfilled copy kept shadowing the update + until its TTL expired — and the auth-time write-back then pushed the stale + copy back into the shared Redis, making the staleness self-sustaining. + + Pins: + 1. After `_cache_team_object` writes an updated team, `get_team_object` + returns the UPDATED model list even though the internal usage cache's + in-memory tier was backfilled with the pre-update team. + 2. The shared Redis still holds the updated team afterwards — the + internal-cache invalidation must happen BEFORE the fresh write, or it + would wipe the value it just wrote. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import _cache_team_object, get_team_object + + team_id = "team-lit-4391" + shared_redis = _SharedFakeRedis() + user_api_key_cache = UserApiKeyCache(redis_cache=shared_redis) + proxy_logging_obj = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache = DualCache( + redis_cache=shared_redis, + default_in_memory_ttl=300, + ) + prisma_client = MagicMock() + + await _cache_team_object( + team_id=team_id, + team_table=LiteLLM_TeamTableCachedObj(team_id=team_id, models=["model-a"]), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + primed = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert primed is not None and primed.models == ["model-a"] + + await _cache_team_object( + team_id=team_id, + team_table=LiteLLM_TeamTableCachedObj( + team_id=team_id, models=["model-a", "model-b"] + ), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + refreshed = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert refreshed is not None and refreshed.models == ["model-a", "model-b"], ( + "get_team_object served a stale team allowlist after _cache_team_object " + f"refreshed it. Got models={refreshed.models if refreshed else None}" + ) + + redis_copy = await shared_redis.async_get_cache(f"team_id:{team_id}") + assert redis_copy is not None and redis_copy["models"] == ["model-a", "model-b"], ( + "The shared Redis lost the refreshed team object — the internal-cache " + "invalidation must run BEFORE the fresh write, not after. " + f"Got: {redis_copy}" + ) + + +@pytest.mark.asyncio +async def test_cache_team_object_tolerates_cache_invalidation_failures(): + """ + Greptile review on the LIT-4391 fix: `_cache_team_object` runs after a + successful DB fetch (inside `get_team_object`) and after every team + mutation's DB write. A cache-backend error during the best-effort + invalidations must NOT fail those operations — otherwise a Redis blip + turns a healthy team lookup into a 404 and a committed /team/update into + a 500. The authoritative team_id-keyed write must still happen. + """ + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import _cache_team_object + + cache = MagicMock() + cache.async_set_cache = AsyncMock() + cache.delete_cache = MagicMock(side_effect=Exception("redis down")) + logging_obj = MagicMock() + logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( + side_effect=Exception("redis down") + ) + + await _cache_team_object( + team_id="team-cache-outage", + team_table=LiteLLM_TeamTableCachedObj( + team_id="team-cache-outage", + team_alias="cache-outage-alias", + models=["model-a"], + ), + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + ) + + written_keys = [ + (c.kwargs.get("key") or c.args[0]) + for c in cache.async_set_cache.await_args_list + ] + assert written_keys == ["team_id:team-cache-outage"] + + MODEL_DISCOVERY_ROUTES = [ "/v1/models", "/models", diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 90f46152837..f53954ad3ae 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2628,6 +2628,123 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_auth_flow_never_persists_fallback_team_object_lit_4391(): + """ + Regression test for LIT-4391 (stale team allowlist poisoning). + + When `get_team_object` fails at the "Check 6" team-auth step (cache miss + inside the DB-throttle window, DB blip, ...), the builder falls back to a + team object reconstructed from the CACHED token's team_* snapshot — which + can be arbitrarily stale (e.g. pre-/team/update models). + + The builder used to write that team object back into `user_api_key_cache` + under "team_id:" after Check 6. Writing a cache-read (or worse, a + token-snapshot) value back into the shared cache re-poisons it — with + enable_redis_auth_cache it clobbered the fresh team `/team/update` had + just written to Redis, making the stale allowlist self-sustaining across + requests. Only authoritative writers (`_cache_team_object` on DB reads and + team mutations) may populate the team cache. + + Pins: the auth flow completes on the fallback path WITHOUT writing any + "team_id:*" cache entry. + """ + from starlette.datastructures import URL + from starlette.requests import Request + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + api_key = "sk-test-lit-4391-no-team-writeback" + valid_token = UserAPIKeyAuth( + api_key=api_key, + token=api_key, + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-lit-4391", + team_models=["model-a"], + models=["all-team-models"], + ) + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(return_value=valid_token) + mock_cache.async_set_cache = AsyncMock(return_value=None) + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _originals = {k: getattr(_proxy_server_mod, k, None) for k in _attrs} + + try: + for k, v in _attrs.items(): + setattr(_proxy_server_mod, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=HTTPException( + status_code=404, + detail={"error": "Team doesn't exist in db."}, + ), + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert result.team_id == "team-lit-4391" + + team_cache_writes = [ + key + for c in mock_cache.async_set_cache.await_args_list + if isinstance(key := (c.kwargs.get("key") if "key" in c.kwargs else c.args[0]), str) + and key.startswith("team_id:") + ] + assert team_cache_writes == [], ( + "The auth flow wrote a team object into the cache. Fallback/" + "cache-read team objects must never be persisted — only " + "_cache_team_object (DB reads and team mutations) may write " + f"'team_id:*' entries. Got writes: {team_cache_writes}" + ) + + finally: + for k, v in _originals.items(): + setattr(_proxy_server_mod, k, v) + + # --------------------------------------------------------------------------- # _run_centralized_common_checks — centralized authz gate @@ -3992,100 +4109,6 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa assert result.is_session_token is True -@pytest.mark.asyncio -async def test_auth_path_caches_team_object_under_canonical_team_id_key(): - """Regression for LIT-4000: the auth builder must cache the team object under - the canonical ``team_id:{id}`` key that ``get_team_object`` and - ``_update_team_cache`` read, never under the raw ``team_id`` (and never under - a ``None`` key, which Redis rejects with a NoneType key error). A raw or None - key is silently dropped by Redis / never served back, so every request - re-hits Postgres for the team object instead of the L2 cache. - - Drives the real builder for a team-scoped key against a real in-memory - ``UserApiKeyCache`` and reads the team object back. Mutating the cache key at - the write site to the raw ``valid_token.team_id`` (or ``None``) makes the - canonical-key read miss and fails this test. - """ - from fastapi import Request - from starlette.datastructures import URL - - import litellm.proxy.proxy_server as _proxy_server_mod - from litellm.proxy._types import LiteLLM_TeamTableCachedObj - from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.proxy_server import hash_token - - team_id = "team-lit-4000" - api_key = "sk-lit-4000-team-key" - cache = UserApiKeyCache() - - team_token = UserAPIKeyAuth(token=hash_token(api_key), team_id=team_id) - team_obj = LiteLLM_TeamTableCachedObj(team_id=team_id) - - proxy_logging_obj = MagicMock() - proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - attrs = { - "prisma_client": MagicMock(), - "user_api_key_cache": cache, - "proxy_logging_obj": proxy_logging_obj, - "master_key": "sk-test-master", - "general_settings": {"allow_requests_on_db_unavailable": False}, - "llm_model_list": [], - "llm_router": None, - "open_telemetry_logger": None, - "model_max_budget_limiter": MagicMock(), - "user_custom_auth": None, - "jwt_handler": None, - "litellm_proxy_admin_name": "admin", - } - originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} - try: - for k, v in attrs.items(): - setattr(_proxy_server_mod, k, v) - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - with ( - patch( - "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", - AsyncMock(return_value=team_token), - ), - patch( - "litellm.proxy.auth.user_api_key_auth.get_team_object", - AsyncMock(return_value=team_obj), - ), - patch( - "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", - new_callable=AsyncMock, - return_value=team_token, - ), - patch( - "litellm.proxy.auth.auth_exception_handler.seed_request_identity", - ), - ): - await _user_api_key_auth_builder( - request=request, - api_key=f"Bearer {api_key}", - azure_api_key_header="", - anthropic_api_key_header=None, - google_ai_studio_api_key_header=None, - azure_apim_header=None, - request_data={}, - ) - finally: - for k, v in originals.items(): - setattr(_proxy_server_mod, k, v) - - served = cache.get_cache( - key=f"team_id:{team_id}", model_type=LiteLLM_TeamTableCachedObj - ) - assert served is not None and served.team_id == team_id - assert cache.get_cache(key=team_id) is None - assert cache.get_cache(key=None) is None - class TestCheckKeyModelBudgetWithFallback: """`_check_key_model_budget_with_fallback` must reroute a request to the diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 59b08a7ec4e..670f967e883 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6418,6 +6418,7 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): return_value=mock_existing_team ) mock_cache.async_set_cache = AsyncMock() + mock_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() # Mock team update mock_updated_team = MagicMock(spec=LiteLLM_TeamTable)