diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 61d7078ae4c..3fa26d3506d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -313,6 +313,17 @@ def model_access_group_cache_key(access_group_name: str) -> str: return f"model_access_group:{access_group_name}" +def live_model_access_group_limits_cache_key(access_group_name: str) -> str: + """Cache key the Live delegation gate stores one access group's full limit row under. + + The gate needs the rpm and tpm columns that ``model_access_group:{name}`` flattens away, so it + keeps its own entry next to the flattened one. Any eviction of the flattened entry must clear + this key too: the gate reads cache-first, and a raised or lowered group limit left cached here + keeps permitting or refusing managed delegation until the entry's TTL expires (LIT-3803). + """ + return f"live:model_access_group_limits:{access_group_name}" + + def model_access_group_registry_cache_key() -> str: """Cache key for the set of model access group names that have a budget row.""" return "model_access_group_registry" diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index e960bdfe337..72e251e8a31 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -26,6 +26,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + live_model_access_group_limits_cache_key, model_access_group_cache_key, model_access_group_registry_cache_key, ) @@ -200,7 +201,14 @@ async def _evict_model_access_group_cache_keys(access_group: str, auth_cache: Us ) await evict_and_broadcast( - cache_keys=(model_access_group_cache_key(access_group), model_access_group_registry_cache_key()), + cache_keys=( + model_access_group_cache_key(access_group), + # The Live delegation gate caches the same group's full limit row next to the flattened + # entry because it needs the rpm and tpm columns; leaving that entry behind keeps the + # old limit deciding managed delegation until its TTL expires. + live_model_access_group_limits_cache_key(access_group), + model_access_group_registry_cache_key(), + ), user_api_key_cache=auth_cache, ) diff --git a/litellm/proxy/realtime_endpoints/live.py b/litellm/proxy/realtime_endpoints/live.py index 70ed3a6422f..596e1e1bb9c 100644 --- a/litellm/proxy/realtime_endpoints/live.py +++ b/litellm/proxy/realtime_endpoints/live.py @@ -13,6 +13,8 @@ from fastapi import APIRouter, HTTPException, Request, Response, WebSocket, WebS from pydantic import BaseModel, Field, JsonValue, TypeAdapter from starlette.types import Message +from litellm._logging import verbose_proxy_logger + if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection @@ -28,10 +30,13 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( + _cache_team_object, # pyright: ignore[reportPrivateUsage] # same cache write the chat path performs + _get_team_object_from_cache, # pyright: ignore[reportPrivateUsage] # same cache read the chat path performs can_key_call_resolved_model, # pyright: ignore[reportUnknownVariableType] # legacy authorization accepts untyped deployment lists can_org_access_model, can_user_call_model, collect_matched_model_access_groups, + get_object_permission, get_org_object, get_project_object, get_team_membership, @@ -40,7 +45,10 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import get_websocket_api_key, user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper -from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl +from litellm.proxy.common_utils.user_api_key_cache import ( + get_management_object_ttl, + live_model_access_group_limits_cache_key, +) from litellm.proxy.hooks.parallel_request_limiter import ( _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage] # limiter class is the existing hook identity ) @@ -62,7 +70,6 @@ _routes: Final = APIRouter() _JSON: Final = TypeAdapter[JsonValue](JsonValue) _EMPTY: Final[Mapping[str, JsonValue]] = MappingProxyType({}) _CACHEABLE_MODEL = TypeVar("_CACHEABLE_MODEL", bound=BaseModel) -_LIVE_GROUP_LIMITS_CACHE_PREFIX: Final = "live:model_access_group_limits:" _MAPPING: Final = TypeAdapter(Mapping[str, object]) _OBJECT: Final = TypeAdapter(Mapping[str, JsonValue]) _DEPLOYMENT: Final = TypeAdapter(LiveDeployment) @@ -727,20 +734,39 @@ async def _live_team(auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: if auth.team_id is None: return None team_id: Final = auth.team_id - - async def load() -> LiteLLM_TeamTableCachedObj | None: - row: Final = await TeamRepository(server.prisma_client).find_by_id(team_id, id_field="team_id") - if row is None: - return None - team: Final = LiteLLM_TeamTableCachedObj.model_validate(row.model_dump()) - team.last_refreshed_at = time.time() - return team - - return await _live_cached_object( + cached: Final = await _get_team_object_from_cache( key=f"team_id:{team_id}", - model_type=LiteLLM_TeamTableCachedObj, - load=load, + user_api_key_cache=server.user_api_key_cache, + parent_otel_span=None, ) + if cached is not None: + return cached + + row: Final = await TeamRepository(server.prisma_client).find_by_id(team_id, id_field="team_id") + if row is None: + return None + team: Final = LiteLLM_TeamTableCachedObj.model_validate(row.model_dump()) + if team.object_permission_id and not team.object_permission: + # The entry is written under the key the chat path reads, so it has to carry the same + # permission relation the chat path caches; a cache hit elsewhere must not see a team + # stripped of the permissions it was about to enforce. + try: + team.object_permission = await get_object_permission( + object_permission_id=team.object_permission_id, + prisma_client=server.prisma_client, + user_api_key_cache=server.user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=server.proxy_logging_obj, + ) + except Exception as exc: # noqa: BLE001 # same degradation as the chat path: cache the team without permissions and log it + verbose_proxy_logger.debug("Failed to load object_permission for Live team %s: %s", team_id, exc) + await _cache_team_object( + team_id=team_id, + team_table=team, + user_api_key_cache=server.user_api_key_cache, + proxy_logging_obj=server.proxy_logging_obj, + ) + return team def _live_team_budget_configured(auth: UserAPIKeyAuth, team: LiteLLM_TeamTable | None) -> bool: @@ -848,7 +874,7 @@ async def _live_fetch_group_limits(groups: tuple[str, ...]) -> tuple[LiteLLM_Bud await asyncio.gather( *( server.user_api_key_cache.async_set_cache( - key=f"{_LIVE_GROUP_LIMITS_CACHE_PREFIX}{group}", + key=live_model_access_group_limits_cache_key(group), value=limit, model_type=LiteLLM_BudgetTable, ttl=get_management_object_ttl(server.user_api_key_cache), @@ -866,7 +892,7 @@ async def _live_model_group_limits(groups: tuple[str, ...]) -> tuple[LiteLLM_Bud cached: Final = await asyncio.gather( *( server.user_api_key_cache.async_get_cache( - key=f"{_LIVE_GROUP_LIMITS_CACHE_PREFIX}{group}", + key=live_model_access_group_limits_cache_key(group), model_type=LiteLLM_BudgetTable, ) for group in groups diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index 59c2921e0d0..cbd286ee4df 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -784,20 +784,22 @@ def _proxy_with_stubbed_reload(prisma): def _eviction_journal(access_group): - """Both auth cache keys, in the order a write path has to evict them.""" + """Every auth cache key that holds this group's limits, in the order a write path has to evict them.""" from litellm.proxy.common_utils.user_api_key_cache import ( + live_model_access_group_limits_cache_key, model_access_group_cache_key, model_access_group_registry_cache_key, ) return [ f"auth_cache.delete:{model_access_group_cache_key(access_group)}", + f"auth_cache.delete:{live_model_access_group_limits_cache_key(access_group)}", f"auth_cache.delete:{model_access_group_registry_cache_key()}", ] def _assert_evicted_after_write(journal, access_group, write_entry): - """Exactly the two keys, in order, after the DB write. Deliberately not a tail slice: what + """Exactly the cached keys, in order, after the DB write. Deliberately not a tail slice: what has to hold is that the eviction follows the write, not that nothing follows the eviction.""" evictions = [entry for entry in journal if entry.startswith("auth_cache.delete:")] assert evictions == _eviction_journal(access_group) @@ -1206,7 +1208,7 @@ async def test_list_access_groups_reports_a_budgetless_group_as_unbudgeted_rathe @pytest.mark.asyncio -async def test_put_access_group_budget_evicts_both_auth_cache_keys(): +async def test_put_access_group_budget_evicts_every_cached_limit_key(): """Auth reads the per-group row and the registry of budgeted groups cache-first with no freshness check, so a PUT that skips either eviction returns 200 and enforces nothing until the TTL expires. Both keys, after the write.""" @@ -1233,7 +1235,7 @@ async def test_put_access_group_budget_evicts_both_auth_cache_keys(): @pytest.mark.asyncio -async def test_delete_access_group_budget_evicts_both_auth_cache_keys(): +async def test_delete_access_group_budget_evicts_every_cached_limit_key(): """Clearing a budget has the same window as setting one: until both keys are dropped, auth keeps enforcing the budget that is already gone.""" from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( @@ -1252,7 +1254,7 @@ async def test_delete_access_group_budget_evicts_both_auth_cache_keys(): @pytest.mark.asyncio -async def test_deleting_the_access_group_evicts_both_auth_cache_keys(): +async def test_deleting_the_access_group_evicts_every_cached_limit_key(): """The group-delete cascade drops the budget row too, so it owes the same two evictions.""" from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( delete_access_group, diff --git a/tests/test_litellm/proxy/realtime_endpoints/test_live.py b/tests/test_litellm/proxy/realtime_endpoints/test_live.py index 1d864fa2919..768412f7570 100644 --- a/tests/test_litellm/proxy/realtime_endpoints/test_live.py +++ b/tests/test_litellm/proxy/realtime_endpoints/test_live.py @@ -1483,6 +1483,32 @@ async def test_managed_budget_fails_closed_when_the_team_row_is_unreadable(monke assert rejected.value.status_code == 503 +@pytest.mark.asyncio +async def test_live_team_caches_the_permission_relation_with_the_team(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + team = LiteLLM_TeamTable(team_id="team", object_permission_id="perm-1") + db = SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team)), + litellm_objectpermissiontable=SimpleNamespace( + find_unique=AsyncMock(return_value=LiteLLM_ObjectPermissionTable(object_permission_id="perm-1")) + ), + ) + cache = _auth_cache() + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + + loaded: Final = await live._live_team(UserAPIKeyAuth(api_key="owner", team_id="team")) + assert loaded is not None and loaded.object_permission is not None + + cached_entries: Final = [ + call.kwargs["value"] for call in cache.async_set_cache.await_args_list if call.kwargs["key"] == "team_id:team" + ] + assert len(cached_entries) == 1, "the team must be cached under the key the chat path reads" + assert cached_entries[0].object_permission is not None + + @pytest.mark.asyncio async def test_managed_budget_fails_closed_when_the_default_budget_is_unreadable(monkeypatch): from litellm.proxy import proxy_server