mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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.
This commit is contained in:
jibanez-staticduo
parent
5a3816a5ad
commit
8e30dd1be5
2 changed files with 152 additions and 134 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue