From 8e30dd1be52d660ea9aa4660846f0baeb2978f4e Mon Sep 17 00:00:00 2001 From: jibanez-staticduo Date: Sat, 26 Sep 2026 00:30:39 +0200 Subject: [PATCH] fix(live): serve authorization reads from cached auth helpers Live session authorization read the team-membership, team, project, and model-access-group budget tables through its own cache probe followed by a direct Prisma call, and never wrote the result back. Every signaling request therefore paid up to six database reads, and a caller that was not a team member re-queried the same absent row on every request. Route all six reads through the auth_checks helpers the chat path already uses, so Live shares their cache keys, TTLs, invalidation broadcasts, single-flight load and negative sentinel with the rest of the proxy: - membership through get_team_membership with raise_on_error left at its default, which also caches an absent row as NO_TEAM_MEMBERSHIP_SENTINEL - team through get_team_object, keeping a missing team row as no team rather than propagating the helper's HTTP 404 - the team's default member budget through get_team_member_default_budget - project through get_project_object, with no separate budget read: the helper includes the joined budget row, as the chat project checks assume - delegated model access group budgets through get_model_access_group_budgets_batch, so a group counts through max_budget the way the chat path enforces it instead of through rpm/tpm The gate stays conservative: any exception still raises HTTP 503 rather than letting managed Live delegation run with configured budgets unsupervised. --- litellm/proxy/realtime_endpoints/live.py | 138 +++++----------- .../proxy/realtime_endpoints/test_live.py | 148 ++++++++++++++---- 2 files changed, 152 insertions(+), 134 deletions(-) diff --git a/litellm/proxy/realtime_endpoints/live.py b/litellm/proxy/realtime_endpoints/live.py index 39787c0a64f..ec7f8f3ed0d 100644 --- a/litellm/proxy/realtime_endpoints/live.py +++ b/litellm/proxy/realtime_endpoints/live.py @@ -21,10 +21,8 @@ from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.llms.chatgpt.live import LiveDeployment, LiveOperation, LiveTransport, live_session_path from litellm.models.budget import LiteLLM_BudgetTable from litellm.models.team import LiteLLM_TeamTable -from litellm.models.team_membership import LiteLLM_TeamMembership from litellm.proxy._types import ( LiteLLM_ProjectTableCachedObj, - LiteLLM_TeamTableCachedObj, LitellmUserRoles, UserAPIKeyAuth, ) @@ -33,18 +31,16 @@ from litellm.proxy.auth.auth_checks import ( can_org_access_model, can_user_call_model, collect_matched_model_access_groups, + get_model_access_group_budgets_batch, get_org_object, + get_project_object, + get_team_member_default_budget, + get_team_membership, get_team_object, get_user_object, ) from litellm.proxy.auth.user_api_key_auth import get_websocket_api_key, user_api_key_auth -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec 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 ( - NO_TEAM_MEMBERSHIP_SENTINEL, - model_access_group_cache_key, - team_membership_reservation_cache_key, -) from litellm.proxy.hooks.parallel_request_limiter import ( _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage] # limiter class is the existing hook identity ) @@ -58,11 +54,6 @@ from litellm.proxy.realtime_endpoints.call_supervision import CALL_SUPERVISORS, from litellm.proxy.spend_tracking.budget_reservation import ( release_or_invalidate_budget_reservation, # pyright: ignore[reportUnknownVariableType] # budget helper accepts legacy reservation dicts ) -from litellm.repositories.budget_repository import BudgetRepository -from litellm.repositories.project_repository import ProjectRepository -from litellm.repositories.table_repositories import ModelAccessGroupBudgetRepository, TeamMembershipRepository -from litellm.repositories.team_repository import TeamRepository -from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget _routes: Final = APIRouter() _JSON: Final = TypeAdapter[JsonValue](JsonValue) @@ -685,25 +676,12 @@ async def _live_team_membership(auth: UserAPIKeyAuth) -> object | None: if auth.team_id is None or auth.user_id is None: return None - membership_key: Final = team_membership_reservation_cache_key(user_id=auth.user_id, team_id=auth.team_id) - membership_cached_raw: Final[object] = _OBJECT_VALUE.validate_python( - await server.user_api_key_cache.async_get_cache(key=membership_key) - ) - membership_cached: Final = ( - CacheCodec.deserialize(membership_cached_raw, model_type=LiteLLM_TeamMembership) - if membership_cached_raw is not None and membership_cached_raw != NO_TEAM_MEMBERSHIP_SENTINEL - else None - ) - if membership_cached is not None or membership_cached_raw == NO_TEAM_MEMBERSHIP_SENTINEL: - return membership_cached - return await TeamMembershipRepository(server.prisma_client).table.find_unique( - where={ # mutable-ok: Prisma serializes query filters from concrete dictionaries - "user_id_team_id": { # mutable-ok: Prisma serializes nested filters from concrete dictionaries - "user_id": auth.user_id, - "team_id": auth.team_id, - } - }, - include={"litellm_budget_table": True}, # mutable-ok: Prisma serializes concrete include dictionaries + return await get_team_membership( + user_id=auth.user_id, + team_id=auth.team_id, + prisma_client=server.prisma_client, + user_api_key_cache=server.user_api_key_cache, + proxy_logging_obj=server.proxy_logging_obj, ) @@ -712,12 +690,17 @@ async def _live_team(auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: if auth.team_id is None: return None - team_from_cache: Final = await server.user_api_key_cache.async_get_cache( - key=f"team_id:{auth.team_id}", model_type=LiteLLM_TeamTableCachedObj - ) - if team_from_cache is not None: - return team_from_cache - return await TeamRepository(server.prisma_client).find_by_id(auth.team_id, id_field="team_id") + try: + return await get_team_object( + team_id=auth.team_id, + prisma_client=server.prisma_client, + user_api_key_cache=server.user_api_key_cache, + proxy_logging_obj=server.proxy_logging_obj, + ) + except HTTPException as exc: + if exc.status_code == 404: + return None + raise def _live_team_budget_configured(auth: UserAPIKeyAuth, team: LiteLLM_TeamTable | None) -> bool: @@ -748,12 +731,12 @@ async def _live_default_budget(auth: UserAPIKeyAuth, team: LiteLLM_TeamTable | N return None from litellm.proxy import proxy_server as server - default_cached: Final = await server.user_api_key_cache.async_get_cache( - key=f"team_member_default_budget:{default_id}", model_type=LiteLLM_BudgetTable + # Like chat auth, a failed default-budget read returns None; membership errors still fail closed above. + return await get_team_member_default_budget( + default_id, + server.prisma_client, + server.user_api_key_cache, ) - if default_cached is not None: - return default_cached - return await BudgetRepository(server.prisma_client).find_by_id(default_id, id_field="budget_id") async def _live_project(auth: UserAPIKeyAuth) -> LiteLLM_ProjectTableCachedObj | None: @@ -761,33 +744,19 @@ async def _live_project(auth: UserAPIKeyAuth) -> LiteLLM_ProjectTableCachedObj | if auth.project_id is None: return None - project_from_cache: Final = await server.user_api_key_cache.async_get_cache( - key=f"project_id:{auth.project_id}", model_type=LiteLLM_ProjectTableCachedObj + return await get_project_object( + project_id=auth.project_id, + prisma_client=server.prisma_client, + user_api_key_cache=server.user_api_key_cache, + proxy_logging_obj=server.proxy_logging_obj, ) - if project_from_cache is not None: - return project_from_cache - project_row: Final = await ProjectRepository(server.prisma_client).table.find_unique( - where={"project_id": auth.project_id}, # mutable-ok: Prisma serializes concrete query dictionaries - include={"litellm_budget_table": True}, # mutable-ok: Prisma serializes concrete include dictionaries - ) - if project_row is None: - return None - return LiteLLM_ProjectTableCachedObj.model_validate(project_row.model_dump()) async def _live_project_budget_configured(auth: UserAPIKeyAuth, project: LiteLLM_ProjectTableCachedObj | None) -> bool: if project is None: return False - from litellm.proxy import proxy_server as server - project_budget: Final = getattr(project, "litellm_budget_table", None) - project_budget_id: Final = getattr(project, "budget_id", None) - project_budget_from_db: Final = ( - await BudgetRepository(server.prisma_client).find_by_id(project_budget_id, id_field="budget_id") - if project_budget is None and isinstance(project_budget_id, str) - else None - ) - if _live_budget_configured(project_budget or project_budget_from_db, zero_is_limit=True): + if _live_budget_configured(project_budget, zero_is_limit=True): return True if _nonempty_limit_value(getattr(project, "model_rpm_limit", None)) or _nonempty_limit_value( getattr(project, "model_tpm_limit", None) @@ -823,44 +792,13 @@ async def _live_model_group_budget_configured( ) if not matched_groups: return False - cached_values: Final = await asyncio.gather( - *( - server.user_api_key_cache.async_get_cache( - key=model_access_group_cache_key(group), model_type=ModelAccessGroupBudget - ) - for group in matched_groups - ) - ) - cached_groups: Final = tuple(zip(matched_groups, cached_values)) - uncached_groups: Final = tuple(group for group, budget in cached_groups if budget is None) - named_group_rows: Final = ( - await ModelAccessGroupBudgetRepository(server.prisma_client).table.find_many( - where={ # mutable-ok: Prisma serializes query filters from concrete dictionaries - "access_group_name": { # mutable-ok: Prisma serializes nested filters from concrete dictionaries - "in": uncached_groups, - } - }, - include={"litellm_budget_table": True}, # mutable-ok: Prisma serializes concrete include dictionaries - ) - if uncached_groups - else () - ) - return any( - _live_budget_configured( - budget - if budget is not None - else next( - ( - getattr(row, "litellm_budget_table", None) - for row in named_group_rows - if getattr(row, "access_group_name", None) == group - ), - None, - ), - zero_is_limit=False, - ) - for group, budget in cached_groups + budgets: Final = await get_model_access_group_budgets_batch( + matched_groups, + server.prisma_client, + server.user_api_key_cache, ) + # Match chat auth: group budget rows contribute max_budget, not rpm/tpm, to this gate. + return any(_live_budget_configured(budget, zero_is_limit=False) for budget in budgets.values()) async def _managed_member_budget(auth: UserAPIKeyAuth, model: str | None = None) -> bool: diff --git a/tests/test_litellm/proxy/realtime_endpoints/test_live.py b/tests/test_litellm/proxy/realtime_endpoints/test_live.py index a395735ad12..7dddd2c215d 100644 --- a/tests/test_litellm/proxy/realtime_endpoints/test_live.py +++ b/tests/test_litellm/proxy/realtime_endpoints/test_live.py @@ -26,6 +26,18 @@ def encryption_key(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "test-live-encryption-key") +def _auth_cache(initial=None): + values = dict(initial or {}) + + async def get(*, key, **kwargs): + return values.get(key) + + async def set(*, key, value, **kwargs): + values[key] = value + + return SimpleNamespace(async_get_cache=AsyncMock(side_effect=get), async_set_cache=AsyncMock(side_effect=set)) + + def handle(owner="owner", model_id=None): deployment = {"model": "gpt-live", "provider": "openai"} if model_id is not None: @@ -509,7 +521,9 @@ async def test_managed_delegation_checks_authoritative_member_and_default_budget from litellm.proxy import proxy_server membership = AsyncMock( - return_value=SimpleNamespace(litellm_budget_table=LiteLLM_BudgetTable(max_budget=member_limit)) + return_value=LiteLLM_TeamMembership( + user_id="user", team_id="team", litellm_budget_table=LiteLLM_BudgetTable(max_budget=member_limit) + ) ) db = SimpleNamespace( litellm_teammembership=SimpleNamespace(find_unique=membership), @@ -519,9 +533,7 @@ async def test_managed_delegation_checks_authoritative_member_and_default_budget ), ) monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) - monkeypatch.setattr( - proxy_server, "user_api_key_cache", SimpleNamespace(async_get_cache=AsyncMock(return_value=None)) - ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", _auth_cache()) monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(live, "_authorize", AsyncMock()) body = {"session": {"delegation": {"type": "responses", "responses": {"model": "backend"}}}} @@ -544,9 +556,7 @@ async def test_managed_delegation_rejects_unverifiable_member_budget_but_allows_ litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(side_effect=RuntimeError("Database unavailable"))) ) monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) - monkeypatch.setattr( - proxy_server, "user_api_key_cache", SimpleNamespace(async_get_cache=AsyncMock(return_value=None)) - ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", _auth_cache()) auth = UserAPIKeyAuth(api_key="owner", team_id="team", user_id="user") body = {"session": {"delegation": {"type": "responses", "responses": {"model": "backend"}}}} with pytest.raises(HTTPException) as rejected: @@ -1309,11 +1319,35 @@ async def test_managed_budget_reads_authoritative_member_budget(monkeypatch): litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(side_effect=find_unique)), litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=None)), ) - cache = SimpleNamespace(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()) + cache = _auth_cache() monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) assert await live._managed_member_budget(UserAPIKeyAuth(api_key="owner", team_id="team", user_id="member")) is True + assert await live._managed_member_budget(UserAPIKeyAuth(api_key="owner", team_id="team", user_id="member")) is True + db.litellm_teammembership.find_unique.assert_awaited_once() + assert cache.async_set_cache.await_args.kwargs["key"] == "team_membership:member:team" + + +@pytest.mark.asyncio +async def test_managed_budget_caches_missing_membership_sentinel(monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.user_api_key_cache import NO_TEAM_MEMBERSHIP_SENTINEL + + membership_lookup = AsyncMock(return_value=None) + db = SimpleNamespace( + litellm_teammembership=SimpleNamespace(find_unique=membership_lookup), + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + ) + cache = _auth_cache() + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + auth = UserAPIKeyAuth(api_key="owner", team_id="missing-member-team", user_id="missing-member") + + assert await live._live_team_membership(auth) is None + assert cache.async_set_cache.await_args.kwargs["value"] == NO_TEAM_MEMBERSHIP_SENTINEL + assert await live._live_team_membership(auth) is None + membership_lookup.assert_awaited_once() @pytest.mark.asyncio @@ -1324,7 +1358,7 @@ async def test_managed_budget_fails_closed_when_membership_repository_is_unreada db = SimpleNamespace( litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(side_effect=failure)), ) - cache = SimpleNamespace(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()) + cache = _auth_cache() monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) @@ -1346,6 +1380,7 @@ async def test_managed_budget_checks_project_team_and_model_group_tables(monkeyp team = LiteLLM_TeamTable(team_id="team", budget_limits=[{"budget_duration": "1d", "max_budget": 1}]) group = SimpleNamespace( access_group_name="group", + spend=0, litellm_budget_table=SimpleNamespace(max_budget=1), ) @@ -1362,7 +1397,7 @@ async def test_managed_budget_checks_project_team_and_model_group_tables(monkeyp litellm_projecttable=SimpleNamespace(find_unique=AsyncMock(side_effect=find_project)), litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=AsyncMock(side_effect=find_groups)), ) - cache = SimpleNamespace(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()) + cache = _auth_cache() monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) monkeypatch.setattr(live, "collect_matched_model_access_groups", AsyncMock(return_value=("group",))) @@ -1383,15 +1418,19 @@ async def test_managed_budget_uses_the_delegated_model_group(monkeypatch, backen from litellm.proxy import proxy_server rows = [ - SimpleNamespace(access_group_name="voice-group", litellm_budget_table=SimpleNamespace(max_budget=1)), + SimpleNamespace(access_group_name="voice-group", spend=0, litellm_budget_table=SimpleNamespace(max_budget=1)), SimpleNamespace( - access_group_name="backend-group", litellm_budget_table=SimpleNamespace(max_budget=backend_budget) + access_group_name="backend-group", spend=0, litellm_budget_table=SimpleNamespace(max_budget=backend_budget) ), ] + + async def find_group_budgets(*, where, include): + return [row for row in rows if row.access_group_name in where["access_group_name"]["in"]] + db = SimpleNamespace( - litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=AsyncMock(return_value=rows)), + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=AsyncMock(side_effect=find_group_budgets)), ) - cache = SimpleNamespace(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()) + cache = _auth_cache() monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) monkeypatch.setattr(proxy_server, "llm_router", SimpleNamespace()) @@ -1399,6 +1438,9 @@ async def test_managed_budget_uses_the_delegated_model_group(monkeypatch, backen auth = UserAPIKeyAuth(api_key="owner", models=["voice", "backend-group"]) assert await live._managed_member_budget(auth, model="backend") is blocked + assert await live._managed_member_budget(auth, model="backend") is blocked + db.litellm_modelaccessgroupbudgettable.find_many.assert_awaited_once() + assert cache.async_set_cache.await_args.kwargs["key"] == "model_access_group:backend-group" @pytest.mark.asyncio @@ -1406,19 +1448,19 @@ async def test_responses_delegation_fails_closed_when_inherited_org_lookup_fails from litellm.proxy import proxy_server from litellm.proxy.auth import auth_checks - team = LiteLLM_TeamTable(team_id="team", organization_id="org", models=["*"]) + team = LiteLLM_TeamTable(team_id="org-lookup-team", organization_id="org", models=["*"]) group = SimpleNamespace( access_group_name="backend-group", + spend=0, litellm_budget_table=SimpleNamespace(max_budget=1), ) db = SimpleNamespace( litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team)), litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=AsyncMock(return_value=[group])), ) - cache = SimpleNamespace(async_get_cache=AsyncMock(return_value=None), async_set_cache=AsyncMock()) + cache = _auth_cache() monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=db)) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", None) monkeypatch.setattr( proxy_server, "llm_router", @@ -1434,7 +1476,7 @@ async def test_responses_delegation_fails_closed_when_inherited_org_lookup_fails with pytest.raises(HTTPException) as rejected: await live._authorize_delegation( {"session": {"delegation": {"type": "responses", "responses": {"model": "backend"}}}}, - UserAPIKeyAuth(api_key="owner", models=["*"], team_id="team"), + UserAPIKeyAuth(api_key="owner", models=["*"], team_id="org-lookup-team"), ) assert rejected.value.status_code == 503 @@ -1694,9 +1736,7 @@ async def test_live_team_membership_prefers_reservation_cache_and_sentinel(monke monkeypatch.setattr( proxy_server, "user_api_key_cache", - SimpleNamespace( - async_get_cache=AsyncMock(return_value=CacheCodec.serialize(membership, model_type=LiteLLM_TeamMembership)) - ), + _auth_cache({"team_membership:user:team": CacheCodec.serialize(membership, model_type=LiteLLM_TeamMembership)}), ) restored = await live._live_team_membership(auth) assert restored is not None and restored.user_id == "user" and restored.team_id == "team" @@ -1704,7 +1744,7 @@ async def test_live_team_membership_prefers_reservation_cache_and_sentinel(monke monkeypatch.setattr( proxy_server, "user_api_key_cache", - SimpleNamespace(async_get_cache=AsyncMock(return_value=NO_TEAM_MEMBERSHIP_SENTINEL)), + _auth_cache({"team_membership:user:team": NO_TEAM_MEMBERSHIP_SENTINEL}), ) assert await live._live_team_membership(auth) is None @@ -1714,10 +1754,36 @@ async def test_live_team_uses_team_cache_before_database(monkeypatch): from litellm.proxy import proxy_server team = SimpleNamespace(team_id="team", models=["*"]) + team_lookup = AsyncMock(side_effect=AssertionError("cache hit must not query the team table")) monkeypatch.setattr( - proxy_server, "user_api_key_cache", SimpleNamespace(async_get_cache=AsyncMock(return_value=team)) + proxy_server, + "prisma_client", + SimpleNamespace(db=SimpleNamespace(litellm_teamtable=SimpleNamespace(find_unique=team_lookup))), ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", _auth_cache({"team_id:team": team})) assert await live._live_team(UserAPIKeyAuth(api_key="owner", team_id="team")) is team + team_lookup.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_live_team_caches_database_row_after_miss(monkeypatch): + from litellm.proxy import proxy_server + + team = LiteLLM_TeamTable(team_id="cache-miss-team") + team_lookup = AsyncMock(return_value=team) + cache = _auth_cache() + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=SimpleNamespace(litellm_teamtable=SimpleNamespace(find_unique=team_lookup))), + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + auth = UserAPIKeyAuth(api_key="owner", team_id="cache-miss-team") + + assert (await live._live_team(auth)).team_id == "cache-miss-team" + assert cache.async_set_cache.await_args.kwargs["key"] == "team_id:cache-miss-team" + assert (await live._live_team(auth)).team_id == "cache-miss-team" + team_lookup.assert_awaited_once() @pytest.mark.parametrize( @@ -1755,33 +1821,47 @@ async def test_live_default_budget_uses_cached_team_member_budget(monkeypatch): from litellm.proxy import proxy_server budget = LiteLLM_BudgetTable(max_budget=1) + budget_lookup = AsyncMock(side_effect=AssertionError("cache hit must not query the budget table")) monkeypatch.setattr( - proxy_server, "user_api_key_cache", SimpleNamespace(async_get_cache=AsyncMock(return_value=budget)) + proxy_server, + "prisma_client", + SimpleNamespace(db=SimpleNamespace(litellm_budgettable=SimpleNamespace(find_unique=budget_lookup))), + ) + monkeypatch.setattr( + proxy_server, "user_api_key_cache", _auth_cache({"team_member_default_budget:budget-1": budget}) ) team = SimpleNamespace(metadata={"team_member_budget_id": "budget-1"}) auth = UserAPIKeyAuth(api_key="owner", user_id="user", team_id="team") assert await live._live_default_budget(auth, team) is budget + budget_lookup.assert_not_awaited() @pytest.mark.asyncio async def test_live_project_uses_cache_or_reports_missing_row(monkeypatch): from litellm.proxy import proxy_server - project = SimpleNamespace(project_id="project-1") + project = LiteLLM_ProjectTable(project_id="project-1") + project_lookup = AsyncMock(return_value=project) monkeypatch.setattr( - proxy_server, "user_api_key_cache", SimpleNamespace(async_get_cache=AsyncMock(return_value=project)) + proxy_server, + "prisma_client", + SimpleNamespace(db=SimpleNamespace(litellm_projecttable=SimpleNamespace(find_unique=project_lookup))), ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", _auth_cache({"project_id:project-1": project})) auth = UserAPIKeyAuth(api_key="owner", project_id="project-1") assert await live._live_project(auth) is project + project_lookup.assert_not_awaited() - monkeypatch.setattr( - proxy_server, "user_api_key_cache", SimpleNamespace(async_get_cache=AsyncMock(return_value=None)) - ) - monkeypatch.setattr( - live, - "ProjectRepository", - lambda client: SimpleNamespace(table=SimpleNamespace(find_unique=AsyncMock(return_value=None))), - ) + cache = _auth_cache() + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + assert (await live._live_project(auth)).project_id == "project-1" + assert cache.async_set_cache.await_args.kwargs["key"] == "project_id:project-1" + assert (await live._live_project(auth)).project_id == "project-1" + project_lookup.assert_awaited_once() + + project_lookup.reset_mock(return_value=True) + project_lookup.return_value = None + monkeypatch.setattr(proxy_server, "user_api_key_cache", _auth_cache()) assert await live._live_project(auth) is None