scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): stop serving stale team model allowlist after /team/update

Browse source 
Mirror of #34266 onto stable/1.93.x (v1.93.0) for an ad hoc release.
Applied from PR head 098cb5dd97 (upstream PR still open against litellm_internal_staging).
This commit is contained in:
mateo-berri 2026-07-22 23:45:30 -07:00
parent 052b5a2169
commit d67a4a6142
5 changed files with 299 additions and 113 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
litellm/proxy/auth/auth_checks.py
View file

@ -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(

11
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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:

160
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -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",

211
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -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:<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

1
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -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)