revert(usage): remove the top-N aggregated key cap and daily global spend rollup code from rc/1.103.0 (#43326)

Reverts the code from #41324 (merge 87694c26ef) and #41293 (merge 2e46b10320).
The LiteLLM_DailyGlobalSpend model and migration stay so databases that applied
rc.1 keep a consistent table and rollup marker when the feature returns.
This commit is contained in:
yuneng-jiang 2026-09-26 11:14:53 -07:00 • committed by GitHub
parent ed9ae4fc8c
commit a78a439c68
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 218 additions and 1877 deletions

View file

@ -2084,16 +2084,12 @@ MCP_SPEND_LOG_MODEL_PREFIX: Final[str] = "MCP: "
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"
# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
# run's cutoff are stamped by different hosts, so clock skew between them must not let
# one run delete a charge another just wrote. A stale row is hours old and a concurrent

View file

@ -3050,18 +3050,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",
@ -3072,18 +3060,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",

View file

@ -9,9 +9,8 @@ from fastapi import HTTPException, status
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_emails,
recover_double_hashed_key_metadata,
@ -149,9 +148,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
@ -169,46 +174,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``
@ -731,8 +702,71 @@ 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"""
def _build_aggregated_sql_query(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
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,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
"""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().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
start_date, end_date, timezone_offset_minutes, include_current_utc_day
)
where_clause, sql_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
)
# 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,
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),
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
SUM(spend)::float AS spend,
{_ptu_flat_cost_select(table_name)},
SUM(prompt_tokens)::bigint AS prompt_tokens,
@ -748,175 +782,27 @@ def _rollup_metric_select(table_name: str) -> str:
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,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
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.
Returns:
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
start_date, end_date, timezone_offset_minutes, include_current_utc_day
)
where_clause, where_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
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)
# 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,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR},
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(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(
@ -961,7 +847,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 (
@ -1063,7 +965,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
@ -1077,7 +978,7 @@ _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
_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
@ -1431,6 +1332,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.
@ -1449,7 +1354,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,
@ -1461,19 +1366,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(**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),
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
)
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)
# 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 [])]
# The grouping-sets dispatcher places each row directly in its bucket
# using the row's GROUPING() bitmask. No Python-side summing needed.
@ -1527,8 +1449,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,
),
)

View file

@ -265,7 +265,6 @@ from litellm.constants import (
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,
@ -696,9 +695,6 @@ from litellm.proxy.route_priority import hot_routes_first
from litellm.proxy.search_endpoints.endpoints import router as search_router
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.daily_global_spend_rollup import (
run_scheduled_daily_global_spend_reconcile,
)
from litellm.proxy.spend_tracking.spend_counter_batch import (
PendingSpendIncrement,
active_spend_counter_batch,
@ -10229,12 +10225,6 @@ class ProxyStartupEvent:
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,
)
### PTU DAILY ROLLUP ###
from litellm.proxy.spend_tracking.ptu_feature_flag import (
is_ptu_cost_attribution_enabled,
@ -10576,39 +10566,6 @@ class ProxyStartupEvent:
"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
async def _initialize_slack_alerting_jobs(
cls,

View file

@ -1,293 +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 pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]:
return (await _scan_pending(prisma_client)).days
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

View file

@ -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):

View file

@ -1,21 +1,13 @@
import pathlib
import re
from collections.abc import Sequence
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 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.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.proxy.management_endpoints.common_daily_activity import (
_adjust_dates_for_timezone,
_build_aggregated_sql_query,
@ -25,12 +17,8 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
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 +169,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 +181,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 +219,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 +232,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)
@ -492,7 +474,9 @@ 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(
@ -869,7 +853,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 +865,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,11 +1309,42 @@ class TestBuildAggregatedSqlQuery:
"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
def test_model_group_rollups_fall_back_to_model_name(self):
"""Aggregated model_groups rollups must fall back to model for group-less rows.
The (date, model_group) grouping level cannot recover the model column
after the fact (it is rolled up), so the fallback has to happen in SQL;
without it, group-less rows silently vanish from the model_groups
breakdown that the usage UI now renders by default. Group-less rows are
stored as empty strings, not NULL (spend_tracking_utils defaults
model_group to ""), so a plain COALESCE is not enough: the fallback must
be NULLIF-wrapped to catch both
"""
sql, _ = _build_aggregated_sql_query(
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
start_date="2026-07-01",
end_date="2026-07-01",
model=None,
api_key=None,
)
normalized = " ".join(sql.split())
fallback = "COALESCE(NULLIF(model_group, ''), model)"
assert f"{fallback} AS model_group" in normalized
assert (
f"GROUPING(date, 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 "(date, model_group)" not in normalized
assert "COALESCE(model_group, model)" not in normalized
class TestAggregatedEmptyEntityFilter:
_BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query)
@ -1351,8 +1364,7 @@ class TestAggregatedEmptyEntityFilter:
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 +1395,7 @@ class TestAggregatedEmptyEntityFilter:
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 +1420,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,
@ -1454,484 +1464,6 @@ 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")
_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
)
"""
def _seed_daily_user_spend(conn: psycopg.Connection, rows: Sequence[tuple[object, ...]]) -> None:
with conn.cursor() as cur:
cur.execute(_DAILY_USER_SPEND_DDL)
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)
""",
rows,
)
conn.commit()
def _psycopg_query_raw(conn: psycopg.Connection, row_counts: list[int]):
"""Run the proxy's $N-parameterized SQL through psycopg, recording each result size."""
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
return query_raw
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_bounds_api_key_rollups(
_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
key_rows: Final = [
(
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(n_keys)
]
sentinel_row: Final = (
"row-ptu",
None,
"2026-06-01",
PTU_SENTINEL_API_KEY,
"gpt-5",
"",
"azure",
None,
0,
1000.0,
0,
0,
)
_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))
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,
)
# 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"}
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 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 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
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both_arms(
_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."""
rows: Final = [
(
f"row-{i}",
f"user-{i}",
"2026-06-01",
f"key-{i}",
"gpt-5",
"",
"openai",
"/v1/chat/completions",
10,
float(i + 1),
1,
1,
)
for i in range(3)
]
_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="key-1",
)
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
assert day.breakdown.models["gpt-5"].metrics.spend == 2.0
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,
):
"""Rows stored with an empty or NULL model_group must land in the model_groups
breakdown under their model name instead of vanishing from the usage UI."""
rows: Final = [
(
"row-0",
"user-0",
"2026-06-01",
"key-0",
"gpt-5",
"gpt-5-eu",
"openai",
"/v1/chat/completions",
10,
7.0,
1,
1,
),
("row-1", "user-1", "2026-06-01", "key-1", "gpt-5", "", "openai", "/v1/chat/completions", 10, 3.0, 1, 1),
("row-2", "user-2", "2026-06-01", "key-2", "claude-x", None, "anthropic", "/v1/messages", 10, 2.0, 1, 1),
]
_seed_daily_user_spend(_aggregated_postgresql, rows)
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
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=[])
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,
)
breakdown: Final = result.results[0].breakdown
assert set(breakdown.model_groups) == {"gpt-5-eu", "gpt-5", "claude-x"}
assert breakdown.model_groups["gpt-5-eu"].metrics.spend == 7.0
assert breakdown.model_groups["gpt-5"].metrics.spend == 3.0
assert breakdown.model_groups["claude-x"].metrics.spend == 2.0
assert set(breakdown.model_groups["gpt-5"].api_key_breakdown) == {"key-1"}
assert set(breakdown.models) == {"gpt-5", "claude-x"}
assert breakdown.models["gpt-5"].metrics.spend == 10.0
def _no_spend_record():
"""A rollup row for a key with no spend, where SUM() returns NULL (None)."""
return SimpleNamespace(
@ -2717,7 +2249,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 +2283,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

View file

@ -28,7 +28,6 @@ from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from fastapi import FastAPI
from pydantic import BaseModel
from typing_extensions import TypedDict
@ -1043,57 +1042,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:

View file

@ -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"}

View file

@ -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();
});
});

View file

@ -81,7 +81,7 @@ const SortableHead = ({
};
const CacheLeakageCard: React.FC<CacheLeakageCardProps> = ({ activity }) => {
const { dateValue, onDateChange, results, loading, isFetchingMore, apiKeyTruncation } = activity;
const { dateValue, onDateChange, results, loading, isFetchingMore } = activity;
const [dimension, setDimension] = useState<CacheLeakageDimension>("key");
const [sort, setSort] = useState<SortState>({ column: "potentialSavings", dir: "desc" });
const leakage = useMemo(() => computeCacheLeakage(results, dimension), [results, dimension]);
@ -123,13 +123,6 @@ const CacheLeakageCard: React.FC<CacheLeakageCardProps> = ({ activity }) => {
</Tabs>
</CardHeader>
<CardContent>
{dimension === "key" && apiKeyTruncation !== undefined && (
<p className="mb-2 text-sm text-muted-foreground" role="note">
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.
</p>
)}
{rows.length > 0 && isFetchingMore && (
<p className="mb-2 text-sm text-muted-foreground">
Data is still loading; rows and totals will update as the rest of the range arrives.

View file

@ -4,13 +4,12 @@ import { describe, expect, it, vi } from "vitest";
const mockUsePaginatedDailyActivity = vi.fn();
const mockCancel = vi.fn();
let mockMetadata: Record<string, number> = {};
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 },
@ -81,18 +80,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();
});
});

View file

@ -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;
}
/**
@ -80,7 +78,6 @@ export const useScopedDailyActivityRange = (
cancelled,
failed,
cancel,
apiKeyTruncation: getApiKeyTruncation(data.metadata?.api_key_limit, data.metadata?.total_api_keys),
};
};

View file

@ -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(<EntityUsage {...defaultProps} entityType="team" />);
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 => {

View file

@ -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<EntityUsageProps> = ({
});
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<EntityUsageProps> = ({
{
key: "keys",
label: "Key Activity",
content: (
<KeyActivityPanel
keyMetrics={keyMetrics}
hidePromptCachingMetrics={entityType === "agent"}
apiKeyTruncation={apiKeyTruncation}
/>
),
content: <KeyActivityPanel keyMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} />,
},
{ key: "endpoints", label: "Endpoint Activity", content: <EndpointUsage userSpendData={spendData} /> },
];
const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation };
const spendFetchState = { coversRange, cancelled, failed };
return (
<div style={{ width: "100%" }} className="relative">

View file

@ -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 {
@ -256,10 +256,6 @@ const UsagePage: React.FC<UsagePageProps> = ({ 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);
@ -865,7 +861,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<ActivityMetrics modelMetrics={modelMetrics} />
</TabsContent>
<TabsContent value="keys" keepMounted>
<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={spendFetchState.apiKeyTruncation} />
<KeyActivityPanel keyMetrics={keyMetrics} />
</TabsContent>
<TabsContent value="mcp" keepMounted>
<ActivityMetrics modelMetrics={mcpServerMetrics} />

View file

@ -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> = {}): 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();
});
});

View file

@ -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;
};

View file

@ -68,14 +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(<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3000 }} />);
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(<KeyActivityPanel keyMetrics={keyMetrics} />);
expect(screen.queryByRole("note")).not.toBeInTheDocument();
});
});

View file

@ -2,7 +2,6 @@ import { Search, X } from "lucide-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,14 +10,9 @@ import type { ModelActivityData } from "../types";
interface KeyActivityPanelProps {
keyMetrics: Record<string, ModelActivityData>;
hidePromptCachingMetrics?: boolean;
apiKeyTruncation?: ApiKeyTruncation;
}
const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
keyMetrics,
hidePromptCachingMetrics = false,
apiKeyTruncation,
}) => {
const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({ keyMetrics, hidePromptCachingMetrics = false }) => {
const [query, setQuery] = useState("");
const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]);
const totalKeys = Object.keys(keyMetrics).length;
@ -49,12 +43,6 @@ const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
<span className="text-sm text-muted-foreground">
Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys
</span>
{apiKeyTruncation !== undefined && (
<span className="text-sm text-muted-foreground" role="note">
Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "}
{apiKeyTruncation.total.toLocaleString()} are loaded
</span>
)}
</div>
{isFiltering && totalKeys > 0 && shownKeys === 0 ? (
<p className="rounded-lg border p-6 text-center text-sm text-muted-foreground">

View file

@ -27961,11 +27961,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
@ -27976,11 +27971,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