From a12070b8b651a425cdfb4062dc0f7963de89f355 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 28 Sep 2026 21:26:40 +0000 Subject: [PATCH 1/2] Revert "feat(usage): search team keys beyond the top-N in the Team usage view (#42857)" --- litellm/proxy/_types.py | 5 - .../management_endpoints/team_endpoints.py | 109 --------- .../management_endpoints/team_endpoints.py | 24 -- .../endpointaudit/coverage_allowlist.txt | 1 - .../unbounded_in_baseline.txt | 6 - .../test_team_daily_activity_key_search.py | 128 ----------- .../management/test_team_daily_activity.py | 15 +- .../proxy/auth/test_route_checks.py | 49 ---- .../test_team_endpoints.py | 212 ------------------ .../components/EntityUsage/EntityUsage.tsx | 14 +- .../src/components/networking.tsx | 25 --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 60 +---- 12 files changed, 7 insertions(+), 641 deletions(-) delete mode 100644 tests/integration/spend/test_team_daily_activity_key_search.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3d8d701d15b5..3bc7cf54db23 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -310,7 +310,6 @@ class KeyManagementRoutes(str, enum.Enum): # team usage routes TEAM_DAILY_ACTIVITY = "/team/daily/activity" TEAM_DAILY_ACTIVITY_AGGREGATED = "/team/daily/activity/aggregated" - TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH = "/team/daily/activity/aggregated/search" # team spend-log viewing SPEND_LOGS = "/spend/logs" @@ -678,7 +677,6 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED.value, - KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH.value, KeyManagementRoutes.SPEND_LOGS.value, KeyManagementRoutes.SPEND_LOGS_V2.value, KeyManagementRoutes.KEY_RESET_SPEND.value, @@ -723,7 +721,6 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_bulk_update", "/team/daily/activity", "/team/daily/activity/aggregated", - "/team/daily/activity/aggregated/search", "/team/spend/by_user", # gateway request counts (SGR); deployment-wide, admin-only "/gateway/daily/activity", @@ -895,7 +892,6 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_update", "/team/daily/activity", "/team/daily/activity/aggregated", - "/team/daily/activity/aggregated/search", "/team/spend/by_user", "/team/{team_id}/members/me", # POST/GET the team's logging callbacks, and DELETE one of them. Every @@ -995,7 +991,6 @@ class LiteLLMRoutes(enum.Enum): "/user/daily/activity", "/team/daily/activity", "/team/daily/activity/aggregated", - "/team/daily/activity/aggregated/search", "/tag/daily/activity", "/tag/list", "/audit", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 493c83c730ae..6cec3e714ec6 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -40,7 +40,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.constants import USAGE_TOP_API_KEYS_LIMIT from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -197,7 +196,6 @@ from litellm.router import Router from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.common_daily_activity import ( - DailySpendMetadata, SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.team_endpoints import ( @@ -206,9 +204,7 @@ BulkUpdateTeamMemberPermissionsRequest, BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, - TeamIdSearchFilter, TeamIdSearchMatch, - TeamKeyActivitySearchWhere, TeamListItem, TeamListResponse, TeamMemberAddResult, @@ -6809,111 +6805,6 @@ async def get_team_daily_activity_aggregated( ) -def _team_key_search_where(*, search: str, scope: _TeamDailyActivityScope) -> TeamKeyActivitySearchWhere: - """Caller scoping lives inside the same Prisma where as the search term so `take` - never trims visible matches in favour of keys the caller is not allowed to see.""" - search_or: Final = ( - {"token": search}, # mutable-ok: prisma where clause leaf - {"key_alias": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf - {"user_id": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf - ) - own_keys: Final = tuple(scope.api_key_filter) if isinstance(scope.api_key_filter, list) else None - team_filter: Final[TeamIdSearchFilter | None] = ( - { # mutable-ok: prisma where clause leaf - "in": tuple(scope.team_ids), - "notIn": tuple(scope.exclude_team_ids), - } - if scope.team_ids is not None and scope.exclude_team_ids is not None - else {"in": tuple(scope.team_ids)} # mutable-ok: prisma where clause leaf - if scope.team_ids is not None - else {"notIn": tuple(scope.exclude_team_ids)} # mutable-ok: prisma where clause leaf - if scope.exclude_team_ids is not None - else None - ) - if team_filter is None and own_keys is None: - return {"OR": search_or} # mutable-ok: prisma where clause root - if team_filter is None and own_keys is not None: - return {"token": {"in": own_keys}, "OR": search_or} # mutable-ok: prisma where clause root - if team_filter is not None and own_keys is None: - return {"team_id": team_filter, "OR": search_or} # mutable-ok: prisma where clause root - assert team_filter is not None and own_keys is not None - return { # mutable-ok: prisma where clause root - "team_id": team_filter, - "token": {"in": own_keys}, # mutable-ok: prisma where clause leaf - "OR": search_or, - } - - -@router.get( - "/team/daily/activity/aggregated/search", - response_model=SpendAnalyticsPaginatedResponse, - tags=["team management"], # mutable-ok: FastAPI route tags shape -) -async def search_team_daily_activity_keys( - user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], - search: str = fastapi.Query( - ..., - min_length=1, - description="Exact token hash, or a case-insensitive substring of the key alias or owning user id", - ), - team_ids: str | None = None, - start_date: str | None = None, - end_date: str | None = None, - exclude_team_ids: str | None = None, - timezone: int | None = None, -) -> SpendAnalyticsPaginatedResponse: - """Aggregated daily team activity for the keys matching `search`, across every key the caller may - see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend.""" - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if prisma_client is None: - raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value) - - range_error: Final = _aggregated_date_range_error(start_date, end_date) - if range_error is not None: - raise _daily_activity_error(status_code=400, message=range_error) - - scope: Final = await _resolve_team_daily_activity_scope( - team_ids=team_ids, - exclude_team_ids=exclude_team_ids, - api_key=None, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - matched_keys: Final = await _tokens_db(prisma_client).find_many( - where=_team_key_search_where(search=search, scope=scope), - take=USAGE_TOP_API_KEYS_LIMIT, - order={"spend": "desc"}, # mutable-ok: prisma serializes order, keep it a plain dict - ) - tokens: Final = [key.token for key in matched_keys] # mutable-ok: get_daily_activity_aggregated takes list[str] - if not tokens: - return SpendAnalyticsPaginatedResponse( - results=[], # mutable-ok: response model field shape - metadata=DailySpendMetadata(api_key_limit=USAGE_TOP_API_KEYS_LIMIT, total_api_keys=0), - ) - - return await get_daily_activity_aggregated( - prisma_client=prisma_client, - table_name="litellm_dailyteamspend", - entity_id_field="team_id", - entity_id=scope.team_ids, - entity_metadata_field=scope.team_alias_metadata, - start_date=start_date, - end_date=end_date, - model=None, - api_key=tokens, - exclude_entity_ids=scope.exclude_team_ids, - timezone_offset_minutes=timezone, - include_entity_breakdown=True, - ) - - def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str: team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count)) user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else "" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index aac2703e918e..4524c47ec389 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,8 +1,6 @@ -from collections.abc import Mapping, Sequence from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.proxy._types import ( KeyManagementRoutes, @@ -13,32 +11,10 @@ MemberDeleteRequest, ) from litellm.proxy.common_utils.timezone_utils import budget_duration_error -from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains from litellm.types.proxy.management_endpoints.management_v1 import ResourceResponse TeamIdSearchMatch = Literal["exact", "prefix"] - -TeamIdSearchFilter = TypedDict( - "TeamIdSearchFilter", - { # mutable-ok: functional TypedDict field map - "in": NotRequired[ReadOnly[Sequence[str]]], - "notIn": NotRequired[ReadOnly[Sequence[str]]], - }, -) - - -class TeamKeyActivitySearchWhere(TypedDict): - """Prisma filter behind `/team/daily/activity/aggregated/search`: exact token hash, or key alias - or user id containing the term, case-insensitive, narrowed to the teams and keys the caller may see.""" - - team_id: NotRequired[ReadOnly[TeamIdSearchFilter]] - token: NotRequired[ReadOnly[Mapping[Literal["in"], Sequence[str]]]] - OR: ReadOnly[ - tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...] - ] - - MAX_BULK_TEAM_MEMBER_DELETES: Final = 500 MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES: Final = 500 diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index c2505977d879..25946eec0098 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -28,7 +28,6 @@ GET /tag/user-agent/per-user-analytics GET /tag/wau GET /team/daily/activity GET /team/daily/activity/aggregated -GET /team/daily/activity/aggregated/search GET /team/spend/by_user GET /team/spend/report GET /user/daily/activity diff --git a/tests/code_coverage_tests/unbounded_in_baseline.txt b/tests/code_coverage_tests/unbounded_in_baseline.txt index b1552d90a91a..17e8d32dde09 100644 --- a/tests/code_coverage_tests/unbounded_in_baseline.txt +++ b/tests/code_coverage_tests/unbounded_in_baseline.txt @@ -102,12 +102,6 @@ litellm/proxy/management_endpoints/team_endpoints.py _resolve_existing_member_us litellm/proxy/management_endpoints/team_endpoints.py _resolve_team_daily_activity_scope prisma team_id.in `list(team_ids_list)` 0 litellm/proxy/management_endpoints/team_endpoints.py _sweep_deleted_team_references prisma team_id.in `tuple(team_ids)` 0 litellm/proxy/management_endpoints/team_endpoints.py _sweep_deleted_team_references_tx prisma team_id.in `tuple(team_ids)` 0 -litellm/proxy/management_endpoints/team_endpoints.py _team_key_search_where prisma ?.in `tuple(scope.team_ids)` 0 -litellm/proxy/management_endpoints/team_endpoints.py _team_key_search_where prisma ?.in `tuple(scope.team_ids)` 1 -litellm/proxy/management_endpoints/team_endpoints.py _team_key_search_where prisma ?.notIn `tuple(scope.exclude_team_ids)` 0 -litellm/proxy/management_endpoints/team_endpoints.py _team_key_search_where prisma ?.notIn `tuple(scope.exclude_team_ids)` 1 -litellm/proxy/management_endpoints/team_endpoints.py _team_key_search_where prisma token.in `own_keys` 0 -litellm/proxy/management_endpoints/team_endpoints.py _team_key_search_where prisma token.in `own_keys` 1 litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(addressed_user_ids)` 0 litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(user_ids_to_delete)` 0 litellm/proxy/management_endpoints/team_endpoints.py _team_member_delete prisma user_id.in `sorted(user_ids_to_delete)` 1 diff --git a/tests/integration/spend/test_team_daily_activity_key_search.py b/tests/integration/spend/test_team_daily_activity_key_search.py deleted file mode 100644 index 2b0395cf4ba2..000000000000 --- a/tests/integration/spend/test_team_daily_activity_key_search.py +++ /dev/null @@ -1,128 +0,0 @@ -import uuid -from datetime import datetime, timedelta, timezone -from hashlib import sha256 -from typing import Final - -import pytest -from integration._support.client import Gateway, eventually, object_value -from integration._support.database import read_rows -from pydantic import JsonValue - -_SEARCH_PATH: Final = "/team/daily/activity/aggregated/search" - - -def _range_around_today() -> dict[str, str]: - today: Final = datetime.now(timezone.utc) - return { - "start_date": (today - timedelta(days=1)).strftime("%Y-%m-%d"), - "end_date": (today + timedelta(days=1)).strftime("%Y-%m-%d"), - "timezone": "0", - } - - -def _team_key_breakdown(body: dict[str, JsonValue], team: str) -> dict[str, JsonValue]: - results: Final = body["results"] - assert isinstance(results, list) and len(results) == 1, body - entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"]) - return object_value(object_value(entities[team])["api_key_breakdown"]) - - -def test_team_key_search_returns_only_the_matching_key_spend_by_alias_and_by_hash(gateway: Gateway) -> None: - with gateway.scenario() as scenario: - model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) - team: Final = scenario.team(models=[model]) - needle_alias: Final = f"needle-{uuid.uuid4().hex}" - needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias) - other: Final = scenario.key(team_id=team, models=[model], key_alias=f"other-{uuid.uuid4().hex}") - needle_digest: Final = sha256(needle.encode()).hexdigest() - other_digest: Final = sha256(other.encode()).hexdigest() - for key in (needle, other): - reply: Final = gateway.chat(model, key=key, text=f"key search {uuid.uuid4().hex}") - assert object_value(reply["usage"])["total_tokens"] == 40, reply - daily: Final = eventually( - lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), - lambda values: sorted(row["api_key"] for row in values) == sorted((needle_digest, other_digest)), - seconds=70, - ) - assert all(float(row["spend"]) == pytest.approx(0.06) for row in daily), daily - for search in (needle_alias.upper(), needle_digest): - response: Final = gateway.request( - "GET", _SEARCH_PATH, params={"team_ids": team, "search": search, **_range_around_today()} - ) - assert response.status_code == 200, response.text - body: Final = object_value(response.json()) - assert object_value(body["metadata"])["total_spend"] == pytest.approx(0.06), response.text - per_key: Final = _team_key_breakdown(body, team) - assert set(per_key) == {needle_digest}, response.text - assert object_value(object_value(per_key[needle_digest])["metrics"])["spend"] == pytest.approx(0.06) - - -def test_team_key_search_is_scoped_to_the_teams_the_caller_belongs_to(gateway: Gateway) -> None: - with gateway.scenario() as scenario: - model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) - team: Final = scenario.team(models=[model]) - needle_alias: Final = f"needle-{uuid.uuid4().hex}" - needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias) - needle_digest: Final = sha256(needle.encode()).hexdigest() - reply: Final = gateway.chat(model, key=needle, text=f"key search {uuid.uuid4().hex}") - assert object_value(reply["usage"])["total_tokens"] == 40, reply - eventually( - lambda: read_rows('SELECT api_key FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), - lambda values: [row["api_key"] for row in values] == [needle_digest], - seconds=70, - ) - outsider: Final = scenario.user(user_role="internal_user") - outsider_team: Final = scenario.team(models=[model], members_with_roles=[{"user_id": outsider, "role": "user"}]) - outsider_key: Final = scenario.key(user_id=outsider, team_id=outsider_team, models=[model]) - params: Final = {"search": needle_alias, **_range_around_today()} - admin_view: Final = gateway.request("GET", _SEARCH_PATH, params={"team_ids": team, **params}) - assert admin_view.status_code == 200, admin_view.text - assert set(_team_key_breakdown(object_value(admin_view.json()), team)) == {needle_digest}, admin_view.text - own_teams_view: Final = gateway.request("GET", _SEARCH_PATH, params=params, key=outsider_key) - assert own_teams_view.status_code == 200, own_teams_view.text - own_teams_body: Final = object_value(own_teams_view.json()) - assert own_teams_body["results"] == [], own_teams_view.text - assert object_value(own_teams_body["metadata"])["total_api_keys"] == 0, own_teams_view.text - foreign_team_view: Final = gateway.request( - "GET", _SEARCH_PATH, params={"team_ids": team, **params}, key=outsider_key - ) - assert foreign_team_view.status_code == 404, foreign_team_view.text - - -def test_team_key_search_excludes_teams_inside_the_where(gateway: Gateway) -> None: - """The dashboard always sends exclude_team_ids; a matching key in an excluded - team with higher spend must not consume a take slot nor appear in the result.""" - with gateway.scenario() as scenario: - model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) - team_keep: Final = scenario.team(models=[model]) - team_drop: Final = scenario.team(models=[model]) - shared_alias: Final = f"needle-{uuid.uuid4().hex}" - keep: Final = scenario.key(team_id=team_keep, models=[model], key_alias=f"{shared_alias}-keep") - drop: Final = scenario.key(team_id=team_drop, models=[model], key_alias=f"{shared_alias}-drop") - keep_digest: Final = sha256(keep.encode()).hexdigest() - drop_digest: Final = sha256(drop.encode()).hexdigest() - for _ in range(2): - reply: Final = gateway.chat(model, key=drop, text=f"key search {uuid.uuid4().hex}") - assert object_value(reply["usage"])["total_tokens"] == 40, reply - reply = gateway.chat(model, key=keep, text=f"key search {uuid.uuid4().hex}") - assert object_value(reply["usage"])["total_tokens"] == 40, reply - eventually( - lambda: read_rows( - 'SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id IN (%s, %s)', - (team_keep, team_drop), - ), - lambda values: sorted(row["api_key"] for row in values) == sorted((keep_digest, drop_digest)), - seconds=70, - ) - response: Final = gateway.request( - "GET", - _SEARCH_PATH, - params={"search": shared_alias, "exclude_team_ids": team_drop, **_range_around_today()}, - ) - assert response.status_code == 200, response.text - body: Final = object_value(response.json()) - results: Final = body["results"] - assert isinstance(results, list) and len(results) == 1, body - entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"]) - assert set(entities) == {team_keep}, response.text - assert set(_team_key_breakdown(body, team_keep)) == {keep_digest}, response.text diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/proxy_behavior/management/test_team_daily_activity.py index 9bbc8fdde290..d84cc4c94af5 100644 --- a/tests/proxy_behavior/management/test_team_daily_activity.py +++ b/tests/proxy_behavior/management/test_team_daily_activity.py @@ -5,9 +5,8 @@ pytestmark = pytest.mark.asyncio(loop_scope="session") -# GET /team/daily/activity, its /aggregated variant, and the key-search -# variant (same shared scope resolver, so the matrix must hold for all -# three). A proxy admin (admin view) sees +# GET /team/daily/activity and its /aggregated variant (same shared scope +# resolver, so the matrix must hold for both). A proxy admin (admin view) sees # activity for any team. A non-admin is scoped to user_info.teams: a bare query # defaults to its own teams (200), and an explicit team_ids filter naming a # team it does not belong to is 404 (the VERIA-43 fix). Org admins have no @@ -44,12 +43,8 @@ def _expected(actor: Actor, team: str) -> int: @pytest.mark.parametrize( "endpoint", - ( - "/team/daily/activity", - "/team/daily/activity/aggregated", - "/team/daily/activity/aggregated/search", - ), - ids=("paginated", "aggregated", "search"), + ("/team/daily/activity", "/team/daily/activity/aggregated"), + ids=("paginated", "aggregated"), ) @pytest.mark.parametrize( "actor,team,expected_status", @@ -59,7 +54,7 @@ def _expected(actor: Actor, team: str) -> int: async def test_team_daily_activity_matrix( actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world ): - query = _DATES + ("&search=x" if endpoint.endswith("/search") else "") + query = _DATES if team == "alpha": query += f"&team_ids={world.team_alpha_id}" elif team == "beta": diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index f76a02e83610..02f5894eb607 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3600,55 +3600,6 @@ def test_user_daily_activity_aggregated_not_covered_by_prefix_match(): ) -@pytest.mark.parametrize( - "route", - [ - "/team/daily/activity", - "/team/daily/activity/aggregated", - "/team/daily/activity/aggregated/search", - ], -) -@pytest.mark.parametrize( - "user_role", - [ - LitellmUserRoles.INTERNAL_USER.value, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, - ], -) -def test_team_daily_activity_routes_reachable_by_non_admin(route, user_role): - """The Team Usage dashboard calls all three team daily-activity routes, and - each handler self-scopes to the caller's teams and own keys - (_resolve_team_daily_activity_scope). self_managed_routes is the only list - granting them to a non-admin, and check_route_access is exact-match, so each - sub-path needs its own entry: dropping one 401s the dashboard before the - handler ever runs. - """ - user_obj = LiteLLM_UserTable( - user_id="test_user", - user_email="test@example.com", - user_role=user_role, - ) - valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) - request = MagicMock(spec=Request) - request.query_params = {} - - def outcome() -> str: - try: - RouteChecks.non_proxy_admin_allowed_routes_check( - user_obj=user_obj, - _user_role=user_role, - route=route, - request=request, - valid_token=valid_token, - request_data={}, - ) - except Exception as exc: - return f"denied: {exc}" - return "allowed" - - assert outcome() == "allowed" - - @pytest.mark.parametrize( "user_role", [ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 38241926f8e0..b066b3b80e62 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14646,218 +14646,6 @@ async def test_get_team_daily_activity_aggregated_rejects_bad_ranges( mock_aggregated.assert_not_called() -def _key_search_team_setup(mock_db_client, user_id: str, team_id: str): - mock_user_info = LiteLLM_UserTable( - user_id=user_id, - teams=[team_id], - max_budget=1000.0, - spend=0.0, - user_email="test@example.com", - user_role="internal_user", - ) - mock_team = MagicMock(spec=LiteLLM_TeamTable) - mock_team.team_id = team_id - mock_team.team_alias = "Test Team" - mock_team.members_with_roles = [Member(user_id=user_id, role="user")] - mock_team.model_dump.return_value = { - "team_id": team_id, - "team_alias": "Test Team", - "members_with_roles": [{"user_id": user_id, "role": "user"}], - } - mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) - return mock_user_info - - -@pytest.mark.asyncio -async def test_search_team_daily_activity_keys_scopes_where_before_take(mock_db_client): - """A member's search must put the team and own-key scoping inside the same - Prisma where as the term, because `take` trims rows before Python sees them: - scoped outside the where, the top-N slice could be spent entirely on keys - the caller is not allowed to see.""" - from litellm.constants import USAGE_TOP_API_KEYS_LIMIT - from litellm.proxy.management_endpoints.team_endpoints import ( - search_team_daily_activity_keys, - ) - - user_id = "test_user_123" - team_id = "test_team_456" - user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) - mock_user_info = _key_search_team_setup(mock_db_client, user_id, team_id) - - user_key_1 = MagicMock() - user_key_1.token = "user_key_1" - matched = MagicMock() - matched.token = "user_key_1" - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=[[user_key_1], [matched]]) - - with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - ) as mock_get_user_object: - mock_get_user_object.return_value = mock_user_info - - with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", - new_callable=AsyncMock, - ) as mock_aggregated: - mock_aggregated.return_value = MagicMock() - - await search_team_daily_activity_keys( - user_api_key_dict=user_api_key_dict, - search="Needle", - team_ids=team_id, - start_date="2024-01-01", - end_date="2024-01-31", - exclude_team_ids=None, - timezone=480, - ) - - token_calls = mock_db_client.db.litellm_verificationtoken.find_many.call_args_list - assert len(token_calls) == 2 - search_kwargs = token_calls[1][1] - assert search_kwargs["where"] == { - "team_id": {"in": (team_id,)}, - "token": {"in": ("user_key_1",)}, - "OR": ( - {"token": "Needle"}, - {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, - {"user_id": {"contains": "Needle", "mode": "insensitive"}}, - ), - } - assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT - assert search_kwargs["order"] == {"spend": "desc"} - - call_kwargs = mock_aggregated.call_args[1] - assert call_kwargs["api_key"] == ["user_key_1"] - assert call_kwargs["entity_id"] == [team_id] - assert call_kwargs["table_name"] == "litellm_dailyteamspend" - assert call_kwargs["include_entity_breakdown"] is True - assert call_kwargs["timezone_offset_minutes"] == 480 - assert call_kwargs["model"] is None - assert call_kwargs["entity_metadata_field"] == {team_id: {"team_alias": "Test Team"}} - - -@pytest.mark.asyncio -async def test_search_team_daily_activity_keys_admin_unscoped_where(mock_db_client): - """An admin's search has no caller scoping, so the where is the bare OR over - token, key alias and user id; every matched hash is passed through to the - aggregation.""" - from litellm.constants import USAGE_TOP_API_KEYS_LIMIT - from litellm.proxy.management_endpoints.team_endpoints import ( - search_team_daily_activity_keys, - ) - - match_1 = MagicMock() - match_1.token = "h1" - match_2 = MagicMock() - match_2.token = "h2" - mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[match_1, match_2]) - - with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", - new_callable=AsyncMock, - ) as mock_aggregated: - mock_aggregated.return_value = MagicMock() - - await search_team_daily_activity_keys( - user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), - search="Needle", - team_ids=None, - start_date="2024-01-01", - end_date="2024-01-31", - exclude_team_ids=None, - timezone=None, - ) - - search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] - assert search_kwargs["where"] == { - "OR": ( - {"token": "Needle"}, - {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, - {"user_id": {"contains": "Needle", "mode": "insensitive"}}, - ) - } - assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT - assert mock_aggregated.call_args[1]["api_key"] == ["h1", "h2"] - - -@pytest.mark.asyncio -async def test_search_team_daily_activity_keys_no_match_returns_empty_without_aggregating( - mock_db_client, -): - """A term matching no key still owes the caller the standard metadata shape - (api_key_limit, total_api_keys), and the aggregated query must not run.""" - from litellm.constants import USAGE_TOP_API_KEYS_LIMIT - from litellm.proxy.management_endpoints.team_endpoints import ( - search_team_daily_activity_keys, - ) - - mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - - with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", - new_callable=AsyncMock, - ) as mock_aggregated: - result = await search_team_daily_activity_keys( - user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), - search="Needle", - team_ids=None, - start_date="2024-01-01", - end_date="2024-01-31", - exclude_team_ids=None, - timezone=None, - ) - - assert result.results == [] - assert result.metadata.total_api_keys == 0 - assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT - mock_aggregated.assert_not_called() - - -@pytest.mark.asyncio -async def test_search_team_daily_activity_keys_excludes_teams_in_where(mock_db_client): - """The dashboard always sends exclude_team_ids=litellm-dashboard; if that - filter stayed out of the where, matching keys in excluded teams could fill - the take=N slice and push visible matches out.""" - from litellm.proxy.management_endpoints.team_endpoints import ( - search_team_daily_activity_keys, - ) - - matched = MagicMock() - matched.token = "h1" - mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) - mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[matched]) - - with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", - new_callable=AsyncMock, - ) as mock_aggregated: - mock_aggregated.return_value = MagicMock() - - await search_team_daily_activity_keys( - user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), - search="Needle", - team_ids=None, - start_date="2024-01-01", - end_date="2024-01-31", - exclude_team_ids="litellm-dashboard", - timezone=None, - ) - - search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] - assert search_kwargs["where"] == { - "team_id": {"notIn": ("litellm-dashboard",)}, - "OR": ( - {"token": "Needle"}, - {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, - {"user_id": {"contains": "Needle", "mode": "insensitive"}}, - ), - } - assert mock_aggregated.call_args[1]["exclude_entity_ids"] == ["litellm-dashboard"] - - def _wire_new_team_prisma(mock_db_client): mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.get_data = AsyncMock(return_value=None) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 16dc41c3ba8d..27460b211086 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -20,7 +20,7 @@ import type { ColumnDef } from "@tanstack/react-table"; import PaginationStatusAlerts from "@/components/shared/PaginationStatusAlerts"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; -import React, { type ReactNode, useCallback, useMemo, useState } from "react"; +import React, { type ReactNode, useMemo, useState } from "react"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; import UserDropdown from "@/components/common_components/UserDropdown"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; @@ -34,7 +34,6 @@ import { tagDailyActivityCall, teamDailyActivityAggregatedCall, teamDailyActivityCall, - teamDailyActivityKeySearchCall, userDailyActivityCall, } from "@/components/networking"; import { Logo } from "@/components/molecules/logo/Logo"; @@ -183,16 +182,6 @@ const EntityUsage: React.FC = ({ const modelBreakdownKey = modelViewType === "groups" ? "model_groups" : "models"; const modelMetrics = processActivityData(spendData, modelBreakdownKey, teams || []); const keyMetrics = processActivityData(spendData, "api_keys", teams || []); - const searchTeamKeys = useCallback( - (query: string) => { - if (!accessToken || !startTime || !endTime) return Promise.resolve({}); - const teamIds = Array.isArray(entityFilterArg) ? entityFilterArg : null; - return teamDailyActivityKeySearchCall(accessToken, startTime, endTime, query, teamIds).then((data) => - processActivityData(data, "api_keys", teams || []), - ); - }, - [accessToken, startTime, endTime, entityFilterArg, teams], - ); const agentMetrics = showAgentBreakdown ? processActivityData(agentSpendData, "entities", teams || []) : {}; const getAllTags = () => { @@ -678,7 +667,6 @@ const EntityUsage: React.FC = ({ keyMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} apiKeyTruncation={apiKeyTruncation} - searchKeys={entityType === "team" ? searchTeamKeys : undefined} /> ), }, diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index aeb416877889..82c446ce28c0 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1467,31 +1467,6 @@ export const teamDailyActivityAggregatedCall = async ( } }; -export const teamDailyActivityKeySearchCall = async ( - accessToken: string, - startTime: Date, - endTime: Date, - ...options: [search: string, teamIds?: string[] | null] -) => { - const [search, teamIds = null] = options; - try { - return await apiClient.get(`/team/daily/activity/aggregated/search`, { - accessToken, - query: { - start_date: formatDate(startTime), - end_date: formatDate(endTime), - timezone: new Date().getTimezoneOffset().toString(), - search, - team_ids: teamIds && teamIds.length > 0 ? teamIds.join(",") : undefined, - exclude_team_ids: "litellm-dashboard", - }, - }); - } catch (error) { - console.error("Failed to search team daily activity keys:", error); - throw error; - } -}; - export type TeamUserSpendResponse = components["schemas"]["TeamUserSpendResponse"]; export const teamSpendByUserCall = async ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 16325970ed19..065e25fc48ab 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15673,27 +15673,6 @@ export interface paths { patch?: never; trace?: never; }; - "/team/daily/activity/aggregated/search": { - parameters: { - query?: never; - header?: never; - path?: never; - cookie?: never; - }; - /** - * Search Team Daily Activity Keys - * @description Aggregated daily team activity for the keys matching `search`, across every key the caller may - * see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend. - */ - get: operations["search_team_daily_activity_keys_team_daily_activity_aggregated_search_get"]; - put?: never; - post?: never; - delete?: never; - options?: never; - head?: never; - patch?: never; - trace?: never; - }; "/team/delete": { parameters: { query?: never; @@ -31596,7 +31575,7 @@ export interface components { * @description Enum for key management routes * @enum {string} */ - KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2"; + KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/spend/logs" | "/spend/logs/v2"; /** * KeyManagementSystem * @enum {string} @@ -66952,43 +66931,6 @@ export interface operations { }; }; }; - search_team_daily_activity_keys_team_daily_activity_aggregated_search_get: { - parameters: { - query: { - /** @description Exact token hash, or a case-insensitive substring of the key alias or owning user id */ - search: string; - team_ids?: string | null; - start_date?: string | null; - end_date?: string | null; - exclude_team_ids?: string | null; - timezone?: number | null; - }; - header?: never; - path?: never; - cookie?: never; - }; - requestBody?: never; - responses: { - /** @description Successful Response */ - 200: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; - }; - }; - /** @description Validation Error */ - 422: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["HTTPValidationError"]; - }; - }; - }; - }; delete_team_team_delete_post: { parameters: { query?: never; From 1903cec22287b327bfec47d62ee15e49ac0e4efc Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 28 Sep 2026 14:40:13 -0700 Subject: [PATCH 2/2] revert: "feat(usage): search keys beyond the top-N usage subset (#42827)" (#43378) * Revert "feat(usage): search keys beyond the top-N usage subset (#42827)" * revert: "feat(proxy): add LiteLLM_DailyGlobalSpend key-free rollup for the usage dashboard (#41324)" (#43595) * Revert "Merge pull request #41324 from BerriAI/litellm_daily_global_spend_table" * Revert "Merge pull request #41293 from BerriAI/litellm_usage_key_free_aggregate_split" (#43596) Co-authored-by: yassin --------- Co-authored-by: yassin Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 4 - litellm/proxy/_lazy_openapi_snapshot.json | 24 - litellm/proxy/_types.py | 2 - .../common_daily_activity.py | 288 ++++------ .../internal_user_endpoints.py | 153 +---- litellm/proxy/proxy_server.py | 43 -- .../daily_global_spend_rollup.py | 289 ---------- .../common_daily_activity.py | 10 - .../internal_user_endpoints.py | 12 +- .../endpointaudit/coverage_allowlist.txt | 1 - .../proxy/auth/test_route_checks.py | 1 - .../test_common_daily_activity.py | 348 ++---------- .../test_internal_user_endpoints.py | 166 ------ .../proxy/proxy_server/test_lifecycle.py | 51 -- .../test_daily_global_spend_rollup.py | 532 ------------------ .../_components/CacheLeakageCard.test.tsx | 24 - .../_components/CacheLeakageCard.tsx | 9 +- .../useDailyActivityRange.test.tsx | 17 +- .../_components/useDailyActivityRange.ts | 3 - .../EntityUsage/EntityUsage.test.tsx | 17 - .../components/EntityUsage/EntityUsage.tsx | 15 +- .../_components/components/UsagePageView.tsx | 22 +- .../exportBlockedReason.test.ts | 26 +- .../EntityUsageExport/exportBlockedReason.ts | 20 +- .../components/KeyActivityPanel.test.tsx | 74 --- .../UsagePage/components/KeyActivityPanel.tsx | 87 +-- .../src/components/networking.tsx | 30 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 75 --- 28 files changed, 181 insertions(+), 2162 deletions(-) delete mode 100644 litellm/proxy/spend_tracking/daily_global_spend_rollup.py delete mode 100644 tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py diff --git a/litellm/constants.py b/litellm/constants.py index 5675a7e1ea5f..fd9812e2cf59 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2140,16 +2140,12 @@ PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__" PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job" PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900 -USAGE_TOP_API_KEYS_LIMIT: Final[int] = int(os.getenv("USAGE_TOP_API_KEYS_LIMIT", "100")) # Furthest back the catch-up pass looks for unpriced PTU days when a deployment # declares no ptu_effective_from, bounding the scan for an open-ended window. PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90 # Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide # expiry cannot produce an alert too large for the channel delivering it. PTU_LAPSED_ALERT_LIMIT: Final[int] = 10 -DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID: Final[str] = "daily_global_spend_reconcile_job" -DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS: Final[int] = 3600 -DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM: Final[str] = "daily_global_spend_reconciled_through" SPEND_CAPTURE_RATE_CHECK_JOB_ID: Final[str] = "spend_capture_rate_check_job" SPEND_CAPTURE_RATE_CHECK_LOCK_TTL_SECONDS: Final[int] = 900 SPEND_CAPTURE_RATE_MAX_RANGE_DAYS: Final[int] = 180 diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index db57aa4f0460..c4ff4b415dd4 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3313,18 +3313,6 @@ }, "DailySpendMetadata": { "properties": { - "api_key_limit": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "description": "When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.", - "title": "Api Key Limit" - }, "has_more": { "default": false, "title": "Has More", @@ -3335,18 +3323,6 @@ "title": "Page", "type": "integer" }, - "total_api_keys": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "description": "Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys.", - "title": "Total Api Keys" - }, "total_api_requests": { "default": 0, "title": "Total Api Requests", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3bc7cf54db23..d9fb053035b1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -703,7 +703,6 @@ class LiteLLMRoutes(enum.Enum): "/user/list", "/user/daily/activity", "/user/daily/activity/aggregated", - "/user/daily/activity/aggregated/search", # team "/team/new", "/team/update", @@ -907,7 +906,6 @@ class LiteLLMRoutes(enum.Enum): "/model/delete", "/user/daily/activity", "/user/daily/activity/aggregated", - "/user/daily/activity/aggregated/search", # Endpoint restricts results to organizations the caller is ORG_ADMIN # of; a caller who administers none gets an empty result set. "/organization/daily/activity", diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 086e8253f313..cecf3e50f5cc 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -9,9 +9,8 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT +from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy._types import CommonProxyErrors -from litellm.proxy.spend_tracking.daily_global_spend_rollup import GLOBAL_SPEND_TABLE_NAME, reconciled_through from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_details, recover_cli_session_key_metadata, @@ -151,9 +150,15 @@ class _AggregatedSpendData(TypedDict): totals: SpendMetrics -class _RollupMetricsRow(SimpleNamespace): +class _GroupingSetsRow(SimpleNamespace): date: str api_key: str | None + model: str | None + model_group: str | None + custom_llm_provider: str | None + mcp_namespaced_tool_name: str | None + endpoint: str | None + group_level: int spend: float | None prompt_tokens: int | None completion_tokens: int | None @@ -171,46 +176,12 @@ class _RollupMetricsRow(SimpleNamespace): timed_requests: int | None -class _GroupingSetsRow(_RollupMetricsRow): - model: str | None - model_group: str | None - custom_llm_provider: str | None - mcp_namespaced_tool_name: str | None - endpoint: str | None - group_level: int - distinct_api_keys: int | None - - -class _EntityRollupRow(_RollupMetricsRow): +class _EntityRollupRow(_GroupingSetsRow): entity_id: str | None api_key_rolled: int -class _AggregatedQueryKwargs(TypedDict): - table_name: ReadOnly[str] - entity_id_field: ReadOnly[str] - entity_id: ReadOnly[str | list[str] | None] - start_date: ReadOnly[str] - end_date: ReadOnly[str] - model: ReadOnly[str | None] - api_key: ReadOnly[str | list[str] | None] - exclude_entity_ids: ReadOnly[list[str] | None] - timezone_offset_minutes: ReadOnly[int | None] - include_current_utc_day: ReadOnly[bool] - - -_SqlQuery = tuple[str, list[str]] - - -async def _query_raw_optional( - prisma_client: PrismaClient, query: _SqlQuery | None -) -> list[dict[str, object]] | None: # mutable-ok: prisma query_raw return shape - if query is None: - return None - return await prisma_client.db.query_raw(query[0], *query[1]) - - -def _reported_flat_cost(record: DailySpendRecord | _RollupMetricsRow) -> float: +def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float: """Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled. Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost`` @@ -734,89 +705,6 @@ def _ptu_flat_cost_select(table_name: str) -> str: return "0::float AS ptu_flat_cost" -def _rollup_metric_select(table_name: str) -> str: - return f""" - SUM(spend)::float AS spend, - {_ptu_flat_cost_select(table_name)}, - SUM(prompt_tokens)::bigint AS prompt_tokens, - SUM(completion_tokens)::bigint AS completion_tokens, - SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens, - SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens, - SUM(compression_saved_tokens)::bigint AS compression_saved_tokens, - SUM(compression_savings_spend)::float AS compression_savings_spend, - SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend, - SUM(gateway_injected_caching_savings_spend)::float AS gateway_injected_caching_savings_spend, - SUM(autorouter_savings_spend)::float AS autorouter_savings_spend, - SUM(api_requests)::bigint AS api_requests, - SUM(successful_requests)::bigint AS successful_requests, - SUM(failed_requests)::bigint AS failed_requests, - SUM(total_response_time_ms)::bigint AS total_response_time_ms, - SUM(timed_requests)::bigint AS timed_requests""" - - -_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" - - -_KEY_FREE_SOURCE_COLUMNS: Final = ( - "date", - "model", - "model_group", - "custom_llm_provider", - "mcp_namespaced_tool_name", - "endpoint", - "spend", - "prompt_tokens", - "completion_tokens", - "cache_read_input_tokens", - "cache_creation_input_tokens", - "compression_saved_tokens", - "compression_savings_spend", - "prompt_caching_savings_spend", - "gateway_injected_caching_savings_spend", - "autorouter_savings_spend", - "api_requests", - "successful_requests", - "failed_requests", - "total_response_time_ms", - "timed_requests", -) - - -async def global_rollup_reconciled_through(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: - """The last day ``LiteLLM_DailyGlobalSpend`` can answer the key-free arm for, or None to - read it all from the per-key table. - - Only an unfiltered read of the user table sums to the same rows as the global table. The - marker read is served from the config cache, so this is not a database round trip per request. - """ - if query["table_name"] != "litellm_dailyuserspend": - return None - if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]: - return None - try: - return await reconciled_through(prisma_client) - except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read - verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc) - return None - - -def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None) -> str: - """The relation the key-free arm aggregates: the per-key table alone, or the global rollup - for days through the marker plus the per-key table for the days still open after it.""" - if marker_param is None: - return f'"{pg_table}"\n WHERE {where_clause}' - columns: Final = ", ".join(_KEY_FREE_SOURCE_COLUMNS) - return f"""( - SELECT {columns} - FROM "{GLOBAL_SPEND_TABLE_NAME}" - WHERE {where_clause} AND date <= {marker_param} - UNION ALL - SELECT {columns} - FROM "{pg_table}" - WHERE {where_clause} AND date > {marker_param} - ) AS key_free_source""" - - def _build_aggregated_sql_query( *, table_name: str, @@ -829,9 +717,13 @@ def _build_aggregated_sql_query( exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path timezone_offset_minutes: int | None = None, include_current_utc_day: bool = False, - global_rollup_through: str | None = None, ) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params - """Build the GROUPING SETS query for aggregated daily activity. + """Build a parameterized SQL GROUP BY query for aggregated daily activity. + + Groups by (date, api_key, model, model_group, custom_llm_provider, + mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns. + The entity_id column is intentionally omitted from GROUP BY to collapse + rows across entities — this is where the biggest row reduction comes from. Returns: Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw(). @@ -844,7 +736,7 @@ def _build_aggregated_sql_query( start_date, end_date, timezone_offset_minutes, include_current_utc_day ) - where_clause, where_params = _build_aggregated_where_clause( + where_clause, sql_params = _build_aggregated_where_clause( entity_id_field=entity_id_field, entity_id=entity_id, adjusted_start=adjusted_start, @@ -853,73 +745,67 @@ def _build_aggregated_sql_query( api_key=api_key, exclude_entity_ids=exclude_entity_ids, ) - sentinel_param: Final = f"${len(where_params) + 1}" - marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}" - metric_select: Final = _rollup_metric_select(table_name) + # Postgres computes every rollup level the response needs — per-date + # totals, per-(date, model), per-(date, model, api_key), per-provider, + # etc. — in a single pass via GROUPING SETS. The GROUPING() bitmask + # encodes which level a row belongs to so Python can dispatch rows + # straight into their buckets without re-summing. The leaf grouping + # is omitted on purpose: nothing in the response shape needs it once + # all the rollups are present. + # # TODO: drop the successful_requests/failed_requests aggregates (and the # total_successful_requests metadata they feed) once the admin UI reads SGR # only from LiteLLM_DailyGatewayRequests. The remaining spend, token and # api_requests rollups are still served from here. sql_query: Final = f""" - (SELECT - date, - NULL::text AS api_key, - model, - {_MODEL_GROUP_EXPR} AS model_group, - custom_llm_provider, - mcp_namespaced_tool_name, - endpoint, - (GROUPING(date) << 6) | {_API_KEY_ROLLED_UP_BIT} - | GROUPING(model, {_MODEL_GROUP_EXPR}, - custom_llm_provider, mcp_namespaced_tool_name, - endpoint) AS group_level, - NULL::bigint AS distinct_api_keys,{metric_select} - FROM {_key_free_source(pg_table, where_clause, marker_param)} - GROUP BY GROUPING SETS ( - (date), - (date, model), - (date, {_MODEL_GROUP_EXPR}), - (date, custom_llm_provider), - (date, mcp_namespaced_tool_name), - (date, endpoint), - () - )) - UNION ALL - (WITH top_api_keys AS ( - SELECT api_key, COUNT(*) OVER () AS distinct_api_keys - FROM "{pg_table}" - WHERE {where_clause} AND api_key <> {sentinel_param} - GROUP BY api_key - ORDER BY SUM(spend) DESC, api_key - LIMIT {USAGE_TOP_API_KEYS_LIMIT} - ) SELECT date, api_key, model, - {_MODEL_GROUP_EXPR} AS model_group, + COALESCE(NULLIF(model_group, ''), model) AS model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint, - GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR}, + GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model), custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level, - MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,{metric_select} - FROM "{pg_table}" JOIN top_api_keys USING (api_key) + SUM(spend)::float AS spend, + {_ptu_flat_cost_select(table_name)}, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens, + SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens, + SUM(compression_saved_tokens)::bigint AS compression_saved_tokens, + SUM(compression_savings_spend)::float AS compression_savings_spend, + SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend, + SUM(gateway_injected_caching_savings_spend)::float AS gateway_injected_caching_savings_spend, + SUM(autorouter_savings_spend)::float AS autorouter_savings_spend, + SUM(api_requests)::bigint AS api_requests, + SUM(successful_requests)::bigint AS successful_requests, + SUM(failed_requests)::bigint AS failed_requests, + SUM(total_response_time_ms)::bigint AS total_response_time_ms, + SUM(timed_requests)::bigint AS timed_requests + FROM "{pg_table}" WHERE {where_clause} GROUP BY GROUPING SETS ( + (date), (date, api_key), + (date, model), (date, model, api_key), - (date, {_MODEL_GROUP_EXPR}, api_key), + (date, COALESCE(NULLIF(model_group, ''), model)), + (date, COALESCE(NULLIF(model_group, ''), model), api_key), + (date, custom_llm_provider), (date, custom_llm_provider, api_key), + (date, mcp_namespaced_tool_name), (date, mcp_namespaced_tool_name, api_key), - (date, endpoint, api_key) - )) + (date, endpoint), + (date, endpoint, api_key), + () + ) """ - marker_params: Final = () if global_rollup_through is None else (global_rollup_through,) - return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params] + return sql_query, sql_params def _build_entity_rollup_sql_query( @@ -964,7 +850,23 @@ def _build_entity_rollup_sql_query( "{entity_id_field}" AS entity_id, date, api_key, - GROUPING(api_key) AS api_key_rolled,{_rollup_metric_select(table_name)} + GROUPING(api_key) AS api_key_rolled, + SUM(spend)::float AS spend, + {_ptu_flat_cost_select(table_name)}, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens, + SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens, + SUM(compression_saved_tokens)::bigint AS compression_saved_tokens, + SUM(compression_savings_spend)::float AS compression_savings_spend, + SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend, + SUM(gateway_injected_caching_savings_spend)::float AS gateway_injected_caching_savings_spend, + SUM(autorouter_savings_spend)::float AS autorouter_savings_spend, + SUM(api_requests)::bigint AS api_requests, + SUM(successful_requests)::bigint AS successful_requests, + SUM(failed_requests)::bigint AS failed_requests, + SUM(total_response_time_ms)::bigint AS total_response_time_ms, + SUM(timed_requests)::bigint AS timed_requests FROM "{pg_table}" WHERE {where_clause} GROUP BY GROUPING SETS ( @@ -1066,7 +968,6 @@ async def _aggregate_spend_records( # current grouping set's key), 0 when the column is part of the key. _GROUP_GRAND_TOTAL: Final = 127 # 0b1111111 — all rolled up _GROUP_DATE: Final = 63 # 0b0111111 — only date kept -_API_KEY_ROLLED_UP_BIT: Final = 32 # 0b0100000 _GROUP_DATE_API_KEY: Final = 31 # 0b0011111 _GROUP_DATE_MODEL: Final = 47 # 0b0101111 _GROUP_DATE_MODEL_API_KEY: Final = 15 # 0b0001111 @@ -1080,7 +981,7 @@ async def _aggregate_spend_records( _GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110 -def _record_to_spend_metrics(record: _RollupMetricsRow) -> SpendMetrics: +def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: """Build a SpendMetrics directly from one already-aggregated rollup row. SUM() over zero rows is SQL NULL, so rollup rows (notably the grand-total @@ -1436,6 +1337,10 @@ async def get_daily_activity_aggregated( ) -> SpendAnalyticsPaginatedResponse: """Aggregated variant that returns the full result set (no pagination). + Uses SQL GROUP BY to aggregate rows in the database rather than fetching + all individual rows into Python. This collapses rows across entities + (users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows. + include_entity_breakdown runs a small companion rollup query and folds `breakdown.entities` onto the response, as entity-scoped views like Team Usage need. @@ -1454,7 +1359,7 @@ async def get_daily_activity_aggregated( ) try: - query_kwargs: Final = _AggregatedQueryKwargs( + sql_query, sql_params = _build_aggregated_sql_query( table_name=table_name, entity_id_field=entity_id_field, entity_id=entity_id, @@ -1466,19 +1371,36 @@ async def get_daily_activity_aggregated( timezone_offset_minutes=timezone_offset_minutes, include_current_utc_day=include_current_utc_day, ) - sql_query, sql_params = _build_aggregated_sql_query( - **query_kwargs, - global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs), + + entity_query: Final = ( + _build_entity_rollup_sql_query( + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + start_date=start_date, + end_date=end_date, + model=model, + api_key=api_key, + exclude_entity_ids=exclude_entity_ids, + timezone_offset_minutes=timezone_offset_minutes, + include_current_utc_day=include_current_utc_day, + ) + if include_entity_breakdown + else None ) - entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None - raw_rows, raw_entity_rows = await asyncio.gather( - prisma_client.db.query_raw(sql_query, *sql_params), - _query_raw_optional(prisma_client, entity_query), + # Execute the GROUPING SETS query (one row per rollup level), alongside + # the per-entity companion rollup when the caller wants entities. + raw_rows, raw_entity_rows = ( + await asyncio.gather( + prisma_client.db.query_raw(sql_query, *sql_params), + prisma_client.db.query_raw(entity_query[0], *entity_query[1]), + ) + if entity_query is not None + else (await prisma_client.db.query_raw(sql_query, *sql_params), None) ) - records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or ())] - total_api_keys: Final = next((r.distinct_api_keys for r in records if r.distinct_api_keys is not None), 0) + records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or [])] # The grouping-sets dispatcher places each row directly in its bucket # using the row's GROUPING() bitmask. No Python-side summing needed. @@ -1532,8 +1454,6 @@ async def get_daily_activity_aggregated( page=1, total_pages=1, has_more=False, - api_key_limit=USAGE_TOP_API_KEYS_LIMIT, - total_api_keys=total_api_keys, ), ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 7b25348aa53c..59d8dd821d85 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -28,7 +28,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.constants import USAGE_TOP_API_KEYS_LIMIT from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( @@ -88,13 +87,11 @@ VerificationTokenRepository, ) from litellm.types.proxy.management_endpoints.common_daily_activity import ( - DailySpendMetadata, SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, - KeyActivitySearchWhere, UserListResponse, UserSearchWhere, UserUpdateResult, @@ -2994,27 +2991,6 @@ async def get_user_daily_activity( ) -def _resolve_user_daily_activity_entity_id( - user_api_key_dict: UserAPIKeyAuth, - user_id: str | None, -) -> str | None: - is_admin: Final = _user_has_admin_view(user_api_key_dict) - - if is_admin: - return user_id - - caller_user_id: Final = require_caller_user_id_for_non_admin(user_api_key_dict) - effective_user_id: Final = user_id if user_id is not None else caller_user_id - if effective_user_id != caller_user_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ # mutable-ok: FastAPI detail payload shape - "error": "Non-admin users can only view their own spend data." - }, - ) - return effective_user_id - - @router.get( "/user/daily/activity/aggregated", tags=["Budget & Spend Tracking", "Internal User management"], @@ -3081,7 +3057,20 @@ async def get_user_daily_activity_aggregated( ) try: - entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id) + is_admin: Final = _user_has_admin_view(user_api_key_dict) + + if is_admin: + entity_id = user_id # None means global view, otherwise filter by user + else: + caller_user_id: Final = require_caller_user_id_for_non_admin(user_api_key_dict) + if user_id is None: + user_id = caller_user_id + if user_id != caller_user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Non-admin users can only view their own spend data."}, + ) + entity_id = user_id return await get_daily_activity_aggregated( prisma_client=prisma_client, @@ -3105,117 +3094,3 @@ async def get_user_daily_activity_aggregated( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Failed to fetch analytics: {e}"}, ) - - -@router.get( - "/user/daily/activity/aggregated/search", - tags=["Budget & Spend Tracking", "Internal User management"], # mutable-ok: FastAPI route tags shape - dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route dependencies shape - response_model=SpendAnalyticsPaginatedResponse, -) -@management_endpoint_wrapper -async def search_user_daily_activity_keys( - search: str = fastapi.Query( - ..., - min_length=1, - description="Matches keys whose hash equals the value, or whose key alias or user ID contains it (case-insensitive)", - ), - start_date: str | None = fastapi.Query( - default=None, - description="Start date in YYYY-MM-DD format", - ), - end_date: str | None = fastapi.Query( - default=None, - description="End date in YYYY-MM-DD format", - ), - user_id: str | None = fastapi.Query( - default=None, - description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", - ), - timezone: int | None = fastapi.Query( - default=None, - description="Timezone offset in minutes from UTC (e.g., 480 for PST). " - "Matches JavaScript's Date.getTimezoneOffset() convention.", - ), - include_current_utc_day: bool = fastapi.Query( - default=False, - description="When the range ends on the caller's current local day, extend it to " - "today's UTC bucket so spend written after the caller's local midnight (in UTC " - "terms) is included. Requires the timezone parameter. Historical ranges are " - "never extended.", - ), - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection -) -> SpendAnalyticsPaginatedResponse: - """ - Search verification tokens by exact token hash or by a case-insensitive substring of - the key alias or owning user ID, then return the aggregated daily activity for the - matches. Lets the Usage page surface keys that fell outside the top-spend subset - the aggregated endpoint loads. - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={ # mutable-ok: FastAPI detail payload shape - "error": CommonProxyErrors.db_not_connected_error.value - }, - ) - - if start_date is None or end_date is None: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={"error": "Please provide start_date and end_date"}, # mutable-ok: FastAPI detail payload shape - ) - - try: - entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id) - - search_or: Final = ( - {"token": search}, # mutable-ok: prisma serializes where clauses, keep plain dicts - {"key_alias": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf - {"user_id": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf - ) - where: Final[KeyActivitySearchWhere] = ( - {"OR": search_or} # mutable-ok: prisma where clause root - if entity_id is None - else {"user_id": entity_id, "OR": search_or} # mutable-ok: prisma where clause root - ) - matched_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many( - where=where, - take=USAGE_TOP_API_KEYS_LIMIT, - order={"spend": "desc"}, # mutable-ok: prisma serializes order, keep it a plain dict - ) - tokens: Final = [key.token for key in matched_keys] # mutable-ok: api_key filter union expects a list - - if not tokens: - return SpendAnalyticsPaginatedResponse( - results=[], # mutable-ok: response model field shape - metadata=DailySpendMetadata( - api_key_limit=USAGE_TOP_API_KEYS_LIMIT, - total_api_keys=0, - ), - ) - - return await get_daily_activity_aggregated( - prisma_client=prisma_client, - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id=entity_id, - entity_metadata_field=None, - start_date=start_date, - end_date=end_date, - model=None, - api_key=tokens, - timezone_offset_minutes=timezone, - include_current_utc_day=include_current_utc_day, - ) - - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.exception("/user/daily/activity/aggregated/search: Exception occured - %s", e) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Failed to fetch analytics: {e}"}, # mutable-ok: FastAPI detail payload shape - ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ae559f30857e..7cd116c23d76 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -278,7 +278,6 @@ def generate_feedback_box(): APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, CLI_SSO_SESSION_TTL_SECONDS, - DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -761,9 +760,6 @@ def generate_feedback_box(): get_budget_window_start, release_unbound_budget_reservation, ) -from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( - run_scheduled_daily_global_spend_reconcile, -) from litellm.proxy.spend_tracking.spend_capture_rate import ( run_scheduled_spend_capture_rate_check, ) @@ -10571,12 +10567,6 @@ async def initialize_scheduled_background_jobs( await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) - cls._initialize_daily_global_spend_reconcile_job( - scheduler=scheduler, - proxy_logging_obj=proxy_logging_obj, - prisma_client=prisma_client, - ) - cls._initialize_spend_capture_rate_check_job( scheduler=scheduler, proxy_logging_obj=proxy_logging_obj, @@ -10934,39 +10924,6 @@ async def _initialize_expired_ui_session_key_cleanup_background_job(cls, schedul "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)" ) - @classmethod - def _initialize_daily_global_spend_reconcile_job( - cls, - scheduler: AsyncIOScheduler, - proxy_logging_obj: ProxyLogging, - prisma_client: PrismaClient, - ) -> None: - async def alert(message: str) -> None: - await proxy_logging_obj.alerting_handler( - message=message, - level="High", - alert_type=AlertType.failed_tracking_spend, - ) - - async def reconcile() -> None: - await run_scheduled_daily_global_spend_reconcile( - prisma_client, - pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, - alert=alert, - ) - - scheduler.add_job( - reconcile, - "cron", - hour=0, - minute=30, - timezone="UTC", - id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2), - ) - @classmethod def _initialize_spend_capture_rate_check_job( cls, diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py deleted file mode 100644 index b81b6c1943e3..000000000000 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ /dev/null @@ -1,289 +0,0 @@ -"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``. - -Only days that are over get rolled up, so a pod still flushing per-key spend for the current -day can never leave the global table short; usage reads serve days through the recorded -marker from the global table and later days live from the per-key table. Per-key rows are -dated by request start, so spend can land on a day that was already rolled up (a flush -straddling midnight, a retry after an outage). Each run therefore also rewrites every closed -day that has rows touched since the previous run's scan, whatever the date. The marker lives -in ``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on -a large deployment the first backfill is minutes of work. -""" - -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from datetime import date, timedelta -from typing import TYPE_CHECKING, Final - -from pydantic import BaseModel, ConfigDict, ValidationError - -from litellm._logging import verbose_proxy_logger -from litellm.constants import ( - DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, - DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, -) - -if TYPE_CHECKING: - from litellm.caching.redis_cache import RedisCache - from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager - from litellm.proxy.utils import PrismaClient - -GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" -# The unique constraint, in constraint order. NULL never matches itself in a unique index, so -# every column is normalized to '' or the same group would be inserted again on every run. -_KEY_COLUMNS: Final = ("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") -_METRIC_COLUMNS: Final = ( - "prompt_tokens", - "completion_tokens", - "cache_read_input_tokens", - "cache_creation_input_tokens", - "compression_saved_tokens", - "api_requests", - "successful_requests", - "failed_requests", - "total_response_time_ms", - "timed_requests", - "compression_savings_spend", - "prompt_caching_savings_spend", - "gateway_injected_caching_savings_spend", - "autorouter_savings_spend", - "spend", -) - - -def _quoted(columns: tuple[str, ...]) -> str: - return ", ".join(f'"{column}"' for column in columns) - - -def _reconcile_day_sql() -> str: - normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in _KEY_COLUMNS) - sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS) - overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS) - return ( - f'INSERT INTO "{GLOBAL_SPEND_TABLE_NAME}" ("id", {_quoted(_KEY_COLUMNS)}, {_quoted(_METRIC_COLUMNS)}, ' - '"updated_at")\n' - f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n" - 'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n' - f"GROUP BY {normalized_keys}\n" - f"ON CONFLICT ({_quoted(_KEY_COLUMNS)}) DO UPDATE SET {overwrite}, " - "\"updated_at\" = (NOW() AT TIME ZONE 'UTC')" - ) - - -RECONCILE_DAY_SQL: Final = _reconcile_day_sql() -_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now, (NOW() AT TIME ZONE 'UTC')::date::text AS today" -_ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"' -# Pod clocks drift from the database clock and from each other, so rows are picked up from a -# little before the previous scan; rewriting a day twice is idempotent. -_PENDING_DAYS_SQL: Final = ( - 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ' - 'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') ' - 'ORDER BY "date"' -) -# Runs can overlap (Redis unreachable, lock expired on a long backfill), so the database keeps the -# later of the stored and the incoming day and scan time in one statement; GREATEST skips NULL. -_ADVANCE_MARKER_SQL: Final = ( - 'INSERT INTO "LiteLLM_Config" ("param_name", "param_value") ' - "VALUES ($1, jsonb_build_object('reconciled_through', $2::text, 'scanned_at', $3::text)) " - 'ON CONFLICT ("param_name") DO UPDATE SET "param_value" = jsonb_build_object(' - "'reconciled_through', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'reconciled_through', " - "EXCLUDED.\"param_value\" ->> 'reconciled_through'), " - "'scanned_at', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'scanned_at', " - "EXCLUDED.\"param_value\" ->> 'scanned_at'))" -) - - -class ReconciledThrough(BaseModel): - """``reconciled_through`` is the last closed UTC day the global table covers. ``scanned_at`` is - the database clock when the scan behind the last fully successful run started: every per-key - row written before it, on any day through the marker, is in the global table.""" - - model_config = ConfigDict(frozen=True, extra="ignore") - - reconciled_through: str - scanned_at: str | None = None - - -class _MarkerRow(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) - - param_value: object = None - - -class _DateRow(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - date: str - - -class _NowRow(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - now: str - today: str - - -@dataclass(frozen=True, slots=True) -class ReconcileResult: - days_reconciled: tuple[str, ...] - reconciled_through: str | None - failed_day: str | None = None - - -@dataclass(frozen=True, slots=True) -class _PendingScan: - marker: ReconciledThrough | None - scanned_at: str - days: tuple[str, ...] - - -def _marker_from_param_value(value: object) -> ReconciledThrough | None: - try: - return ( - ReconciledThrough.model_validate_json(value) - if isinstance(value, str) - else ReconciledThrough.model_validate(value) - ) - except ValidationError: - return None - - -async def read_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: - from litellm.proxy.utils import get_config_param - - row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value) - - -async def reconciled_through(prisma_client: "PrismaClient") -> str | None: - """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" - marker: Final = await read_marker(prisma_client) - return None if marker is None else marker.reconciled_through - - -async def _advance_marker(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: - """Move the stored marker to the last of ``days`` and to ``scanned_at`` where those are later - than what is stored, so a slower overlapping run can only add to a faster run's marker.""" - from litellm.proxy.utils import invalidate_config_param - - await prisma_client.db.execute_raw( - _ADVANCE_MARKER_SQL, - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, - max(days) if days else None, - scanned_at, - ) - await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -async def _db_now(prisma_client: "PrismaClient") -> _NowRow: - rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) - return _NowRow.model_validate(rows[0]) - - -async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan: - """Every closed UTC day (strictly before the database's today) still to roll up, oldest first: - days past the marker, plus any day with per-key rows written since the scan behind the marker. - Before a run has fully succeeded there is no such scan, so every closed day is rolled up.""" - marker: Final = await read_marker(prisma_client) - db_now: Final = await _db_now(prisma_client) - last_closed_day: Final = (date.fromisoformat(db_now.today) - timedelta(days=1)).isoformat() - rows: Final = ( - await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day) - if marker is None or marker.scanned_at is None - else await prisma_client.db.query_raw( - _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at - ) - ) - return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows)) - - -async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: - """Rewrite one day of the global table from the per-key sums. Idempotent: a rerun - overwrites every group with the same totals.""" - await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day) - - -async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> ReconcileResult: - """Roll up every pending day, advancing the marker after each; a failing day stops the run - with the marker on the last good day so the next run resumes there. The scan time is only - recorded once every pending day is done, so late rows a failed run saw are found again.""" - scan: Final = await _scan_pending(prisma_client) - done: Final = await _reconcile_until_failure(prisma_client, scan) - if len(done) < len(scan.days): - marker: Final = await reconciled_through(prisma_client) - return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) - if scan.marker is not None or done: - await _advance_marker(prisma_client, done, scanned_at=scan.scanned_at) - return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) - - -async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: - for index, day in enumerate(scan.days): - if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]): - return scan.days[:index] - return scan.days - - -async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: tuple[str, ...]) -> bool: - day: Final = done_with_this[-1] - try: - await reconcile_day(prisma_client, day) - await _advance_marker(prisma_client, done_with_this, scanned_at=None) - except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done - verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) - return False - return True - - -async def run_scheduled_daily_global_spend_reconcile( - prisma_client: "PrismaClient", - pod_lock_manager: "PodLockManager | None" = None, - alert: Callable[[str], Awaitable[None]] | None = None, -) -> ReconcileResult | None: - """Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves - effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping.""" - redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache - if pod_lock_manager is None or redis_cache is None: - return await _run_and_alert(prisma_client, alert=alert) - - acquired: Final = await pod_lock_manager.acquire_lock( - cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS - ) - if not acquired and await _lock_is_held(pod_lock_manager, redis_cache): - verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run") - return None - try: - return await _run_and_alert(prisma_client, alert=alert) - finally: - if acquired: - await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) - - -async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool: - try: - lock_key: Final = pod_lock_manager.get_redis_lock_key(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) - return bool(await redis_cache.async_get_cache(lock_key)) - except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the run - verbose_proxy_logger.warning("Daily global spend reconcile: could not read the lock: %s", exc) - return False - - -async def _run_and_alert( - prisma_client: "PrismaClient", - *, - alert: Callable[[str], Awaitable[None]] | None, -) -> ReconcileResult: - result: Final = await run_daily_global_spend_reconcile(prisma_client) - if result.days_reconciled: - verbose_proxy_logger.info( - "Daily global spend reconcile: rolled up %d day(s), reconciled through %s", - len(result.days_reconciled), - result.reconciled_through, - ) - if result.failed_day is not None and alert is not None: - await alert( - f"Daily global spend reconcile stopped at {result.failed_day}; usage totals keep reading the per-key " - f"table for ranges past {result.reconciled_through or 'the beginning'} until the next run succeeds." - ) - return result diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 2a4f6b2944a6..28488ba7de6b 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -101,16 +101,6 @@ class DailySpendMetadata(BaseModel): page: int = Field(default=1) total_pages: int = Field(default=1) has_more: bool = Field(default=False) - api_key_limit: int | None = Field( - default=None, - description="When set, api_keys and every api_key_breakdown list at most this many keys, " - "ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.", - ) - total_api_keys: int | None = Field( - default=None, - description="Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key " - "lists are truncated to the highest-spend keys.", - ) class SpendAnalyticsPaginatedResponse(BaseModel): diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index 05cbd4507a27..43e3899d5237 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -2,7 +2,7 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import ( LiteLLM_UserTableWithKeyCount, @@ -28,16 +28,6 @@ class UserSearchWhere(TypedDict): OR: ReadOnly[tuple[Mapping[Literal["user_id", "user_email"], InsensitiveContains], ...]] -class KeyActivitySearchWhere(TypedDict): - """Prisma filter behind `/user/daily/activity/aggregated/search`: exact token hash, or key alias - or user id containing the term, case-insensitive.""" - - user_id: NotRequired[ReadOnly[str]] - OR: ReadOnly[ - tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...] - ] - - class UserListResponse(BaseModel): """ Response model for the user list endpoint diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 25946eec0098..99ed82dcad8f 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -32,7 +32,6 @@ GET /team/spend/by_user GET /team/spend/report GET /user/daily/activity GET /user/daily/activity/aggregated -GET /user/daily/activity/aggregated/search GET /user/spend/report # Admin UI helper endpoints; serve UI forms and caller-scoped views, not desired state diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 02f5894eb607..d55316ca4299 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3517,7 +3517,6 @@ def test_internal_user_still_blocked_from_another_users_info(): [ "/user/daily/activity", "/user/daily/activity/aggregated", - "/user/daily/activity/aggregated/search", ], ) @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 1274d69cdf2b..32856a3bee9f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1,4 +1,3 @@ -import pathlib import re from collections.abc import Sequence from datetime import datetime, timedelta, timezone @@ -11,11 +10,7 @@ from psycopg.rows import dict_row from pytest_postgresql import factories -from litellm.constants import ( - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, - PTU_SENTINEL_API_KEY, - USAGE_TOP_API_KEYS_LIMIT, -) +from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy.management_endpoints.common_daily_activity import ( _adjust_dates_for_timezone, _build_aggregated_sql_query, @@ -25,12 +20,9 @@ get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, - global_rollup_reconciled_through, update_metrics, ) -from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.proxy.utils import evict_config_param from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, SpendMetrics, @@ -181,7 +173,6 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": "/v1/chat/completions", "api_key": None, "group_level": 62, - "distinct_api_keys": None, "spend": 15.0, "prompt_tokens": 150, "completion_tokens": 75, @@ -194,7 +185,31 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": "/v1/embeddings", "api_key": None, "group_level": 62, - "distinct_api_keys": None, + "spend": 3.0, + "prompt_tokens": 30, + "completion_tokens": 0, + "api_requests": 1, + "successful_requests": 1, + }, + # (date, endpoint, api_key) — populates the per-key sub-bucket + { + **base, + "date": "2024-01-01", + "endpoint": "/v1/chat/completions", + "api_key": "key-1", + "group_level": 30, + "spend": 15.0, + "prompt_tokens": 150, + "completion_tokens": 75, + "api_requests": 2, + "successful_requests": 2, + }, + { + **base, + "date": "2024-01-01", + "endpoint": "/v1/embeddings", + "api_key": "key-2", + "group_level": 30, "spend": 3.0, "prompt_tokens": 30, "completion_tokens": 0, @@ -208,7 +223,6 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": None, "api_key": None, "group_level": 63, - "distinct_api_keys": None, "spend": 18.0, "prompt_tokens": 180, "completion_tokens": 75, @@ -222,40 +236,12 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": None, "api_key": None, "group_level": 127, - "distinct_api_keys": None, "spend": 18.0, "prompt_tokens": 180, "completion_tokens": 75, "api_requests": 3, "successful_requests": 3, }, - # (date, endpoint, api_key) — populates the per-key sub-bucket - { - **base, - "date": "2024-01-01", - "endpoint": "/v1/chat/completions", - "api_key": "key-1", - "group_level": 30, - "distinct_api_keys": 2, - "spend": 15.0, - "prompt_tokens": 150, - "completion_tokens": 75, - "api_requests": 2, - "successful_requests": 2, - }, - { - **base, - "date": "2024-01-01", - "endpoint": "/v1/embeddings", - "api_key": "key-2", - "group_level": 30, - "distinct_api_keys": 2, - "spend": 3.0, - "prompt_tokens": 30, - "completion_tokens": 0, - "api_requests": 1, - "successful_requests": 1, - }, ] mock_prisma.db.query_raw = AsyncMock(return_value=mock_rows) @@ -869,7 +855,6 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): "endpoint": "/v1/chat/completions", "api_key": None, "group_level": 62, - "distinct_api_keys": None, "spend": 10.0, "prompt_tokens": 100, "completion_tokens": 50, @@ -882,7 +867,6 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): "endpoint": "/v1/chat/completions", "api_key": "deleted-key-hash", "group_level": 30, - "distinct_api_keys": 1, "spend": 10.0, "prompt_tokens": 100, "completion_tokens": 50, @@ -1327,7 +1311,6 @@ def test_optional_filters_appear_in_params_in_order(self): "user-1", "bedrock/global.anthropic.claude-opus-4-8", "sk-test", - PTU_SENTINEL_API_KEY, ] assert "model = $4" in sql assert "api_key = $5" in sql @@ -1351,8 +1334,7 @@ def test_empty_entity_list_emits_no_degenerate_in_clause(self, build): normalized = " ".join(sql.split()) assert "IN ()" not in normalized assert '"team_id" IN' not in normalized - sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else [] - assert params == ["2026-08-01", "2026-08-19", *sentinel_params] + assert params == ["2026-08-01", "2026-08-19"] @pytest.mark.parametrize("build", _BUILDERS) def test_empty_entity_list_matches_nothing_rather_than_everything(self, build): @@ -1383,8 +1365,7 @@ def test_populated_entity_list_still_filters_on_its_ids(self, build): normalized = " ".join(sql.split()) assert '"team_id" IN ($3, $4)' in normalized assert "FALSE" not in normalized - sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else [] - assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta", *sentinel_params] + assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta"] @pytest.mark.asyncio @@ -1409,7 +1390,6 @@ async def test_get_daily_activity_aggregated_empty_result_set(): "mcp_namespaced_tool_name": None, "endpoint": None, "group_level": 127, - "distinct_api_keys": None, "spend": None, "prompt_tokens": None, "completion_tokens": None, @@ -1520,18 +1500,10 @@ async def query_raw(sql: str, *params: str) -> list[dict[str, object]]: @pytest.mark.asyncio -async def test_get_daily_activity_aggregated_bounds_api_key_rollups( +async def test_get_daily_activity_aggregated_returns_every_api_key( _aggregated_postgresql: psycopg.Connection, ): - """Run the GROUPING SETS statement against real Postgres with more keys than the cap. - - key-004 and key-005 tie on spend exactly at the USAGE_TOP_API_KEYS_LIMIT - cutoff; the api_key tiebreaker must keep key-004 and drop key-005. The PTU - sentinel outspends every key but must not take a slot. Excluded keys and the - sentinel still count toward the totals and the model rollup, which come from - the key-free arm. - """ - n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 5 + n_keys: Final = 105 key_rows: Final = [ ( f"row-{i:03d}", @@ -1565,11 +1537,11 @@ async def test_get_daily_activity_aggregated_bounds_api_key_rollups( ) _seed_daily_user_spend(_aggregated_postgresql, [*key_rows, sentinel_row]) key_spend: Final = sum(6.0 if i == 4 else float(i + 1) for i in range(n_keys)) + expected_api_keys: Final = {f"key-{i:03d}" for i in range(n_keys)} - row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts) + mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) @@ -1585,37 +1557,24 @@ async def test_get_daily_activity_aggregated_bounds_api_key_rollups( api_key=None, ) - # Key-free arm: (), (date), (date, model), (date, model_group), two providers, - # one mcp NULL bucket, endpoint plus its NULL bucket = 9 rows regardless of key count. - # Per-key arm: six per-key grouping sets, each capped at the limit. - assert row_counts == [9 + 6 * USAGE_TOP_API_KEYS_LIMIT] - assert result.metadata.total_spend == pytest.approx(key_spend + 1000.0) - assert result.metadata.total_api_requests == n_keys - assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT - assert result.metadata.total_api_keys == n_keys - - expected_top: Final = {f"key-{i:03d}" for i in range(6, n_keys)} | {"key-004"} + assert result.metadata.total_api_requests == 105 day: Final = result.results[0] assert day.metrics.spend == pytest.approx(key_spend + 1000.0) - assert set(day.breakdown.api_keys) == expected_top - assert day.breakdown.api_keys["key-004"].metrics.spend == 6.0 - assert "key-005" not in day.breakdown.api_keys + assert set(day.breakdown.api_keys) == expected_api_keys assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys - assert day.breakdown.models["gpt-5"].metrics.spend == pytest.approx(key_spend + 1000.0) - assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_top + assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_api_keys assert day.breakdown.providers["openai"].metrics.spend == pytest.approx(key_spend) - assert set(day.breakdown.providers["openai"].api_key_breakdown) == expected_top - assert day.breakdown.endpoints["/v1/chat/completions"].metrics.api_requests == n_keys + assert set(day.breakdown.providers["openai"].api_key_breakdown) == expected_api_keys + assert day.breakdown.endpoints["/v1/chat/completions"].metrics.api_requests == 105 @pytest.mark.asyncio -async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both_arms( +async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_results( _aggregated_postgresql: psycopg.Connection, ): - """An explicit api_key filter must scope the key-free totals and the per-key - rollups to that key alone, so the two arms never disagree.""" + """An explicit api_key filter must scope the results to that key alone.""" rows: Final = [ ( f"row-{i}", @@ -1635,10 +1594,9 @@ async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both ] _seed_daily_user_spend(_aggregated_postgresql, rows) - row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts) + mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) @@ -1655,7 +1613,6 @@ async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both ) assert result.metadata.total_spend == 2.0 - assert result.metadata.total_api_keys == 1 day: Final = result.results[0] assert set(day.breakdown.api_keys) == {"key-1"} assert day.breakdown.api_keys["key-1"].metrics.spend == 2.0 @@ -1663,221 +1620,6 @@ async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == {"key-1"} -def _prisma_with_marker(marker: str | None) -> MagicMock: - prisma = MagicMock() - prisma.db = MagicMock() - prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - row = ( - None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') - ) - prisma.get_generic_data = AsyncMock(return_value=row) - return prisma - - -def _unfiltered_user_query(**overrides): - return { - "table_name": "litellm_dailyuserspend", - "entity_id_field": "user_id", - "entity_id": None, - "start_date": "2026-06-01", - "end_date": "2026-06-02", - "model": None, - "api_key": None, - "exclude_entity_ids": None, - "timezone_offset_minutes": None, - "include_current_utc_day": False, - **overrides, - } - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("marker", "overrides", "expected"), - [ - ("2026-06-02", {}, "2026-06-02"), - ("2026-06-02", {"model": "gpt-5"}, "2026-06-02"), - ("2026-05-01", {}, "2026-05-01"), - (None, {}, None), - ("2026-06-02", {"api_key": "sk-1"}, None), - ("2026-06-02", {"api_key": []}, None), - ("2026-06-02", {"entity_id": "u-1"}, None), - ("2026-06-02", {"exclude_entity_ids": ["u-1"]}, None), - ("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None), - ], -) -async def test_global_rollup_marker_is_used_only_for_unfiltered_user_reads(marker, overrides, expected): - """Anything that filters by key or entity has no counterpart in the global table; the - SQL splits the range at the marker itself, so the marker passes through unchanged.""" - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - prisma = _prisma_with_marker(marker) - - assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query(**overrides)) == expected - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -@pytest.mark.asyncio -async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table(): - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - prisma = _prisma_with_marker(None) - prisma.get_generic_data = AsyncMock(side_effect=RuntimeError("db down")) - - assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query()) is None - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -_GLOBAL_SPEND_MIGRATION: Final = ( - pathlib.Path(__file__).resolve().parents[4] - / "litellm-proxy-extras" - / "litellm_proxy_extras" - / "migrations" - / "20260915000000_add_daily_global_spend" - / "migration.sql" -) - - -@pytest.mark.asyncio -async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_table_and_open_days_live( - _aggregated_postgresql: psycopg.Connection, -): - """Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must - give the same response as reading everything per-key: day 1 from the global table, day 2 - live, one grand total across both. Per-key rows that land after the rollup then tell the - two sources apart: a late day 1 row is invisible to totals until the next reconcile while a - late day 2 row shows up at once, and both keys rank in the key breakdown, which stays - per-key throughout.""" - n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 - rows: Final = [ - ( - f"row-{day}-{i:03d}", - f"user-{i % 7}", - day, - f"key-{i:03d}", - "gpt-5" if i % 2 else "claude", - "" if i % 3 else "gpt-5", - "openai" if i % 2 else None, - "/v1/chat/completions" if i % 5 else None, - 10, - float(i + 1), - 1, - 1, - ) - for day in ("2026-06-01", "2026-06-02") - for i in range(n_keys) - ] - _seed_daily_user_spend(_aggregated_postgresql, rows) - with _aggregated_postgresql.cursor() as cur: - cur.execute( - 'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms = prompt_tokens * 25, ' - "timed_requests = api_requests" - ) - cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - cur.execute( - re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg - {"p1": "2026-06-01"}, - ) - _aggregated_postgresql.commit() - - async def read(marker: str | None): - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - prisma = _prisma_with_marker(marker) - prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) - return await get_daily_activity_aggregated( - prisma_client=prisma, - entity_metadata_field=None, - **_unfiltered_user_query(), - ) - - from_per_key = await read(None) - from_global = await read("2026-06-01") - - assert from_global.model_dump() == from_per_key.model_dump() - seeded_spend: Final = 2 * sum(float(i + 1) for i in range(n_keys)) - assert from_global.metadata.total_spend == pytest.approx(seeded_spend) - assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25 - assert from_global.metadata.total_timed_requests == 2 * n_keys - assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} - assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT - assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} - - with _aggregated_postgresql.cursor() as cur: - cur.executemany( - """ - INSERT INTO "LiteLLM_DailyUserSpend" - (id, user_id, date, api_key, model, model_group, custom_llm_provider, - endpoint, prompt_tokens, spend, api_requests, successful_requests) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - """, - [ - ("late-1", "user-late", "2026-06-01", "key-late-1", "gpt-5", "", "openai", None, 10, 1000.0, 1, 1), - ("late-2", "user-late", "2026-06-02", "key-late-2", "gpt-5", "", "openai", None, 10, 500.0, 1, 1), - ], - ) - _aggregated_postgresql.commit() - - late_per_key = await read(None) - late_global = await read("2026-06-01") - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - assert late_per_key.metadata.total_spend == pytest.approx(seeded_spend + 1000.0 + 500.0) - assert late_global.metadata.total_spend == pytest.approx(seeded_spend + 500.0) - by_day: Final = {day.date.isoformat(): day for day in late_global.results} - assert by_day["2026-06-01"].metrics.spend == pytest.approx(seeded_spend / 2) - assert by_day["2026-06-02"].metrics.spend == pytest.approx(seeded_spend / 2 + 500.0) - assert by_day["2026-06-01"].breakdown.api_keys["key-late-1"].metrics.spend == pytest.approx(1000.0) - assert by_day["2026-06-02"].breakdown.api_keys["key-late-2"].metrics.spend == pytest.approx(500.0) - assert late_global.metadata.total_api_keys == n_keys + 2 - - -@pytest.mark.asyncio -async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete( - _aggregated_postgresql: psycopg.Connection, -): - """With exactly USAGE_TOP_API_KEYS_LIMIT keys nothing is dropped, and the - response must say so: total_api_keys equals the limit rather than exceeding it.""" - rows: Final = [ - ( - f"row-{i:03d}", - f"user-{i:03d}", - "2026-06-01", - f"key-{i:03d}", - "gpt-5", - "", - "openai", - "/v1/chat/completions", - 10, - float(i + 1), - 1, - 1, - ) - for i in range(USAGE_TOP_API_KEYS_LIMIT) - ] - _seed_daily_user_spend(_aggregated_postgresql, rows) - - row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts) - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - - result = await get_daily_activity_aggregated( - prisma_client=mock_prisma, - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id=None, - entity_metadata_field=None, - start_date="2026-06-01", - end_date="2026-06-01", - model=None, - api_key=None, - ) - - assert result.metadata.total_api_keys == USAGE_TOP_API_KEYS_LIMIT - assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT - assert len(result.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT - - @pytest.mark.asyncio async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_model_name( _aggregated_postgresql: psycopg.Connection, @@ -2717,7 +2459,7 @@ def test_entity_rollup_sql_query_and_api_key_list_filter(): api_key=[], ) assert "FALSE" in empty_sql - assert empty_params == ["2024-01-01", "2024-01-31", PTU_SENTINEL_API_KEY] + assert empty_params == ["2024-01-01", "2024-01-31"] @pytest.mark.asyncio @@ -2751,10 +2493,10 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown(): "successful_requests": 0, } main_rows = [ - {**base, "date": None, "group_level": 127, "distinct_api_keys": None, "spend": 18.0}, - {**base, "date": "2024-01-01", "group_level": 63, "distinct_api_keys": None, "spend": 18.0}, - {**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "distinct_api_keys": None, "spend": 18.0}, - {**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "distinct_api_keys": 1, "spend": 12.0}, + {**base, "date": None, "group_level": 127, "spend": 18.0}, + {**base, "date": "2024-01-01", "group_level": 63, "spend": 18.0}, + {**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "spend": 18.0}, + {**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "spend": 12.0}, ] entity_base = { key: value diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 8260aec9326a..c663e63414c6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2659,172 +2659,6 @@ async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_us assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "regular-user-123" -@pytest.mark.asyncio -async def test_search_user_daily_activity_keys_passes_matched_tokens_to_aggregation(monkeypatch): - """The search endpoint resolves matching verification tokens by hash, alias, or - user id, then aggregates daily spend for exactly those tokens. This is what lets - the Usage page find keys outside the top-spend subset the aggregated endpoint caps.""" - from types import SimpleNamespace - from unittest.mock import AsyncMock, MagicMock - - from litellm.constants import USAGE_TOP_API_KEYS_LIMIT - from litellm.proxy.management_endpoints.internal_user_endpoints import ( - search_user_daily_activity_keys, - ) - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[SimpleNamespace(token="tok-a"), SimpleNamespace(token="tok-b")] - ) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - mock_response = MagicMock() - mock_get_daily_agg = AsyncMock(return_value=mock_response) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", - mock_get_daily_agg, - ) - - admin_key_dict = UserAPIKeyAuth( - user_id="admin-user-001", - user_role=LitellmUserRoles.PROXY_ADMIN, - ) - - result = await search_user_daily_activity_keys( - search="gamma", - start_date="2025-02-01", - end_date="2025-02-28", - user_id=None, - timezone=480, - include_current_utc_day=False, - user_api_key_dict=admin_key_dict, - ) - - assert result is mock_response - - find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs - assert find_many_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT - assert find_many_kwargs["where"]["OR"] == ( - {"token": "gamma"}, - {"key_alias": {"contains": "gamma", "mode": "insensitive"}}, - {"user_id": {"contains": "gamma", "mode": "insensitive"}}, - ) - assert "user_id" not in find_many_kwargs["where"] - - mock_get_daily_agg.assert_called_once_with( - prisma_client=mock_prisma_client, - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id=None, - entity_metadata_field=None, - start_date="2025-02-01", - end_date="2025-02-28", - model=None, - api_key=["tok-a", "tok-b"], - timezone_offset_minutes=480, - include_current_utc_day=False, - ) - - -@pytest.mark.asyncio -async def test_search_user_daily_activity_keys_no_match_returns_empty_without_aggregating(monkeypatch): - from unittest.mock import AsyncMock, MagicMock - - from litellm.constants import USAGE_TOP_API_KEYS_LIMIT - from litellm.proxy.management_endpoints.internal_user_endpoints import ( - search_user_daily_activity_keys, - ) - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - mock_get_daily_agg = AsyncMock() - monkeypatch.setattr( - "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", - mock_get_daily_agg, - ) - - admin_key_dict = UserAPIKeyAuth( - user_id="admin-user-001", - user_role=LitellmUserRoles.PROXY_ADMIN, - ) - - result = await search_user_daily_activity_keys( - search="nothing-matches", - start_date="2025-02-01", - end_date="2025-02-28", - user_id=None, - timezone=None, - include_current_utc_day=False, - user_api_key_dict=admin_key_dict, - ) - - assert result.results == [] - assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT - assert result.metadata.total_api_keys == 0 - mock_get_daily_agg.assert_not_called() - - -@pytest.mark.asyncio -async def test_search_user_daily_activity_keys_non_admin_scoped_to_caller(monkeypatch): - """Same scoping contract as the aggregated route: a non-admin with no user_id - is scoped to their own rows, and any other user_id is a 403.""" - from types import SimpleNamespace - from unittest.mock import AsyncMock, MagicMock - - from fastapi import HTTPException - - from litellm.proxy.management_endpoints.internal_user_endpoints import ( - search_user_daily_activity_keys, - ) - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[SimpleNamespace(token="tok-a")]) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - - non_admin_key_dict = UserAPIKeyAuth( - user_id="user-1", - user_role=LitellmUserRoles.INTERNAL_USER, - ) - - mock_response = MagicMock() - mock_get_daily_agg = AsyncMock(return_value=mock_response) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", - mock_get_daily_agg, - ) - - result = await search_user_daily_activity_keys( - search="gamma", - start_date="2025-02-01", - end_date="2025-02-28", - user_id=None, - timezone=None, - include_current_utc_day=False, - user_api_key_dict=non_admin_key_dict, - ) - - assert result is mock_response - assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "user-1" - find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs - assert find_many_kwargs["where"]["user_id"] == "user-1" - - with pytest.raises(HTTPException) as exc_info: - await search_user_daily_activity_keys( - search="gamma", - start_date="2025-02-01", - end_date="2025-02-28", - user_id="user-2", - timezone=None, - include_current_utc_day=False, - user_api_key_dict=non_admin_key_dict, - ) - - assert exc_info.value.status_code == 403 - assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail) - - @pytest.mark.asyncio async def test_delete_user_cleans_up_created_by_invitation_links(mocker): """ diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 4812135e4e1a..4047473e9d71 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -1232,57 +1232,6 @@ async def test_spend_report_locks_are_never_released(): proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited() -def _init_daily_global_spend_reconcile_job() -> tuple[AsyncIOScheduler, MagicMock, MagicMock]: - scheduler = AsyncIOScheduler() - proxy_logging_obj = MagicMock() - proxy_logging_obj.alerting_handler = AsyncMock() - prisma_client = MagicMock() - ProxyStartupEvent._initialize_daily_global_spend_reconcile_job( - scheduler=scheduler, - proxy_logging_obj=proxy_logging_obj, - prisma_client=prisma_client, - ) - return scheduler, proxy_logging_obj, prisma_client - - -def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run(): - """Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a - fresh deploy switches usage reads to the global table without waiting for the nightly - run, and after that it fires once a day at 00:30 UTC, when the previous UTC day is closed.""" - from datetime import datetime, timedelta, timezone - - from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID - - scheduler, _, _ = _init_daily_global_spend_reconcile_job() - job = scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) - assert job is not None - - assert timedelta(0) < job.next_run_time - datetime.now(timezone.utc) <= timedelta(minutes=2) - after_catch_up = datetime(2026, 9, 16, 12, 0, tzinfo=timezone.utc) - assert job.trigger.get_next_fire_time(None, after_catch_up) == datetime(2026, 9, 17, 0, 30, tzinfo=timezone.utc) - just_after_a_run = datetime(2026, 9, 17, 0, 30, 1, tzinfo=timezone.utc) - assert job.trigger.get_next_fire_time(None, just_after_a_run) == datetime(2026, 9, 18, 0, 30, tzinfo=timezone.utc) - - -@pytest.mark.asyncio -async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch): - from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID - - scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job() - run = AsyncMock() - monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run) - - await scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func() - - run.assert_awaited_once() - assert run.await_args.args == (prisma_client,) - assert run.await_args.kwargs["pod_lock_manager"] is proxy_logging_obj.db_spend_update_writer.pod_lock_manager - await run.await_args.kwargs["alert"]("day 2026-09-01 failed") - proxy_logging_obj.alerting_handler.assert_awaited_once() - assert proxy_logging_obj.alerting_handler.await_args.kwargs["message"] == "day 2026-09-01 failed" - assert proxy_logging_obj.alerting_handler.await_args.kwargs["level"] == "High" - - @pytest.mark.asyncio async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch): """The boot-time send goes through the same gate, so a losing pod sends nothing at all: diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py deleted file mode 100644 index 3da587435ad7..000000000000 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ /dev/null @@ -1,532 +0,0 @@ -"""Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818).""" - -import json -import pathlib -import re -from datetime import date -from typing import Final -from unittest.mock import AsyncMock, MagicMock - -import psycopg -import pytest -from psycopg.rows import dict_row -from pytest_postgresql import factories - -from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM -from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key -from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( - _ADVANCE_MARKER_SQL, - RECONCILE_DAY_SQL, - read_marker, - reconciled_through, - run_daily_global_spend_reconcile, - run_scheduled_daily_global_spend_reconcile, -) -from litellm.proxy.utils import evict_config_param - -USER_TABLE: Final = DAILY_SPEND_TABLES["user"] -TODAY: Final = date(2026, 9, 15) - - -class _FakeConfigRow: - def __init__(self, param_name: str, param_value: object) -> None: - self.param_name = param_name - self.param_value = param_value - - -class _FakeConfigTable: - def __init__(self) -> None: - self.rows: dict[str, object] = {} - - def advance(self, param_name: str, through: str | None, scanned_at: str | None) -> None: - """What ``_ADVANCE_MARKER_SQL`` does in Postgres: keep the later of stored and incoming per field.""" - stored = self.rows.get(param_name) - current: dict[str, str | None] = json.loads(stored) if isinstance(stored, str) else {} - self.rows[param_name] = json.dumps( - { - "reconciled_through": _greatest(current.get("reconciled_through"), through), - "scanned_at": _greatest(current.get("scanned_at"), scanned_at), - } - ) - - -def _greatest(stored: str | None, incoming: str | None) -> str | None: - present = [value for value in (stored, incoming) if value is not None] - return max(present) if present else None - - -class _FakeDb: - """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, - so "rows written since the last scan" behaves like Postgres would. The database's own - date decides which day is still open, never the pod's clock.""" - - def __init__(self, prisma: "_FakePrisma") -> None: - self._prisma = prisma - self.litellm_config = _FakeConfigTable() - - async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: - if sql.startswith("SELECT (NOW()"): - self._prisma.clock += 1 - return [{"now": f"clock-{self._prisma.clock:04d}", "today": self._prisma.today.isoformat()}] - rows = self._prisma.user_rows - if len(params) == 1: - (last,) = params - return [{"date": d} for d in sorted(rows) if d <= last] - last, marker, scanned_at = params - return [ - {"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at) - ] - - async def execute_raw(self, sql: str, *params: str | None) -> int: - if sql == _ADVANCE_MARKER_SQL: - param_name, through, scanned_at = params - assert param_name is not None - self.litellm_config.advance(param_name, through, scanned_at) - return 1 - (day,) = params - if day is None or day in self._prisma.failing_days: - raise RuntimeError(f"day {day} exploded") - self._prisma.reconciled.append(day) - landing = self._prisma.marker_landing_on_day.get(day) - if landing is not None: - self.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = landing - return 1 - - -class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw. - ``marker_landing_on_day`` stores another pod's marker the moment this run rewrites that day.""" - - def __init__( - self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY - ) -> None: - self.clock = 0 - self.today = today - self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} - self.failing_days = failing_days - self.marker_landing_on_day: dict[str, str] = {} - self.reconciled: list[str] = [] - self.db = _FakeDb(self) - - def write_late_row(self, day: str) -> None: - """A per-key row for ``day`` lands now, after whatever scans already happened.""" - self.clock += 1 - self.user_rows[day] = f"clock-{self.clock:04d}" - - async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: - stored = self.db.litellm_config.rows.get(value) - return None if stored is None else _FakeConfigRow(value, stored) - - -@pytest.fixture(autouse=True) -async def _fresh_marker_cache(): - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - yield - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -@pytest.mark.asyncio -async def test_first_run_rolls_up_every_closed_day_and_never_the_database_s_today(): - """Before any marker exists every closed day with per-key rows is rolled up. Today is left - out: pods are still flushing it, so it is served live from the per-key table until it closes. - The database clock says which day that is; a pod booting with its clock a day ahead must not - roll the open day up and mark it reconciled.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15")) - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14") - assert result.failed_day is None - assert result.reconciled_through == "2026-09-14" - assert await reconciled_through(prisma) == "2026-09-14" - assert "2026-09-15" not in prisma.reconciled - - -@pytest.mark.asyncio -async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14"), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.reconciled.clear() - prisma.today = TODAY - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-14",) - assert await reconciled_through(prisma) == "2026-09-14" - - -@pytest.mark.asyncio -async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run(): - """Per-key rows carry the request start date, so a delayed flush or retry can add spend to a - day far behind the marker. That day is rewritten, and the marker never moves back for it.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13"), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.reconciled.clear() - prisma.today = TODAY - prisma.write_late_row("2026-09-01") - prisma.write_late_row("2026-09-03") - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-03") - assert "2026-09-05" not in prisma.reconciled - assert await reconciled_through(prisma) == "2026-09-13" - - -@pytest.mark.asyncio -async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one(): - """The scan time only advances when every pending day was rewritten, otherwise a late row - found by the failed run would be counted as handled.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13"), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.today = TODAY - prisma.write_late_row("2026-09-01") - prisma.failing_days = frozenset({"2026-09-01"}) - failed = await run_daily_global_spend_reconcile(prisma) - prisma.failing_days = frozenset() - prisma.reconciled.clear() - - result = await run_daily_global_spend_reconcile(prisma) - - assert failed.failed_day == "2026-09-01" - assert failed.reconciled_through == "2026-09-13" - assert result.days_reconciled == ("2026-09-01",) - assert result.failed_day is None - - -@pytest.mark.asyncio -async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) - prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}' - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-13") - marker = await read_marker(prisma) - assert marker is not None and marker.reconciled_through == "2026-09-13" and marker.scanned_at is not None - - -@pytest.mark.asyncio -async def test_a_run_with_no_new_closed_days_keeps_the_marker(): - prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.reconciled.clear() - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == () - assert result.reconciled_through == "2026-09-13" - - -@pytest.mark.asyncio -async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_good_day(): - """The marker may never claim a day that was not rewritten: reads past it would then trust - a global table missing that day's spend.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01",) - assert result.failed_day == "2026-09-02" - assert result.reconciled_through == "2026-09-01" - assert prisma.reconciled == ["2026-09-01"] - assert await reconciled_through(prisma) == "2026-09-01" - - -@pytest.mark.asyncio -async def test_the_next_run_resumes_from_the_failed_day(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - await run_daily_global_spend_reconcile(prisma) - prisma.failing_days = frozenset() - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03") - assert await reconciled_through(prisma) == "2026-09-03" - - -@pytest.mark.asyncio -async def test_a_slower_overlapping_run_never_rewinds_the_marker_a_faster_run_stored(): - """Two pods can reconcile at once (Redis unreachable, or the lock expired on a long backfill). - When the faster one has already stored a later marker, the slower one may only add to it. Putting - its own older prefix back, or dropping the scan time, would send usage reads for every day in - between back to the per-key table until the next run.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-03"})) - prisma.marker_landing_on_day = { - "2026-09-02": '{"reconciled_through": "2026-09-14", "scanned_at": "clock-0009"}', - } - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-02") - assert result.reconciled_through == "2026-09-14" - marker = await read_marker(prisma) - assert marker is not None and (marker.reconciled_through, marker.scanned_at) == ("2026-09-14", "clock-0009") - - -@pytest.mark.asyncio -async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" - prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.today = TODAY - prisma.write_late_row("2026-09-12") - prisma.failing_days = frozenset({"2026-09-12"}) - alert = AsyncMock() - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) - - assert result is not None - assert result.days_reconciled == () - assert result.failed_day == "2026-09-12" - assert result.reconciled_through == "2026-09-13" - alert.assert_awaited_once() - assert "2026-09-12" in alert.await_args.args[0] - - -@pytest.mark.asyncio -async def test_a_clean_run_does_not_alert(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - alert = AsyncMock() - - await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) - - alert.assert_not_awaited() - - -def _pod_lock(acquired: bool) -> MagicMock: - lock = MagicMock() - lock.redis_cache = MagicMock() - lock.redis_cache.async_get_cache = AsyncMock(return_value="other-pod") - lock.get_redis_lock_key = MagicMock(return_value="lock-key") - lock.acquire_lock = AsyncMock(return_value=acquired) - lock.release_lock = AsyncMock() - return lock - - -@pytest.mark.asyncio -async def test_scheduled_run_skips_when_another_pod_holds_the_lock(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - lock = _pod_lock(acquired=False) - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) - - assert result is None - assert prisma.reconciled == [] - lock.release_lock.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - lock = _pod_lock(acquired=True) - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) - - assert result is not None and result.days_reconciled == ("2026-09-13",) - lock.release_lock.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read(): - """A Redis outage must not stall the backfill: the day rewrite is idempotent, so running - twice is only wasted effort while skipping forever leaves usage on the slow path.""" - prisma = _FakePrisma(user_days=("2026-09-13",)) - lock = _pod_lock(acquired=False) - lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) - - assert result is not None and result.days_reconciled == ("2026-09-13",) - lock.release_lock.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_marker_is_read_back_from_the_json_string_the_config_table_stores(): - prisma = _FakePrisma(user_days=()) - prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-10"}' - - assert await reconciled_through(prisma) == "2026-09-10" - - -@pytest.mark.asyncio -async def test_an_unparseable_marker_reads_as_never_reconciled(): - prisma = _FakePrisma(user_days=()) - prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"something_else": 1}' - - assert await reconciled_through(prisma) is None - - -_rollup_postgresql_proc: Final = factories.postgresql_proc() -_rollup_postgresql: Final = factories.postgresql("_rollup_postgresql_proc") - -_MIGRATIONS_DIR: Final = ( - pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" -) -_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" - -_DAILY_USER_SPEND_DDL: Final = """ - CREATE TABLE "LiteLLM_DailyUserSpend" ( - id TEXT PRIMARY KEY, - user_id TEXT, - date TEXT NOT NULL, - api_key TEXT NOT NULL, - model TEXT, - model_group TEXT, - custom_llm_provider TEXT, - mcp_namespaced_tool_name TEXT, - endpoint TEXT, - prompt_tokens BIGINT DEFAULT 0, - completion_tokens BIGINT DEFAULT 0, - cache_read_input_tokens BIGINT DEFAULT 0, - cache_creation_input_tokens BIGINT DEFAULT 0, - compression_saved_tokens BIGINT DEFAULT 0, - compression_savings_spend DOUBLE PRECISION DEFAULT 0, - prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, - spend DOUBLE PRECISION DEFAULT 0, - api_requests BIGINT DEFAULT 0, - successful_requests BIGINT DEFAULT 0, - failed_requests BIGINT DEFAULT 0, - total_response_time_ms BIGINT DEFAULT 0, - timed_requests BIGINT DEFAULT 0, - created_at TIMESTAMP DEFAULT now(), - updated_at TIMESTAMP, - UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) - ) -""" - -_PER_KEY_SUMS_SQL: Final = """ - SELECT COALESCE(model, '') AS model, COALESCE(model_group, '') AS model_group, - COALESCE(custom_llm_provider, '') AS custom_llm_provider, - SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests, - SUM(total_response_time_ms) AS total_response_time_ms, SUM(timed_requests) AS timed_requests - FROM "LiteLLM_DailyUserSpend" WHERE date = %s - GROUP BY 1, 2, 3 ORDER BY 1, 2, 3 -""" -_GLOBAL_ROWS_SQL: Final = """ - SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests, - total_response_time_ms, timed_requests - FROM "LiteLLM_DailyGlobalSpend" WHERE date = %s ORDER BY 1, 2, 3 -""" - - -def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: - converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) - conn.execute( - converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query - {f"p{i}": v for i, v in enumerate(params, start=1)}, - ) - conn.commit() - - -def _user_txn(**overrides): - return { - "user_id": "u-1", - "date": "2026-09-14", - "api_key": "sk-1", - "model": "gpt-5", - "model_group": "gpt-5", - "custom_llm_provider": "openai", - "mcp_namespaced_tool_name": "", - "endpoint": "/chat/completions", - "prompt_tokens": 10, - "completion_tokens": 20, - "spend": 1.0, - "api_requests": 1, - "successful_requests": 1, - "failed_requests": 0, - "total_response_time_ms": 800, - "timed_requests": 1, - **overrides, - } - - -def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: - return [ - ( - r["model"], - r["model_group"], - r["custom_llm_provider"], - float(r["spend"]), - int(r["prompt_tokens"]), - int(r["api_requests"]), - int(r["total_response_time_ms"]), - int(r["timed_requests"]), - ) # pyright: ignore[reportArgumentType] # dict_row values are untyped - for r in rows - ] - - -def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: writer-shaped rows and legacy rows - (NULL and '' dimension spellings) fold into one global day, running the day twice changes - nothing, and other days are left alone.""" - conn: Final = _rollup_postgresql - conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - conn.commit() - - written_batch = merge_by_conflict_key( - USER_TABLE, - (_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)), - ) - _execute_dollar_sql(conn, *build_bulk_upsert(USER_TABLE, written_batch)) - - conn.execute( - """ - INSERT INTO "LiteLLM_DailyUserSpend" - (id, user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, - endpoint, prompt_tokens, spend, api_requests) - VALUES - ('legacy-1', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', NULL, 'openai', NULL, NULL, 5, 4.0, 1), - ('legacy-2', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', '', 'openai', '', '', 5, 8.0, 1), - ('legacy-3', 'u-9', '2026-09-13', 'sk-9', 'claude', '', 'anthropic', '', '', 7, 16.0, 1) - """ - ) - conn.commit() - - _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) - _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) - - with conn.cursor(row_factory=dict_row) as cur: - global_rows = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-14",)).fetchall() - per_key = cur.execute(_PER_KEY_SUMS_SQL, ("2026-09-14",)).fetchall() - untouched = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-13",)).fetchall() - - assert _normalized(global_rows) == _normalized(per_key) - assert sum(float(r["spend"]) for r in global_rows) == pytest.approx(15.0) # pyright: ignore[reportArgumentType] # dict_row values are untyped - assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped - assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] - assert untouched == [] - - -_CONFIG_DDL: Final = 'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB)' -_MARKER_SQL: Final = 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s' - - -def test_advance_marker_sql_only_ever_moves_the_stored_marker_forward(_rollup_postgresql: psycopg.Connection): - """Against real Postgres: the statement a slower overlapping run issues after the faster run - already stored a later marker leaves that marker alone, whether it carries an older scan time or - none at all, while a run that is further along moves both fields on.""" - conn: Final = _rollup_postgresql - conn.execute(_CONFIG_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.commit() - param: Final = DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM - - def stored() -> object: - with conn.cursor(row_factory=dict_row) as cur: - row = cur.execute(_MARKER_SQL, (param,)).fetchone() - return None if row is None else row["param_value"] - - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-01", None)) - assert stored() == {"reconciled_through": "2026-09-01", "scanned_at": None} - - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-14", "2026-09-15 00:30:02.5")) - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-02", None)) - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-03", "2026-09-15 00:30:01.25")) - assert stored() == {"reconciled_through": "2026-09-14", "scanned_at": "2026-09-15 00:30:02.5"} - - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-15", "2026-09-16 00:30:00.75")) - assert stored() == {"reconciled_through": "2026-09-15", "scanned_at": "2026-09-16 00:30:00.75"} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx index b06986d01f07..8d328c4c329d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx @@ -178,28 +178,4 @@ describe("CacheLeakageCard", () => { screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); - - it("says which keys are missing from the key ranking when the proxy capped the per-key lists", () => { - const day = dayWithKeys("2026-07-12", { - "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), - }); - renderWith([day], { apiKeyTruncation: { limit: 100, total: 3000 } }); - - expect(screen.getByRole("note")).toHaveTextContent( - "Only the 100 highest-spend keys of 3,000 are loaded, so a lower-spend key that leaks more is not listed here.", - ); - - fireEvent.click(screen.getByRole("tab", { name: "By model" })); - - expect(screen.queryByRole("note")).not.toBeInTheDocument(); - }); - - it("keeps the key ranking note off when every key was loaded", () => { - const day = dayWithKeys("2026-07-12", { - "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), - }); - renderWith([day]); - - expect(screen.queryByRole("note")).not.toBeInTheDocument(); - }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx index f5b71a00061c..cfac77788b7a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx @@ -80,7 +80,7 @@ const SortableHead = ({ }; const CacheLeakageCard: React.FC = ({ activity }) => { - const { results, loading, isFetchingMore, apiKeyTruncation } = activity; + const { results, loading, isFetchingMore } = activity; const [dimension, setDimension] = useState("key"); const [sort, setSort] = useState({ column: "potentialSavings", dir: "desc" }); const leakage = useMemo(() => computeCacheLeakage(results, dimension), [results, dimension]); @@ -119,13 +119,6 @@ const CacheLeakageCard: React.FC = ({ activity }) => { - {dimension === "key" && apiKeyTruncation !== undefined && ( -

- Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "} - {apiKeyTruncation.total.toLocaleString()} are loaded, so a lower-spend key that leaks more is not listed - here. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys. -

- )} {rows.length > 0 && isFetchingMore && (

Data is still loading; rows and totals will update as the rest of the range arrives. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx index e501cf00b905..b94fa45ecb70 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx @@ -4,13 +4,12 @@ import { describe, expect, it, vi } from "vitest"; const mockUsePaginatedDailyActivity = vi.fn(); const mockCancel = vi.fn(); -let mockMetadata: Record = {}; vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", () => ({ usePaginatedDailyActivity: (args: unknown) => { mockUsePaginatedDailyActivity(args); return { - data: { results: [], metadata: mockMetadata }, + data: { results: [] }, loading: false, isFetchingMore: false, progress: { currentPage: 4, totalPages: 9 }, @@ -83,18 +82,4 @@ describe("useDailyActivityRange", () => { expect(mockUsePaginatedDailyActivity).toHaveBeenLastCalledWith(expect.objectContaining({ enabled: false })); }); - - it("reports how many keys the proxy left out of the per-key lists", () => { - mockMetadata = { api_key_limit: 100, total_api_keys: 3000 }; - const { result } = renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin")); - - expect(result.current.apiKeyTruncation).toEqual({ limit: 100, total: 3000 }); - }); - - it("reports no key truncation when every key fit under the proxy limit", () => { - mockMetadata = { api_key_limit: 100, total_api_keys: 100 }; - const { result } = renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin")); - - expect(result.current.apiKeyTruncation).toBeUndefined(); - }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts index 4eb9f257d30b..605926132e95 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts @@ -1,7 +1,6 @@ import { useMemo, useState } from "react"; import { userDailyActivityAggregatedCall, userDailyActivityCall } from "@/components/networking"; -import { ApiKeyTruncation, getApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason"; import { DailyData } from "@/components/UsagePage/types"; import { spendScopeUserId } from "@/utils/roles"; import { usePaginatedDailyActivity } from "@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity"; @@ -23,7 +22,6 @@ export interface DailyActivityRange { cancelled: boolean; failed: boolean; cancel: () => void; - apiKeyTruncation?: ApiKeyTruncation; } /** @@ -82,7 +80,6 @@ export const useScopedDailyActivityRange = ( cancelled, failed, cancel, - apiKeyTruncation: getApiKeyTruncation(data.metadata?.api_key_limit, data.metadata?.total_api_keys), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index 6bd16095351a..8f5bb9455d0d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -650,23 +650,6 @@ describe("EntityUsage", () => { expect(screen.getAllByText("Activity Metrics")[1]).toBeInTheDocument(); }); - it("tells the team view how many keys the proxy left out of the per-key lists", async () => { - mockTeamDailyActivityAggregatedCall.mockResolvedValue({ - ...mockSpendData, - metadata: { ...mockSpendData.metadata, api_key_limit: 100, total_api_keys: 3000 }, - }); - render(); - - await waitFor(() => { - expect(mockTeamDailyActivityAggregatedCall).toHaveBeenCalled(); - }); - act(() => { - fireEvent.click(screen.getByText("Key Activity")); - }); - - expect(await screen.findByRole("note")).toHaveTextContent("Only the 100 highest-spend keys of 3,000 are loaded"); - }); - // An inactive tab panel is marked aria-selected="false" by one tab library and hidden by the // other, so treat either as "not on screen" and the assertion holds whichever one is rendering. const isShowing = (element: HTMLElement): boolean => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 27460b211086..6687bd4df038 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -25,7 +25,7 @@ import TeamMultiSelect from "@/components/common_components/team_multi_select"; import UserDropdown from "@/components/common_components/UserDropdown"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import { UsageExportHeader } from "@/components/EntityUsageExport"; -import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; +import { getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; import type { EntityType } from "@/components/EntityUsageExport/types"; import { agentDailyActivityCall, @@ -71,8 +71,6 @@ interface EntitySpendData { total_successful_requests: number; total_failed_requests: number; total_tokens: number; - api_key_limit?: number | null; - total_api_keys?: number | null; }; } @@ -162,7 +160,6 @@ const EntityUsage: React.FC = ({ }); const spendData = spendDataRaw as unknown as EntitySpendData; - const apiKeyTruncation = getApiKeyTruncation(spendData.metadata?.api_key_limit, spendData.metadata?.total_api_keys); const { data: agentSpendDataRaw, @@ -662,18 +659,12 @@ const EntityUsage: React.FC = ({ { key: "keys", label: "Key Activity", - content: ( - - ), + content: , }, { key: "endpoints", label: "Endpoint Activity", content: }, ]; - const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation }; + const spendFetchState = { coversRange, cancelled, failed }; return (

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index e0171d974236..018dfc847408 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -30,7 +30,7 @@ import { ActivityMetrics, processActivityData } from "@/components/activity_metr import CloudZeroExportModal from "@/components/cloudzero_export_modal"; import UserDropdown from "@/components/common_components/UserDropdown"; import EntityUsageExportModal from "@/components/EntityUsageExport"; -import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; +import { getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel"; import { Team } from "@/components/key_team_helpers/key_list"; import { @@ -39,7 +39,6 @@ import { tagListCall, userDailyActivityAggregatedCall, userDailyActivityCall, - userDailyActivityKeySearchCall, } from "@/components/networking"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { ChartLoader } from "@/components/shared/chart_loader"; @@ -257,10 +256,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { coversRange: activeAggregated !== null || paginatedResult.coversRange, cancelled: paginatedResult.cancelled, failed: paginatedResult.failed, - apiKeyTruncation: getApiKeyTruncation( - userSpendData.metadata?.api_key_limit, - userSpendData.metadata?.total_api_keys, - ), }; const exportBlockedReason = getExportBlockedReason(spendFetchState); @@ -438,15 +433,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { [userSpendData, modelViewType, teams], ); const keyMetrics = useMemo(() => processActivityData(userSpendData, "api_keys", teams), [userSpendData, teams]); - const searchKeys = useCallback( - (q: string) => { - if (!accessToken || !startTime || !endTime) return Promise.resolve({}); - return userDailyActivityKeySearchCall(accessToken, startTime, endTime, q, effectiveUserId).then((data) => - processActivityData(data, "api_keys", teams), - ); - }, - [accessToken, startTime, endTime, effectiveUserId, teams], - ); const mcpServerMetrics = useMemo( () => processActivityData(userSpendData, "mcp_servers", teams), [userSpendData, teams], @@ -875,11 +861,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { - + diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts index 8491b31f5f98..e39b01a5dea0 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts @@ -1,12 +1,11 @@ import { describe, expect, it } from "vitest"; -import { getApiKeyTruncation, getExportBlockedReason, type UsageFetchState } from "./exportBlockedReason"; +import { getExportBlockedReason, type UsageFetchState } from "./exportBlockedReason"; const state = (overrides: Partial = {}): UsageFetchState => ({ coversRange: true, cancelled: false, failed: false, - apiKeyTruncation: undefined, ...overrides, }); @@ -32,27 +31,4 @@ describe("getExportBlockedReason", () => { expect(reason).toMatch(/failed to load/i); expect(reason).not.toMatch(/stopped/i); }); - - it("blocks when the aggregated endpoint dropped keys, since a per-team CSV would miss them", () => { - const reason = getExportBlockedReason(state({ apiKeyTruncation: { limit: 100, total: 3000 } })); - - expect(reason).toMatch(/100 highest-spend keys of 3000/); - expect(reason).toMatch(/USAGE_TOP_API_KEYS_LIMIT/); - }); -}); - -describe("getApiKeyTruncation", () => { - it("reports truncation once the proxy saw more keys than it returned", () => { - expect(getApiKeyTruncation(100, 101)).toEqual({ limit: 100, total: 101 }); - }); - - it("stays quiet when exactly the cap exists, since every key is on screen", () => { - expect(getApiKeyTruncation(100, 100)).toBeUndefined(); - expect(getApiKeyTruncation(100, 7)).toBeUndefined(); - }); - - it("stays quiet when the response carries no cap, as the paginated fallback does", () => { - expect(getApiKeyTruncation(undefined, undefined)).toBeUndefined(); - expect(getApiKeyTruncation(100, null)).toBeUndefined(); - }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts index 6c5a5f83231c..71408ba8f3ff 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts @@ -1,31 +1,13 @@ -export interface ApiKeyTruncation { - limit: number; - total: number; -} - export interface UsageFetchState { coversRange: boolean; cancelled: boolean; failed: boolean; - apiKeyTruncation: ApiKeyTruncation | undefined; } -export const getApiKeyTruncation = (apiKeyLimit: unknown, totalApiKeys: unknown): ApiKeyTruncation | undefined => { - if (typeof apiKeyLimit !== "number" || typeof totalApiKeys !== "number") return undefined; - return totalApiKeys > apiKeyLimit ? { limit: apiKeyLimit, total: totalApiKeys } : undefined; -}; - -export const getExportBlockedReason = ({ - coversRange, - cancelled, - failed, - apiKeyTruncation, -}: UsageFetchState): string | undefined => { +export const getExportBlockedReason = ({ coversRange, cancelled, failed }: UsageFetchState): string | undefined => { if (failed) return "Some spend data failed to load, so an export would under-report. Reload the page to try again."; if (cancelled) return "Loading was stopped before the whole range arrived, so an export would under-report. Reload the page to load it all."; if (!coversRange) return "Spend data is still loading, so an export would under-report. Wait for it to finish."; - if (apiKeyTruncation !== undefined) - return `Only the ${apiKeyTruncation.limit} highest-spend keys of ${apiKeyTruncation.total} were loaded, so a per-team export would under-report. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys.`; return undefined; }; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx index f8a7d07633b4..693ac20a3601 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx @@ -68,78 +68,4 @@ describe("KeyActivityPanel", () => { expect(screen.getByLabelText("Search keys")).toHaveValue(""); expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alicehash-bob"); }); - - it("says how many keys the proxy left out when only the top spenders were loaded", () => { - render(); - expect(screen.getByRole("note")).toHaveTextContent("Only the 2 highest-spend keys of 3,000 are loaded"); - }); - - it("shows no truncation note when every key is loaded", () => { - render(); - expect(screen.queryByRole("note")).not.toBeInTheDocument(); - }); - - it("finds keys outside the loaded top-spend subset via server search", async () => { - const searchKeys = vi - .fn<(query: string) => Promise>>() - .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); - render( - , - ); - - fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); - - expect(await screen.findByText("hash-gamma")).toBeInTheDocument(); - expect(searchKeys).toHaveBeenCalledWith("gamma"); - expect(screen.getByText("Showing 1 of 3 keys")).toBeInTheDocument(); - }); - - it("never calls the server search when every key is already loaded", async () => { - const searchKeys = vi - .fn<(query: string) => Promise>>() - .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); - render(); - - fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); - - expect(await screen.findByText('No keys match "gamma" in this date range')).toBeInTheDocument(); - await new Promise((resolve) => setTimeout(resolve, 400)); - expect(searchKeys).not.toHaveBeenCalled(); - }); - - it("drops stale server results as soon as the search callback is rebuilt", async () => { - const searchKeysA = vi - .fn<(query: string) => Promise>>() - .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); - const searchKeysB = vi - .fn<(query: string) => Promise>>() - .mockReturnValue(new Promise(() => {})); - const { rerender } = render( - , - ); - - fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); - expect(await screen.findByText("hash-gamma")).toBeInTheDocument(); - - rerender( - , - ); - - expect(screen.getByRole("status")).toHaveTextContent("Searching all keys"); - expect(screen.queryByText("hash-gamma")).not.toBeInTheDocument(); - }); - - it("reports a failed server search but keeps the local matches", async () => { - const searchKeys = vi - .fn<(query: string) => Promise>>() - .mockRejectedValue(new Error("boom")); - render( - , - ); - - fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "alice" } }); - - expect(await screen.findByRole("alert")).toHaveTextContent("Key search failed"); - expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alice"); - }); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx index 3467ba61b220..8287a04d0c78 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx @@ -1,8 +1,7 @@ import { Search, X } from "lucide-react"; -import React, { useEffect, useMemo, useState } from "react"; +import React, { useMemo, useState } from "react"; import { ActivityMetrics } from "@/components/activity_metrics"; -import type { ApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { filterKeyActivity } from "../keyActivityFilter"; @@ -11,68 +10,14 @@ import type { ModelActivityData } from "../types"; interface KeyActivityPanelProps { keyMetrics: Record; hidePromptCachingMetrics?: boolean; - apiKeyTruncation?: ApiKeyTruncation; - searchKeys?: SearchKeys; } -type SearchKeys = (query: string) => Promise>; - -type RemoteSearch = - | { status: "idle" } - | { status: "loading"; query: string; searchKeys: SearchKeys } - | { status: "done"; query: string; searchKeys: SearchKeys; keys: Record } - | { status: "error"; query: string; searchKeys: SearchKeys }; - -const REMOTE_SEARCH_DEBOUNCE_MS = 300; - -const KeyActivityPanel: React.FC = ({ - keyMetrics, - hidePromptCachingMetrics = false, - apiKeyTruncation, - searchKeys, -}) => { +const KeyActivityPanel: React.FC = ({ keyMetrics, hidePromptCachingMetrics = false }) => { const [query, setQuery] = useState(""); - const [remote, setRemote] = useState({ status: "idle" }); const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]); - const trimmedQuery = query.trim(); - const remoteEnabled = searchKeys !== undefined && apiKeyTruncation !== undefined && trimmedQuery !== ""; - - useEffect(() => { - if (!remoteEnabled) return; - let cancelled = false; - const timer = setTimeout(() => { - setRemote({ status: "loading", query: trimmedQuery, searchKeys }); - searchKeys(trimmedQuery) - .then((keys) => { - if (!cancelled) setRemote({ status: "done", query: trimmedQuery, searchKeys, keys }); - }) - .catch(() => { - if (!cancelled) setRemote({ status: "error", query: trimmedQuery, searchKeys }); - }); - }, REMOTE_SEARCH_DEBOUNCE_MS); - return () => { - cancelled = true; - clearTimeout(timer); - }; - }, [remoteEnabled, trimmedQuery, searchKeys]); - - const remoteMatchesSearch = - "searchKeys" in remote && remote.searchKeys === searchKeys && remote.query === trimmedQuery; - const remoteCurrent = remoteEnabled && remoteMatchesSearch; - const remoteLoading = remoteEnabled && (remote.status === "loading" || !remoteCurrent); - const remoteFailed = remoteCurrent && remote.status === "error"; - - const extraRemoteKeys = useMemo(() => { - const remoteKeys = remoteCurrent && remote.status === "done" ? remote.keys : {}; - return Object.fromEntries(Object.entries(remoteKeys).filter(([hash]) => !(hash in keyMetrics))); - }, [remoteCurrent, remote, keyMetrics]); - const displayed = useMemo(() => ({ ...extraRemoteKeys, ...filtered }), [extraRemoteKeys, filtered]); - const totalKeys = Object.keys(keyMetrics).length; - const shownKeys = Object.keys(displayed).length; - const totalShown = totalKeys + Object.keys(extraRemoteKeys).length; - const isFiltering = trimmedQuery !== ""; - const noMatches = isFiltering && !remoteLoading && totalKeys > 0 && shownKeys === 0; + const shownKeys = Object.keys(filtered).length; + const isFiltering = query.trim() !== ""; return (
@@ -96,31 +41,15 @@ const KeyActivityPanel: React.FC = ({ )} - Showing {shownKeys.toLocaleString()} of {totalShown.toLocaleString()} keys + Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys - {remoteLoading && ( - - Searching all keys... - - )} - {remoteFailed && ( - - Key search failed - - )} - {apiKeyTruncation !== undefined && ( - - Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "} - {apiKeyTruncation.total.toLocaleString()} are loaded - - )}
- {noMatches ? ( + {isFiltering && totalKeys > 0 && shownKeys === 0 ? (

- No keys match "{trimmedQuery}" in this date range + No keys match "{query.trim()}" in this date range

) : ( - + )}
); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 82c446ce28c0..42dc5350b495 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2556,36 +2556,6 @@ export const userDailyActivityAggregatedCall = async ( } }; -export const userDailyActivityKeySearchCall = async ( - accessToken: string, - startTime: Date, - endTime: Date, - ...options: [search: string, userId?: string | null] -) => { - const [search, userId = null] = options; - try { - const formatDate = (date: Date) => { - const year = date.getFullYear(); - const month = String(date.getMonth() + 1).padStart(2, "0"); - const day = String(date.getDate()).padStart(2, "0"); - return `${year}-${month}-${day}`; - }; - return await apiClient.get(`/user/daily/activity/aggregated/search`, { - accessToken, - query: { - start_date: formatDate(startTime), - end_date: formatDate(endTime), - timezone: new Date().getTimezoneOffset().toString(), - search, - user_id: userId || undefined, - }, - }); - } catch (error) { - console.error("Failed to search user daily activity keys:", error); - throw error; - } -}; - export const gatewayDailyActivityCall = async (accessToken: string, startTime: Date, endTime: Date) => { /** * Get gateway request counts (SGR) recorded by the proxy middleware. diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 065e25fc48ab..4e64a98dd658 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -17354,29 +17354,6 @@ export interface paths { patch?: never; trace?: never; }; - "/user/daily/activity/aggregated/search": { - parameters: { - query?: never; - header?: never; - path?: never; - cookie?: never; - }; - /** - * Search User Daily Activity Keys - * @description Search verification tokens by exact token hash or by a case-insensitive substring of - * the key alias or owning user ID, then return the aggregated daily activity for the - * matches. Lets the Usage page surface keys that fell outside the top-spend subset - * the aggregated endpoint loads. - */ - get: operations["search_user_daily_activity_keys_user_daily_activity_aggregated_search_get"]; - put?: never; - post?: never; - delete?: never; - options?: never; - head?: never; - patch?: never; - trace?: never; - }; "/user/delete": { parameters: { query?: never; @@ -29461,11 +29438,6 @@ export interface components { }; /** DailySpendMetadata */ DailySpendMetadata: { - /** - * Api Key Limit - * @description When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key. - */ - api_key_limit?: number | null; /** * Has More * @default false @@ -29476,11 +29448,6 @@ export interface components { * @default 1 */ page: number; - /** - * Total Api Keys - * @description Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys. - */ - total_api_keys?: number | null; /** * Total Api Requests * @default 0 @@ -69061,48 +69028,6 @@ export interface operations { }; }; }; - search_user_daily_activity_keys_user_daily_activity_aggregated_search_get: { - parameters: { - query: { - /** @description Matches keys whose hash equals the value, or whose key alias or user ID contains it (case-insensitive) */ - search: string; - /** @description Start date in YYYY-MM-DD format */ - start_date?: string | null; - /** @description End date in YYYY-MM-DD format */ - end_date?: string | null; - /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ - user_id?: string | null; - /** @description Timezone offset in minutes from UTC (e.g., 480 for PST). Matches JavaScript's Date.getTimezoneOffset() convention. */ - timezone?: number | null; - /** @description When the range ends on the caller's current local day, extend it to today's UTC bucket so spend written after the caller's local midnight (in UTC terms) is included. Requires the timezone parameter. Historical ranges are never extended. */ - include_current_utc_day?: boolean; - }; - header?: never; - path?: never; - cookie?: never; - }; - requestBody?: never; - responses: { - /** @description Successful Response */ - 200: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; - }; - }; - /** @description Validation Error */ - 422: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["HTTPValidationError"]; - }; - }; - }; - }; delete_user_user_delete_post: { parameters: { query?: never;