From c8b0530c30c678f27f0b70359566925d303aca99 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 25 Jul 2026 17:38:58 -0700 Subject: [PATCH 1/7] fix(proxy): roll up tool spend daily instead of scanning SpendLogs GET /v1/tool/spend served the Cost Optimization card with two raw queries over LiteLLM_SpendLogToolIndex x LiteLLM_SpendLogs on every dashboard load; the totals query's driving scan was all of SpendLogs in the window. Both per-request tables reach 1M+ rows at customer scale, so the card cost O(traffic) per view and had to be capped at 30 days. The index writer also mined proxy_server_request.tools, i.e. tools DECLARED in the request body, attributing each request's full spend to tools that never ran; and all non-MCP mining ran against payload fields that are '{}' unless store_prompts_in_spend_logs is enabled, so non-MCP coverage silently depended on a privacy setting. Now the spend writer builds a ToolUsageTransaction at request time from invoked tools only, resolved by the shared get_tool_calls_from_response normalizer so every response surface (chat completions, Responses API, Anthropic Messages) is covered; the tool registry's response arm delegates to the same owner. Transactions queue beside the spend-log queue and the flush job writes index rows plus a new LiteLLM_DailyToolSpend rollup (date, tool_name PK) in one transaction, retrying connection errors with backoff (a failed batch commits nothing, so the retry cannot double-count) and dropping the batch with an error log on anything else. The endpoint aggregates in SQL: by_tool is the top TOOL_SPEND_TOP_TOOLS tools by spend via group_by and daily covers only those tools, so the response is bounded by days x TOOL_SPEND_TOP_TOOLS regardless of range or tool-name cardinality; the 30-day clamp is gone. total_spend is dropped from the response; it was never rendered and its deduplicated semantics are not computable from a rollup. Spend-log retention deliberately does not touch the rollup, so tool spend history outlives per-request rows. --- db_scripts/backfill_daily_tool_spend.sql | 44 +++ .../migration.sql | 12 + .../litellm_proxy_extras/schema.prisma | 13 + litellm/constants.py | 2 +- .../prompt_templates/factory.py | 5 +- litellm/proxy/_lazy_openapi_snapshot.json | 10 +- litellm/proxy/db/db_spend_update_writer.py | 55 +++- litellm/proxy/db/spend_log_tool_index.py | 261 ++++++++------- .../tool_management_endpoints.py | 170 ++++------ litellm/proxy/schema.prisma | 13 + litellm/proxy/utils.py | 42 ++- litellm/repositories/__init__.py | 2 + litellm/repositories/table_repositories.py | 4 + litellm/types/tool_management.py | 7 - schema.prisma | 13 + tests/proxy_unit_tests/test_update_spend.py | 4 +- .../proxy/db/test_db_spend_update_writer.py | 138 ++++++++ .../proxy/db/test_spend_log_tool_index.py | 309 ++++++++++++++++++ .../test_tool_management_endpoints.py | 220 ++++++------- .../proxy/test_spend_log_cleanup.py | 6 + .../proxy/utils/prisma_and_spend/conftest.py | 2 + .../prisma_and_spend/test_spend_functions.py | 31 +- ui/litellm-dashboard/eslint-suppressions.json | 7 +- .../CostOptimizationView.activity.test.tsx | 2 +- .../_components/UsageTab.test.tsx | 30 +- .../_components/UsageTab.tsx | 13 +- .../src/components/ToolDetail.tsx | 2 +- .../src/components/networking.tsx | 1 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 27 +- 29 files changed, 981 insertions(+), 464 deletions(-) create mode 100644 db_scripts/backfill_daily_tool_spend.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql create mode 100644 tests/test_litellm/proxy/db/test_spend_log_tool_index.py diff --git a/db_scripts/backfill_daily_tool_spend.sql b/db_scripts/backfill_daily_tool_spend.sql new file mode 100644 index 00000000000..358ebf1f23f --- /dev/null +++ b/db_scripts/backfill_daily_tool_spend.sql @@ -0,0 +1,44 @@ +-- One-shot backfill of the LiteLLM_DailyToolSpend rollup from the per-request +-- LiteLLM_SpendLogToolIndex x LiteLLM_SpendLogs tables. +-- +-- This is an opt-in, manual operation. New deployments do not need it: the +-- rollup is written at request time from the moment the release is deployed. +-- Run it only if you want the Cost Optimization "Spend by tool" card to show +-- history from before the deploy, and only once. +-- +-- IMPORTANT caveats before running: +-- +-- 1. Pre-deploy index rows may include tools that were merely DECLARED in a +-- request body but never invoked (the release this ships with stops +-- recording those). For agentic clients that declare many tools per +-- request, backfilled history attributes each request's full spend to +-- every declared tool, overstating per-tool spend. Post-deploy rows do not +-- have this problem. If your traffic is mostly such clients, consider not +-- backfilling. +-- +-- 2. Coverage is bounded by spend-log retention: rows older than +-- maximum_spend_logs_retention_period are already gone. +-- +-- 3. Replace the cutover timestamp below with the time you deployed the +-- release, so backfilled per-request rows cannot double-count on top of +-- rollup rows the new writer already created. ON CONFLICT DO NOTHING is a +-- second guard for (date, tool_name) buckets the writer already touched: +-- such buckets keep the writer's numbers and skip the backfill's. +-- +-- Usage: +-- psql "$DATABASE_URL" -v cutover="'2026-07-25T00:00:00Z'" -f db_scripts/backfill_daily_tool_spend.sql + +INSERT INTO "LiteLLM_DailyToolSpend" (date, tool_name, spend, total_tokens, request_count, created_at, updated_at) +SELECT + to_char(ti.start_time, 'YYYY-MM-DD') AS date, + ti.tool_name, + COALESCE(SUM(sl.spend), 0) AS spend, + COALESCE(SUM(sl.total_tokens), 0) AS total_tokens, + COUNT(*) AS request_count, + now() AS created_at, + now() AS updated_at +FROM "LiteLLM_SpendLogToolIndex" ti +JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id +WHERE ti.start_time < :cutover::timestamptz +GROUP BY 1, 2 +ON CONFLICT (date, tool_name) DO NOTHING; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql new file mode 100644 index 00000000000..e02ed01a554 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql @@ -0,0 +1,12 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyToolSpend" ( + "date" TEXT NOT NULL, + "tool_name" TEXT NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "total_tokens" BIGINT NOT NULL DEFAULT 0, + "request_count" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyToolSpend_pkey" PRIMARY KEY ("date","tool_name") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 6713b212314..37ea55f8c13 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1097,6 +1097,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/litellm/constants.py b/litellm/constants.py index a9edf135731..1014b472c61 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1457,7 +1457,7 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEA SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float( os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) ) -TOOL_SPEND_MAX_WINDOW_DAYS = 30 +TOOL_SPEND_TOP_TOOLS = 100 SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index f7ff4d6b16f..c13cf0817b5 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -6,6 +6,7 @@ import mimetypes import re import xml.etree.ElementTree as ET from enum import Enum +from collections.abc import Mapping from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict, Union, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -5350,7 +5351,9 @@ def prompt_factory( def get_attribute_or_key(tool_or_function, attribute, default=None): if hasattr(tool_or_function, attribute): return getattr(tool_or_function, attribute) - return tool_or_function.get(attribute, default) + if isinstance(tool_or_function, Mapping): + return tool_or_function.get(attribute, default) + return default class NormalizedToolCall(TypedDict): diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 972c831073f..96f6ee89d56 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -26858,12 +26858,6 @@ } ], "title": "Start Date" - }, - "total_spend": { - "default": 0.0, - "description": "Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist", - "title": "Total Spend", - "type": "number" } }, "title": "ToolSpendResponse", @@ -27417,7 +27411,7 @@ }, "/v1/tool/spend": { "get": { - "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nJoins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to\n``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools\ncounts its full spend toward each of those tools, so per-tool numbers are\nattributions. ``total_spend`` is the deduplicated spend of every request that\ncalled at least one tool in the window, so it never double counts.\n\n``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to\n31 calendar dates inclusive, the same width as the endpoint's default window):\na wider requested range is clamped, and the response's ``start_date`` reflects\nthe effective window actually served.", + "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nReads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked\ntools only (MCP tool calls and response tool_calls; declaring a tool without\ninvoking it does not count). A request that invoked multiple tools counts its\nfull spend toward each of them, so per-tool numbers are attributions and do not\nsum to a deduplicated total.\n\n``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in\nSQL, and ``daily`` covers only those tools, so the response is bounded by\ndays x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many\ndistinct tool names exist.", "operationId": "get_tool_spend_v1_tool_spend_get", "parameters": [ { @@ -27588,7 +27582,7 @@ }, "/v1/tool/{tool_name}/logs": { "get": { - "description": "Return paginated spend logs for requests that used this tool (from SpendLogToolIndex).", + "description": "Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex).\nDeclaring a tool in a request body without the model invoking it does not create an entry.", "operationId": "get_tool_usage_logs_v1_tool__tool_name__logs_get", "parameters": [ { diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2262141f426..ebdb08a681a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -182,6 +182,12 @@ class DBSpendUpdateWriter: payload=payload, prisma_client=prisma_client, ) + await self._enqueue_tool_usage_transaction( + payload=payload, + completion_response=completion_response, + prisma_client=prisma_client, + kwargs=kwargs, + ) else: verbose_proxy_logger.debug( "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." @@ -223,6 +229,36 @@ class DBSpendUpdateWriter: end_user_id, ) + async def _enqueue_tool_usage_transaction( + self, + payload: SpendLogsPayload, + completion_response: "litellm.ModelResponse | Any | Exception | None", + prisma_client: "PrismaClient | None", + kwargs: "dict | None" = None, + ) -> None: + try: + if prisma_client is None: + return + from litellm.proxy.db.spend_log_tool_index import ( + build_tool_usage_transaction, + ) + + transaction = build_tool_usage_transaction( + request_id=payload["request_id"], + start_time_iso=str(payload["startTime"]), + mcp_namespaced_tool_name=payload.get("mcp_namespaced_tool_name"), + spend=payload["spend"], + total_tokens=payload["total_tokens"], + completion_response=completion_response, + realtime_tool_calls=(kwargs or {}).get("realtime_tool_calls"), + ) + if transaction is None: + return + async with prisma_client._tool_usage_transactions_lock: + prisma_client.tool_usage_transactions.append(transaction) + except Exception as e: + verbose_proxy_logger.debug("_enqueue_tool_usage_transaction error (non-blocking): %s", e) + def _enqueue_tool_registry_upsert( self, kwargs: Optional[dict], @@ -299,21 +335,10 @@ class DBSpendUpdateWriter: _enqueue(name) # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) --- - if completion_response is not None and hasattr(completion_response, "choices"): - for choice in completion_response.choices or []: - message = getattr(choice, "message", None) - if message is None: - continue - tool_calls = getattr(message, "tool_calls", None) - if not tool_calls: - continue - for tc in tool_calls: - fn = getattr(tc, "function", None) - if fn is None: - continue - tool_name = getattr(fn, "name", None) - if tool_name: - _enqueue(tool_name) + from litellm.proxy.db.spend_log_tool_index import response_tool_call_names + + for tool_name in response_tool_call_names(completion_response): + _enqueue(tool_name) except Exception as e: verbose_proxy_logger.debug("_enqueue_tool_registry_upsert error (non-blocking): %s", e) diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 80036e235f7..064d08acb59 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -1,140 +1,147 @@ """ -Track tool usage for the dashboard: insert into SpendLogToolIndex when spend logs -are written, so "last N requests for tool X" and "how is this tool called in production" -queries are fast. +Tool usage tracking for the dashboard. + +At request time the spend writer builds one ToolUsageTransaction per request that +invoked tools (MCP namespaced tool name plus response tool_calls; declared-but-not- +invoked tools are excluded) and queues it on the prisma client. The spend-log flush +job drains the queue into LiteLLM_SpendLogToolIndex (per-request drill-down) and +LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads) in a +single transaction, so a failed flush never leaves a partial rollup increment. """ +from __future__ import annotations + +import asyncio +import random +from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Dict, List, Set +from itertools import groupby +from typing import TYPE_CHECKING, Any, Sequence -from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy.utils import PrismaClient -from litellm.repositories.table_repositories import SpendLogToolIndexRepository +from litellm.proxy._types import DB_CONNECTION_ERROR_TYPES + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient -def _add_tool_calls_to_set(tool_calls: Any, out: Set[str]) -> None: - """Extract tool names from OpenAI-style tool_calls list into out.""" - if not isinstance(tool_calls, list): - return - for tc in tool_calls: - if not isinstance(tc, dict): - continue - fn = tc.get("function") - if isinstance(fn, dict): - name = fn.get("name") - if name and isinstance(name, str) and name.strip(): - out.add(name.strip()) +@dataclass(frozen=True, slots=True) +class ToolUsageTransaction: + request_id: str + date: str + start_time: datetime + tool_names: tuple[str, ...] + spend: float + total_tokens: int -def _parse_tool_names_from_payload(payload: Dict[str, Any]) -> Set[str]: - """ - Extract deduplicated tool names from a spend log payload. - Sources: mcp_namespaced_tool_name, response (tool_calls), proxy_server_request (tools). - """ - tool_names: Set[str] = set() +def response_tool_call_names(completion_response: Any) -> tuple[str, ...]: + """Tool names invoked in a completion response, in call order, for any response + surface get_tool_calls_from_response understands (chat completions, Responses + API output items, Anthropic Messages tool_use blocks).""" + if completion_response is None or isinstance(completion_response, Exception): + return () + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) - # Top-level MCP tool name (single tool per request for that flow) - mcp_name = payload.get("mcp_namespaced_tool_name") - if mcp_name and isinstance(mcp_name, str) and mcp_name.strip(): - tool_names.add(mcp_name.strip()) - - # Response: OpenAI-style tool_calls[].function.name or choices[0].message.tool_calls - response_raw = payload.get("response") - if response_raw: - response_obj = safe_json_loads(response_raw, default=None) if isinstance(response_raw, str) else response_raw - if isinstance(response_obj, dict): - _add_tool_calls_to_set(response_obj.get("tool_calls"), tool_names) - choices = response_obj.get("choices") - if isinstance(choices, list) and choices: - msg = choices[0].get("message") if isinstance(choices[0], dict) else None - if isinstance(msg, dict): - _add_tool_calls_to_set(msg.get("tool_calls"), tool_names) - - # Request body: tools[].function.name - request_raw = payload.get("proxy_server_request") - if request_raw: - request_obj = safe_json_loads(request_raw, default=None) if isinstance(request_raw, str) else request_raw - if isinstance(request_obj, dict): - body = request_obj.get("body", request_obj) - if isinstance(body, dict): - request_obj = body - if isinstance(request_obj, dict): - tools = request_obj.get("tools") - if isinstance(tools, list): - for t in tools: - if isinstance(t, dict): - fn = t.get("function") - if isinstance(fn, dict): - name = fn.get("name") - if name and isinstance(name, str) and name.strip(): - tool_names.add(name.strip()) - - return tool_names + return tuple( + stripped + for tool_call in get_tool_calls_from_response(completion_response) + if isinstance(name := tool_call.get("name"), str) and (stripped := name.strip()) + ) -async def process_spend_logs_tool_usage( - prisma_client: PrismaClient, - logs_to_process: List[Dict[str, Any]], -) -> None: - """ - After spend logs are written: insert SpendLogToolIndex rows from each payload. - Extracts tool names from mcp_namespaced_tool_name, response tool_calls, and - proxy_server_request tools. - """ - if not logs_to_process: - return - - index_rows: List[Dict[str, Any]] = [] - - for payload in logs_to_process: - request_id = payload.get("request_id") - start_time = payload.get("startTime") - if not request_id or not start_time: - continue - if isinstance(start_time, str): - try: - start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue - if start_time.tzinfo is None: - start_time = start_time.replace(tzinfo=timezone.utc) - - tool_names = _parse_tool_names_from_payload(payload) - for tool_name in tool_names: - index_rows.append( - { - "request_id": request_id, - "tool_name": tool_name, - "start_time": start_time, - } - ) - - if not index_rows: - return - +def build_tool_usage_transaction( + request_id: str, + start_time_iso: str, + mcp_namespaced_tool_name: str | None, + spend: float, + total_tokens: int, + completion_response: Any, + realtime_tool_calls: Any = None, +) -> ToolUsageTransaction | None: + """None when the request invoked no tools. Realtime sessions carry invoked + tools in kwargs["realtime_tool_calls"] (OpenAI tool_calls shape) rather than + on a response object, so they are normalized through the same owner by + wrapping them in the chat-completion shape. Date derivation must match the + daily spend writer's ``startTime.split("T")[0]`` so rollup rows land in the + same UTC day bucket as LiteLLM_DailyUserSpend.""" + mcp_names = ( + (mcp_namespaced_tool_name.strip(),) if mcp_namespaced_tool_name and mcp_namespaced_tool_name.strip() else () + ) + realtime_names = ( + response_tool_call_names({"choices": [{"message": {"tool_calls": realtime_tool_calls}}]}) + if realtime_tool_calls + else () + ) + tool_names = tuple(dict.fromkeys(mcp_names + response_tool_call_names(completion_response) + realtime_names)) + if not tool_names: + return None try: - index_data = [] - for r in index_rows: - st = r["start_time"] - if isinstance(st, str): - try: - st = datetime.fromisoformat(st.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue - if st.tzinfo is None: - st = st.replace(tzinfo=timezone.utc) - index_data.append( - { - "request_id": r["request_id"], - "tool_name": r["tool_name"], - "start_time": st, - } - ) - if index_data: - await SpendLogToolIndexRepository(prisma_client).table.create_many( - data=index_data, - skip_duplicates=True, - ) - except Exception as e: - verbose_proxy_logger.warning("Tool usage tracking (SpendLogToolIndex) failed (non-fatal): %s", e) + start_time = datetime.fromisoformat(start_time_iso.replace("Z", "+00:00")) + except ValueError: + return None + return ToolUsageTransaction( + request_id=request_id, + date=start_time_iso.split("T")[0], + start_time=start_time if start_time.tzinfo else start_time.replace(tzinfo=timezone.utc), + tool_names=tool_names, + spend=spend, + total_tokens=total_tokens, + ) + + +async def flush_tool_usage_transactions( + prisma_client: PrismaClient, + transactions: Sequence[ToolUsageTransaction], + n_retry_times: int = 3, +) -> None: + """Write index rows and rollup upserts for a drained queue batch in one + transaction. Connection errors are retried with backoff, which cannot + double-count because a failed batch commits nothing; every other error + propagates so the caller drops the batch. Callers must not add their own + retry around this function: a batch that DID commit must never run again, + since the rollup update increments counters.""" + if not transactions: + return + + index_rows = [ + {"request_id": txn.request_id, "tool_name": tool_name, "start_time": txn.start_time} + for txn in transactions + for tool_name in txn.tool_names + ] + per_tool_day = sorted( + ((txn.date, tool_name, txn.spend, txn.total_tokens) for txn in transactions for tool_name in txn.tool_names), + key=lambda entry: (entry[0], entry[1]), + ) + + for attempt in range(n_retry_times + 1): + try: + async with prisma_client.db.batch_() as batcher: + batcher.litellm_spendlogtoolindex.create_many(data=index_rows, skip_duplicates=True) + for (date_key, tool_name), grouped in groupby(per_tool_day, key=lambda entry: (entry[0], entry[1])): + entries = tuple(grouped) + spend = sum(entry[2] for entry in entries) + total_tokens = sum(entry[3] for entry in entries) + batcher.litellm_dailytoolspend.upsert( + where={"date_tool_name": {"date": date_key, "tool_name": tool_name}}, + data={ + "create": { + "date": date_key, + "tool_name": tool_name, + "spend": spend, + "total_tokens": total_tokens, + "request_count": len(entries), + }, + "update": { + "spend": {"increment": spend}, + "total_tokens": {"increment": total_tokens}, + "request_count": {"increment": len(entries)}, + }, + }, + ) + return + except DB_CONNECTION_ERROR_TYPES: + if attempt >= n_retry_times: + raise + await asyncio.sleep(2**attempt + random.uniform(0, 1)) diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index b6a445ef327..ad0b4f7444c 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -11,21 +11,21 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a import uuid from datetime import datetime, timedelta, timezone -from itertools import groupby from typing import TYPE_CHECKING, Annotated, Any, List, Optional from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient from litellm._logging import verbose_proxy_logger -from litellm.constants import TOOL_SPEND_MAX_WINDOW_DAYS +from litellm.constants import TOOL_SPEND_TOP_TOOLS from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ( + DailyToolSpendRepository, SpendLogsRepository, SpendLogToolIndexRepository, ) @@ -142,53 +142,18 @@ def _parse_day_start(value: str | None) -> datetime | None: ) -class _ToolSpendRow(BaseModel): - date: str +class _ToolSpendSums(BaseModel): + spend: float = 0.0 + total_tokens: int = 0 + request_count: int = 0 + + +class _TopToolRow(BaseModel): tool_name: str - call_count: int - spend: float - total_tokens: int + sums: _ToolSpendSums = Field(alias="_sum") -class _RequestTotalRow(BaseModel): - total_spend: float - - -_TOOL_SPEND_ROWS = TypeAdapter(list[_ToolSpendRow]) -_REQUEST_TOTAL_ROWS = TypeAdapter(list[_RequestTotalRow]) - - -def _summarize_tool(name: str, grp: tuple[_ToolSpendRow, ...]) -> ToolSpendEntry: - return ToolSpendEntry( - tool_name=name, - spend=sum(r.spend for r in grp), - call_count=sum(r.call_count for r in grp), - total_tokens=sum(r.total_tokens for r in grp), - ) - - -def _build_tool_spend_response( - rows: list[_ToolSpendRow], - total_spend: float, - start_date: str, - end_date: str, -) -> ToolSpendResponse: - daily = [ - ToolSpendDailyEntry(date=r.date, tool_name=r.tool_name, spend=r.spend, call_count=r.call_count) for r in rows - ] - grouped = groupby(sorted(rows, key=lambda r: r.tool_name), key=lambda r: r.tool_name) - by_tool = sorted( - (_summarize_tool(name, tuple(grp)) for name, grp in grouped), - key=lambda e: e.spend, - reverse=True, - ) - return ToolSpendResponse( - by_tool=by_tool, - daily=daily, - total_spend=total_spend, - start_date=start_date, - end_date=end_date, - ) +_TOP_TOOL_ROWS = TypeAdapter(list[_TopToolRow]) @router.get( @@ -205,16 +170,16 @@ async def get_tool_spend( """ Spend attributed to each tool over a date range, for the Cost Optimization dashboard. - Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to - ``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools - counts its full spend toward each of those tools, so per-tool numbers are - attributions. ``total_spend`` is the deduplicated spend of every request that - called at least one tool in the window, so it never double counts. + Reads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked + tools only (MCP tool calls and response tool_calls; declaring a tool without + invoking it does not count). A request that invoked multiple tools counts its + full spend toward each of them, so per-tool numbers are attributions and do not + sum to a deduplicated total. - ``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to - 31 calendar dates inclusive, the same width as the endpoint's default window): - a wider requested range is clamped, and the response's ``start_date`` reflects - the effective window actually served. + ``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in + SQL, and ``daily`` covers only those tools, so the response is bounded by + days x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many + distinct tool names exist. """ from litellm.proxy.proxy_server import prisma_client @@ -230,64 +195,46 @@ async def get_tool_spend( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - now = datetime.now(timezone.utc) - end_day = _parse_day_start(end_date) - # Anchor the floor to a midnight so the clamp compares dates with dates: - # parsed start_dates are midnight-aligned, and a floor carrying now's - # time-of-day would invisibly truncate an explicit start_date to mid-day. - today = now.replace(hour=0, minute=0, second=0, microsecond=0) - window_floor = (end_day or today) - timedelta(days=TOOL_SPEND_MAX_WINDOW_DAYS) - start_dt = _parse_day_start(start_date) or window_floor - if start_dt < window_floor: - start_dt = window_floor - end_exclusive = (end_day + timedelta(days=1)) if end_day else now + end_day = _parse_day_start(end_date) or datetime.now(timezone.utc) + start_day = _parse_day_start(start_date) or end_day - timedelta(days=30) + start_str = start_day.strftime("%Y-%m-%d") + end_str = end_day.strftime("%Y-%m-%d") + date_window = {"date": {"gte": start_str, "lte": end_str}} - # ti.start_time defines the window in both queries; the sl."startTime" bounds - # exist only so the planner can use the SpendLogs startTime index, and carry a - # 1s margin because the two writers can disagree by ~1ms on the same request. - rows = await prisma_client.db.query_raw( - """ - SELECT to_char(ti.start_time, 'YYYY-MM-DD') AS date, - ti.tool_name AS tool_name, - COUNT(*)::int AS call_count, - COALESCE(SUM(sl.spend), 0)::double precision AS spend, - COALESCE(SUM(sl.total_tokens), 0)::bigint AS total_tokens - FROM "LiteLLM_SpendLogToolIndex" ti - JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id - WHERE ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') - AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') - AND sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') - interval '1 second' - AND sl."startTime" < ($2::timestamptz AT TIME ZONE 'UTC') + interval '1 second' - GROUP BY date, ti.tool_name - ORDER BY date ASC, spend DESC - """, - start_dt.isoformat(), - end_exclusive.isoformat(), - ) - totals = await prisma_client.db.query_raw( - """ - SELECT COALESCE(SUM(sl.spend), 0)::double precision AS total_spend - FROM "LiteLLM_SpendLogs" sl - WHERE sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') - interval '1 second' - AND sl."startTime" < ($2::timestamptz AT TIME ZONE 'UTC') + interval '1 second' - AND EXISTS ( - SELECT 1 - FROM "LiteLLM_SpendLogToolIndex" ti - WHERE ti.request_id = sl.request_id - AND ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') - AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') + table = DailyToolSpendRepository(prisma_client).table + top_tools = _TOP_TOOL_ROWS.validate_python( + await table.group_by( + by=["tool_name"], + sum={"spend": True, "total_tokens": True, "request_count": True}, + where=date_window, + order={"_sum": {"spend": "desc"}}, + take=TOOL_SPEND_TOP_TOOLS, ) - """, - start_dt.isoformat(), - end_exclusive.isoformat(), + or [] ) - total_rows = _REQUEST_TOTAL_ROWS.validate_python(totals or []) - return _build_tool_spend_response( - rows=_TOOL_SPEND_ROWS.validate_python(rows or []), - total_spend=total_rows[0].total_spend if total_rows else 0.0, - start_date=start_dt.strftime("%Y-%m-%d"), - end_date=(end_day or now).strftime("%Y-%m-%d"), + by_tool = [ + ToolSpendEntry( + tool_name=row.tool_name, + spend=row.sums.spend, + call_count=row.sums.request_count, + total_tokens=row.sums.total_tokens, + ) + for row in top_tools + ] + + daily_rows = ( + await table.find_many( + where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}}, + order=[{"date": "asc"}, {"spend": "desc"}], + ) + if top_tools + else [] ) + daily = [ + ToolSpendDailyEntry(date=row.date, tool_name=row.tool_name, spend=row.spend, call_count=row.request_count) + for row in daily_rows + ] + return ToolSpendResponse(by_tool=by_tool, daily=daily, start_date=start_str, end_date=end_str) @router.get( @@ -388,7 +335,8 @@ async def get_tool_usage_logs( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Return paginated spend logs for requests that used this tool (from SpendLogToolIndex). + Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex). + Declaring a tool in a request body without the model invoking it does not create an entry. """ from litellm.proxy.proxy_server import prisma_client diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 6713b212314..37ea55f8c13 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1097,6 +1097,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e85ccf150d2..d7a95284818 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -175,6 +175,7 @@ if TYPE_CHECKING: from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction Span = Union[_Span, Any] else: @@ -2917,6 +2918,8 @@ async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> class PrismaClient: spend_log_transactions: List = [] _spend_log_transactions_lock = asyncio.Lock() + tool_usage_transactions: List["ToolUsageTransaction"] = [] + _tool_usage_transactions_lock = asyncio.Lock() def __init__( self, @@ -5473,12 +5476,15 @@ async def update_spend( queue_size = len(prisma_client.spend_log_transactions) verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size)) + async with prisma_client._tool_usage_transactions_lock: + tool_usage_queue_size = len(prisma_client.tool_usage_transactions) + # Process spend log transactions when called directly. # This keeps backwards compatibility with the old behavior. # See update_spend_logs_job and _monitor_spend_logs_queue for the new behavior. # Safe to keep: under high concurrency this can take up to ~30s to run, # so it's unlikely to overlap with monitor_spend_logs_queue. - if queue_size > 0: + if queue_size > 0 or tool_usage_queue_size > 0: await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, @@ -5545,10 +5551,14 @@ async def update_spend_logs_job( n_retry_times = 3 MAX_LOGS_PER_INTERVAL = 10000 - # Atomically pop batch from queue + # Atomically pop batch from queue. The tool usage queue counts toward the + # emptiness check: a spend-log write failure aborts a run before the tool + # drain below, and those entries must not strand once the spend queue drains. async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - if queue_size == 0: + async with prisma_client._tool_usage_transactions_lock: + tool_queue_size = len(prisma_client.tool_usage_transactions) + if queue_size == 0 and tool_queue_size == 0: return async with prisma_client._spend_log_transactions_lock: @@ -5579,17 +5589,23 @@ async def update_spend_logs_job( guardrail_tracking_err, ) - # Tool usage tracking (same batch): SpendLogToolIndex for "last N requests for tool X" + # Tool usage tracking: drain the request-time queue into the tool index and the + # LiteLLM_DailyToolSpend rollup. Never retried; a dropped batch is permanently + # absent from the rollup, so failures log at error. + async with prisma_client._tool_usage_transactions_lock: + tool_usage_to_process = prisma_client.tool_usage_transactions[:MAX_LOGS_PER_INTERVAL] + prisma_client.tool_usage_transactions = prisma_client.tool_usage_transactions[len(tool_usage_to_process) :] try: - from litellm.proxy.db.spend_log_tool_index import process_spend_logs_tool_usage + from litellm.proxy.db.spend_log_tool_index import flush_tool_usage_transactions - await process_spend_logs_tool_usage( + await flush_tool_usage_transactions( prisma_client=prisma_client, - logs_to_process=logs_to_process, + transactions=tool_usage_to_process, ) except Exception as tool_tracking_err: - verbose_proxy_logger.warning( - "Spend tracking - tool usage tracking failed (non-fatal): %s", + verbose_proxy_logger.error( + "Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s", + len(tool_usage_to_process), tool_tracking_err, ) @@ -5625,9 +5641,13 @@ async def _monitor_spend_logs_queue( while True: try: - # Check queue size with lock protection + # Check queue sizes with lock protection; the tool usage queue keeps + # the monitor firing when a prior failed run left it nonempty. async with prisma_client._spend_log_transactions_lock: - queue_size = len(prisma_client.spend_log_transactions) + spend_queue_size = len(prisma_client.spend_log_transactions) + async with prisma_client._tool_usage_transactions_lock: + tool_queue_size = len(prisma_client.tool_usage_transactions) + queue_size = spend_queue_size + tool_queue_size if queue_size > 0: if queue_size >= threshold: diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index 4451f0865da..29c953e06cf 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -23,6 +23,7 @@ from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, DailyPolicyMetricsRepository, DailyTagSpendRepository, + DailyToolSpendRepository, DeletedTeamRepository, DeletedVerificationTokenRepository, DeprecatedVerificationTokenRepository, @@ -104,6 +105,7 @@ __all__ = [ "ManagedVectorStoreIndexRepository", "WorkflowMessageRepository", "DailyTagSpendRepository", + "DailyToolSpendRepository", "SpendLogToolIndexRepository", "SpendLogGuardrailIndexRepository", "UserNotificationsRepository", diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index dc2a7d25259..54008c0950c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -181,6 +181,10 @@ class SpendLogToolIndexRepository(PrismaTableRepository): table_name = "litellm_spendlogtoolindex" +class DailyToolSpendRepository(PrismaTableRepository): + table_name = "litellm_dailytoolspend" + + class SpendLogGuardrailIndexRepository(PrismaTableRepository): table_name = "litellm_spendlogguardrailindex" diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 71ec412e8ef..ccf4b7dbc9f 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -124,12 +124,5 @@ class ToolSpendDailyEntry(BaseModel): class ToolSpendResponse(BaseModel): by_tool: List[ToolSpendEntry] = Field(default_factory=list) daily: List[ToolSpendDailyEntry] = Field(default_factory=list) - total_spend: float = Field( - 0.0, - description=( - "Deduplicated spend of every request that called at least one tool in the window; " - "less than the sum of per-tool attributed spend whenever multi-tool requests exist" - ), - ) start_date: str | None = None end_date: str | None = None diff --git a/schema.prisma b/schema.prisma index 6713b212314..37ea55f8c13 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1097,6 +1097,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index e2dca0a0f81..131f46a3e21 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -28,11 +28,13 @@ class MockPrismaClient: # Initialize transaction lists self.spend_log_transactions = [] self.daily_user_spend_transactions = {} + self.tool_usage_transactions = [] - # Add lock for spend_log_transactions (matches real PrismaClient) + # Add locks for the transaction queues (matches real PrismaClient) import asyncio self._spend_log_transactions_lock = asyncio.Lock() + self._tool_usage_transactions_lock = asyncio.Lock() def jsonify_object(self, obj): return obj diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 8149cf90e70..8759b008549 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -76,6 +76,144 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert call_args["payload"]["custom_llm_provider"] == "openai" +def _tool_call_response(*names: str) -> object: + from types import SimpleNamespace + + tool_calls = [SimpleNamespace(function=SimpleNamespace(name=name)) for name in names] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=tool_calls))]) + + +def _tool_usage_prisma() -> MagicMock: + prisma = MagicMock() + prisma.tool_usage_transactions = [] + prisma._tool_usage_transactions_lock = asyncio.Lock() + prisma.spend_log_transactions = [] + prisma._spend_log_transactions_lock = asyncio.Lock() + return prisma + + +def _minimal_spend_payload() -> dict: + return { + "request_id": "req-tool-1", + "startTime": datetime(2026, 7, 25, 10, 0, tzinfo=timezone.utc), + "endTime": datetime(2026, 7, 25, 10, 0, 1, tzinfo=timezone.utc), + "spend": 0.0, + "total_tokens": 42, + "mcp_namespaced_tool_name": None, + } + + +@pytest.mark.asyncio +async def test_update_database_enqueues_tool_usage_for_invoked_tools(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=_tool_call_response("get_weather"), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert len(prisma.tool_usage_transactions) == 1 + transaction = prisma.tool_usage_transactions[0] + assert transaction.request_id == "req-tool-1" + assert transaction.tool_names == ("get_weather",) + assert transaction.spend == 0.1 + assert transaction.total_tokens == 42 + assert transaction.date == "2026-07-25" + + +@pytest.mark.asyncio +async def test_update_database_enqueues_realtime_tool_usage(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={ + "model": "gpt-realtime", + "realtime_tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "rt_tool", "arguments": "{}"}} + ], + }, + completion_response=None, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.2, + ) + await asyncio.sleep(0) + + assert len(prisma.tool_usage_transactions) == 1 + assert prisma.tool_usage_transactions[0].tool_names == ("rt_tool",) + + +@pytest.mark.asyncio +async def test_update_database_skips_tool_usage_when_spend_logs_disabled(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", True), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=_tool_call_response("get_weather"), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert prisma.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_update_daily_spend_with_null_entity_id(): """ diff --git a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py new file mode 100644 index 00000000000..3b6acaa1eb3 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py @@ -0,0 +1,309 @@ +""" +Tests for the tool usage writer: ToolUsageTransaction construction (invoked tools +only) and the flush that writes LiteLLM_SpendLogToolIndex plus the +LiteLLM_DailyToolSpend rollup in one transaction. +""" + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.spend_log_tool_index import ( + ToolUsageTransaction, + build_tool_usage_transaction, + flush_tool_usage_transactions, + response_tool_call_names, +) + + +def _response_with_tool_calls(*names: str) -> SimpleNamespace: + tool_calls = [SimpleNamespace(function=SimpleNamespace(name=name)) for name in names] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=tool_calls))]) + + +class _FakeBatcher: + def __init__(self) -> None: + self.litellm_spendlogtoolindex = MagicMock() + self.litellm_dailytoolspend = MagicMock() + + async def __aenter__(self) -> "_FakeBatcher": + return self + + async def __aexit__(self, *args: Any) -> None: + return None + + +def _prisma_with_batcher() -> tuple[MagicMock, _FakeBatcher]: + batcher = _FakeBatcher() + prisma = MagicMock() + prisma.db.batch_ = MagicMock(return_value=batcher) + return prisma, batcher + + +class TestBuildToolUsageTransaction: + def test_declared_tools_never_reach_the_transaction(self): + # Regression for the inflation bug: the builder's only non-MCP source is + # the response's tool_calls, so a request declaring N tools while the + # model invokes one produces exactly one attribution. + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("get_weather"), + ) + assert transaction is not None + assert transaction.tool_names == ("get_weather",) + + def test_no_invoked_tools_returns_none(self): + assert ( + build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=None))]), + ) + is None + ) + + def test_mcp_name_and_response_names_dedupe(self): + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("srv/tool_a", "tool_b", "tool_b"), + ) + assert transaction is not None + assert transaction.tool_names == ("srv/tool_a", "tool_b") + + def test_date_matches_daily_spend_writer_derivation(self): + # The daily spend writer derives its date bucket as + # payload["startTime"].split("T")[0] (db_spend_update_writer.py), i.e. the + # timestamp's own calendar date, NOT the astimezone-UTC date. A non-UTC + # isoformat pins the difference: 2026-07-25T22:00:00-07:00 is 2026-07-26 + # in UTC but must bucket as 2026-07-25 to match LiteLLM_DailyUserSpend. + start_time_iso = "2026-07-25T22:00:00-07:00" + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso=start_time_iso, + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=None, + ) + assert transaction is not None + assert transaction.date == start_time_iso.split("T")[0] == "2026-07-25" + + def test_realtime_tool_calls_reach_the_transaction(self): + # Realtime sessions carry invoked tools in kwargs["realtime_tool_calls"] + # (OpenAI tool_calls dict shape, built in realtime_streaming.py), not on a + # response object; they must land in the rollup like any other invocation. + realtime_tool_calls = [ + {"id": "call_1", "type": "function", "function": {"name": "rt_get_weather", "arguments": "{}"}}, + ] + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=None, + realtime_tool_calls=realtime_tool_calls, + ) + assert transaction is not None + assert transaction.tool_names == ("rt_get_weather",) + + def test_realtime_names_dedupe_against_response_names(self): + realtime_tool_calls = [{"type": "function", "function": {"name": "get_weather", "arguments": "{}"}}] + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("get_weather"), + realtime_tool_calls=realtime_tool_calls, + ) + assert transaction is not None + assert transaction.tool_names == ("get_weather",) + + def test_unparseable_start_time_returns_none(self): + assert ( + build_tool_usage_transaction( + request_id="r1", + start_time_iso="not-a-timestamp", + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=None, + ) + is None + ) + + +class TestResponseToolCallNames: + def test_unrecognized_shapes_yield_nothing(self): + assert response_tool_call_names(None) == () + assert response_tool_call_names(SimpleNamespace()) == () + assert response_tool_call_names(ValueError("boom")) == () + + def test_blank_names_are_dropped(self): + assert response_tool_call_names(_response_with_tool_calls(" ", "real_tool")) == ("real_tool",) + + def test_responses_api_output_function_calls(self): + # Regression: /v1/responses carries invocations in output[] items of + # type function_call, not in choices; they must reach the rollup. + response = SimpleNamespace( + output=[ + SimpleNamespace(type="function_call", name="get_weather", call_id="c1", arguments="{}"), + SimpleNamespace(type="message", name=None, call_id=None, arguments=None), + ] + ) + assert response_tool_call_names(response) == ("get_weather",) + + def test_anthropic_messages_tool_use_blocks(self): + response = { + "content": [ + {"type": "text", "text": "checking"}, + {"type": "tool_use", "id": "t1", "name": "ant_get_weather", "input": {"city": "Paris"}}, + ] + } + assert response_tool_call_names(response) == ("ant_get_weather",) + + +def _transaction( + request_id: str, + date: str = "2026-07-25", + tool_names: tuple = ("tool_a",), + spend: float = 1.0, + total_tokens: int = 10, +) -> ToolUsageTransaction: + from datetime import datetime, timezone + + return ToolUsageTransaction( + request_id=request_id, + date=date, + start_time=datetime(2026, 7, 25, 10, 0, tzinfo=timezone.utc), + tool_names=tool_names, + spend=spend, + total_tokens=total_tokens, + ) + + +class TestFlushToolUsageTransactions: + @pytest.mark.asyncio + async def test_multi_tool_request_attributes_full_spend_to_each_tool(self): + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[_transaction("r1", tool_names=("tool_a", "tool_b"), spend=0.10, total_tokens=100)], + ) + index_rows = batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["data"] + assert [(r["request_id"], r["tool_name"]) for r in index_rows] == [("r1", "tool_a"), ("r1", "tool_b")] + assert batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["skip_duplicates"] is True + + upserts = { + c.kwargs["where"]["date_tool_name"]["tool_name"]: c.kwargs["data"] + for c in batcher.litellm_dailytoolspend.upsert.call_args_list + } + assert set(upserts) == {"tool_a", "tool_b"} + for data in upserts.values(): + assert data["create"]["spend"] == 0.10 + assert data["create"]["request_count"] == 1 + assert data["update"]["spend"] == {"increment": 0.10} + assert data["update"]["request_count"] == {"increment": 1} + + @pytest.mark.asyncio + async def test_same_day_same_tool_aggregates_within_batch(self): + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[ + _transaction("r1", spend=0.10, total_tokens=100), + _transaction("r2", spend=0.30, total_tokens=200), + ], + ) + assert batcher.litellm_dailytoolspend.upsert.call_count == 1 + data = batcher.litellm_dailytoolspend.upsert.call_args.kwargs["data"] + assert data["create"] == { + "date": "2026-07-25", + "tool_name": "tool_a", + "spend": pytest.approx(0.40), + "total_tokens": 300, + "request_count": 2, + } + assert data["update"]["spend"] == {"increment": pytest.approx(0.40)} + assert data["update"]["total_tokens"] == {"increment": 300} + assert data["update"]["request_count"] == {"increment": 2} + + @pytest.mark.asyncio + async def test_index_rows_and_rollup_share_one_transaction(self): + # Both writes go through the same batch_() so a failed flush cannot leave + # index rows without their rollup increments (or vice versa); increments + # are not idempotent, so partial states must be unreachable. + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[_transaction("r1")], + ) + prisma.db.batch_.assert_called_once() + batcher.litellm_spendlogtoolindex.create_many.assert_called_once() + batcher.litellm_dailytoolspend.upsert.assert_called_once() + + @pytest.mark.asyncio + async def test_empty_batch_touches_nothing(self): + prisma, _ = _prisma_with_batcher() + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[]) + prisma.db.batch_.assert_not_called() + + @pytest.mark.asyncio + async def test_connection_errors_retry_and_succeed(self, monkeypatch): + # A failed batch commits nothing, so retrying a connection error cannot + # double-count; the flush must retry rather than drop the batch. + import httpx + + batcher = _FakeBatcher() + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), batcher]) + sleeps: list[float] = [] + + async def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep) + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + assert prisma.db.batch_.call_count == 2 + assert len(sleeps) == 1 + batcher.litellm_dailytoolspend.upsert.assert_called_once() + + @pytest.mark.asyncio + async def test_connection_errors_exhaust_retries_then_raise(self, monkeypatch): + import httpx + + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=httpx.ConnectError("down")) + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep) + with pytest.raises(httpx.ConnectError): + await flush_tool_usage_transactions( + prisma_client=prisma, transactions=[_transaction("r1")], n_retry_times=2 + ) + assert prisma.db.batch_.call_count == 3 + + @pytest.mark.asyncio + async def test_non_connection_errors_do_not_retry(self): + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=ValueError("bad data")) + with pytest.raises(ValueError): + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + prisma.db.batch_.assert_called_once() diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py index c908250fa64..45c3c6c2466 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -19,11 +19,7 @@ from fastapi.testclient import TestClient sys.path.insert(0, os.path.abspath("../../..")) -from litellm.proxy.management_endpoints.tool_management_endpoints import ( - _build_tool_spend_response, - _ToolSpendRow, - router, -) +from litellm.proxy.management_endpoints.tool_management_endpoints import router from litellm.types.tool_management import LiteLLM_ToolTableRow # --- helpers --- @@ -64,6 +60,30 @@ def _override_auth(): _MOCK_PRISMA = MagicMock() +def _rollup_row(date: str, tool_name: str, spend: float, request_count: int, total_tokens: int) -> MagicMock: + row = MagicMock() + row.date = date + row.tool_name = tool_name + row.spend = spend + row.request_count = request_count + row.total_tokens = total_tokens + return row + + +def _group_row(tool_name: str, spend: float, request_count: int, total_tokens: int) -> dict: + return {"tool_name": tool_name, "_sum": {"spend": spend, "total_tokens": total_tokens, "request_count": request_count}} + + +def _rollup_prisma(group_rows: list, daily_rows: list | None = None) -> MagicMock: + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_spendlogtoolindex.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_dailytoolspend.group_by = AsyncMock(return_value=group_rows) + prisma.db.litellm_dailytoolspend.find_many = AsyncMock(return_value=daily_rows or []) + return prisma + + # --- test class --- @@ -154,21 +174,23 @@ class TestToolManagementEndpoints: assert resp.status_code == 422 def test_tool_spend_route_not_shadowed_by_get_tool(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend") assert resp.status_code == 200 assert resp.json()["by_tool"] == [] - def test_tool_spend_aggregates_and_sorts(self): - rows = [ - {"date": "2026-07-01", "tool_name": "search", "call_count": 2, "spend": 1.0, "total_tokens": 100}, - {"date": "2026-07-02", "tool_name": "search", "call_count": 1, "spend": 4.0, "total_tokens": 50}, - {"date": "2026-07-01", "tool_name": "read_file", "call_count": 3, "spend": 2.0, "total_tokens": 300}, + def test_tool_spend_serves_sql_aggregates_and_daily_series(self): + group_rows = [ + _group_row("search", spend=5.0, request_count=3, total_tokens=150), + _group_row("read_file", spend=2.0, request_count=3, total_tokens=300), ] - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(side_effect=[rows, [{"total_spend": 5.5}]]) + daily_rows = [ + _rollup_row("2026-07-01", "search", spend=1.0, request_count=2, total_tokens=100), + _rollup_row("2026-07-01", "read_file", spend=2.0, request_count=3, total_tokens=300), + _rollup_row("2026-07-02", "search", spend=4.0, request_count=1, total_tokens=50), + ] + prisma = _rollup_prisma(group_rows, daily_rows) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") assert resp.status_code == 200 @@ -179,94 +201,89 @@ class TestToolManagementEndpoints: assert search["call_count"] == 3 assert search["total_tokens"] == 150 assert len(body["daily"]) == 3 + assert body["daily"][0]["call_count"] == 2 assert body["start_date"] == "2026-07-01" assert body["end_date"] == "2026-07-02" - assert body["total_spend"] == 5.5 + + def test_tool_spend_coerces_bigint_string_sums(self): + # prisma group_by returns BigInt sums as strings ("808"); the response + # must coerce them to ints rather than 500 on validation. + group_rows = [{"tool_name": "search", "_sum": {"spend": 0.5, "total_tokens": "808", "request_count": "3"}}] + prisma = _rollup_prisma(group_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + assert resp.json()["by_tool"][0]["total_tokens"] == 808 + assert resp.json()["by_tool"][0]["call_count"] == 3 + + def test_tool_spend_daily_restricted_to_top_tools_and_capped(self): + from litellm.constants import TOOL_SPEND_TOP_TOOLS + + group_rows = [_group_row("search", spend=5.0, request_count=1, total_tokens=10)] + prisma = _rollup_prisma(group_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + group_kwargs = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs + assert group_kwargs["take"] == TOOL_SPEND_TOP_TOOLS + assert group_kwargs["order"] == {"_sum": {"spend": "desc"}} + daily_where = prisma.db.litellm_dailytoolspend.find_many.await_args.kwargs["where"] + assert daily_where["tool_name"] == {"in": ["search"]} + + def test_tool_spend_skips_daily_query_when_no_tools(self): + prisma = _rollup_prisma([]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + prisma.db.litellm_dailytoolspend.find_many.assert_not_awaited() @patch("litellm.proxy.proxy_server.prisma_client", None) def test_tool_spend_no_db_returns_500(self): resp = self.client.get("/v1/tool/spend") assert resp.status_code == 500 - def test_tool_spend_end_date_is_inclusive_via_exclusive_next_day_bound(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_reads_rollup_only_never_spendlogs(self): + # Regression for the GA blocker: the dashboard aggregate must be served + # entirely from LiteLLM_DailyToolSpend; any query_raw or SpendLogs table + # access on this path reintroduces the per-request scan. + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") assert resp.status_code == 200 - expected_binds = ( - datetime(2026, 7, 1, tzinfo=timezone.utc).isoformat(), - datetime(2026, 7, 3, tzinfo=timezone.utc).isoformat(), - ) - assert prisma.db.query_raw.await_count == 2 - for call in prisma.db.query_raw.await_args_list: - assert tuple(call.args[1:]) == expected_binds + prisma.db.query_raw.assert_not_awaited() + prisma.db.litellm_spendlogs.find_many.assert_not_awaited() + prisma.db.litellm_spendlogtoolindex.find_many.assert_not_awaited() + prisma.db.litellm_dailytoolspend.group_by.assert_awaited_once() + + def test_tool_spend_windows_rollup_by_inclusive_date_strings(self): + prisma = _rollup_prisma([]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + where = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs["where"] + assert where == {"date": {"gte": "2026-07-01", "lte": "2026-07-02"}} assert resp.json()["end_date"] == "2026-07-02" - def test_tool_spend_start_clamped_to_30_days_before_end(self): - # Clamped floor is end_date minus 30 days, serving up to 31 calendar dates - # inclusive: deliberately the same width as the endpoint's default window, - # so the dashboard's default range never triggers the clamp. - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_wide_range_served_fully(self): + # Regression: the 30-day clamp is gone; a 182-day request is served as + # requested because the rollup read is O(tools x dates). + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-01-01&end_date=2026-07-01") assert resp.status_code == 200 - expected_binds = ( - datetime(2026, 6, 1, tzinfo=timezone.utc).isoformat(), - datetime(2026, 7, 2, tzinfo=timezone.utc).isoformat(), - ) - assert prisma.db.query_raw.await_count == 2 - for call in prisma.db.query_raw.await_args_list: - assert tuple(call.args[1:]) == expected_binds - assert resp.json()["start_date"] == "2026-06-01" + where = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs["where"] + assert where == {"date": {"gte": "2026-01-01", "lte": "2026-07-01"}} + assert resp.json()["start_date"] == "2026-01-01" assert resp.json()["end_date"] == "2026-07-01" - def test_tool_spend_range_within_cap_is_not_clamped(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_defaults_to_trailing_30_days(self): + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2026-06-25&end_date=2026-07-01") + resp = self.client.get("/v1/tool/spend") assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == datetime(2026, 6, 25, tzinfo=timezone.utc).isoformat() - assert resp.json()["start_date"] == "2026-06-25" - - def test_tool_spend_start_honored_when_end_date_omitted(self): - # Regression: with end_date omitted the floor anchors to today's UTC - # midnight, not now's time-of-day, so an explicit start_date exactly 30 - # days back is served from midnight rather than truncated to mid-day. - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - floor_day = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=30) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get(f"/v1/tool/spend?start_date={floor_day.strftime('%Y-%m-%d')}") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == floor_day.isoformat() - assert resp.json()["start_date"] == floor_day.strftime("%Y-%m-%d") - - def test_tool_spend_clamp_without_end_date_lands_on_midnight(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - floor_day = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=30) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2020-01-01") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == floor_day.isoformat() - assert resp.json()["start_date"] == floor_day.strftime("%Y-%m-%d") - - def test_tool_spend_total_query_bounds_outer_spendlogs_scan(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - sql = call.args[0] - assert 'sl."startTime" >=' in sql - assert 'sl."startTime" <' in sql + today = datetime.now(timezone.utc) + assert resp.json()["end_date"] == today.strftime("%Y-%m-%d") + assert resp.json()["start_date"] == (today - timedelta(days=30)).strftime("%Y-%m-%d") @pytest.mark.parametrize( "query", @@ -279,13 +296,12 @@ class TestToolManagementEndpoints: ], ) def test_tool_spend_malformed_date_returns_400(self, query: str): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get(f"/v1/tool/spend?{query}") assert resp.status_code == 400 assert "Invalid date format" in resp.json()["detail"] - prisma.db.query_raw.assert_not_awaited() + prisma.db.litellm_dailytoolspend.group_by.assert_not_awaited() def test_tool_spend_non_admin_returns_403(self): from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -296,38 +312,8 @@ class TestToolManagementEndpoints: api_key="sk-user", user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER ) client = TestClient(app, raise_server_exceptions=True) - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = client.get("/v1/tool/spend") assert resp.status_code == 403 - prisma.db.query_raw.assert_not_awaited() - - -def _spend_row(date: str, tool_name: str, spend: float, call_count: int = 1, total_tokens: int = 10) -> _ToolSpendRow: - return _ToolSpendRow(date=date, tool_name=tool_name, call_count=call_count, spend=spend, total_tokens=total_tokens) - - -class TestBuildToolSpendResponse: - def test_multi_tool_attribution_double_counts_per_tool_but_not_total(self): - rows = [ - _spend_row("2026-07-01", "a", spend=3.0), - _spend_row("2026-07-01", "b", spend=3.0), - ] - resp = _build_tool_spend_response(rows, total_spend=3.0, start_date="2026-07-01", end_date="2026-07-01") - by_tool = {t.tool_name: t.spend for t in resp.by_tool} - assert by_tool == {"a": 3.0, "b": 3.0} - assert resp.total_spend == 3.0 - - def test_groups_across_days_and_sorts_by_spend(self): - rows = [ - _spend_row("2026-07-01", "b", spend=1.0, call_count=2, total_tokens=100), - _spend_row("2026-07-02", "b", spend=4.0, call_count=1, total_tokens=50), - _spend_row("2026-07-01", "a", spend=2.0, call_count=3, total_tokens=300), - ] - resp = _build_tool_spend_response(rows, total_spend=7.0, start_date="2026-07-01", end_date="2026-07-02") - assert [(t.tool_name, t.spend, t.call_count, t.total_tokens) for t in resp.by_tool] == [ - ("b", 5.0, 3, 150), - ("a", 2.0, 3, 300), - ] - assert len(resp.daily) == 3 + prisma.db.litellm_dailytoolspend.group_by.assert_not_awaited() diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index f969b040a0d..2ba9257e1da 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -193,6 +193,12 @@ async def test_cleanup_old_spend_logs_batch_deletion(): tool_index_sql = mock_db.execute_raw.call_args_list[3][0][0] assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in tool_index_sql + # The LiteLLM_DailyToolSpend rollup must outlive spend-log retention: it is + # the only copy of tool spend history once its per-request sources expire, + # so spend-log cleanup must never touch it. + for call in mock_db.execute_raw.call_args_list: + assert "LiteLLM_DailyToolSpend" not in call[0][0] + @pytest.mark.asyncio async def test_cleanup_old_spend_logs_retention_period_cutoff(): diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py index af62b7eef62..74c9abd9978 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -128,6 +128,8 @@ def mock_prisma_client() -> MagicMock: client.proxy_logging_obj.failure_handler = AsyncMock() client.spend_log_transactions = [] client._spend_log_transactions_lock = asyncio.Lock() + client.tool_usage_transactions = [] + client._tool_usage_transactions_lock = asyncio.Lock() client.jsonify_object = lambda data: dict(data) client.db.is_connected = MagicMock(return_value=False) client.db.connect = AsyncMock() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a0b3af54750..d9eeb168611 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -188,6 +188,35 @@ async def test_update_spend_logs_job_skips_when_queue_empty( assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 0 +@pytest.mark.asyncio +async def test_update_spend_logs_job_drains_tool_queue_when_spend_queue_empty( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # Regression: a spend-log write failure aborts a run before the tool drain, + # so tool transactions can outlive the spend queue; the job must still run + # for them instead of early-returning on the empty spend queue. + import litellm.proxy.db.spend_log_tool_index as tool_mod + import litellm.proxy.guardrails.usage_tracking as guard_mod + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [MagicMock()] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock() + monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False) + flush_stub = AsyncMock() + monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", flush_stub, raising=False) + + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert len(flush_stub.await_args.kwargs["transactions"]) == 1 + assert mock_prisma_client.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_update_spend_logs_job_processes_and_clears_queue( mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch @@ -208,7 +237,7 @@ async def test_update_spend_logs_job_processes_and_clears_queue( guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False ) monkeypatch.setattr( - tool_mod, "process_spend_logs_tool_usage", AsyncMock(), raising=False + tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False ) await update_spend_logs_job( diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index ae09a893194..26f19f1b35c 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -225,11 +225,6 @@ "count": 1 } }, - "src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx": { "local/filename-pascal-case": { "count": 1 @@ -4350,4 +4345,4 @@ "count": 1 } } -} +} \ No newline at end of file diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 27261768b8d..305a65ec5a2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -5,7 +5,7 @@ const mockUserDailyActivityCall = vi.fn(); vi.mock("@/components/networking", () => ({ userDailyActivityCall: (...args: unknown[]) => mockUserDailyActivityCall(...args), - getToolSpend: vi.fn().mockResolvedValue({ by_tool: [], daily: [], total_spend: 0, start_date: null, end_date: null }), + getToolSpend: vi.fn().mockResolvedValue({ by_tool: [], daily: [], start_date: null, end_date: null }), getGeneralSettingsCall: vi.fn().mockResolvedValue([]), })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index f84167f5a82..125c8dff694 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -34,7 +34,7 @@ vi.mock("@/components/shared/charts", () => ({ import UsageTab from "./UsageTab"; -const emptyToolSpend: ToolSpendResponse = { by_tool: [], daily: [], total_spend: 0, start_date: null, end_date: null }; +const emptyToolSpend: ToolSpendResponse = { by_tool: [], daily: [], start_date: null, end_date: null }; const baseMetrics = (overrides: Partial): SpendMetrics => ({ spend: 0, @@ -216,7 +216,6 @@ describe("UsageTab", () => { { tool_name: "read_file", spend: 1.0, call_count: 2, total_tokens: 50 }, ], daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 5.0, start_date: "2026-07-12", end_date: "2026-07-12", }; @@ -226,31 +225,4 @@ describe("UsageTab", () => { const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); }); - - it("notes the 30-day cap when the server clamps the tool spend window", async () => { - const toolSpend = { - by_tool: [{ tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }], - daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 4.0, - start_date: "2026-07-05", - end_date: "2026-07-14", - }; - const { findByText } = renderWith([day("2026-07-12", {})], { toolSpend }); - - expect(await findByText(/capped at 30 days before the end of the selected range/)).toBeInTheDocument(); - }); - - it("shows no cap note when the served window matches the request", async () => { - const toolSpend = { - by_tool: [{ tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }], - daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 4.0, - start_date: "2026-07-01", - end_date: "2026-07-14", - }; - const { findAllByTestId, queryByText } = renderWith([day("2026-07-12", {})], { toolSpend }); - - await findAllByTestId("bar-chart"); - expect(queryByText(/capped at 30 days before the end of the selected range/)).not.toBeInTheDocument(); - }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index 15ce84b8445..508bc13496c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -34,7 +34,6 @@ interface UsageTabProps { const EMPTY_TOOL_SPEND: ToolSpendResponse = { by_tool: [], daily: [], - total_spend: 0, start_date: null, end_date: null, }; @@ -103,7 +102,6 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { const toolSpend = toolSpendState?.key === rangeKey ? toolSpendState.data : null; const toolSpendLoading = toolSpendEnabled && toolSpend === null; - const toolSpendWindowClamped = !!toolSpend?.start_date && !!startTime && toolSpend.start_date > isoDay(startTime); const compressionTotal = useMemo(() => results.reduce((sum, d) => sum + compressionOf(d.metrics), 0), [results]); const cachingTotal = useMemo(() => results.reduce((sum, d) => sum + cachingOf(d.metrics), 0), [results]); @@ -262,15 +260,10 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { Spend by tool

- Spend on requests that called each tool (MCP and client-side tools). A request that used multiple tools - counts its full spend toward each, so this attributes rather than partitions spend. + Spend on requests that invoked each tool (MCP and client-side tools); declaring a tool without invoking it + does not count. A request that invoked multiple tools counts its full spend toward each, so this attributes + rather than partitions spend.

- {toolSpendWindowClamped && ( -

- Tool spend is capped at 30 days before the end of the selected range; showing spend since{" "} - {toolSpend?.start_date}. -

- )} {topTools.length === 0 ? ( diff --git a/ui/litellm-dashboard/src/components/ToolDetail.tsx b/ui/litellm-dashboard/src/components/ToolDetail.tsx index 6a8457a559f..06f14638141 100644 --- a/ui/litellm-dashboard/src/components/ToolDetail.tsx +++ b/ui/litellm-dashboard/src/components/ToolDetail.tsx @@ -430,7 +430,7 @@ export function ToolDetail({ toolName, onBack, accessToken }: ToolDetailProps) {

- Recent logs + Recent invocations

Date: Sun, 26 Jul 2026 04:58:28 +0000 Subject: [PATCH 2/7] fix(ui): keep the spend-by-tool legend from overlapping the charts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/UsageTab.test.tsx | 41 ++++++++++++++++++- .../_components/UsageTab.tsx | 5 ++- .../shared/charts/bar_chart.test.tsx | 34 +++++++++++++++ .../components/shared/charts/bar_chart.tsx | 13 ++++-- .../shared/charts/chart_legend.test.tsx | 8 ++++ .../components/shared/charts/chart_legend.tsx | 2 +- .../src/components/ui/chart.tsx | 6 ++- 7 files changed, 101 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index 125c8dff694..4d1f1c182db 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -23,8 +23,24 @@ vi.mock("@/components/shared/charts", () => ({ DonutChart: ({ data, label }: { data: unknown; label: string }) => (
), - BarChart: ({ data, categories }: { data: unknown; categories: string[] }) => ( -
+ BarChart: ({ + data, + categories, + colors, + showLegend, + }: { + data: unknown; + categories: string[]; + colors?: readonly string[]; + showLegend?: boolean; + }) => ( +
), CustomLegend: ({ categories }: { categories: readonly string[] }) => (
{categories.join(",")}
@@ -225,4 +241,25 @@ describe("UsageTab", () => { const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); }); + + it("renders the tool legend once outside the charts, with both charts sharing the tool colors", async () => { + const toolSpend = { + by_tool: [ + { tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }, + { tool_name: "read_file", spend: 1.0, call_count: 2, total_tokens: 50 }, + ], + daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], + start_date: "2026-07-12", + end_date: "2026-07-12", + }; + const { findAllByTestId, getAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); + + const bars = await findAllByTestId("bar-chart"); + const [totalByTool, dailyByTool] = bars.slice(-2); + expect(dailyByTool.getAttribute("data-show-legend")).toBe("false"); + expect(totalByTool.getAttribute("data-colors")).toBe(dailyByTool.getAttribute("data-colors")); + + const toolLegends = getAllByTestId("chart-legend").filter((legend) => legend.textContent === "search,read_file"); + expect(toolLegends).toHaveLength(1); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index 508bc13496c..68f9c1d0ba4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -278,7 +278,8 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { data={topToolsChart} index="tool_name" categories={["spend"]} - colors={["emerald"]} + colors={toolColors} + colorByDatum layout="vertical" yAxisWidth={140} showLegend={false} @@ -287,6 +288,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => {

Daily spend by tool

+ = ({ accessToken, activity }) => { colors={toolColors} stack valueFormatter={usd} + showLegend={false} />
diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx index b30a252659f..cb0d5c603a4 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx @@ -113,6 +113,40 @@ describe("BarChart", () => { expect(container.querySelector("style")).toBeNull(); }); + it("colors each bar by its datum when colorByDatum is set, instead of one fill for the series", () => { + const singleCategory = [ + { tool: "alpha", spend: 3 }, + { tool: "beta", spend: 2 }, + { tool: "gamma", spend: 1 }, + ]; + + const { container, rerender } = render( + , + ); + const sharedFills = Array.from(container.querySelectorAll("path.recharts-rectangle")).map((rect) => + rect.getAttribute("fill"), + ); + expect(new Set(sharedFills).size).toBe(1); + + rerender( + , + ); + const perDatumFills = Array.from(container.querySelectorAll("path.recharts-rectangle")).map((rect) => + rect.getAttribute("fill"), + ); + expect(perDatumFills).toEqual([ + "var(--color-blue-500, #3b82f6)", + "var(--color-cyan-500, #06b6d4)", + "var(--color-violet-500, #8b5cf6)", + ]); + }); + it("stacks bars into a single column per index when stack is set", () => { const { container } = render( , diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx index 7069ececb70..6bfcf14c2a0 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx @@ -1,17 +1,20 @@ "use client"; import * as React from "react"; -import { Bar, BarChart as RechartsBarChart, CartesianGrid, XAxis, YAxis } from "recharts"; +import { Bar, BarChart as RechartsBarChart, CartesianGrid, Cell, XAxis, YAxis } from "recharts"; import { ChartContainer, ChartLegend, ChartLegendContent, ChartTooltip, type ChartConfig } from "@/components/ui/chart"; import { cn } from "@/lib/cva.config"; import { ValueTooltip, type ChartTooltipComponent } from "./chart_tooltip"; import { categoryFills, type ChartColor } from "./colors"; +const MAX_BAR_SIZE = 64; + export type BarChartProps> = { data: readonly TDatum[]; index: string; categories: readonly string[]; colors?: readonly ChartColor[]; + colorByDatum?: boolean; valueFormatter?: (value: number) => string; stack?: boolean; layout?: "horizontal" | "vertical"; @@ -32,6 +35,7 @@ export function BarChart>({ index, categories, colors, + colorByDatum = false, valueFormatter, stack = false, layout = "horizontal", @@ -57,7 +61,7 @@ export function BarChart>({ ); } - const fills = categoryFills(categories.length, colors); + const fills = categoryFills(colorByDatum ? data.length : categories.length, colors); const config: ChartConfig = Object.fromEntries(categories.map((category) => [category, { label: category }])); const vertical = layout === "vertical"; const TooltipContent = customTooltip ?? ValueTooltip; @@ -115,6 +119,7 @@ export function BarChart>({ fill={fills[i]} stackId={stack ? "stack" : undefined} isAnimationActive={false} + maxBarSize={MAX_BAR_SIZE} onClick={ onValueChange ? (item: { payload?: TDatum }) => { @@ -122,7 +127,9 @@ export function BarChart>({ } : undefined } - /> + > + {colorByDatum && data.map((_, dataIndex) => )} + ))} diff --git a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx index 889927aca43..28afe5faf9c 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx @@ -21,6 +21,14 @@ describe("CustomLegend", () => { expect(dots[1]?.getAttribute("style")).toContain("--color-green-500"); }); + it("wraps onto multiple lines instead of overflowing when there are many categories", () => { + const { container } = render( + `metrics.tool_${i}`)} colors={["blue", "green"]} />, + ); + + expect(container.firstElementChild?.className).toContain("flex-wrap"); + }); + it("cycles colors when there are more categories than colors", () => { const { container } = render( , diff --git a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx index da252d8bf63..1551f3d0e39 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx @@ -11,7 +11,7 @@ export const CustomLegend = ({ categories: readonly string[]; colors: readonly ChartColor[]; }) => ( -
+
{categories.map((category, idx) => (
{payload .filter((item) => item.type !== "none") From 5d77c39bbba17dc37e8683f33017b5e9f3be0733 Mon Sep 17 00:00:00 2001 From: tin Date: Sun, 26 Jul 2026 05:23:57 +0000 Subject: [PATCH 3/7] fix(ui): color spend-by-tool charts with an ordered ramp Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../CostOptimizationView.activity.test.tsx | 2 +- .../cost-optimization/_components/UsageTab.test.tsx | 2 +- .../cost-optimization/_components/UsageTab.tsx | 4 ++-- .../src/components/shared/charts/colors.ts | 11 +++++++++++ .../src/components/shared/charts/index.ts | 9 ++++++++- 5 files changed, 23 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 305a65ec5a2..363525c48af 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -19,7 +19,7 @@ vi.mock("@/components/shared/charts", () => ({ DonutChart: () =>
, BarChart: () =>
, CustomLegend: () =>
, - DEFAULT_COLOR_CYCLE: ["emerald"], + SEQUENTIAL_COLOR_RAMP: ["indigo"], })); vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () => ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index 4d1f1c182db..beeb9466b1d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -45,7 +45,7 @@ vi.mock("@/components/shared/charts", () => ({ CustomLegend: ({ categories }: { categories: readonly string[] }) => (
{categories.join(",")}
), - DEFAULT_COLOR_CYCLE: ["emerald", "blue", "violet", "amber"], + SEQUENTIAL_COLOR_RAMP: ["indigo", "blue", "sky", "cyan"], })); import UsageTab from "./UsageTab"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index 68f9c1d0ba4..829a79d7638 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -3,7 +3,7 @@ import React, { useEffect, useMemo, useState } from "react"; import { Info } from "lucide-react"; -import { AreaChart, BarChart, CustomLegend, DonutChart, DEFAULT_COLOR_CYCLE } from "@/components/shared/charts"; +import { AreaChart, BarChart, CustomLegend, DonutChart, SEQUENTIAL_COLOR_RAMP } from "@/components/shared/charts"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; @@ -166,7 +166,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { })), [toolSpend, topToolNames], ); - const toolColors = useMemo(() => DEFAULT_COLOR_CYCLE.slice(0, Math.max(topToolNames.length, 1)), [topToolNames]); + const toolColors = useMemo(() => SEQUENTIAL_COLOR_RAMP.slice(0, Math.max(topToolNames.length, 1)), [topToolNames]); return (
diff --git a/ui/litellm-dashboard/src/components/shared/charts/colors.ts b/ui/litellm-dashboard/src/components/shared/charts/colors.ts index c30f58e9e4d..3efbd54cfd2 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/colors.ts +++ b/ui/litellm-dashboard/src/components/shared/charts/colors.ts @@ -50,6 +50,17 @@ export const DEFAULT_COLOR_CYCLE: readonly ChartColor[] = [ "rose", ]; +export const SEQUENTIAL_COLOR_RAMP: readonly ChartColor[] = [ + "indigo", + "blue", + "sky", + "cyan", + "teal", + "emerald", + "green", + "lime", +]; + export const chartColorValue = (color: ChartColor): string => `var(--color-${color}-500, ${CHART_COLOR_HEX[color]})`; export const categoryFills = (count: number, colors?: readonly ChartColor[]): readonly string[] => { diff --git a/ui/litellm-dashboard/src/components/shared/charts/index.ts b/ui/litellm-dashboard/src/components/shared/charts/index.ts index 8383c767064..69edd3fb13f 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/index.ts +++ b/ui/litellm-dashboard/src/components/shared/charts/index.ts @@ -8,6 +8,13 @@ export { type ChartTooltipComponent, type ChartTooltipProps, } from "./chart_tooltip"; -export { CHART_COLOR_HEX, DEFAULT_COLOR_CYCLE, categoryFills, chartColorValue, type ChartColor } from "./colors"; +export { + CHART_COLOR_HEX, + DEFAULT_COLOR_CYCLE, + SEQUENTIAL_COLOR_RAMP, + categoryFills, + chartColorValue, + type ChartColor, +} from "./colors"; export { DonutChart, type DonutChartProps } from "./donut_chart"; export { LineChart, type LineChartCurveType, type LineChartProps } from "./line_chart"; From 708a3a19df8d9f803fcbad4f6c365438033a39c8 Mon Sep 17 00:00:00 2001 From: tin Date: Sun, 26 Jul 2026 06:53:26 +0000 Subject: [PATCH 4/7] fix(ui): use a single muted blue ramp for the tool charts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/shared/charts/colors.ts | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/charts/colors.ts b/ui/litellm-dashboard/src/components/shared/charts/colors.ts index 3efbd54cfd2..8b5717cc58e 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/colors.ts +++ b/ui/litellm-dashboard/src/components/shared/charts/colors.ts @@ -23,7 +23,7 @@ export const CHART_COLOR_HEX = { rose: "#f43f5e", } as const; -export type ChartColor = keyof typeof CHART_COLOR_HEX; +export type ChartColor = keyof typeof CHART_COLOR_HEX | `#${string}`; export const DEFAULT_COLOR_CYCLE: readonly ChartColor[] = [ "blue", @@ -51,17 +51,20 @@ export const DEFAULT_COLOR_CYCLE: readonly ChartColor[] = [ ]; export const SEQUENTIAL_COLOR_RAMP: readonly ChartColor[] = [ - "indigo", - "blue", - "sky", - "cyan", - "teal", - "emerald", - "green", - "lime", + "#1e3a8a", + "#1d4ed8", + "#2563eb", + "#3b82f6", + "#60a5fa", + "#93c5fd", + "#bfdbfe", + "#dbeafe", ]; -export const chartColorValue = (color: ChartColor): string => `var(--color-${color}-500, ${CHART_COLOR_HEX[color]})`; +const NAMED_COLOR_HEX: Readonly> = CHART_COLOR_HEX; + +export const chartColorValue = (color: ChartColor): string => + color in NAMED_COLOR_HEX ? `var(--color-${color}-500, ${NAMED_COLOR_HEX[color]})` : color; export const categoryFills = (count: number, colors?: readonly ChartColor[]): readonly string[] => { const cycle = colors && colors.length > 0 ? colors : DEFAULT_COLOR_CYCLE; From 1240c1a76d12b8d9643af799a755b57b09d7b5ec Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sun, 26 Jul 2026 01:55:23 -0700 Subject: [PATCH 5/7] fix(proxy): close the adversarial-review findings on the tool spend rollup Three fixes from an adversarial review of this branch, each at the owning seam rather than the report site. The flush retried DB_CONNECTION_ERROR_TYPES, which includes ReadTimeout. A ReadTimeout is the committed-but-unacked case: the review reproduced the engine abandoning the transaction open on the pooled connection, the retry stacking its statements into it, and one commit applying both increment sets while the flush reports success. The retry now covers only ConnectError, the one failure that proves the statements never reached the database; post-send failures drop the batch with an error log. The docstring no longer claims an idempotency the pattern does not have. The same hazard exists in the untouched daily spend writer and is left for its own change. get_tool_calls_from_response read choices[0] only, so a tool invoked in a later choice of an n>1 response earned spend but never reached the rollup, the index, or the registry. Choice scope is now an explicit parameter: accounting passes include_all_choices=True because every choice costs money; guardrails keep the primary-choice default because they rebuild the primary assistant message. First multi-choice fixtures in the suite pin both scopes. maxBarSize=64 had been added to the shared BarChart unconditionally, resizing every existing consumer. It is now a prop; only the tool spend charts opt in. The legend flex-wrap changes stay global because clipping overflow was a defect, not a preference. --- .../prompt_templates/factory.py | 27 +++++++++---- litellm/proxy/db/spend_log_tool_index.py | 21 +++++----- ...llm_core_utils_prompt_templates_factory.py | 31 +++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 18 +++++++++ .../proxy/db/test_spend_log_tool_index.py | 39 +++++++++++++++++++ .../_components/UsageTab.test.tsx | 6 +++ .../_components/UsageTab.tsx | 2 + .../components/shared/charts/bar_chart.tsx | 6 +-- 8 files changed, 131 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index c13cf0817b5..4e3d94e2ab3 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5382,14 +5382,18 @@ def _parse_tool_call_arguments(raw: Any, tool_name: Optional[str], context: str) return parsed if isinstance(parsed, dict) else {} -def _tool_calls_from_chat_completion_response(response: Any) -> list[NormalizedToolCall]: +def _tool_calls_from_chat_completion_response( + response: Any, include_all_choices: bool = False +) -> list[NormalizedToolCall]: choices = get_attribute_or_key(response, "choices", None) if not (isinstance(choices, list) and choices): return [] - message = get_attribute_or_key(choices[0], "message", None) - tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None - if not isinstance(tool_calls, list): - return [] + tool_calls: list[Any] = [] + for choice in choices if include_all_choices else choices[:1]: + message = get_attribute_or_key(choice, "message", None) + choice_tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None + if isinstance(choice_tool_calls, list): + tool_calls.extend(choice_tool_calls) result: list[NormalizedToolCall] = [] for tc in tool_calls: fn = get_attribute_or_key(tc, "function", None) @@ -5452,7 +5456,7 @@ def _tool_calls_from_anthropic_messages_response(response: Any) -> list[Normaliz return result -def get_tool_calls_from_response(response: Any) -> list[NormalizedToolCall]: +def get_tool_calls_from_response(response: Any, include_all_choices: bool = False) -> list[NormalizedToolCall]: """ Extract tool/function calls from a response object into a normalized ``{"id", "name", "arguments"}`` shape, regardless of which API surface @@ -5460,11 +5464,20 @@ def get_tool_calls_from_response(response: Any) -> list[NormalizedToolCall]: the Responses API (``output`` items of type ``function_call``), or the Anthropic Messages API (``content`` blocks of type ``tool_use``). + ``include_all_choices`` decides the chat-completions scope: the default + reads only ``choices[0]``, which is what consumers that act on THE reply + (e.g. guardrails rebuilding the primary assistant message) want; usage + accounting passes True because every choice of an ``n>1`` request costs + money and its tool calls really ran. The other surfaces have a single + output, so the flag has no effect on them. + Callers that only care about a specific tool should filter the result by ``name`` themselves -- this returns every tool call found. """ + chat_tool_calls = _tool_calls_from_chat_completion_response(response, include_all_choices=include_all_choices) + if chat_tool_calls: + return chat_tool_calls for extractor in ( - _tool_calls_from_chat_completion_response, _tool_calls_from_responses_api_response, _tool_calls_from_anthropic_messages_response, ): diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 064d08acb59..a478248b0fa 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -18,7 +18,7 @@ from datetime import datetime, timezone from itertools import groupby from typing import TYPE_CHECKING, Any, Sequence -from litellm.proxy._types import DB_CONNECTION_ERROR_TYPES +import httpx if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -37,7 +37,8 @@ class ToolUsageTransaction: def response_tool_call_names(completion_response: Any) -> tuple[str, ...]: """Tool names invoked in a completion response, in call order, for any response surface get_tool_calls_from_response understands (chat completions, Responses - API output items, Anthropic Messages tool_use blocks).""" + API output items, Anthropic Messages tool_use blocks). Reads every choice of + an ``n>1`` chat response: each choice cost money and its tool calls ran.""" if completion_response is None or isinstance(completion_response, Exception): return () from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -46,7 +47,7 @@ def response_tool_call_names(completion_response: Any) -> tuple[str, ...]: return tuple( stripped - for tool_call in get_tool_calls_from_response(completion_response) + for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) if isinstance(name := tool_call.get("name"), str) and (stripped := name.strip()) ) @@ -97,11 +98,13 @@ async def flush_tool_usage_transactions( n_retry_times: int = 3, ) -> None: """Write index rows and rollup upserts for a drained queue batch in one - transaction. Connection errors are retried with backoff, which cannot - double-count because a failed batch commits nothing; every other error - propagates so the caller drops the batch. Callers must not add their own - retry around this function: a batch that DID commit must never run again, - since the rollup update increments counters.""" + transaction. Retries only ConnectError, the one failure that proves the + statements never reached the database. Post-send failures (Read timeouts + and errors) are ambiguous and are NOT retried: the engine can abandon the + transaction open on the pooled connection, so a retry's statements stack + into the same transaction and one commit applies both increment sets. + Ambiguous failures drop the batch; the caller logs it at error. Callers + must not add their own retry around this function.""" if not transactions: return @@ -141,7 +144,7 @@ async def flush_tool_usage_transactions( }, ) return - except DB_CONNECTION_ERROR_TYPES: + except httpx.ConnectError: if attempt >= n_retry_times: raise await asyncio.sleep(2**attempt + random.uniform(0, 1)) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index bcda88ea609..9565de1139c 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3166,3 +3166,34 @@ async def test_bedrock_converse_message_level_cache_point_preserves_ttl_async(): ) assert _collect_cache_points(result) == [{"type": "default", "ttl": "1h"}] + + +def _n_choices_response(*names_per_choice): + from types import SimpleNamespace + + choices = [ + SimpleNamespace( + message=SimpleNamespace( + tool_calls=[SimpleNamespace(id=f"c{i}", function=SimpleNamespace(name=name, arguments="{}"))] + ) + ) + for i, name in enumerate(names_per_choice) + ] + return SimpleNamespace(choices=choices) + + +def test_get_tool_calls_from_response_defaults_to_primary_choice_only(): + from litellm.litellm_core_utils.prompt_templates.factory import get_tool_calls_from_response + + response = _n_choices_response("tool_alpha", "tool_beta") + + assert [tc["name"] for tc in get_tool_calls_from_response(response)] == ["tool_alpha"] + + +def test_get_tool_calls_from_response_include_all_choices_reads_every_choice(): + from litellm.litellm_core_utils.prompt_templates.factory import get_tool_calls_from_response + + response = _n_choices_response("tool_alpha", "tool_beta") + + names = [tc["name"] for tc in get_tool_calls_from_response(response, include_all_choices=True)] + assert names == ["tool_alpha", "tool_beta"] diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 8759b008549..cd293325c15 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -181,6 +181,24 @@ async def test_update_database_enqueues_realtime_tool_usage(): assert prisma.tool_usage_transactions[0].tool_names == ("rt_tool",) +def test_enqueue_tool_registry_upsert_reads_every_choice(): + from types import SimpleNamespace as NS + + db_writer = DBSpendUpdateWriter() + db_writer.tool_discovery_queue = MagicMock() + response = NS( + choices=[ + NS(message=NS(tool_calls=[NS(function=NS(name="tool_alpha"))])), + NS(message=NS(tool_calls=[NS(function=NS(name="tool_beta"))])), + ] + ) + + db_writer._enqueue_tool_registry_upsert(kwargs={}, completion_response=response) + + enqueued = [call.args[0]["tool_name"] for call in db_writer.tool_discovery_queue.add_update.call_args_list] + assert enqueued == ["tool_alpha", "tool_beta"] + + @pytest.mark.asyncio async def test_update_database_skips_tool_usage_when_spend_logs_disabled(): db_writer = DBSpendUpdateWriter() diff --git a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py index 3b6acaa1eb3..71073fd216e 100644 --- a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py +++ b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py @@ -134,6 +134,29 @@ class TestBuildToolUsageTransaction: assert transaction is not None assert transaction.tool_names == ("get_weather",) + def test_n_greater_than_one_tools_from_every_choice_reach_the_transaction(self): + # Regression: an n>1 request pays for every choice, and a tool invoked + # only in a later choice really ran; it must not be dropped because the + # extractor read choices[0] alone. + from types import SimpleNamespace as NS + + response = NS( + choices=[ + NS(message=NS(tool_calls=[NS(function=NS(name="tool_alpha"))])), + NS(message=NS(tool_calls=[NS(function=NS(name="tool_beta"))])), + ] + ) + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=response, + ) + assert transaction is not None + assert transaction.tool_names == ("tool_alpha", "tool_beta") + def test_unparseable_start_time_returns_none(self): assert ( build_tool_usage_transaction( @@ -307,3 +330,19 @@ class TestFlushToolUsageTransactions: with pytest.raises(ValueError): await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) prisma.db.batch_.assert_called_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("ambiguous_error", ["ReadTimeout", "ReadError"]) + async def test_post_send_ambiguous_errors_drop_without_retry(self, ambiguous_error): + # A ReadTimeout means the statements were sent and the outcome is + # unknown; the engine can leave the transaction open on the pooled + # connection, so a retry's statements would stack into it and one + # commit would apply both increment sets. These must never retry. + import httpx + + error = getattr(httpx, ambiguous_error)("ambiguous") + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=error) + with pytest.raises((httpx.ReadTimeout, httpx.ReadError)): + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + prisma.db.batch_.assert_called_once() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index beeb9466b1d..5c26ac30477 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -28,17 +28,20 @@ vi.mock("@/components/shared/charts", () => ({ categories, colors, showLegend, + maxBarSize, }: { data: unknown; categories: string[]; colors?: readonly string[]; showLegend?: boolean; + maxBarSize?: number; }) => (
), @@ -240,6 +243,9 @@ describe("UsageTab", () => { const bars = await findAllByTestId("bar-chart"); const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); + // The 64px bar cap is this card's opt-in; the shared BarChart must not cap + // by default (other consumers keep their pre-existing geometry). + expect(bars[0].getAttribute("data-max-bar-size")).toBe("64"); }); it("renders the tool legend once outside the charts, with both charts sharing the tool colors", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index 829a79d7638..ec37418e0b5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -282,6 +282,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { colorByDatum layout="vertical" yAxisWidth={140} + maxBarSize={64} showLegend={false} valueFormatter={usd} /> @@ -295,6 +296,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { categories={topToolNames} colors={toolColors} stack + maxBarSize={64} valueFormatter={usd} showLegend={false} /> diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx index 6bfcf14c2a0..ab2cc66eaf4 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx @@ -7,14 +7,13 @@ import { cn } from "@/lib/cva.config"; import { ValueTooltip, type ChartTooltipComponent } from "./chart_tooltip"; import { categoryFills, type ChartColor } from "./colors"; -const MAX_BAR_SIZE = 64; - export type BarChartProps> = { data: readonly TDatum[]; index: string; categories: readonly string[]; colors?: readonly ChartColor[]; colorByDatum?: boolean; + maxBarSize?: number; valueFormatter?: (value: number) => string; stack?: boolean; layout?: "horizontal" | "vertical"; @@ -36,6 +35,7 @@ export function BarChart>({ categories, colors, colorByDatum = false, + maxBarSize, valueFormatter, stack = false, layout = "horizontal", @@ -119,7 +119,7 @@ export function BarChart>({ fill={fills[i]} stackId={stack ? "stack" : undefined} isAnimationActive={false} - maxBarSize={MAX_BAR_SIZE} + maxBarSize={maxBarSize} onClick={ onValueChange ? (item: { payload?: TDatum }) => { From 2e12614a5b1f175d602e33ab079a4c0fb8d36c94 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 27 Jul 2026 11:13:49 -0700 Subject: [PATCH 6/7] fix(proxy): stop retrying post-send ambiguous DB errors in every spend writer Resolves LIT-4823. An adversarial review reproduced against real Postgres that a batched increment upsert stalling past the prisma engine timeout leaves its transaction open on the pooled connection; the retry draws the same connection, its statements stack into the still-open transaction, and one commit applies both increment sets while the writer reports success. httpx.ReadTimeout is exactly that post-send case and every spend writer retried it. DB_RETRY_SAFE_ERROR_TYPES (ConnectError only, the failure that proves the statements never reached the database) is now the single owner of what a non-idempotent writer may retry. All seven entity and daily spend writer retry arms and the tool usage flush consume it. DB_CONNECTION_ERROR_TYPES is unchanged for the idempotent spend-log writer, whose create_many with skip_duplicates may safely retry the full tuple. The corruption was reproduced on update_daily_user_spend (seeded 10|100|1, expected 11|110|2, observed 12|120|3); the new policy tests pin that a ReadTimeout drops the batch loudly on the first attempt and a ConnectError still retries. --- litellm/proxy/_types.py | 7 ++ litellm/proxy/db/db_spend_update_writer.py | 16 ++-- litellm/proxy/db/spend_log_tool_index.py | 4 +- litellm/proxy/utils.py | 3 +- .../proxy/db/test_db_spend_update_writer.py | 78 +++++++++++++++++++ .../test_proxy_update_spend.py | 35 +++++++-- 6 files changed, 126 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7575091be54..b94a34fa14c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3638,6 +3638,13 @@ DB_CONNECTION_ERROR_TYPES = ( httpx.ReadTimeout, ) +# What a NON-IDEMPOTENT write (increment upsert) may retry: only ConnectError +# proves the statements never reached the database. Post-send errors are +# ambiguous; a stalled statement can leave its transaction open on the pooled +# connection, where a retry stacks a second increment set into the same commit. +# Idempotent writes (create_many with skip_duplicates) may retry the full tuple. +DB_RETRY_SAFE_ERROR_TYPES = (httpx.ConnectError,) + class SSOUserDefinedValues(TypedDict): models: List[str] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ebdb08a681a..fd8132fef22 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,7 +34,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, DailyEndUserSpendTransaction, @@ -1121,7 +1121,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1164,7 +1164,7 @@ class DBSpendUpdateWriter: }, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1197,7 +1197,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1244,7 +1244,7 @@ class DBSpendUpdateWriter: ) # Transaction succeeded, break out of retry loop break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1286,7 +1286,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1372,7 +1372,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, @@ -1669,7 +1669,7 @@ class DBSpendUpdateWriter: break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index a478248b0fa..802e893d473 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -18,7 +18,7 @@ from datetime import datetime, timezone from itertools import groupby from typing import TYPE_CHECKING, Any, Sequence -import httpx +from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -144,7 +144,7 @@ async def flush_tool_usage_transactions( }, ) return - except httpx.ConnectError: + except DB_RETRY_SAFE_ERROR_TYPES: if attempt >= n_retry_times: raise await asyncio.sleep(2**attempt + random.uniform(0, 1)) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d7a95284818..924189fed4b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -41,6 +41,7 @@ from litellm.constants import ( ) from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, CommonProxyErrors, ProxyErrorTypes, ProxyException, @@ -5337,7 +5338,7 @@ class ProxyUpdateSpend: ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index cd293325c15..191080e3a48 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -308,6 +308,84 @@ async def test_update_daily_spend_with_null_entity_id(): assert create_data["failed_requests"] == 0 +def _daily_txn(user_id: str = "user1") -> dict: + return { + "user_id": user_id, + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + +@pytest.mark.asyncio +async def test_update_daily_spend_does_not_retry_post_send_ambiguous_errors(): + # Regression for the double-apply hazard: a ReadTimeout means the batch was + # sent and its outcome is unknown; the engine can leave the transaction open + # on the pooled connection, so retrying stacks a second set of increments + # into it and one commit applies both. Post-send failures must drop the + # batch (loudly), never retry it. + import httpx + + mock_prisma_client = MagicMock() + mock_prisma_client.db.batch_ = MagicMock(side_effect=httpx.ReadTimeout("ambiguous")) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(httpx.ReadTimeout): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions={"k1": _daily_txn()}, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", + ) + + mock_prisma_client.db.batch_.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_daily_spend_retries_connect_errors(monkeypatch): + # ConnectError proves the statements never reached the database, so it is + # the one failure the writer may retry. + import httpx + + mock_batcher = MagicMock() + good_ctx = MagicMock() + good_ctx.__aenter__ = AsyncMock(return_value=mock_batcher) + good_ctx.__aexit__ = AsyncMock(return_value=None) + mock_prisma_client = MagicMock() + mock_prisma_client.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), good_ctx]) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", fake_sleep) + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions={"k1": _daily_txn()}, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", + ) + + assert mock_prisma_client.db.batch_.call_count == 2 + + @pytest.mark.asyncio async def test_update_daily_spend_sorting(): """ diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index d5d4de7f2cf..f075acc7307 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -68,12 +68,12 @@ async def test_update_end_user_spend_upserts_each_end_user( @pytest.mark.asyncio -async def test_update_end_user_spend_retries_on_connection_error( +async def test_update_end_user_spend_retries_on_connect_error( mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """``DB_CONNECTION_ERROR_TYPES`` failures should be retried with backoff; - once retries are exhausted, ``_raise_failed_update_spend_exception`` is - invoked and the original exception bubbles up. + """``DB_RETRY_SAFE_ERROR_TYPES`` (ConnectError, statements provably never + sent) retries with backoff; once retries are exhausted the original + exception bubbles up via ``_raise_failed_update_spend_exception``. """ import httpx import litellm.proxy.utils as utils_mod @@ -85,11 +85,11 @@ async def test_update_end_user_spend_retries_on_connection_error( monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) - err = httpx.ReadError("conn reset") + err = httpx.ConnectError("down") mock_prisma_client.db.tx = MagicMock(side_effect=err) proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() - with pytest.raises(httpx.ReadError): + with pytest.raises(httpx.ConnectError): await ProxyUpdateSpend.update_end_user_spend( n_retry_times=1, prisma_client=mock_prisma_client, @@ -99,6 +99,29 @@ async def test_update_end_user_spend_retries_on_connection_error( assert sleeps == [1.0] +@pytest.mark.asyncio +@pytest.mark.parametrize("ambiguous_error_name", ["ReadTimeout", "ReadError"]) +async def test_update_end_user_spend_does_not_retry_post_send_ambiguous_errors( + mock_prisma_client: Any, ambiguous_error_name: str +) -> None: + """Post-send errors are ambiguous and retrying can double-apply increments + (see DB_RETRY_SAFE_ERROR_TYPES); they must raise on the first attempt.""" + import httpx + + err = getattr(httpx, ambiguous_error_name)("ambiguous") + mock_prisma_client.db.tx = MagicMock(side_effect=err) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + with pytest.raises((httpx.ReadTimeout, httpx.ReadError)): + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + end_user_list_transactions={"u": 1.0}, + ) + mock_prisma_client.db.tx.assert_called_once() + + @pytest.mark.asyncio async def test_update_end_user_spend_non_connection_error_raises_immediately( mock_prisma_client: Any, From a87754d7fc6d0a42bd2c822de94e7c042a84c2fd Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 27 Jul 2026 12:29:19 -0700 Subject: [PATCH 7/7] fix(db_scripts): pin the tool spend backfill session to UTC The backfill compares the naive start_time column against a timestamptz cutover, and that coercion follows the session time zone, so a non-UTC session shifts the cutover boundary by the offset. Pinning the session makes the whole script timezone-independent. The date bucketing itself was already safe: to_char on a timestamp without time zone ignores the session time zone and the stored values are UTC --- db_scripts/backfill_daily_tool_spend.sql | 2 ++ 1 file changed, 2 insertions(+) diff --git a/db_scripts/backfill_daily_tool_spend.sql b/db_scripts/backfill_daily_tool_spend.sql index 358ebf1f23f..309b9dbe0ff 100644 --- a/db_scripts/backfill_daily_tool_spend.sql +++ b/db_scripts/backfill_daily_tool_spend.sql @@ -28,6 +28,8 @@ -- Usage: -- psql "$DATABASE_URL" -v cutover="'2026-07-25T00:00:00Z'" -f db_scripts/backfill_daily_tool_spend.sql +SET TIME ZONE 'UTC'; + INSERT INTO "LiteLLM_DailyToolSpend" (date, tool_name, spend, total_tokens, request_count, created_at, updated_at) SELECT to_char(ti.start_time, 'YYYY-MM-DD') AS date,