From 49fde3761b99177697543ba6fa02843805fafa78 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 15:52:30 +0000 Subject: [PATCH 1/5] fix(proxy): bound api_key rollups in aggregated usage query to top keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../common_daily_activity.py | 42 ++++++++++++---- .../test_common_daily_activity.py | 49 ++++++++++++++++++- 2 files changed, 80 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index ce6a97708ab..e6e06accdb6 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -710,20 +710,33 @@ def _build_aggregated_sql_query( # is omitted on purpose: nothing in the response shape needs it once # all the rollups are present. # + # api_key appears as tk.top_api_key from the top_api_keys CTE, bounding + # the api_key-keyed sets to the top _MAX_API_KEYS_IN_BREAKDOWN keys by + # spend instead of every distinct key in the window. The same + # where_clause/$N params run in both the CTE and the outer query. + # # 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""" + WITH top_api_keys AS ( + SELECT api_key AS top_api_key + FROM "{pg_table}" + WHERE {where_clause} + GROUP BY api_key + ORDER BY SUM(spend) DESC + LIMIT {_MAX_API_KEYS_IN_BREAKDOWN} + ) SELECT date, - api_key, + tk.top_api_key AS api_key, model, COALESCE(NULLIF(model_group, ''), model) AS model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint, - GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model), + GROUPING(date, tk.top_api_key, model, COALESCE(NULLIF(model_group, ''), model), custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level, SUM(spend)::float AS spend, @@ -740,21 +753,22 @@ def _build_aggregated_sql_query( SUM(api_requests)::bigint AS api_requests, SUM(successful_requests)::bigint AS successful_requests, SUM(failed_requests)::bigint AS failed_requests - FROM "{pg_table}" + FROM "{pg_table}" t + LEFT JOIN top_api_keys tk ON tk.top_api_key = t.api_key WHERE {where_clause} GROUP BY GROUPING SETS ( (date), - (date, api_key), + (date, tk.top_api_key), (date, model), - (date, model, api_key), + (date, model, tk.top_api_key), (date, COALESCE(NULLIF(model_group, ''), model)), - (date, COALESCE(NULLIF(model_group, ''), model), api_key), + (date, COALESCE(NULLIF(model_group, ''), model), tk.top_api_key), (date, custom_llm_provider), - (date, custom_llm_provider, api_key), + (date, custom_llm_provider, tk.top_api_key), (date, mcp_namespaced_tool_name), - (date, mcp_namespaced_tool_name, api_key), + (date, mcp_namespaced_tool_name, tk.top_api_key), (date, endpoint), - (date, endpoint, api_key), + (date, endpoint, tk.top_api_key), () ) """ @@ -930,6 +944,16 @@ _GROUP_DATE_MCP_API_KEY: Final = 29 # 0b0011101 _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110 _GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110 +# Cap on distinct api_keys carried into the api_key-keyed grouping sets of +# _build_aggregated_sql_query. Six of the thirteen sets include api_key, so +# result rows scale with distinct-key count; on large deployments the +# prisma-query-engine buffers the whole result and gets OOM-killed. The UI +# only renders up to 50 top keys (TOP_KEYS_LIMITS), so 100 is generous. +# Keys outside the top N group into a NULL api_key bucket that the +# dispatcher skips, and non-keyed totals are unaffected (LEFT JOIN keeps +# every row). +_MAX_API_KEYS_IN_BREAKDOWN: Final = 100 + def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: """Build a SpendMetrics directly from one already-aggregated rollup row. 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 37a54c4901a..173c4b4cccd 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 @@ -9,6 +9,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_E from litellm.proxy.management_endpoints.common_daily_activity import ( + _MAX_API_KEYS_IN_BREAKDOWN, _adjust_dates_for_timezone, _build_aggregated_sql_query, _build_entity_rollup_sql_query, @@ -1149,13 +1150,57 @@ class TestBuildAggregatedSqlQuery: fallback = "COALESCE(NULLIF(model_group, ''), model)" assert f"{fallback} AS model_group" in normalized assert ( - f"GROUPING(date, api_key, model, {fallback}, " + f"GROUPING(date, tk.top_api_key, model, {fallback}, " "custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level" in normalized ) - assert f"(date, {fallback}), (date, {fallback}, api_key)," in normalized + assert f"(date, {fallback}), (date, {fallback}, tk.top_api_key)," in normalized assert "(date, model_group)" not in normalized assert "COALESCE(model_group, model)" not in normalized + def test_api_key_rollups_are_bounded_to_top_keys(self): + """api_key-keyed grouping sets must group on the bounded top-N CTE, not the raw column. + + Six of the thirteen grouping sets include api_key, so result rows used + to scale with the total distinct-key count in the window; the + prisma-query-engine buffered the whole result and got OOM-killed. The + top_api_keys CTE bounds the dimension to the top + _MAX_API_KEYS_IN_BREAKDOWN keys by spend; every api_key grouping must + reference tk.top_api_key, and grouping on the raw api_key column is the + regression this guards. + """ + sql, params = _build_aggregated_sql_query( + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id="user-1", + start_date="2026-05-29", + end_date="2026-06-02", + model="bedrock/global.anthropic.claude-opus-4-8", + api_key="sk-test", + timezone_offset_minutes=-330, + ) + + normalized = " ".join(sql.split()) + assert "WITH top_api_keys AS ( SELECT api_key AS top_api_key" in normalized + assert f"ORDER BY SUM(spend) DESC LIMIT {_MAX_API_KEYS_IN_BREAKDOWN}" in normalized + + grouping_block = normalized.split("GROUP BY GROUPING SETS (", 1)[1] + grouping_sets = {part.strip() for part in grouping_block.split("),")} + for grouping_set in grouping_sets: + assert "api_key" not in grouping_set.replace("tk.top_api_key", ""), ( + f"grouping set uses the unbounded raw api_key column: {grouping_set}" + ) + assert grouping_block.count("tk.top_api_key") == 6 + + # The CTE and the outer WHERE reuse the same $N placeholders, so the + # params list is unchanged: date bounds, entity, model, api_key filter. + assert params == [ + "2026-05-29", + "2026-06-02", + "user-1", + "bedrock/global.anthropic.claude-opus-4-8", + "sk-test", + ] + class TestAggregatedEmptyEntityFilter: _BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query) From 570251e935ad51d47ea139dd6ca9083ae4ab08e4 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 15:53:24 +0000 Subject: [PATCH 2/5] refactor(proxy): trim comments on api_key rollup cap Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../common_daily_activity.py | 16 +++------------- .../test_common_daily_activity.py | 11 +---------- 2 files changed, 4 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index e6e06accdb6..c982fa2de64 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -710,11 +710,6 @@ def _build_aggregated_sql_query( # is omitted on purpose: nothing in the response shape needs it once # all the rollups are present. # - # api_key appears as tk.top_api_key from the top_api_keys CTE, bounding - # the api_key-keyed sets to the top _MAX_API_KEYS_IN_BREAKDOWN keys by - # spend instead of every distinct key in the window. The same - # where_clause/$N params run in both the CTE and the outer query. - # # 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 @@ -944,14 +939,9 @@ _GROUP_DATE_MCP_API_KEY: Final = 29 # 0b0011101 _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110 _GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110 -# Cap on distinct api_keys carried into the api_key-keyed grouping sets of -# _build_aggregated_sql_query. Six of the thirteen sets include api_key, so -# result rows scale with distinct-key count; on large deployments the -# prisma-query-engine buffers the whole result and gets OOM-killed. The UI -# only renders up to 50 top keys (TOP_KEYS_LIMITS), so 100 is generous. -# Keys outside the top N group into a NULL api_key bucket that the -# dispatcher skips, and non-keyed totals are unaffected (LEFT JOIN keeps -# every row). +# Six of the thirteen grouping sets in _build_aggregated_sql_query are keyed on api_key, so +# without a cap the result grows with every distinct key and the prisma query engine OOMs +# buffering it. Keys outside the top N fall into a NULL api_key bucket the dispatcher skips. _MAX_API_KEYS_IN_BREAKDOWN: Final = 100 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 173c4b4cccd..415157960c0 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 @@ -1158,16 +1158,7 @@ class TestBuildAggregatedSqlQuery: assert "COALESCE(model_group, model)" not in normalized def test_api_key_rollups_are_bounded_to_top_keys(self): - """api_key-keyed grouping sets must group on the bounded top-N CTE, not the raw column. - - Six of the thirteen grouping sets include api_key, so result rows used - to scale with the total distinct-key count in the window; the - prisma-query-engine buffered the whole result and got OOM-killed. The - top_api_keys CTE bounds the dimension to the top - _MAX_API_KEYS_IN_BREAKDOWN keys by spend; every api_key grouping must - reference tk.top_api_key, and grouping on the raw api_key column is the - regression this guards. - """ + """Grouping on the raw api_key column makes result rows scale with every distinct key.""" sql, params = _build_aggregated_sql_query( table_name="litellm_dailyuserspend", entity_id_field="user_id", From 31cedbad81648040c27fd4169c013cfef7bf7a28 Mon Sep 17 00:00:00 2001 From: jesus Date: Tue, 8 Sep 2026 15:53:52 +0000 Subject: [PATCH 3/5] style(tests): format daily activity test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_common_daily_activity.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 415157960c0..9ad0753f914 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 @@ -473,9 +473,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] ) mock_prisma.db.query_raw = AsyncMock( - return_value=[ - {"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"} - ] + return_value=[{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}] ) result = await get_api_key_metadata( From 9f25236425f8bc9ed0ca5adfef893f2347ed5da1 Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 12 Sep 2026 16:49:50 +0000 Subject: [PATCH 4/5] fix(proxy): break top api_key ties deterministically Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/common_daily_activity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index c982fa2de64..d37e32c37bb 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -720,7 +720,7 @@ def _build_aggregated_sql_query( FROM "{pg_table}" WHERE {where_clause} GROUP BY api_key - ORDER BY SUM(spend) DESC + ORDER BY SUM(spend) DESC, api_key LIMIT {_MAX_API_KEYS_IN_BREAKDOWN} ) SELECT From d43af0ece0c8dbc900a31fe87029e071c8acb4e0 Mon Sep 17 00:00:00 2001 From: jesus Date: Sat, 12 Sep 2026 16:49:50 +0000 Subject: [PATCH 5/5] test(proxy): exercise bounded api_key rollups against postgres Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_common_daily_activity.py | 165 ++++++++++++++---- 1 file changed, 127 insertions(+), 38 deletions(-) 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 9ad0753f914..b91e51c3946 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,12 +1,13 @@ +import re from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock +import psycopg import pytest - -from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR - +from psycopg.rows import dict_row +from pytest_postgresql import factories from litellm.proxy.management_endpoints.common_daily_activity import ( _MAX_API_KEYS_IN_BREAKDOWN, @@ -20,6 +21,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, update_metrics, ) +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, SpendMetrics, @@ -1155,41 +1157,6 @@ class TestBuildAggregatedSqlQuery: assert "(date, model_group)" not in normalized assert "COALESCE(model_group, model)" not in normalized - def test_api_key_rollups_are_bounded_to_top_keys(self): - """Grouping on the raw api_key column makes result rows scale with every distinct key.""" - sql, params = _build_aggregated_sql_query( - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id="user-1", - start_date="2026-05-29", - end_date="2026-06-02", - model="bedrock/global.anthropic.claude-opus-4-8", - api_key="sk-test", - timezone_offset_minutes=-330, - ) - - normalized = " ".join(sql.split()) - assert "WITH top_api_keys AS ( SELECT api_key AS top_api_key" in normalized - assert f"ORDER BY SUM(spend) DESC LIMIT {_MAX_API_KEYS_IN_BREAKDOWN}" in normalized - - grouping_block = normalized.split("GROUP BY GROUPING SETS (", 1)[1] - grouping_sets = {part.strip() for part in grouping_block.split("),")} - for grouping_set in grouping_sets: - assert "api_key" not in grouping_set.replace("tk.top_api_key", ""), ( - f"grouping set uses the unbounded raw api_key column: {grouping_set}" - ) - assert grouping_block.count("tk.top_api_key") == 6 - - # The CTE and the outer WHERE reuse the same $N placeholders, so the - # params list is unchanged: date bounds, entity, model, api_key filter. - assert params == [ - "2026-05-29", - "2026-06-02", - "user-1", - "bedrock/global.anthropic.claude-opus-4-8", - "sk-test", - ] - class TestAggregatedEmptyEntityFilter: _BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query) @@ -1307,6 +1274,128 @@ async def test_get_daily_activity_aggregated_empty_result_set(): assert result.metadata.total_compression_saved_tokens == 0 +_aggregated_postgresql_proc: Final = factories.postgresql_proc() +_aggregated_postgresql: Final = factories.postgresql("_aggregated_postgresql_proc") + + +@pytest.mark.asyncio +async def test_get_daily_activity_aggregated_bounds_api_key_rollups( + _aggregated_postgresql: psycopg.Connection, +): + """Exercise the GROUPING SETS query against real Postgres with more keys than the cap. + + key-004 and key-005 tie on spend exactly at the _MAX_API_KEYS_IN_BREAKDOWN + cutoff; the deterministic api_key tiebreaker must keep key-004 and drop + key-005. Excluded keys still count toward the totals via the grand-total + and date-level rollup rows. + """ + conn: Final = _aggregated_postgresql + with conn.cursor() as cur: + cur.execute( + """ + 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 + ) + """ + ) + 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) + """, + [ + ( + f"row-{i:03d}", + f"user-{i:03d}", + "2026-06-01", + f"key-{i:03d}", + "gpt-5", + "", + "openai", + "/v1/chat/completions", + 10, + 6.0 if i == 4 else float(i + 1), + 1, + 1, + ) + for i in range(_MAX_API_KEYS_IN_BREAKDOWN + 5) + ], + ) + conn.commit() + + row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim + + async def query_raw(sql: str, *params: str) -> list[dict[str, object]]: + converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) + with conn.cursor(row_factory=dict_row) as cur: + cur.execute( + converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + rows: Final = cur.fetchall() + row_counts.append(len(rows)) + return rows + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = query_raw + 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, + ) + + # 13 grouping sets: 7 key-less sets emit one row each; the 6 api_key sets + # emit the top N keys plus one NULL bucket for the excluded keys. + assert row_counts == [7 + 6 * (_MAX_API_KEYS_IN_BREAKDOWN + 1)] + + # Totals still cover every key, capped or not. + assert result.metadata.total_spend == pytest.approx(5566.0) + assert result.metadata.total_api_requests == _MAX_API_KEYS_IN_BREAKDOWN + 5 + + expected_top: Final = {f"key-{i:03d}" for i in range(6, 105)} | {"key-004"} + day: Final = result.results[0] + assert len(day.breakdown.api_keys) == _MAX_API_KEYS_IN_BREAKDOWN + 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 day.breakdown.models["gpt-5"].metrics.spend == pytest.approx(5566.0) + assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_top + + def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace(