fix(usage): keep spend when the team, organization, customer, or agent name lookup drops

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}

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)

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

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

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: