From 5d66f2d8ff55378ee01c0cb04f3193330dc09e83 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:28:42 -0800 Subject: [PATCH] fix(usage): keep spend when metadata lookups hit database failures (#50253) * fix(usage): tolerate entity metadata database failures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(usage): keep spend when metadata lookups hit database failures Treat transport errors and the query engine crash on a closed connection as database failures in the key, user and guardrail lookups of the aggregated, keys, users, model_top_keys and cache_leakage_keys routes, and keep prisma out of module-level imports so a proxy install without prisma still imports the usage endpoints. * fix(usage): fall back only on database outages, on every daily activity route Data-layer Prisma errors such as a rejected query or a missing table return 500 again, a key whose details failed to load reports key_exists null instead of false, and the non-aggregated daily activity routes keep spend when the key or entity metadata lookup hits an outage * fix(proxy): skip the key metadata lookup when the daily rows carry no keys * fix(usage): keep the guardrail breakdown fallback for any Prisma error A missing guardrail usage table or column made the Usage call return 500 after the outage-only narrowing. The guardrail lookup now falls back on outages and on any Prisma error, as it did before this PR * fix(usage): keep spend when the team, organization, customer, or agent name lookup drops --------- Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/endpoints.py | 10 +- .../common_daily_activity.py | 90 +++--- .../customer_endpoints.py | 10 +- .../daily_activity_routes.py | 15 +- .../organization_endpoints.py | 10 +- .../management_endpoints/team_endpoints.py | 8 +- ruff.toml | 5 +- .../test_common_daily_activity.py | 298 +++++++++++++++++- .../test_daily_activity_routes.py | 201 +++++++++++- 9 files changed, 579 insertions(+), 68 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/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index bbd6a1c9021..0ac3c06cbbb 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, replace from datetime import date, datetime, timedelta from itertools import groupby from types import MappingProxyType -from typing import Final, Literal, NoReturn, Protocol +from typing import Final, Literal, NoReturn, Protocol, TypeVar from fastapi import HTTPException, status from typing_extensions import ReadOnly, TypedDict, assert_never @@ -14,6 +14,7 @@ from litellm import constants from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY, SUBSCRIPTION_BACKED_PROVIDERS from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_details, recover_cli_session_key_metadata, @@ -56,6 +57,8 @@ from litellm.types.repositories.daily_activity import ( SpendLogsWindow, ) +_T: Final = TypeVar("_T") + @dataclass(frozen=True, slots=True) class ScopeDenied: @@ -188,9 +191,11 @@ class _AggregatedSpendData(TypedDict): totals: ReadOnly[SpendMetrics] -def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str | None) -> KeyMetadata: +def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow] | None, api_key: str | None) -> KeyMetadata: if not api_key: return KeyMetadata(key_alias="Unattributed", key_exists=False) + if api_key_metadata is None: + return KeyMetadata() meta: Final = api_key_metadata.get(api_key) return KeyMetadata( key_alias=meta.key_alias if meta is not None else None, @@ -306,7 +311,7 @@ def update_breakdown_metrics( record: DailySpendRecord, model_metadata: Mapping[str, dict[str, object]], provider_metadata: Mapping[str, dict[str, object]], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, entity_id_field: str | None = None, entity_metadata_field: Mapping[str, dict[str, object]] | None = None, ) -> BreakdownMetrics: @@ -626,7 +631,7 @@ async def get_api_key_metadata( def _aggregate_spend_records_sync( *, records: Sequence[DailySpendRecord], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, entity_id_field: str | None, entity_metadata_field: Mapping[str, dict[str, object]] | None, ) -> _AggregatedSpendData: @@ -685,22 +690,10 @@ async def _aggregate_spend_records( The per-row loop is offloaded to a worker thread via asyncio.to_thread so a large result set doesn't peg the event loop. """ - api_keys: Final[set[str]] = { - record.api_key for record in records if record.api_key and record.api_key != PTU_SENTINEL_API_KEY - } - - api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = ( - await repository.key_metadata( - frozenset(api_keys), spend_logs_window(frozenset(record.date for record in records)) - ) - if api_keys - else MappingProxyType({}) - ) - return await asyncio.to_thread( _aggregate_spend_records_sync, records=records, - api_key_metadata=api_key_metadata, + api_key_metadata=await _key_metadata_unless_database_fails(repository, records), entity_id_field=entity_id_field, entity_metadata_field=entity_metadata_field, ) @@ -783,7 +776,7 @@ def _record_to_spend_metrics(record: RollupMetricsRow) -> SpendMetrics: def _aggregate_grouping_sets_records_sync( *, records: Sequence[GroupingSetsRow], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, ) -> _AggregatedSpendData: """Build the response from rollup rows produced by the GROUPING SETS query. @@ -919,24 +912,40 @@ def _aggregate_grouping_sets_records_sync( return {"results": results, "totals": total_metrics} +async def unless_database_fails(lookup: Awaitable[_T], fallback: _T, left_out: str) -> _T: + try: + return await lookup + except Exception as exc: + if not PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): + raise + verbose_proxy_logger.warning("Returning daily activity without %s: %s", left_out, exc) + return fallback + + +async def _key_metadata_unless_database_fails( + repository: DailyActivityRepository, + rows: Sequence[GroupingSetsRow] | Sequence[EntityRollupRow] | Sequence[DailySpendRecord], +) -> Mapping[str, KeyMetadataRow] | None: + api_keys: Final = frozenset(row.api_key for row in rows if row.api_key and row.api_key != PTU_SENTINEL_API_KEY) + if not api_keys: + return MappingProxyType({}) + return await unless_database_fails( + repository.key_metadata(api_keys, spend_logs_window(frozenset(row.date for row in rows))), + None, + "its key metadata", + ) + + async def _aggregate_grouping_sets_records( *, repository: DailyActivityRepository, records: Sequence[GroupingSetsRow], ) -> _AggregatedSpendData: """Async wrapper: fetch api_key_metadata, then dispatch on a worker thread.""" - api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY} - - api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = ( - await repository.key_metadata(frozenset(api_keys), spend_logs_window(frozenset(r.date for r in records))) - if api_keys - else MappingProxyType({}) - ) - return await asyncio.to_thread( _aggregate_grouping_sets_records_sync, records=records, - api_key_metadata=api_key_metadata, + api_key_metadata=await _key_metadata_unless_database_fails(repository, records), ) @@ -991,7 +1000,13 @@ async def get_daily_activity( if resolve_entity_metadata is not None: resolved_entity_metadata = { **(entity_metadata_field or {}), - **(await resolve_entity_metadata(daily_spend_data)), + **( + await unless_database_fails( + resolve_entity_metadata(daily_spend_data), + MappingProxyType[str, dict[str, object]]({}), + "its entity metadata", + ) + ), } aggregated: Final = await _aggregate_spend_records( repository=repository, @@ -1038,7 +1053,7 @@ def _fold_entity_rollups_sync( *, results: Sequence[DailySpendData], entity_rows: Sequence[EntityRollupRow], - api_key_metadata: Mapping[str, KeyMetadataRow], + api_key_metadata: Mapping[str, KeyMetadataRow] | None, entity_metadata_field: Mapping[str, Mapping[str, object]] | None, ) -> None: """Write breakdown.entities onto the already-built per-day results.""" @@ -1107,7 +1122,7 @@ async def _guardrail_spend_or_empty( from prisma.errors import PrismaError try: - return await repository.guardrail_spend(scope) + return await unless_database_fails(repository.guardrail_spend(scope), (), "its guardrail breakdown") except PrismaError as exc: verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc) return () @@ -1117,7 +1132,9 @@ async def _user_entity_metadata( repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow] ) -> Mapping[str, Mapping[str, object]]: user_ids: Final = frozenset(row.entity_id for row in entity_rows if row.entity_id) - users: Final = await repository.user_metadata(user_ids) + users: Final = await unless_database_fails( + repository.user_metadata(user_ids), MappingProxyType({}), "its entity user metadata" + ) return MappingProxyType( { user_id: MappingProxyType({"user_email": user.user_email, "user_alias": user.user_alias}) @@ -1160,16 +1177,7 @@ async def get_daily_activity_aggregated( ) if aggregated_rows.entity_rows: entity_records: Final = aggregated_rows.entity_rows - entity_api_keys: Final = frozenset( - row.api_key for row in entity_records if row.api_key and row.api_key != PTU_SENTINEL_API_KEY - ) - entity_key_metadata: Final[Mapping[str, KeyMetadataRow]] = ( - await repository.key_metadata( - entity_api_keys, spend_logs_window(frozenset(row.date for row in entity_records)) - ) - if entity_api_keys - else MappingProxyType({}) - ) + entity_key_metadata: Final = await _key_metadata_unless_database_fails(repository, entity_records) user_entity_metadata: Final[Mapping[str, Mapping[str, object]]] = ( await _user_entity_metadata(repository, entity_records) if scope.table is DailyActivityTable.USER 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/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 1a3445aae52..b5af673585c 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -4,6 +4,7 @@ import json from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import asdict, fields, replace from datetime import date, datetime +from types import MappingProxyType from typing import Annotated, Final, Literal from fastapi import APIRouter, Depends, HTTPException, Query @@ -23,6 +24,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( raise_public, spend_logs_window, sum_spend_metrics, + unless_database_fails, ) from litellm.proxy.management_endpoints.daily_activity_scopes import ( AGENT_RESOLVER, @@ -149,9 +151,10 @@ async def _key_activity_rows( spend_window: Final[tuple[datetime, datetime] | None] = spend_logs_window( frozenset((resolved_scope.scope.start_date, resolved_scope.scope.end_date)) ) - metadata: Final[Mapping[str, KeyMetadataRow]] = await repository.key_metadata( - frozenset(row.api_key for row in rows if row.api_key), - spend_window, + metadata: Final[Mapping[str, KeyMetadataRow]] = await unless_database_fails( + repository.key_metadata(frozenset(row.api_key for row in rows if row.api_key), spend_window), + MappingProxyType({}), + "its key metadata", ) return [_key_activity_row(row, metadata) for row in rows] @@ -176,8 +179,10 @@ async def _user_activity_rows( repository: DailyActivityRepository, rows: Sequence[UserSpendRow], ) -> Sequence[UserActivityRow]: - metadata: Final[Mapping[str, UserMetadataRow]] = await repository.user_metadata( - frozenset(user_id for row in rows if (user_id := row.user_id) is not None) + metadata: Final[Mapping[str, UserMetadataRow]] = await unless_database_fails( + repository.user_metadata(frozenset(user_id for row in rows if (user_id := row.user_id) is not None)), + MappingProxyType({}), + "its user metadata", ) return [_user_activity_row(row, metadata) for row in rows] 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 cf2ae392677..59fa8a18ac7 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 @@ -7303,7 +7304,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/ruff.toml b/ruff.toml index 715e5551fad..037b15dbfe9 100644 --- a/ruff.toml +++ b/ruff.toml @@ -9,7 +9,7 @@ lint.extend-select = [ "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "SIM201", "SIM211", "SIM222", "TC005", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045", - "C404", "C419", + "C404", "C419", "TID253", ] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external @@ -35,6 +35,9 @@ format.exclude = ["**/enterprise/**"] lint.exclude = ["litellm/__init__.py", "litellm/proxy/example_config_yaml/*", "tests/*"] +[lint.flake8-tidy-imports] +banned-module-level-imports = ["prisma"] + [lint.per-file-ignores] "litellm/main.py" = ["F401"] "litellm/types/caching.py" = ["F401"] diff --git a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py index 5fe4cd2edcd..1cfe295b15d 100644 --- a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py @@ -1,11 +1,14 @@ from collections.abc import Mapping, Sequence +from dataclasses import asdict, fields from datetime import date, datetime -from types import SimpleNamespace +from types import FunctionType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from fastapi import HTTPException +from prisma.errors import DataError, PrismaError, TableNotFoundError import litellm.proxy.management_endpoints.common_daily_activity as common_daily_activity_module from litellm.constants import USAGE_TOP_API_KEYS_DEFAULT @@ -32,10 +35,301 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_E from litellm.proxy.utils import PrismaClient, hash_token from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, + KeyMetadata, SpendAnalyticsPaginatedResponse, SpendMetrics, ) -from litellm.types.repositories.daily_activity import GroupingSetsRow, KeyMetadataRow +from litellm.types.repositories.daily_activity import ( + AggregatedRows, + DailyActivityScope, + EntityRollupRow, + GroupingSetsRow, + KeyMetadataRow, +) + + +def _raise_missing_error_meta() -> None: + raise AttributeError("'NoneType' object has no attribute 'get'") + + +def _query_engine_crash_on_closed_connection() -> AttributeError: + raise_inside_query_engine: Final = FunctionType( + _raise_missing_error_meta.__code__, {"__name__": "prisma.engine.utils"} + ) + with pytest.raises(AttributeError) as crash: + raise_inside_query_engine() + return crash.value + + +DATABASE_FAILURES: Final = pytest.mark.parametrize( + "failure", + [ + DataError( + {"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server at `db:5432`"}} + ), + PrismaError("metadata unavailable"), + httpx.ReadTimeout("the query engine stopped answering"), + httpx.ConnectError("the query engine refused the connection"), + _query_engine_crash_on_closed_connection(), + ], + ids=[ + "unreachable-database", + "prisma-error", + "engine-read-timeout", + "engine-connect-error", + "engine-crash-on-closed-connection", + ], +) + +CODE_BUGS: Final = ( + pytest.param(TypeError("not a database failure"), id="type-error"), + pytest.param(AttributeError("not a database failure"), id="attribute-error"), +) +REJECTED_QUERIES: Final = ( + pytest.param( + DataError({"user_facing_error": {"error_code": "P2010", "message": "not a database failure"}}), + id="rejected-query", + ), + pytest.param( + TableNotFoundError( + { + "user_facing_error": { + "error_code": "P2021", + "message": "not a database failure", + "meta": {"table": "LiteLLM_VerificationToken"}, + } + } + ), + id="missing-table", + ), +) +NOT_DATABASE_OUTAGES: Final = pytest.mark.parametrize("bug", (*CODE_BUGS, *REJECTED_QUERIES)) + + +@pytest.mark.asyncio +@DATABASE_FAILURES +@pytest.mark.parametrize( + ("key_failure", "user_failure", "empty"), + [(False, False, False), (True, False, False), (False, True, False), (True, True, False), (False, False, True)], +) +async def test_entity_metadata_outages_preserve_daily_spend_and_other_metadata( + failure: Exception, key_failure: bool, user_failure: bool, empty: bool +) -> None: + prisma: Final = MagicMock() + key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None) + user: Final = SimpleNamespace(user_id="user", user_email="user@example.com", user_alias="User") + main_rows: Final = tuple( + asdict(_grouping_row(level, api_key=api_key, spend=2.0)) + for level, api_key in ((127, None), (63, None), (31, "key")) + ) + entity_rows: Final = tuple( + {**asdict(_grouping_row(63, api_key=api_key, spend=2.0)), "entity_id": "user", "api_key_rolled": rolled} + for rolled, api_key in ((1, None), (0, "key")) + ) + prisma.db.query_raw = AsyncMock(side_effect=[main_rows, entity_rows]) + prisma.db.litellm_verificationtoken.find_many = AsyncMock( + side_effect=[[key], failure if key_failure else ([] if empty else [key])] + ) + prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_usertable.find_many = AsyncMock( + side_effect=failure if user_failure else None, + return_value=[] if empty else [user], + ) + + result: Final = await get_daily_activity_aggregated( + prisma_client=prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + include_entity_breakdown=True, + ) + + day: Final = result.results[0] + entity: Final = day.breakdown.entities["user"] + assert result.metadata.total_spend == day.metrics.spend == entity.metrics.spend == 2.0 + assert day.breakdown.api_keys["key"].metadata.key_alias == "Alias" + assert entity.metadata == ( + {} if user_failure or empty else {"user_email": "user@example.com", "user_alias": "User"} + ) + assert entity.api_key_breakdown["key"].metrics.spend == 2.0 + assert entity.api_key_breakdown["key"].metadata == ( + KeyMetadata() if key_failure else KeyMetadata(key_alias=None if empty else "Alias", key_exists=not empty) + ) + + +def _usage_scope() -> DailyActivityScope: + return daily_activity_scope( + "litellm_dailyuserspend", "user_id", None, None, None, "2024-01-01", "2024-01-01", None, None + ) + + +def _entity_row() -> EntityRollupRow: + grouping: Final = asdict(_grouping_row(63, api_key="key", spend=2.0)) + return EntityRollupRow( + **{ + **{field.name: grouping[field.name] for field in fields(EntityRollupRow) if field.name in grouping}, + "entity_id": "user", + "api_key_rolled": 0, + } + ) + + +def _usage_repository() -> MagicMock: + repository: Final = MagicMock() + repository.aggregated = AsyncMock( + return_value=AggregatedRows( + grouping_rows=tuple( + _grouping_row(level, api_key=api_key, spend=2.0) + for level, api_key in ((127, None), (63, None), (31, "key")) + ), + entity_rows=(_entity_row(),), + distinct_api_keys=1, + ) + ) + repository.key_metadata = AsyncMock(return_value={}) + repository.user_metadata = AsyncMock(return_value={}) + repository.guardrail_spend = AsyncMock(return_value=()) + return repository + + +@pytest.mark.asyncio +@DATABASE_FAILURES +async def test_usage_keeps_spend_when_the_first_key_lookup_hits_a_database_failure(failure: Exception) -> None: + repository: Final = _usage_repository() + repository.key_metadata = AsyncMock(side_effect=failure) + + result: Final = await _get_daily_activity_aggregated(repository, _usage_scope()) + + day: Final = result.results[0] + assert result.metadata.total_spend == day.metrics.spend == 2.0 + assert day.breakdown.api_keys["key"].metrics.spend == 2.0 + assert day.breakdown.api_keys["key"].metadata == KeyMetadata() + + +@pytest.mark.asyncio +@DATABASE_FAILURES +async def test_usage_keeps_spend_when_the_guardrail_lookup_hits_a_database_failure(failure: Exception) -> None: + repository: Final = _usage_repository() + repository.guardrail_spend = AsyncMock(side_effect=failure) + + result: Final = await _get_daily_activity_aggregated(repository, _usage_scope(), include_guardrail_breakdown=True) + + assert result.metadata.total_spend == result.results[0].metrics.spend == 2.0 + assert result.results[0].breakdown.guardrails == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", REJECTED_QUERIES) +async def test_usage_keeps_spend_when_the_guardrail_table_rejects_the_query(failure: Exception) -> None: + repository: Final = _usage_repository() + repository.guardrail_spend = AsyncMock(side_effect=failure) + + result: Final = await _get_daily_activity_aggregated(repository, _usage_scope(), include_guardrail_breakdown=True) + + assert result.metadata.total_spend == result.results[0].metrics.spend == 2.0 + assert result.results[0].breakdown.guardrails == {} + + +async def _assert_usage_reports_the_bug(repository: MagicMock) -> None: + with pytest.raises(HTTPException) as error: + await _get_daily_activity_aggregated( + repository, _usage_scope(), include_entity_breakdown=True, include_guardrail_breakdown=True + ) + + assert error.value.status_code == 500 + assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["key_metadata", "user_metadata"]) +@NOT_DATABASE_OUTAGES +async def test_usage_still_reports_a_lookup_bug_that_is_not_a_database_failure(lookup: str, bug: Exception) -> None: + repository: Final = _usage_repository() + setattr(repository, lookup, AsyncMock(side_effect=bug)) + + await _assert_usage_reports_the_bug(repository) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bug", CODE_BUGS) +async def test_usage_still_reports_a_guardrail_lookup_bug_that_is_not_a_database_error(bug: Exception) -> None: + repository: Final = _usage_repository() + repository.guardrail_spend = AsyncMock(side_effect=bug) + + await _assert_usage_reports_the_bug(repository) + + +def _daily_user_spend_prisma(key_lookup: AsyncMock) -> MagicMock: + prisma: Final = MagicMock() + records: Final = [_daily_user_spend_record(user_id="user", api_key="key", spend=7.0)] + prisma.db.litellm_dailyuserspend.count = AsyncMock(return_value=len(records)) + prisma.db.litellm_dailyuserspend.find_many = AsyncMock(return_value=records) + prisma.db.litellm_verificationtoken.find_many = key_lookup + prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + return prisma + + +async def _daily_user_activity(prisma: MagicMock, user_emails: AsyncMock) -> SpendAnalyticsPaginatedResponse: + return await get_daily_activity( + prisma_client=prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + page=1, + page_size=1000, + resolve_entity_metadata=user_emails, + ) + + +@pytest.mark.asyncio +@DATABASE_FAILURES +@pytest.mark.parametrize(("key_failure", "email_failure"), [(True, False), (False, True)]) +async def test_daily_activity_keeps_spend_when_a_metadata_lookup_hits_a_database_failure( + failure: Exception, key_failure: bool, email_failure: bool +) -> None: + key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None) + prisma: Final = _daily_user_spend_prisma( + AsyncMock(side_effect=failure if key_failure else None, return_value=[key]) + ) + user_emails: Final = AsyncMock( + side_effect=failure if email_failure else None, return_value={"user": {"user_email": "user@example.com"}} + ) + + result: Final = await _daily_user_activity(prisma, user_emails) + + day: Final = result.results[0] + assert result.metadata.total_spend == day.metrics.spend == 7.0 + assert day.breakdown.api_keys["key"].metadata == ( + KeyMetadata() if key_failure else KeyMetadata(key_alias="Alias", key_exists=True) + ) + assert day.breakdown.entities["user"].metadata == ({} if email_failure else {"user_email": "user@example.com"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lookup", ["key metadata", "user emails"]) +@NOT_DATABASE_OUTAGES +async def test_daily_activity_still_reports_a_metadata_lookup_bug_that_is_not_a_database_failure( + lookup: str, bug: Exception +) -> None: + prisma: Final = _daily_user_spend_prisma( + AsyncMock(side_effect=bug if lookup == "key metadata" else None, return_value=[]) + ) + user_emails: Final = AsyncMock(side_effect=bug if lookup == "user emails" else None, return_value={}) + + with pytest.raises(HTTPException) as error: + await _daily_user_activity(prisma, user_emails) + + assert error.value.status_code == 500 + assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"} async def _run_aggregated_daily_activity( 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 d658e59b76c..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,15 +1,18 @@ 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 from typing import Final from unittest.mock import AsyncMock +import httpx import pytest from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from prisma.errors import DataError from litellm import constants from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY @@ -520,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"] @@ -529,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", @@ -552,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"),)), @@ -581,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) @@ -607,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 @@ -1015,6 +1032,178 @@ def test_user_page_route_ranks_users_by_spend_and_attaches_metadata( assert repository.user_metadata.call_args.args[0] == frozenset(("user-a", "user-b")) +_LOOKUP_DATABASE_FAILURES: Final = pytest.mark.parametrize( + "failure", + [ + DataError({"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server"}}), + httpx.ConnectError("the query engine refused the connection"), + ], + ids=["unreachable-database", "engine-connect-error"], +) +_KEY_ROWS_ROUTES: Final = pytest.mark.parametrize( + ("suffix", "extra_params"), + [ + ("keys", {}), + ("search", {"search": "key"}), + ("model_top_keys", {"model_group": "rare-group"}), + ("cache_leakage_keys", {}), + ], + ids=["keys", "search", "model_top_keys", "cache_leakage_keys"], +) + + +@_LOOKUP_DATABASE_FAILURES +@_KEY_ROWS_ROUTES +def test_key_routes_keep_spend_when_the_key_lookup_hits_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], + suffix: str, + extra_params: Mapping[str, str], + failure: Exception, +) -> None: + client, repository = daily_activity_client + repository.key_metadata.side_effect = failure + + response: Final = client.get( + f"/user/daily/activity/aggregated/{suffix}", + params={**_entity_params("user_id", "user-a"), **extra_params}, + ) + + assert response.status_code == 200, response.text + rows: Final = response.json()["api_keys"] + assert rows + assert all(row["metrics"]["spend"] > 0 for row in rows) + assert {row["metadata"]["key_alias"] for row in rows} == {None} + + +@_KEY_ROWS_ROUTES +def test_key_routes_still_report_a_lookup_bug_that_is_not_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], + suffix: str, + extra_params: Mapping[str, str], +) -> None: + client, repository = daily_activity_client + repository.key_metadata.side_effect = AttributeError("not a database failure") + + response: Final = client.get( + f"/user/daily/activity/aggregated/{suffix}", + params={**_entity_params("user_id", "user-a"), **extra_params}, + ) + + assert response.status_code == 500, response.text + assert "not a database failure" in response.text + + +@_LOOKUP_DATABASE_FAILURES +def test_user_page_route_keeps_spend_when_the_user_lookup_hits_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], + failure: Exception, +) -> None: + client, repository = daily_activity_client + repository.user_metadata.side_effect = failure + + response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS) + + assert response.status_code == 200, response.text + assert [(user["user_id"], user["spend"], user["user_email"]) for user in response.json()["users"]] == [ + ("user-b", 100.0, None), + ("user-a", 14.5, None), + ] + + +def test_user_page_route_still_reports_a_lookup_bug_that_is_not_a_database_failure( + daily_activity_client: tuple[TestClient, _FakeRepository], +) -> None: + client, repository = daily_activity_client + repository.user_metadata.side_effect = AttributeError("not a database failure") + + response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS) + + assert response.status_code == 500, response.text + 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: