From abd3f901013c8b1ee36beb8baa8a2e4ffaf51dfc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 11:49:50 -0800 Subject: [PATCH] fix(usage): keep spend when the team, organization, customer, or agent name lookup drops --- litellm/proxy/agent_endpoints/endpoints.py | 10 +- .../customer_endpoints.py | 10 +- .../organization_endpoints.py | 10 +- .../management_endpoints/team_endpoints.py | 8 +- .../test_daily_activity_routes.py | 109 +++++++++++++++++- 5 files changed, 128 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 42cf44a3c95..2bd9b4357d7 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -69,7 +69,7 @@ from litellm.proxy.agent_endpoints.kill_switch import ( from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails from litellm.proxy.utils import PrismaClient, get_custom_url from litellm.repositories.chunked_in import find_many_in from litellm.types.agents import ( @@ -394,10 +394,12 @@ async def resolve_agent_daily_activity_scope( prisma_client=prisma_client, ) - agent_records: Final = ( - await agents_table(prisma_client).find_many(where={}) + agent_records: Final = await unless_database_fails( + agents_table(prisma_client).find_many(where={}) if resolved_agent_ids is None - else await find_many_in(agents_table(prisma_client), "agent_id", resolved_agent_ids) + else find_many_in(agents_table(prisma_client), "agent_id", resolved_agent_ids), + (), + "its agent names", ) agent_metadata: Final[Mapping[str, dict[str, object]]] = MappingProxyType( {agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records} diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index b426369adea..0640bf0fa4b 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -35,7 +35,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, end_user_restricted_registry_cache_key, ) -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails from litellm.proxy.management_endpoints.common_utils import validate_budget_duration from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports _set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -522,10 +522,12 @@ async def resolve_customer_daily_activity_scope( prisma_client: "PrismaClient", ) -> _CustomerDailyActivityScope: end_user_table: Final = _typed_table(EndUserRepository(prisma_client)) - end_user_aliases: Final = ( - await find_many_in(end_user_table, "user_id", end_user_ids) + end_user_aliases: Final = await unless_database_fails( + find_many_in(end_user_table, "user_id", end_user_ids) if end_user_ids is not None - else await end_user_table.find_many(where={}) + else end_user_table.find_many(where={}), + (), + "its customer aliases", ) metadata: Final = MappingProxyType({end_user.user_id: {"alias": end_user.alias} for end_user in end_user_aliases}) return _CustomerDailyActivityScope(end_user_ids, metadata) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 4e511aee24a..15d05633679 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -48,7 +48,7 @@ from litellm.proxy.management_endpoints.budget_management_endpoints import ( new_budget, update_budget, ) -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity, unless_database_fails from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports _set_object_metadata_field, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _user_has_admin_view, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -655,10 +655,12 @@ async def resolve_organization_daily_activity_scope( ) organization_table: Final = _table(OrganizationRepository(prisma_client)) - organization_rows: Final = ( - await find_many_in(organization_table, "organization_id", resolved_organization_ids) + organization_rows: Final = await unless_database_fails( + find_many_in(organization_table, "organization_id", resolved_organization_ids) if resolved_organization_ids is not None - else await organization_table.find_many(where={}) + else organization_table.find_many(where={}), + (), + "its organization aliases", ) metadata: Final = MappingProxyType( { diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 54b9b26123e..b3fe8685148 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -155,6 +155,7 @@ from litellm.proxy.management.teams.dependencies import get_team_access from litellm.proxy.management_endpoints.common_daily_activity import ( InvalidDateRange, parse_canonical_date_range, + unless_database_fails, ) from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # legacy module exports _check_disable_global_guardrails_caller_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -7302,7 +7303,12 @@ async def resolve_team_daily_activity_scope( where_condition: Final[_TeamIdInFilter] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition) + team_alias_lookup: Final = _team_db(prisma_client).find_many(where=where_condition) + team_aliases: Final = ( + await unless_database_fails(team_alias_lookup, [], "its team aliases") + if user_api_key_has_admin_view(user_api_key_dict) + else await team_alias_lookup + ) team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index 90690111867..4ad17272fad 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -1,6 +1,7 @@ import csv import io from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from contextlib import contextmanager from dataclasses import dataclass, fields from itertools import chain from types import MappingProxyType, SimpleNamespace @@ -522,6 +523,14 @@ class _PrismaTable: return next((row for row in self._rows if _matches(row, where)), None) +class _FailingPrismaTable: + def __init__(self, failure: Exception) -> None: + self._failure: Final = failure + + async def find_many(self, **kwargs: object) -> tuple[object, ...]: + raise self._failure + + def _matches(row: object, where: Mapping[str, object]) -> bool: return all( getattr(row, field_name, None) in value["in"] @@ -531,7 +540,7 @@ def _matches(row: object, where: Mapping[str, object]) -> bool: ) -def _prisma_client() -> object: +def _prisma_client(**table_overrides: object) -> object: user: Final = LiteLLM_UserTable( user_id="user-a", user_email="user@example.test", @@ -554,7 +563,7 @@ def _prisma_client() -> object: team_alias="Team B", members_with_roles=[Member(user_id="user-b", role="user")], ) - db: Final = SimpleNamespace( + default_tables: Final = SimpleNamespace( litellm_usertable=_PrismaTable((user, other_user)), litellm_teamtable=_PrismaTable((team, other_team)), litellm_verificationtoken=_PrismaTable((SimpleNamespace(token="key-alpha", user_id="user-a"),)), @@ -583,13 +592,12 @@ def _prisma_client() -> object: ) ), ) + db: Final = SimpleNamespace(**{**vars(default_tables), **table_overrides}) return SimpleNamespace(db=db, writer_db=db) -@pytest.fixture -def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]: - repository: Final = _FakeRepository(_seeded_activity()) - prisma_client: Final = _prisma_client() +@contextmanager +def _daily_activity_test_client(repository: _FakeRepository, prisma_client: object) -> Iterator[TestClient]: app: Final = FastAPI() app.include_router(router) @@ -609,6 +617,13 @@ def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]: app.dependency_overrides[get_daily_activity_llm_router] = lambda: None app.dependency_overrides[user_api_key_auth] = resolve_auth with TestClient(app) as client: + yield client + + +@pytest.fixture +def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]: + repository: Final = _FakeRepository(_seeded_activity()) + with _daily_activity_test_client(repository, _prisma_client()) as client: yield client, repository @@ -1107,6 +1122,88 @@ def test_user_page_route_still_reports_a_lookup_bug_that_is_not_a_database_failu assert "not a database failure" in response.text +_ENTITY_ALIAS_LOOKUPS: Final = pytest.mark.parametrize( + ("prefix", "query_name", "entity_id", "alias_table", "alias"), + [ + ("/team", "team_ids", "team-a", "litellm_teamtable", {"team_alias": "Team A"}), + ("/organization", "organization_ids", "org-a", "litellm_organizationtable", {"organization_alias": "Org A"}), + ("/customer", "end_user_ids", "customer-a", "litellm_endusertable", {"alias": "Customer A"}), + ("/agent", "agent_ids", "agent-a", "litellm_agentstable", {"agent_name": "Agent A"}), + ], + ids=["team", "organization", "customer", "agent"], +) + + +def _aggregated_entity_metadata( + prefix: str, query_name: str, entity_id: str, prisma_client: object +) -> tuple[httpx.Response, tuple[Mapping[str, object], ...]]: + with _daily_activity_test_client(_FakeRepository(_seeded_activity()), prisma_client) as client: + response: Final = client.get(f"{prefix}/daily/activity/aggregated", params=_entity_params(query_name, entity_id)) + if response.status_code != 200: + return response, () + entities: Final = chain.from_iterable( + result["breakdown"]["entities"].values() for result in response.json()["results"] + ) + return response, tuple(entity["metadata"] for entity in entities) + + +@_LOOKUP_DATABASE_FAILURES +@_ENTITY_ALIAS_LOOKUPS +def test_entity_routes_keep_spend_and_blank_the_alias_when_the_alias_lookup_hits_a_database_failure( + prefix: str, + query_name: str, + entity_id: str, + alias_table: str, + alias: Mapping[str, object], + failure: Exception, +) -> None: + healthy, healthy_metadata = _aggregated_entity_metadata(prefix, query_name, entity_id, _prisma_client()) + failed, failed_metadata = _aggregated_entity_metadata( + prefix, query_name, entity_id, _prisma_client(**{alias_table: _FailingPrismaTable(failure)}) + ) + + assert failed.status_code == 200, failed.text + assert failed.json()["metadata"] == healthy.json()["metadata"] + assert healthy_metadata and set(map(str, healthy_metadata)) == {str(alias)} + assert failed_metadata and set(map(str, failed_metadata)) == {str({})} + + +@_ENTITY_ALIAS_LOOKUPS +def test_entity_routes_still_report_an_alias_lookup_bug_that_is_not_a_database_failure( + prefix: str, + query_name: str, + entity_id: str, + alias_table: str, + alias: Mapping[str, object], +) -> None: + failed, _ = _aggregated_entity_metadata( + prefix, + query_name, + entity_id, + _prisma_client(**{alias_table: _FailingPrismaTable(AttributeError("not a database failure"))}), + ) + + assert failed.status_code == 500, failed.text + assert "not a database failure" in failed.text + + +@_LOOKUP_DATABASE_FAILURES +def test_team_route_fails_closed_for_a_team_member_when_the_team_lookup_hits_a_database_failure( + failure: Exception, +) -> None: + repository: Final = _FakeRepository(_seeded_activity()) + prisma_client: Final = _prisma_client(litellm_teamtable=_FailingPrismaTable(failure)) + with _daily_activity_test_client(repository, prisma_client) as client: + response: Final = client.get( + "/team/daily/activity/aggregated", + params=_entity_params("team_ids", "team-a"), + headers={"x-user-role": LitellmUserRoles.INTERNAL_USER.value, "x-user-id": "user-a"}, + ) + + assert response.status_code == 500, response.text + repository.aggregated.assert_not_called() + + def test_user_page_route_forwards_offset_and_limit( daily_activity_client: tuple[TestClient, _FakeRepository], ) -> None: