From f1dea17be1153cb50c1e51e2cd6f57b03216575b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:37:27 -0700 Subject: [PATCH] fix(spend_logs): store litellm_call_id and match it in request_id lookups Success spend rows are keyed by the upstream provider response id, so the x-litellm-call-id response header value never found them. Add a nullable indexed litellm_call_id column to LiteLLM_SpendLogs, populate it at write time, and widen every request_id lookup surface (/spend/logs, /spend/logs/ui, request details, ownership check) to match either id. --- .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/proxy/_types.py | 1 + litellm/proxy/schema.prisma | 2 + .../spend_management_endpoints.py | 40 +++++-- .../spend_tracking/spend_tracking_utils.py | 1 + schema.prisma | 2 + .../test_spend_management_endpoints.py | 107 ++++++++++++++++-- .../test_spend_tracking_utils.py | 27 +++++ 9 files changed, 165 insertions(+), 22 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql new file mode 100644 index 00000000000..b3bcad738ee --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "litellm_call_id" TEXT; + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 01a607b68a9..a9c468c3e9d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -656,12 +656,14 @@ model LiteLLM_SpendLogs { mcp_namespaced_tool_name String? agent_id String? proxy_server_request Json? @default("{}") + litellm_call_id String? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([litellm_call_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c549f48126e..428736da96d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3644,6 +3644,7 @@ class SpendLogsPayload(TypedDict): session_id: str | None request_duration_ms: int | None status: Literal["success", "failure"] + litellm_call_id: ReadOnly[str | None] class SpanAttributes(str, enum.Enum): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 01a607b68a9..a9c468c3e9d 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -656,12 +656,14 @@ model LiteLLM_SpendLogs { mcp_namespaced_tool_name String? agent_id String? proxy_server_request Json? @default("{}") + litellm_call_id String? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([litellm_call_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 41c65b1d5c5..e7e8e0c5341 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -229,10 +229,24 @@ async def _find_spend_logs( return rows +class _RequestIdEquals(TypedDict): + request_id: ReadOnly[str] + + +class _LitellmCallIdEquals(TypedDict): + litellm_call_id: ReadOnly[str] + + +def _request_id_or_call_id_clause(request_id: str) -> tuple[_RequestIdEquals, _LitellmCallIdEquals]: + request_id_clause: Final[_RequestIdEquals] = {"request_id": request_id} + call_id_clause: Final[_LitellmCallIdEquals] = {"litellm_call_id": request_id} + return (request_id_clause, call_id_clause) + + async def _find_spend_log_row(prisma_client: PrismaClient, request_id: str) -> _SpendLogOwnershipRow | None: - """Read the single spend log row identified by ``request_id``.""" - return await _spend_logs_table(prisma_client).find_unique( - where={"request_id": request_id}, + """Read the single spend log row identified by ``request_id`` or ``litellm_call_id``.""" + return await _spend_logs_table(prisma_client).find_first( + where={"OR": _request_id_or_call_id_clause(request_id)}, include=None, ) @@ -2543,7 +2557,6 @@ async def ui_view_spend_logs( ("team_id", "team_id"), ('"user"', "user"), ("api_key", "api_key"), - ("request_id", "request_id"), ("model", "model"), ("model_id", "model_id"), ("model_group", "model_group"), @@ -2555,6 +2568,12 @@ async def ui_view_spend_logs( sql_params.append(val) p += 1 + request_id_filter: Final = where_conditions.get("request_id") + if isinstance(request_id_filter, str): + sql_conditions.append(f"(request_id = ${p} OR litellm_call_id = ${p})") + sql_params.append(request_id_filter) + p += 1 + # Multi-team OR filter: (user = $X OR team_id = ANY($Y)) if permitted_team_ids: or_clause: Final = f'("user" = ${p} OR team_id = ANY(${p + 1}::text[]))' @@ -2662,6 +2681,7 @@ async def ui_view_spend_logs( cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_id, + litellm_call_id, COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms FROM "LiteLLM_SpendLogs" WHERE {" AND ".join(sql_conditions)} @@ -2735,7 +2755,7 @@ def _hydrate_spend_log_metadata(rows: Sequence[Mapping[str, object]]) -> None: def _cold_storage_object_key_from_metadata( - metadata: str | dict | None, + metadata: str | Mapping[str, object] | None, ) -> str | None: if isinstance(metadata, str): try: @@ -2870,7 +2890,7 @@ async def ui_view_request_response_for_request_id( sql_query: Final = """ SELECT messages, response, proxy_server_request, metadata FROM "LiteLLM_SpendLogs" - WHERE request_id = $1 + WHERE request_id = $1 OR litellm_call_id = $1 LIMIT 1 """ db_result: Final[Sequence[Mapping[str, object]] | None] = await _query_raw_or_none( @@ -2989,7 +3009,7 @@ async def view_spend_logs( start_date_iso: Final = start_date_obj.isoformat() end_date_iso: Final = end_date_obj.isoformat() - filter_query: Final = { + filter_query: Final[dict[str, object]] = { "startTime": { "gte": start_date_iso, # Greater than or equal to Start Date "lte": end_date_iso, # Less than or equal to End Date @@ -3002,7 +3022,7 @@ async def view_spend_logs( else: filter_query["api_key"] = api_key if request_id is not None and isinstance(request_id, str): - filter_query["request_id"] = request_id + filter_query["OR"] = _request_id_or_call_id_clause(request_id) if user_id is not None and isinstance(user_id, str): filter_query["user"] = user_id @@ -3073,7 +3093,7 @@ async def view_spend_logs( return response else: - scoped_filter: Final[dict[str, str]] = {} + scoped_filter: Final[dict[str, object]] = {} if api_key is not None and isinstance(api_key, str): if api_key.startswith("sk-"): hashed_token = prisma_client.hash_token(token=api_key) @@ -3081,7 +3101,7 @@ async def view_spend_logs( hashed_token = api_key scoped_filter["api_key"] = hashed_token if request_id is not None and isinstance(request_id, str): - scoped_filter["request_id"] = request_id + scoped_filter["OR"] = _request_id_or_call_id_clause(request_id) if user_id is not None and isinstance(user_id, str): scoped_filter["user"] = user_id diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 9f718b7d20d..da5ce2f727e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -565,6 +565,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs status=_get_status_for_spend_log( metadata=metadata, ), + litellm_call_id=litellm_call_id, ) verbose_proxy_logger.debug( diff --git a/schema.prisma b/schema.prisma index 01a607b68a9..a9c468c3e9d 100644 --- a/schema.prisma +++ b/schema.prisma @@ -656,12 +656,14 @@ model LiteLLM_SpendLogs { mcp_namespaced_tool_name String? agent_id String? proxy_server_request Json? @default("{}") + litellm_call_id String? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([litellm_call_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index a0dcbf802ef..6d392fc5b3b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -98,7 +98,10 @@ def _reconstruct_ui_where_from_sql(sql_query, params): sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond) status = re.fullmatch(r"status = \$(\d+)", cond) api_key_not_in = re.fullmatch(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", cond) - if gte: + req_or_call = re.fullmatch(r"\(request_id = \$(\d+) OR litellm_call_id = \$\1\)", cond) + if req_or_call: + where["request_id_or_call_id"] = params[int(req_or_call.group(1)) - 1] + elif gte: date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1]) elif lte: date_bounds["lte"] = _iso(params[int(lte.group(1)) - 1]) @@ -410,7 +413,7 @@ async def test_assert_user_can_view_request_id_rejects_both_users_none(): team_id = None class MockSpendLogs: - async def find_unique(self, where, include=None): + async def find_first(self, where=None, include=None): return MockRow() class MockDB: @@ -453,6 +456,7 @@ def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypat ignored_keys = [ "request_id", + "litellm_call_id", "metadata.litellm_call_id", "session_id", "startTime", @@ -2162,7 +2166,10 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window( def filter_fn(where): captured["where"] = where rows = _filter_logs_by_date_range(mock_spend_logs, where) - if where.get("request_id"): + rid_either = where.get("request_id_or_call_id") + if rid_either: + rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))] + elif where.get("request_id"): rows = [r for r in rows if r["request_id"] == where["request_id"]] return rows @@ -2192,9 +2199,82 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window( data = response.json() assert data["total"] == 1 assert data["data"][0]["request_id"] == "req-old" - # Query dropped the time window and scoped solely by the primary key. + # Query dropped the time window and scoped solely by the id lookup. assert "startTime" not in captured["where"] - assert captured["where"]["request_id"] == "req-old" + assert captured["where"]["request_id_or_call_id"] == "req-old" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_lookup_matches_litellm_call_id( + client, monkeypatch +): + """ + LIT-6302: success rows are keyed by the upstream provider response id, so a + lookup with the x-litellm-call-id response header value found nothing. The id + lookup now matches request_id OR litellm_call_id, resolving the header value. + """ + today = datetime.datetime.now(timezone.utc) + mock_spend_logs = [ + { + "id": "log_provider_keyed", + "request_id": "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm", + "litellm_call_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd", + "api_key": "sk-test-key", + "user": "test_user_1", + "team_id": "team1", + "spend": 0.05, + "startTime": today.isoformat(), + "model": "gpt-4", + }, + { + "id": "log_other", + "request_id": "chatcmpl-other", + "litellm_call_id": "11111111-2222-3333-4444-555555555555", + "api_key": "sk-test-key", + "user": "test_user_1", + "team_id": "team1", + "spend": 0.01, + "startTime": today.isoformat(), + "model": "gpt-4", + }, + ] + + def filter_fn(where): + rid_either = where.get("request_id_or_call_id") + if rid_either: + return [ + r + for r in mock_spend_logs + if rid_either in (r["request_id"], r.get("litellm_call_id")) + ] + if where.get("request_id"): + return [ + r for r in mock_spend_logs if r["request_id"] == where["request_id"] + ] + return list(mock_spend_logs) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/ui", + params={"request_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert ( + data["data"][0]["request_id"] == "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm" + ) finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -2254,7 +2334,7 @@ async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatc team_id = None class _SpendLogs: - async def find_unique(self, where, include=None): + async def find_first(self, where=None, include=None): return _ForeignRow() class _DB: @@ -2307,7 +2387,10 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( def filter_fn(where): captured["where"] = where rows = _filter_logs_by_date_range(mock_spend_logs, where) - if where.get("request_id"): + rid_either = where.get("request_id_or_call_id") + if rid_either: + rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))] + elif where.get("request_id"): rows = [r for r in rows if r["request_id"] == where["request_id"]] return rows @@ -2317,10 +2400,10 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( user = "user_1" team_id = "team1" - async def _find_unique(where, include=None): + async def _find_first(where=None, include=None): return _OwnedRow() - mock_prisma.db.find_unique = _find_unique + mock_prisma.db.find_first = _find_first monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) # A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends. @@ -2345,7 +2428,7 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( assert data["total"] == 1 assert data["data"][0]["request_id"] == "req-old" assert "startTime" not in captured["where"] - assert captured["where"]["request_id"] == "req-old" + assert captured["where"]["request_id_or_call_id"] == "req-old" assert "user" not in captured["where"] assert "OR" not in captured["where"] finally: @@ -4435,7 +4518,7 @@ async def test_view_spend_logs_internal_user_combines_user_with_request_id( where = mock_client.db.captured_where assert where is not None assert where["user"] == "internal-user-2" - assert where["request_id"] == "req-abc" + assert where["OR"] == ({"request_id": "req-abc"}, {"litellm_call_id": "req-abc"}) finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -4462,7 +4545,7 @@ async def test_view_spend_logs_non_date_range_combines_user_with_request_id( where = mock_client.db.captured_where assert where is not None assert where["user"] == "internal-user-3" - assert where["request_id"] == "req-xyz" + assert where["OR"] == ({"request_id": "req-xyz"}, {"litellm_call_id": "req-xyz"}) finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9e5917637a8..488ccd8b81b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1071,6 +1071,33 @@ def test_get_logging_payload_includes_agent_id_from_kwargs(): ), f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'" +def test_get_logging_payload_populates_litellm_call_id_alongside_provider_request_id(): + """ + LIT-6302: request_id stays the provider response id, so clients holding the + x-litellm-call-id header value could never find their row. The payload now + also carries litellm_call_id as its own column for lookups by either id. + """ + call_id = "b980eea9-5cd9-4099-93cd-8291e46c76fd" + + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_call_id": call_id, + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=litellm.ModelResponse( + id="chatcmpl-provider-id", + choices=[], + usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["request_id"] == "chatcmpl-provider-id" + assert payload["litellm_call_id"] == call_id + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_includes_overhead_in_spend_logs_metadata():