From 2285640eeab7206f609d54286df0bac4130987f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:02:39 -0700 Subject: [PATCH 1/9] fix(spend-tracking): recover key alias for session tokens from spend logs CLI session tokens are in-memory only and never get a LiteLLM_VerificationToken row, so the usage APIs could not resolve key_alias, team_id, or user_email for their spend rows: the exact join and the reverse-hash recovery both miss. The owner is written to LiteLLM_SpendLogs.metadata at request time under the same hashed api_key, so read it back from there for keys still unresolved after the token-table passes. The lookup is sha256-gated like the existing reverse-hash recovery and bounded to the records' startTime window (min date minus one day, max date plus two) so it stays on the startTime index. No migration. Also guard the window parser against the date=None rollup rows GROUPING SETS aggregation emits, which raised TypeError from strptime and turned the aggregated usage endpoints into HTTP 500s. --- .../common_daily_activity.py | 51 ++++++++-- .../spend_tracking/key_metadata_recovery.py | 48 ++++++++++ .../test_common_daily_activity.py | 53 +++++++++++ .../test_key_metadata_recovery.py | 95 ++++++++++++++++++- 4 files changed, 238 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index ce6a97708ab..c6c5aad5926 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -14,6 +14,7 @@ from litellm.proxy._types import CommonProxyErrors from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, + recover_key_metadata_from_spend_logs, ) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import PrismaClient @@ -433,15 +434,37 @@ def update_breakdown_metrics( return breakdown +def _spend_logs_window(dates: AbstractSet[str | None]) -> tuple[datetime, datetime] | None: + parsed: Final = sorted(day for day in (_parse_spend_date(raw) for raw in dates) if day is not None) + if not parsed: + return None + return (parsed[0] - timedelta(days=1), parsed[-1] + timedelta(days=2)) + + +def _parse_spend_date(raw: str | None) -> datetime | None: + if not isinstance(raw, str): + return None + try: + return datetime.fromisoformat(raw) + except ValueError: + return None + + +_EMPTY_KEY_METADATA: Final[Mapping[str, _KeyMetadataDict]] = MappingProxyType({}) + + async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], + spend_logs_window: tuple[datetime, datetime] | None = None, ) -> Mapping[str, _KeyMetadataDict]: """Get api key metadata, falling back to deleted keys table for keys not found in active table. This ensures that key_alias and team_id are preserved in historical activity logs even after a key is deleted or regenerated. Also recovers aliases for api_key - values that were double-hashed by the v1.99 spend-log provenance gate. + values that were double-hashed by the v1.99 spend-log provenance gate, and, when + spend_logs_window is given, for keys never written to either token table (CLI + session tokens) from the spend-log rows those requests wrote in that window. """ key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(api_keys)}} @@ -481,11 +504,17 @@ async def get_api_key_metadata( ) still_missing: Final = api_keys - frozenset(result) - combined: Final = ( - result - if not still_missing - else MappingProxyType({**result, **(await recover_double_hashed_key_metadata(prisma_client, still_missing))}) + from_reverse_hash: Final = ( + await recover_double_hashed_key_metadata(prisma_client, still_missing) if still_missing else _EMPTY_KEY_METADATA ) + after_token_recovery: Final = MappingProxyType({**result, **from_reverse_hash}) + unresolved: Final = api_keys - frozenset(after_token_recovery) + from_spend_logs: Final = ( + await recover_key_metadata_from_spend_logs(prisma_client, unresolved, spend_logs_window) + if unresolved and spend_logs_window is not None + else _EMPTY_KEY_METADATA + ) + combined: Final = MappingProxyType({**after_token_recovery, **from_spend_logs}) return await attach_user_emails(prisma_client, combined) @@ -898,7 +927,9 @@ async def _aggregate_spend_records( api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: - api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) + api_key_metadata = await get_api_key_metadata( + prisma_client, api_keys, _spend_logs_window(frozenset(record.date for record in records)) + ) return await asyncio.to_thread( _aggregate_spend_records_sync, @@ -1094,7 +1125,9 @@ async def _aggregate_grouping_sets_records( api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: - api_key_metadata = await get_api_key_metadata(prisma_client, api_keys) + api_key_metadata = await get_api_key_metadata( + prisma_client, api_keys, _spend_logs_window(frozenset(r.date for r in records)) + ) return await asyncio.to_thread( _aggregate_grouping_sets_records_sync, @@ -1357,7 +1390,9 @@ async def get_daily_activity_aggregated( r.api_key for r in entity_records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY ) entity_key_metadata: Final = ( - await get_api_key_metadata(prisma_client, entity_api_keys) + await get_api_key_metadata( + prisma_client, entity_api_keys, _spend_logs_window(frozenset(r.date for r in entity_records)) + ) if entity_api_keys else {} # mutable-ok: matches the helper's dict return ) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 7de18521edd..81c1f4cf119 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,5 +1,6 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet +from datetime import datetime from types import MappingProxyType from typing import Final, TypeVar @@ -27,6 +28,19 @@ WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[]) ORDER BY token, deleted_at DESC """ +_SPEND_LOG_ALIAS_SQL: Final = """ +SELECT DISTINCT ON (api_key) + api_key AS digest, + metadata->>'user_api_key_alias' AS key_alias, + COALESCE(NULLIF(team_id, ''), metadata->>'user_api_key_team_id') AS team_id, + COALESCE(NULLIF("user", ''), metadata->>'user_api_key_user_id') AS user_id +FROM "LiteLLM_SpendLogs" +WHERE api_key = ANY($1::text[]) + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp +ORDER BY api_key, (metadata->>'user_api_key_alias') IS NULL, "startTime" DESC +""" + class KeyMetadataDict(TypedDict, total=False): key_alias: ReadOnly[str | None] @@ -168,6 +182,40 @@ async def recover_double_hashed_key_metadata( return MappingProxyType({**from_active, **from_deleted}) +async def recover_key_metadata_from_spend_logs( + prisma_client: PrismaClient, + missing_keys: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Mapping[str, KeyMetadataDict]: + """ + Recover key_alias/team_id/user_id for hashed api_key values absent from both + verification-token tables, e.g. in-memory CLI session tokens that never get a + token row. Their owner is written to LiteLLM_SpendLogs metadata at request + time under the same hashed api_key, so it is the only surviving source. Only + sha256 digests are looked up, matching the reverse-hash recovery gate, since + every current api_key value in spend logs is a token hash. The [start, end) + bound keeps the lookup on the startTime index instead of scanning the table. + """ + sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) + if not sha_missing: + return _EMPTY_KEY_METADATA + start, end = window + rows: Final = await _db_or_empty( + lambda: prisma_client.db.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(sha_missing), start, end), + "Failed spend-log alias recovery for %d missing keys: %s", + len(sha_missing), + ) + if rows is None: + return _EMPTY_KEY_METADATA + return MappingProxyType( + { + row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=row.user_id) + for row in _TOKEN_DIGEST_ROWS.validate_python(rows) + if row.digest in sha_missing and (row.key_alias or row.user_id or row.team_id) + } + ) + + def _row_with_recovered_fields( row: Mapping[str, object], recovered: Mapping[str, KeyMetadataDict], diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 37a54c4901a..4d3dc320d8f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -2105,3 +2105,56 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown(): # Rollups with the entity bit set must still land in their usual buckets assert daily.breakdown.models["gpt-4o"].metrics.spend == 18.0 assert daily.breakdown.api_keys["key-1"].metrics.spend == 12.0 + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_resolves_session_key_via_spend_log_window(): + """ + A CLI session token has no verification-token row, so the active/deleted lookups + and reverse-hash all miss. Given a spend-log window, its alias and owner are + recovered from the spend-log metadata and its email is filled from the user table. + """ + from litellm.proxy.utils import hash_token + + session_digest = hash_token("cli-session-user-42") + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="user-42", user_email="user42@example.com")] + ) + + async def query_raw(sql, *params): + if "LiteLLM_SpendLogs" in sql: + return [{"digest": session_digest, "key_alias": "cli-session-user-42", "team_id": None, "user_id": "user-42"}] + return [] + + mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={session_digest}, + spend_logs_window=(datetime(2026, 9, 7), datetime(2026, 9, 10)), + ) + + assert result[session_digest]["key_alias"] == "cli-session-user-42" + assert result[session_digest]["user_id"] == "user-42" + assert result[session_digest]["user_email"] == "user42@example.com" + spend_log_calls = [call.args for call in mock_prisma.db.query_raw.call_args_list if "LiteLLM_SpendLogs" in call.args[0]] + ((_, digests, start, end),) = spend_log_calls + assert digests == [session_digest] + assert (start, end) == (datetime(2026, 9, 7), datetime(2026, 9, 10)) + + +def test_spend_logs_window_pads_min_minus_one_day_and_max_plus_two_days(): + from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window + + window = _spend_logs_window({"2026-09-08", "2026-09-05", "not-a-date"}) + + assert window == (datetime(2026, 9, 4), datetime(2026, 9, 10)) + + +def test_spend_logs_window_is_none_when_no_date_parses(): + from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window + + assert _spend_logs_window({"garbage", ""}) is None diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 7a80319239d..9e42f017106 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,4 +1,5 @@ from collections.abc import Sequence +from datetime import datetime from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -8,14 +9,24 @@ from prisma.errors import PrismaError from litellm.proxy.spend_tracking.key_metadata_recovery import ( fill_missing_api_key_aliases, recover_double_hashed_key_metadata, + recover_key_metadata_from_spend_logs, ) from litellm.proxy.utils import hash_token -def _digest_row(digest: str, key_alias: str, team_id: str | None, user_id: str | None) -> dict[str, str | None]: +def _digest_row(digest: str, key_alias: str | None, team_id: str | None, user_id: str | None) -> dict[str, str | None]: return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id} +def _query_raw_spend_logs(rows: Sequence[dict[str, str | None]]) -> AsyncMock: + async def query_raw(sql: str, *params: object) -> list[dict[str, str | None]]: + if '"LiteLLM_SpendLogs"' in sql: + return list(rows) + raise AssertionError(f"unexpected query: {sql}") + + return AsyncMock(side_effect=query_raw) + + def _query_raw_by_table( active_rows: Sequence[dict[str, str | None]], deleted_rows: Sequence[dict[str, str | None]], @@ -218,3 +229,85 @@ async def test_fill_missing_api_key_aliases_skips_named_keys_that_have_no_email( assert filled == rows mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_resolves_session_token_from_metadata(): + """ + CLI session tokens never get a verification-token row, so both the exact join and + the reverse-hash lookup miss them. Their owner survives only in the spend-log + metadata written at request time, keyed by the same hashed api_key. + """ + session_digest = hash_token("cli-session-repro-user-6852") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs( + [_digest_row(session_digest, "cli-session-repro-user-6852", None, "repro-user-6852")] + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {session_digest}, window) + + assert result[session_digest]["key_alias"] == "cli-session-repro-user-6852" + assert result[session_digest]["user_id"] == "repro-user-6852" + ((_, digests, start, end),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + assert digests == [session_digest] + assert (start, end) == window + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_skips_query_when_no_missing_keys(): + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, set(), window) + + assert result == {} + mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_returns_empty_on_prisma_error(): + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down")) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {hash_token("cli-session-x")}, window) + + assert result == {} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_ignores_foreign_and_all_null_rows(): + wanted = hash_token("cli-session-wanted") + all_null = hash_token("cli-session-null") + foreign = hash_token("cli-session-foreign") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs( + [ + _digest_row(wanted, "kept-alias", None, "owner-1"), + _digest_row(all_null, None, None, None), + _digest_row(foreign, "foreign-alias", None, "owner-2"), + ] + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {wanted, all_null}, window) + + assert set(result) == {wanted} + assert result[wanted]["key_alias"] == "kept-alias" + assert result[wanted]["user_id"] == "owner-1" + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_skips_non_sha256_keys(): + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {"cli-session-raw-1798", "key-hash-short"}, window + ) + + assert result == {} + mock_prisma.db.query_raw.assert_not_called() From cf0488316e23b3f9859691d77043e49eb97f6eb6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:40:28 -0700 Subject: [PATCH 2/9] fix(spend-tracking): take each unresolved key's newest spend row and drop docstrings --- .../common_daily_activity.py | 4 +- .../spend_tracking/key_metadata_recovery.py | 42 +++++++------------ .../test_common_daily_activity.py | 5 --- .../test_key_metadata_recovery.py | 5 --- 4 files changed, 16 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index c6c5aad5926..a8aef30107c 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -462,9 +462,7 @@ async def get_api_key_metadata( This ensures that key_alias and team_id are preserved in historical activity logs even after a key is deleted or regenerated. Also recovers aliases for api_key - values that were double-hashed by the v1.99 spend-log provenance gate, and, when - spend_logs_window is given, for keys never written to either token table (CLI - session tokens) from the spend-log rows those requests wrote in that window. + values that were double-hashed by the v1.99 spend-log provenance gate. """ key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(api_keys)}} diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 81c1f4cf119..e207f226255 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -29,16 +29,21 @@ ORDER BY token, deleted_at DESC """ _SPEND_LOG_ALIAS_SQL: Final = """ -SELECT DISTINCT ON (api_key) - api_key AS digest, - metadata->>'user_api_key_alias' AS key_alias, - COALESCE(NULLIF(team_id, ''), metadata->>'user_api_key_team_id') AS team_id, - COALESCE(NULLIF("user", ''), metadata->>'user_api_key_user_id') AS user_id -FROM "LiteLLM_SpendLogs" -WHERE api_key = ANY($1::text[]) - AND "startTime" >= $2::timestamp - AND "startTime" < $3::timestamp -ORDER BY api_key, (metadata->>'user_api_key_alias') IS NULL, "startTime" DESC +SELECT newest.digest, newest.key_alias, newest.team_id, newest.user_id +FROM unnest($1::text[]) AS missing(digest) +CROSS JOIN LATERAL ( + SELECT + api_key AS digest, + metadata->>'user_api_key_alias' AS key_alias, + COALESCE(NULLIF(team_id, ''), metadata->>'user_api_key_team_id') AS team_id, + COALESCE(NULLIF("user", ''), metadata->>'user_api_key_user_id') AS user_id + FROM "LiteLLM_SpendLogs" + WHERE api_key = missing.digest + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp + ORDER BY "startTime" DESC + LIMIT 1 +) AS newest """ @@ -152,14 +157,6 @@ async def recover_double_hashed_key_metadata( prisma_client: PrismaClient, missing_keys: AbstractSet[str], ) -> Mapping[str, KeyMetadataDict]: - """ - Recover key_alias/team_id/user_id for DailyUserSpend.api_key values that - were double-hashed by the v1.99 spend-log provenance gate. - - Those rows store hash(VerificationToken.token) instead of the token, so the - exact join misses. Postgres hashes the token column itself, one pass over - active keys and one over deleted keys, so no key row crosses the wire. - """ sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) if not sha_missing: return _EMPTY_KEY_METADATA @@ -187,15 +184,6 @@ async def recover_key_metadata_from_spend_logs( missing_keys: AbstractSet[str], window: tuple[datetime, datetime], ) -> Mapping[str, KeyMetadataDict]: - """ - Recover key_alias/team_id/user_id for hashed api_key values absent from both - verification-token tables, e.g. in-memory CLI session tokens that never get a - token row. Their owner is written to LiteLLM_SpendLogs metadata at request - time under the same hashed api_key, so it is the only surviving source. Only - sha256 digests are looked up, matching the reverse-hash recovery gate, since - every current api_key value in spend logs is a token hash. The [start, end) - bound keeps the lookup on the startTime index instead of scanning the table. - """ sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) if not sha_missing: return _EMPTY_KEY_METADATA diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 4d3dc320d8f..4c8be415204 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -2109,11 +2109,6 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown(): @pytest.mark.asyncio async def test_get_api_key_metadata_resolves_session_key_via_spend_log_window(): - """ - A CLI session token has no verification-token row, so the active/deleted lookups - and reverse-hash all miss. Given a spend-log window, its alias and owner are - recovered from the spend-log metadata and its email is filled from the user table. - """ from litellm.proxy.utils import hash_token session_digest = hash_token("cli-session-user-42") diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 9e42f017106..05a90137e3b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -233,11 +233,6 @@ async def test_fill_missing_api_key_aliases_skips_named_keys_that_have_no_email( @pytest.mark.asyncio async def test_recover_key_metadata_from_spend_logs_resolves_session_token_from_metadata(): - """ - CLI session tokens never get a verification-token row, so both the exact join and - the reverse-hash lookup miss them. Their owner survives only in the spend-log - metadata written at request time, keyed by the same hashed api_key. - """ session_digest = hash_token("cli-session-repro-user-6852") window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() From 550e733bf35ac862970cfebeed872e45f539bc7d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:09:43 -0700 Subject: [PATCH 3/9] perf(spend-tracking): recover key identity in one bounded spend-log pass with a per-worker cache The spend-log lookup for permanently unresolvable digests is back to a single DISTINCT ON scan over the requested window, keeping only rows that carry an alias, user, or team so a newer nameless row cannot hide an older named one. Results and misses are cached per worker for ten minutes keyed by digest and window, failed queries are not cached, and JWT rows keyed hashed-jwt- now pass the digest gate. Tests cover the JWT gate, cache reuse and partial misses, window changes, error handling, the with-window guard, and the daily activity wiring --- litellm/constants.py | 2 + .../spend_tracking/key_metadata_recovery.py | 105 +++++++++++++----- .../test_common_daily_activity.py | 74 +++++++++++- .../test_key_metadata_recovery.py | 99 ++++++++++++++++- 4 files changed, 248 insertions(+), 32 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 8ef9523a60a..c1896573eb5 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1763,6 +1763,8 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) +SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = int(os.getenv("SPEND_LOG_KEY_METADATA_CACHE_TTL", "600")) +SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = int(os.getenv("SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS", "10000")) # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding # staleness so a transient DB error (which surfaces as an empty list) cannot diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index e207f226255..7212d8b6acb 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -8,6 +8,8 @@ from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, SPEND_LOG_KEY_METADATA_CACHE_TTL from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient from litellm.repositories.user_repository import UserRepository @@ -29,23 +31,27 @@ ORDER BY token, deleted_at DESC """ _SPEND_LOG_ALIAS_SQL: Final = """ -SELECT newest.digest, newest.key_alias, newest.team_id, newest.user_id -FROM unnest($1::text[]) AS missing(digest) -CROSS JOIN LATERAL ( - SELECT - api_key AS digest, - metadata->>'user_api_key_alias' AS key_alias, - COALESCE(NULLIF(team_id, ''), metadata->>'user_api_key_team_id') AS team_id, - COALESCE(NULLIF("user", ''), metadata->>'user_api_key_user_id') AS user_id - FROM "LiteLLM_SpendLogs" - WHERE api_key = missing.digest - AND "startTime" >= $2::timestamp - AND "startTime" < $3::timestamp - ORDER BY "startTime" DESC - LIMIT 1 -) AS newest +SELECT DISTINCT ON (api_key) + api_key AS digest, + metadata->>'user_api_key_alias' AS key_alias, + COALESCE(NULLIF(team_id, ''), metadata->>'user_api_key_team_id') AS team_id, + COALESCE(NULLIF("user", ''), metadata->>'user_api_key_user_id') AS user_id +FROM "LiteLLM_SpendLogs" +WHERE api_key = ANY($1::text[]) + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp + AND COALESCE( + metadata->>'user_api_key_alias', + NULLIF("user", ''), + metadata->>'user_api_key_user_id', + NULLIF(team_id, ''), + metadata->>'user_api_key_team_id' + ) IS NOT NULL +ORDER BY api_key, "startTime" DESC """ +_HASHED_JWT_PREFIX: Final = "hashed-jwt-" + class KeyMetadataDict(TypedDict, total=False): key_alias: ReadOnly[str | None] @@ -62,6 +68,11 @@ class _TokenDigestRow(BaseModel): _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) +_CACHED_KEY_METADATA: Final = TypeAdapter(KeyMetadataDict) +_SPEND_LOG_METADATA_CACHE: Final = InMemoryCache( + max_size_in_memory=SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, + default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, +) _EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({}) _EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({}) @@ -179,31 +190,73 @@ async def recover_double_hashed_key_metadata( return MappingProxyType({**from_active, **from_deleted}) -async def recover_key_metadata_from_spend_logs( - prisma_client: PrismaClient, - missing_keys: AbstractSet[str], +def _is_spend_log_digest(key: str) -> bool: + return is_valid_sha256_hash(key.removeprefix(_HASHED_JWT_PREFIX)) + + +def _spend_log_cache_key(digest: str, window: tuple[datetime, datetime]) -> str: + start, end = window + return f"spend_log_key_metadata:{digest}:{start.isoformat()}:{end.isoformat()}" + + +def _cached_spend_log_metadata( + cache: InMemoryCache, + digest: str, window: tuple[datetime, datetime], -) -> Mapping[str, KeyMetadataDict]: - sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) - if not sha_missing: - return _EMPTY_KEY_METADATA +) -> KeyMetadataDict | None: + cached: Final[object] = cache.get_cache(_spend_log_cache_key(digest, window)) + return None if cached is None else _CACHED_KEY_METADATA.validate_python(cached) + + +async def _query_spend_log_metadata( + prisma_client: PrismaClient, + digests: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Mapping[str, KeyMetadataDict] | None: start, end = window rows: Final = await _db_or_empty( - lambda: prisma_client.db.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(sha_missing), start, end), + lambda: prisma_client.db.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(digests), start, end), "Failed spend-log alias recovery for %d missing keys: %s", - len(sha_missing), + len(digests), ) if rows is None: - return _EMPTY_KEY_METADATA + return None return MappingProxyType( { row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=row.user_id) for row in _TOKEN_DIGEST_ROWS.validate_python(rows) - if row.digest in sha_missing and (row.key_alias or row.user_id or row.team_id) + if row.digest in digests and (row.key_alias or row.user_id or row.team_id) } ) +async def recover_key_metadata_from_spend_logs( + prisma_client: PrismaClient, + missing_keys: AbstractSet[str], + window: tuple[datetime, datetime], + cache: InMemoryCache = _SPEND_LOG_METADATA_CACHE, +) -> Mapping[str, KeyMetadataDict]: + digests: Final = frozenset(key for key in missing_keys if _is_spend_log_digest(key)) + if not digests: + return _EMPTY_KEY_METADATA + cached: Final = MappingProxyType( + { + digest: meta + for digest in digests + for meta in (_cached_spend_log_metadata(cache, digest, window),) + if meta is not None + } + ) + uncached: Final = digests - frozenset(cached) + fresh: Final = await _query_spend_log_metadata(prisma_client, uncached, window) if uncached else _EMPTY_KEY_METADATA + if fresh is not None: + for digest in uncached: + cache.set_cache(_spend_log_cache_key(digest, window), fresh.get(digest, KeyMetadataDict())) + return MappingProxyType( + {digest: meta for digest, meta in (*cached.items(), *(fresh or _EMPTY_KEY_METADATA).items()) if meta} + ) + + def _row_with_recovered_fields( row: Mapping[str, object], recovered: Mapping[str, KeyMetadataDict], diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 4c8be415204..d214ab6b5ac 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -492,7 +492,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( @pytest.mark.asyncio async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_spend_logs(): - """A dirty key no table can explain costs two digest lookups, never a token page walk or a SpendLogs scan.""" + """Without a spend-log window a dirty key no table can explain costs two digest lookups and never a token page walk.""" from litellm.proxy.utils import hash_token double_hashed = hash_token("b" * 64) @@ -518,6 +518,78 @@ async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_s assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups) +@pytest.mark.asyncio +async def test_get_api_key_metadata_permanent_miss_with_a_window_reads_spend_logs_once_within_it(): + from litellm.proxy.utils import hash_token + + double_hashed = hash_token("permanent-miss-with-window-6852") + window = (datetime(2024, 1, 1), datetime(2024, 1, 4)) + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={double_hashed}, spend_logs_window=window) + + assert double_hashed not in result + assert mock_prisma.db.query_raw.await_count == 3 + ((_, digests, start, end),) = [ + call.args for call in mock_prisma.db.query_raw.call_args_list if "LiteLLM_SpendLogs" in call.args[0] + ] + assert digests == [double_hashed] + assert (start, end) == window + + +@pytest.mark.asyncio +async def test_get_daily_activity_recovers_a_session_key_alias_from_spend_logs_around_the_page_dates(): + from litellm.proxy.utils import hash_token + + session_digest = hash_token("cli-session-daily-activity-6852") + records = [_daily_user_spend_record(user_id="session-user", api_key=session_digest, spend=1.5)] + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=len(records)) + mock_table.find_many = AsyncMock(return_value=records) + mock_prisma.db.litellm_dailyuserspend = mock_table + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="session-user", user_email="session@example.com")] + ) + + async def query_raw(sql, *params): + if "LiteLLM_SpendLogs" in sql: + return [{"digest": session_digest, "key_alias": "cli-session-alias", "team_id": None, "user_id": "session-user"}] + return [] + + mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw) + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + page=1, + page_size=1000, + ) + + key_metadata = result.results[0].breakdown.api_keys[session_digest].metadata + assert key_metadata.key_alias == "cli-session-alias" + assert key_metadata.user_email == "session@example.com" + ((_, digests, start, end),) = [ + call.args for call in mock_prisma.db.query_raw.call_args_list if "LiteLLM_SpendLogs" in call.args[0] + ] + assert digests == [session_digest] + assert (start, end) == (datetime(2023, 12, 31), datetime(2024, 1, 3)) + + def test_key_metadata_includes_recovered_user_email(): from litellm.proxy.management_endpoints.common_daily_activity import _key_metadata diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 05a90137e3b..3d7362b9de0 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from prisma.errors import PrismaError +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy.spend_tracking.key_metadata_recovery import ( fill_missing_api_key_aliases, recover_double_hashed_key_metadata, @@ -240,7 +241,7 @@ async def test_recover_key_metadata_from_spend_logs_resolves_session_token_from_ [_digest_row(session_digest, "cli-session-repro-user-6852", None, "repro-user-6852")] ) - result = await recover_key_metadata_from_spend_logs(mock_prisma, {session_digest}, window) + result = await recover_key_metadata_from_spend_logs(mock_prisma, {session_digest}, window, cache=InMemoryCache()) assert result[session_digest]["key_alias"] == "cli-session-repro-user-6852" assert result[session_digest]["user_id"] == "repro-user-6852" @@ -255,7 +256,7 @@ async def test_recover_key_metadata_from_spend_logs_skips_query_when_no_missing_ mock_prisma = MagicMock() mock_prisma.db.query_raw = AsyncMock(return_value=[]) - result = await recover_key_metadata_from_spend_logs(mock_prisma, set(), window) + result = await recover_key_metadata_from_spend_logs(mock_prisma, set(), window, cache=InMemoryCache()) assert result == {} mock_prisma.db.query_raw.assert_not_called() @@ -267,7 +268,9 @@ async def test_recover_key_metadata_from_spend_logs_returns_empty_on_prisma_erro mock_prisma = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down")) - result = await recover_key_metadata_from_spend_logs(mock_prisma, {hash_token("cli-session-x")}, window) + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {hash_token("cli-session-x")}, window, cache=InMemoryCache() + ) assert result == {} @@ -287,7 +290,7 @@ async def test_recover_key_metadata_from_spend_logs_ignores_foreign_and_all_null ] ) - result = await recover_key_metadata_from_spend_logs(mock_prisma, {wanted, all_null}, window) + result = await recover_key_metadata_from_spend_logs(mock_prisma, {wanted, all_null}, window, cache=InMemoryCache()) assert set(result) == {wanted} assert result[wanted]["key_alias"] == "kept-alias" @@ -301,8 +304,94 @@ async def test_recover_key_metadata_from_spend_logs_skips_non_sha256_keys(): mock_prisma.db.query_raw = AsyncMock(return_value=[]) result = await recover_key_metadata_from_spend_logs( - mock_prisma, {"cli-session-raw-1798", "key-hash-short"}, window + mock_prisma, {"cli-session-raw-1798", "key-hash-short"}, window, cache=InMemoryCache() ) assert result == {} mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_accepts_hashed_jwt_digests(): + jwt_digest = f"hashed-jwt-{hash_token('jwt-subject-1')}" + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(jwt_digest, None, "team-jwt", "jwt-user")]) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {jwt_digest}, window, cache=InMemoryCache()) + + assert result[jwt_digest]["team_id"] == "team-jwt" + assert result[jwt_digest]["user_id"] == "jwt-user" + ((_, digests, _, _),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + assert digests == [jwt_digest] + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_serves_repeat_lookups_from_the_cache(): + found = hash_token("cli-session-found") + unknown = hash_token("cli-session-unknown") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache() + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(found, "found-alias", None, "owner-1")]) + + first = await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) + second = await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) + + assert first == second + assert set(first) == {found} + assert first[found]["key_alias"] == "found-alias" + assert mock_prisma.db.query_raw.await_count == 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_only_queries_digests_the_cache_has_not_seen(): + cached_digest = hash_token("cli-session-cached") + new_digest = hash_token("cli-session-new") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache() + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(cached_digest, "cached-alias", None, None)]) + await recover_key_metadata_from_spend_logs(mock_prisma, {cached_digest}, window, cache=cache) + mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(new_digest, "new-alias", None, None)]) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {cached_digest, new_digest}, window, cache=cache) + + assert result[cached_digest]["key_alias"] == "cached-alias" + assert result[new_digest]["key_alias"] == "new-alias" + ((_, digests, _, _),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + assert digests == [new_digest] + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_rescans_when_the_window_changes(): + digest = hash_token("cli-session-windowed") + cache = InMemoryCache() + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs([]) + await recover_key_metadata_from_spend_logs( + mock_prisma, {digest}, (datetime(2026, 9, 1), datetime(2026, 9, 4)), cache=cache + ) + mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(digest, "later-alias", None, None)]) + + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {digest}, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=cache + ) + + assert result[digest]["key_alias"] == "later-alias" + assert mock_prisma.db.query_raw.await_count == 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_does_not_cache_a_failed_query(): + digest = hash_token("cli-session-retry") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache() + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down")) + assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} + mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(digest, "back-online", None, None)]) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) + + assert result[digest]["key_alias"] == "back-online" From f4aef5a1db67d35f606d09f5e8bfe008181af015 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:19:35 -0700 Subject: [PATCH 4/9] fix(spend-tracking): make the spend-log metadata cache limits plain constants The ten minute TTL and the 10000 item ceiling need no env override, and the documentation env-key check flags any os.getenv read that the docs do not list --- litellm/constants.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c1896573eb5..9654c84c412 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1763,8 +1763,8 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) -SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = int(os.getenv("SPEND_LOG_KEY_METADATA_CACHE_TTL", "600")) -SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = int(os.getenv("SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS", "10000")) +SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 +SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding # staleness so a transient DB error (which surfaces as an empty list) cannot From 0c2d0f4777f4113b42510e2daef2330ca0bb87a8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:28:13 -0700 Subject: [PATCH 5/9] fix(spend-tracking): forget an empty spend-log lookup after thirty seconds The daily spend rows and the spend logs of one batch are written a moment apart, so a usage read landing between them used to remember the session as nameless for ten minutes on that worker. Found identities keep the ten minute entry --- litellm/constants.py | 1 + .../spend_tracking/key_metadata_recovery.py | 19 ++++++++++++++++-- .../test_key_metadata_recovery.py | 20 +++++++++++++++++++ 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 9654c84c412..cfc5b6b86a7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1764,6 +1764,7 @@ SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 +SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 7212d8b6acb..354c6479d34 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -9,7 +9,11 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, SPEND_LOG_KEY_METADATA_CACHE_TTL +from litellm.constants import ( + SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, + SPEND_LOG_KEY_METADATA_CACHE_TTL, + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, +) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient from litellm.repositories.user_repository import UserRepository @@ -230,6 +234,17 @@ async def _query_spend_log_metadata( ) +def _remember_spend_log_metadata( + cache: InMemoryCache, digest: str, window: tuple[datetime, datetime], meta: KeyMetadataDict | None +) -> None: + if meta is None: + cache.set_cache( + _spend_log_cache_key(digest, window), KeyMetadataDict(), ttl=SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + ) + return + cache.set_cache(_spend_log_cache_key(digest, window), meta) + + async def recover_key_metadata_from_spend_logs( prisma_client: PrismaClient, missing_keys: AbstractSet[str], @@ -251,7 +266,7 @@ async def recover_key_metadata_from_spend_logs( fresh: Final = await _query_spend_log_metadata(prisma_client, uncached, window) if uncached else _EMPTY_KEY_METADATA if fresh is not None: for digest in uncached: - cache.set_cache(_spend_log_cache_key(digest, window), fresh.get(digest, KeyMetadataDict())) + _remember_spend_log_metadata(cache, digest, window, fresh.get(digest)) return MappingProxyType( {digest: meta for digest, meta in (*cached.items(), *(fresh or _EMPTY_KEY_METADATA).items()) if meta} ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 3d7362b9de0..4f065070330 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,3 +1,4 @@ +import time from collections.abc import Sequence from datetime import datetime from types import SimpleNamespace @@ -7,6 +8,7 @@ import pytest from prisma.errors import PrismaError from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL from litellm.proxy.spend_tracking.key_metadata_recovery import ( fill_missing_api_key_aliases, recover_double_hashed_key_metadata, @@ -395,3 +397,21 @@ async def test_recover_key_metadata_from_spend_logs_does_not_cache_a_failed_quer result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) assert result[digest]["key_alias"] == "back-online" + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_forgets_a_miss_long_before_a_hit(): + found = hash_token("cli-session-found") + unknown = hash_token("cli-session-unknown") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(found, "found-alias", None, None)]) + started = time.time() + + await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) + + hit_expires = next(deadline for key, deadline in cache.ttl_dict.items() if found in key) + miss_expires = next(deadline for key, deadline in cache.ttl_dict.items() if unknown in key) + assert miss_expires - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 + assert hit_expires - started >= SPEND_LOG_KEY_METADATA_CACHE_TTL - 1 From 907c200b3162368cfaa0e363411997330b2da4f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:00:56 -0700 Subject: [PATCH 6/9] fix(spend-tracking): run one spend-log scan at a time and back off repeated misses Concurrent usage reads on one worker now share a single spend-log query instead of each scanning the same window, and a digest that comes back nameless a second time is remembered for the full ten minutes rather than thirty seconds, so a key that never resolves costs at most two scans per worker per window per ten minutes. The first miss still expires after thirty seconds so a read that lands between the daily spend flush and the spend-log flush recovers on the next read --- .../spend_tracking/key_metadata_recovery.py | 71 +++++++++++++------ .../test_key_metadata_recovery.py | 44 ++++++++++++ 2 files changed, 92 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 354c6479d34..319050b87cd 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime @@ -77,6 +78,7 @@ _SPEND_LOG_METADATA_CACHE: Final = InMemoryCache( max_size_in_memory=SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, ) +_SPEND_LOG_QUERY_LOCK: Final = asyncio.Lock() _EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({}) _EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({}) @@ -205,11 +207,17 @@ def _spend_log_cache_key(digest: str, window: tuple[datetime, datetime]) -> str: def _cached_spend_log_metadata( cache: InMemoryCache, - digest: str, + digests: AbstractSet[str], window: tuple[datetime, datetime], -) -> KeyMetadataDict | None: - cached: Final[object] = cache.get_cache(_spend_log_cache_key(digest, window)) - return None if cached is None else _CACHED_KEY_METADATA.validate_python(cached) +) -> Mapping[str, KeyMetadataDict]: + return MappingProxyType( + { + digest: _CACHED_KEY_METADATA.validate_python(cached) + for digest in digests + for cached in (cache.get_cache(_spend_log_cache_key(digest, window)),) + if cached is not None + } + ) async def _query_spend_log_metadata( @@ -237,12 +245,36 @@ async def _query_spend_log_metadata( def _remember_spend_log_metadata( cache: InMemoryCache, digest: str, window: tuple[datetime, datetime], meta: KeyMetadataDict | None ) -> None: - if meta is None: - cache.set_cache( - _spend_log_cache_key(digest, window), KeyMetadataDict(), ttl=SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL - ) + key: Final = _spend_log_cache_key(digest, window) + if meta is not None: + cache.set_cache(key, meta) return - cache.set_cache(_spend_log_cache_key(digest, window), meta) + missed_before: Final = f"{key}:missed-before" + if cache.get_cache(missed_before) is not None: + cache.set_cache(key, KeyMetadataDict()) + return + cache.set_cache(key, KeyMetadataDict(), ttl=SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL) + cache.set_cache(missed_before, True) + + +async def _spend_log_metadata_one_query_at_a_time( + prisma_client: PrismaClient, + cache: InMemoryCache, + lock: asyncio.Lock, + digests: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Mapping[str, KeyMetadataDict]: + async with lock: + settled: Final = _cached_spend_log_metadata(cache, digests, window) + pending: Final = digests - frozenset(settled) + fresh: Final = ( + await _query_spend_log_metadata(prisma_client, pending, window) if pending else _EMPTY_KEY_METADATA + ) + if fresh is None: + return settled + for digest in pending: + _remember_spend_log_metadata(cache, digest, window, fresh.get(digest)) + return MappingProxyType({**settled, **fresh}) async def recover_key_metadata_from_spend_logs( @@ -250,26 +282,19 @@ async def recover_key_metadata_from_spend_logs( missing_keys: AbstractSet[str], window: tuple[datetime, datetime], cache: InMemoryCache = _SPEND_LOG_METADATA_CACHE, + lock: asyncio.Lock = _SPEND_LOG_QUERY_LOCK, ) -> Mapping[str, KeyMetadataDict]: digests: Final = frozenset(key for key in missing_keys if _is_spend_log_digest(key)) if not digests: return _EMPTY_KEY_METADATA - cached: Final = MappingProxyType( - { - digest: meta - for digest in digests - for meta in (_cached_spend_log_metadata(cache, digest, window),) - if meta is not None - } - ) + cached: Final = _cached_spend_log_metadata(cache, digests, window) uncached: Final = digests - frozenset(cached) - fresh: Final = await _query_spend_log_metadata(prisma_client, uncached, window) if uncached else _EMPTY_KEY_METADATA - if fresh is not None: - for digest in uncached: - _remember_spend_log_metadata(cache, digest, window, fresh.get(digest)) - return MappingProxyType( - {digest: meta for digest, meta in (*cached.items(), *(fresh or _EMPTY_KEY_METADATA).items()) if meta} + settled: Final = ( + await _spend_log_metadata_one_query_at_a_time(prisma_client, cache, lock, uncached, window) + if uncached + else _EMPTY_KEY_METADATA ) + return MappingProxyType({digest: meta for digest, meta in (*cached.items(), *settled.items()) if meta}) def _row_with_recovered_fields( diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 4f065070330..5b45777502a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,3 +1,4 @@ +import asyncio import time from collections.abc import Sequence from datetime import datetime @@ -415,3 +416,46 @@ async def test_recover_key_metadata_from_spend_logs_forgets_a_miss_long_before_a miss_expires = next(deadline for key, deadline in cache.ttl_dict.items() if unknown in key) assert miss_expires - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 assert hit_expires - started >= SPEND_LOG_KEY_METADATA_CACHE_TTL - 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_runs_one_query_for_concurrent_lookups(): + digest = hash_token("cli-session-shared") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache() + lock = asyncio.Lock() + mock_prisma = MagicMock() + + async def slow_query_raw(sql: str, *params: object) -> list[dict[str, str | None]]: + await asyncio.sleep(0.01) + return [_digest_row(digest, "shared-alias", None, None)] + + mock_prisma.db.query_raw = AsyncMock(side_effect=slow_query_raw) + + results = await asyncio.gather( + *( + recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache, lock=lock) + for _ in range(9) + ) + ) + + assert all(result[digest]["key_alias"] == "shared-alias" for result in results) + assert mock_prisma.db.query_raw.await_count == 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_keeps_a_repeated_miss_as_long_as_a_hit(): + unknown = hash_token("cli-session-never-named") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs([]) + await recover_key_metadata_from_spend_logs(mock_prisma, {unknown}, window, cache=cache) + first_miss_key = next(key for key in cache.ttl_dict if unknown in key and not key.endswith(":missed-before")) + cache.ttl_dict[first_miss_key] = time.time() - 1 + started = time.time() + + await recover_key_metadata_from_spend_logs(mock_prisma, {unknown}, window, cache=cache) + + assert mock_prisma.db.query_raw.await_count == 2 + assert cache.ttl_dict[first_miss_key] - started >= SPEND_LOG_KEY_METADATA_CACHE_TTL - 1 From 7a6c0cbf08d9c564697c4df7d70d661b0e6241b9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:43:28 -0700 Subject: [PATCH 7/9] fix(spend-tracking): drop the owner of a digest shared by several users and back off failed scans --- .../spend_tracking/key_metadata_recovery.py | 54 +++++++++++-------- .../test_key_metadata_recovery.py | 49 +++++++++++++++-- 2 files changed, 79 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 319050b87cd..d555da99397 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -38,20 +38,23 @@ ORDER BY token, deleted_at DESC _SPEND_LOG_ALIAS_SQL: Final = """ SELECT DISTINCT ON (api_key) api_key AS digest, - metadata->>'user_api_key_alias' AS key_alias, - COALESCE(NULLIF(team_id, ''), metadata->>'user_api_key_team_id') AS team_id, - COALESCE(NULLIF("user", ''), metadata->>'user_api_key_user_id') AS user_id -FROM "LiteLLM_SpendLogs" -WHERE api_key = ANY($1::text[]) - AND "startTime" >= $2::timestamp - AND "startTime" < $3::timestamp - AND COALESCE( - metadata->>'user_api_key_alias', - NULLIF("user", ''), - metadata->>'user_api_key_user_id', - NULLIF(team_id, ''), - metadata->>'user_api_key_team_id' - ) IS NOT NULL + key_alias, + team_id, + user_id, + MIN(user_id) OVER (PARTITION BY api_key) AS first_owner, + MAX(user_id) OVER (PARTITION BY api_key) AS last_owner +FROM ( + SELECT api_key, + "startTime", + NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, + COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id, + COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id + FROM "LiteLLM_SpendLogs" + WHERE api_key = ANY($1::text[]) + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp +) named +WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL ORDER BY api_key, "startTime" DESC """ @@ -72,7 +75,16 @@ class _TokenDigestRow(BaseModel): user_id: str | None = None +class _SpendLogDigestRow(_TokenDigestRow): + first_owner: str | None = None + last_owner: str | None = None + + def unanimous_owner(self) -> str | None: + return self.user_id if self.first_owner == self.last_owner else None + + _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) +_SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...]) _CACHED_KEY_METADATA: Final = TypeAdapter(KeyMetadataDict) _SPEND_LOG_METADATA_CACHE: Final = InMemoryCache( max_size_in_memory=SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, @@ -235,9 +247,10 @@ async def _query_spend_log_metadata( return None return MappingProxyType( { - row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=row.user_id) - for row in _TOKEN_DIGEST_ROWS.validate_python(rows) - if row.digest in digests and (row.key_alias or row.user_id or row.team_id) + row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=owner) + for row in _SPEND_LOG_DIGEST_ROWS.validate_python(rows) + for owner in (row.unanimous_owner(),) + if row.digest in digests and (row.key_alias or owner or row.team_id) } ) @@ -270,11 +283,10 @@ async def _spend_log_metadata_one_query_at_a_time( fresh: Final = ( await _query_spend_log_metadata(prisma_client, pending, window) if pending else _EMPTY_KEY_METADATA ) - if fresh is None: - return settled + found: Final = fresh if fresh is not None else _EMPTY_KEY_METADATA for digest in pending: - _remember_spend_log_metadata(cache, digest, window, fresh.get(digest)) - return MappingProxyType({**settled, **fresh}) + _remember_spend_log_metadata(cache, digest, window, found.get(digest)) + return MappingProxyType({**settled, **found}) async def recover_key_metadata_from_spend_logs( diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 5b45777502a..682b11f5bb7 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -386,20 +386,63 @@ async def test_recover_key_metadata_from_spend_logs_rescans_when_the_window_chan @pytest.mark.asyncio -async def test_recover_key_metadata_from_spend_logs_does_not_cache_a_failed_query(): +async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl(): digest = hash_token("cli-session-retry") window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) - cache = InMemoryCache() + cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) mock_prisma = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down")) + mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("statement timeout")) + started = time.time() assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(digest, "back-online", None, None)]) + assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} + mock_prisma.db.query_raw.assert_not_awaited() + miss_key = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before")) + assert cache.ttl_dict[miss_key] - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 + cache.ttl_dict[miss_key] = time.time() - 1 + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) assert result[digest]["key_alias"] == "back-online" +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_drops_the_owner_of_a_digest_shared_by_several_users(): + shared_ui_digest = hash_token("ui-token") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs( + [ + { + **_digest_row(shared_ui_digest, "ui-token", "litellm-dashboard", "bob"), + "first_owner": "alice", + "last_owner": "bob", + } + ] + ) + + result = await recover_key_metadata_from_spend_logs( + mock_prisma, {shared_ui_digest}, window, cache=InMemoryCache() + ) + + assert result[shared_ui_digest] == {"key_alias": "ui-token", "team_id": "litellm-dashboard", "user_id": None} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_keeps_the_owner_when_every_named_row_agrees(): + digest = hash_token("cli-session-one-owner") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + mock_prisma.db.query_raw = _query_raw_spend_logs( + [{**_digest_row(digest, None, None, "carol"), "first_owner": "carol", "last_owner": "carol"}] + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert result[digest]["user_id"] == "carol" + + @pytest.mark.asyncio async def test_recover_key_metadata_from_spend_logs_forgets_a_miss_long_before_a_hit(): found = hash_token("cli-session-found") From 268b944081e7139e54dc88dda6e311152274dc68 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:24:22 -0700 Subject: [PATCH 8/9] fix(spend-tracking): bound the spend-log scan with a statement timeout and name only unanimous alias, team, and owner --- litellm/constants.py | 1 + .../spend_tracking/key_metadata_recovery.py | 62 ++++-- .../test_key_metadata_recovery.py | 176 +++++++++++++----- 3 files changed, 176 insertions(+), 63 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cfc5b6b86a7..e599a5dd3f9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1766,6 +1766,7 @@ DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACH SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 +SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding # staleness so a transient DB error (which surfaces as an empty list) cannot diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index d555da99397..29688b61b3d 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,7 +1,7 @@ import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet -from datetime import datetime +from datetime import datetime, timedelta from types import MappingProxyType from typing import Final, TypeVar @@ -14,6 +14,7 @@ from litellm.constants import ( SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS, SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, + SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, ) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient @@ -36,16 +37,15 @@ ORDER BY token, deleted_at DESC """ _SPEND_LOG_ALIAS_SQL: Final = """ -SELECT DISTINCT ON (api_key) - api_key AS digest, - key_alias, - team_id, - user_id, - MIN(user_id) OVER (PARTITION BY api_key) AS first_owner, - MAX(user_id) OVER (PARTITION BY api_key) AS last_owner +SELECT api_key AS digest, + MIN(key_alias) AS first_alias, + MAX(key_alias) AS last_alias, + MIN(team_id) AS first_team, + MAX(team_id) AS last_team, + MIN(user_id) AS first_owner, + MAX(user_id) AS last_owner FROM ( SELECT api_key, - "startTime", NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id, COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id @@ -55,9 +55,12 @@ FROM ( AND "startTime" < $3::timestamp ) named WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL -ORDER BY api_key, "startTime" DESC +GROUP BY api_key """ +_SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" +_SPEND_LOG_TRANSACTION_TIMEOUT: Final = timedelta(milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS) + _HASHED_JWT_PREFIX: Final = "hashed-jwt-" @@ -75,12 +78,25 @@ class _TokenDigestRow(BaseModel): user_id: str | None = None -class _SpendLogDigestRow(_TokenDigestRow): +def _unanimous(first: str | None, last: str | None) -> str | None: + return first if first == last else None + + +class _SpendLogDigestRow(BaseModel): + digest: str + first_alias: str | None = None + last_alias: str | None = None + first_team: str | None = None + last_team: str | None = None first_owner: str | None = None last_owner: str | None = None - def unanimous_owner(self) -> str | None: - return self.user_id if self.first_owner == self.last_owner else None + def metadata(self) -> KeyMetadataDict: + return KeyMetadataDict( + key_alias=_unanimous(self.first_alias, self.last_alias), + team_id=_unanimous(self.first_team, self.last_team), + user_id=_unanimous(self.first_owner, self.last_owner), + ) _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) @@ -232,14 +248,24 @@ def _cached_spend_log_metadata( ) +async def _spend_log_rows_within_the_statement_timeout( + prisma_client: PrismaClient, + digests: AbstractSet[str], + window: tuple[datetime, datetime], +) -> Sequence[Mapping[str, object]]: + start, end = window + async with prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction: + await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) + return await transaction.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(digests), start, end) + + async def _query_spend_log_metadata( prisma_client: PrismaClient, digests: AbstractSet[str], window: tuple[datetime, datetime], ) -> Mapping[str, KeyMetadataDict] | None: - start, end = window rows: Final = await _db_or_empty( - lambda: prisma_client.db.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(digests), start, end), + lambda: _spend_log_rows_within_the_statement_timeout(prisma_client, digests, window), "Failed spend-log alias recovery for %d missing keys: %s", len(digests), ) @@ -247,10 +273,10 @@ async def _query_spend_log_metadata( return None return MappingProxyType( { - row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=owner) + row.digest: meta for row in _SPEND_LOG_DIGEST_ROWS.validate_python(rows) - for owner in (row.unanimous_owner(),) - if row.digest in digests and (row.key_alias or owner or row.team_id) + for meta in (row.metadata(),) + if row.digest in digests and any(meta.values()) } ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 682b11f5bb7..89be341c87b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,7 +1,7 @@ import asyncio import time from collections.abc import Sequence -from datetime import datetime +from datetime import datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -9,7 +9,11 @@ import pytest from prisma.errors import PrismaError from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL +from litellm.constants import ( + SPEND_LOG_KEY_METADATA_CACHE_TTL, + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, + SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, +) from litellm.proxy.spend_tracking.key_metadata_recovery import ( fill_missing_api_key_aliases, recover_double_hashed_key_metadata, @@ -31,6 +35,26 @@ def _query_raw_spend_logs(rows: Sequence[dict[str, str | None]]) -> AsyncMock: return AsyncMock(side_effect=query_raw) +def _spend_log_row(digest: str, key_alias: str | None, team_id: str | None, user_id: str | None) -> dict[str, str | None]: + return { + "digest": digest, + "first_alias": key_alias, + "last_alias": key_alias, + "first_team": team_id, + "last_team": team_id, + "first_owner": user_id, + "last_owner": user_id, + } + + +def _spend_log_transaction(mock_prisma: MagicMock, query_raw: AsyncMock) -> AsyncMock: + transaction = MagicMock() + transaction.execute_raw = AsyncMock(return_value=0) + transaction.query_raw = query_raw + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return query_raw + + def _query_raw_by_table( active_rows: Sequence[dict[str, str | None]], deleted_rows: Sequence[dict[str, str | None]], @@ -240,15 +264,18 @@ async def test_recover_key_metadata_from_spend_logs_resolves_session_token_from_ session_digest = hash_token("cli-session-repro-user-6852") window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs( - [_digest_row(session_digest, "cli-session-repro-user-6852", None, "repro-user-6852")] + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [_spend_log_row(session_digest, "cli-session-repro-user-6852", None, "repro-user-6852")] + ), ) result = await recover_key_metadata_from_spend_logs(mock_prisma, {session_digest}, window, cache=InMemoryCache()) assert result[session_digest]["key_alias"] == "cli-session-repro-user-6852" assert result[session_digest]["user_id"] == "repro-user-6852" - ((_, digests, start, end),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + ((_, digests, start, end),) = [call.args for call in query_raw.call_args_list] assert digests == [session_digest] assert (start, end) == window @@ -257,19 +284,19 @@ async def test_recover_key_metadata_from_spend_logs_resolves_session_token_from_ async def test_recover_key_metadata_from_spend_logs_skips_query_when_no_missing_keys(): window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = AsyncMock(return_value=[]) + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(return_value=[])) result = await recover_key_metadata_from_spend_logs(mock_prisma, set(), window, cache=InMemoryCache()) assert result == {} - mock_prisma.db.query_raw.assert_not_called() + query_raw.assert_not_called() @pytest.mark.asyncio async def test_recover_key_metadata_from_spend_logs_returns_empty_on_prisma_error(): window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down")) + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("db down"))) result = await recover_key_metadata_from_spend_logs( mock_prisma, {hash_token("cli-session-x")}, window, cache=InMemoryCache() @@ -285,12 +312,15 @@ async def test_recover_key_metadata_from_spend_logs_ignores_foreign_and_all_null foreign = hash_token("cli-session-foreign") window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs( - [ - _digest_row(wanted, "kept-alias", None, "owner-1"), - _digest_row(all_null, None, None, None), - _digest_row(foreign, "foreign-alias", None, "owner-2"), - ] + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [ + _spend_log_row(wanted, "kept-alias", None, "owner-1"), + _spend_log_row(all_null, None, None, None), + _spend_log_row(foreign, "foreign-alias", None, "owner-2"), + ] + ), ) result = await recover_key_metadata_from_spend_logs(mock_prisma, {wanted, all_null}, window, cache=InMemoryCache()) @@ -304,14 +334,14 @@ async def test_recover_key_metadata_from_spend_logs_ignores_foreign_and_all_null async def test_recover_key_metadata_from_spend_logs_skips_non_sha256_keys(): window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = AsyncMock(return_value=[]) + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(return_value=[])) result = await recover_key_metadata_from_spend_logs( mock_prisma, {"cli-session-raw-1798", "key-hash-short"}, window, cache=InMemoryCache() ) assert result == {} - mock_prisma.db.query_raw.assert_not_called() + query_raw.assert_not_called() @pytest.mark.asyncio @@ -319,13 +349,13 @@ async def test_recover_key_metadata_from_spend_logs_accepts_hashed_jwt_digests() jwt_digest = f"hashed-jwt-{hash_token('jwt-subject-1')}" window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(jwt_digest, None, "team-jwt", "jwt-user")]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(jwt_digest, None, "team-jwt", "jwt-user")])) result = await recover_key_metadata_from_spend_logs(mock_prisma, {jwt_digest}, window, cache=InMemoryCache()) assert result[jwt_digest]["team_id"] == "team-jwt" assert result[jwt_digest]["user_id"] == "jwt-user" - ((_, digests, _, _),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + ((_, digests, _, _),) = [call.args for call in query_raw.call_args_list] assert digests == [jwt_digest] @@ -336,7 +366,7 @@ async def test_recover_key_metadata_from_spend_logs_serves_repeat_lookups_from_t window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) cache = InMemoryCache() mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(found, "found-alias", None, "owner-1")]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(found, "found-alias", None, "owner-1")])) first = await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) second = await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) @@ -344,7 +374,7 @@ async def test_recover_key_metadata_from_spend_logs_serves_repeat_lookups_from_t assert first == second assert set(first) == {found} assert first[found]["key_alias"] == "found-alias" - assert mock_prisma.db.query_raw.await_count == 1 + assert query_raw.await_count == 1 @pytest.mark.asyncio @@ -354,15 +384,15 @@ async def test_recover_key_metadata_from_spend_logs_only_queries_digests_the_cac window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) cache = InMemoryCache() mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(cached_digest, "cached-alias", None, None)]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(cached_digest, "cached-alias", None, None)])) await recover_key_metadata_from_spend_logs(mock_prisma, {cached_digest}, window, cache=cache) - mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(new_digest, "new-alias", None, None)]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(new_digest, "new-alias", None, None)])) result = await recover_key_metadata_from_spend_logs(mock_prisma, {cached_digest, new_digest}, window, cache=cache) assert result[cached_digest]["key_alias"] == "cached-alias" assert result[new_digest]["key_alias"] == "new-alias" - ((_, digests, _, _),) = [call.args for call in mock_prisma.db.query_raw.call_args_list] + ((_, digests, _, _),) = [call.args for call in query_raw.call_args_list] assert digests == [new_digest] @@ -371,18 +401,18 @@ async def test_recover_key_metadata_from_spend_logs_rescans_when_the_window_chan digest = hash_token("cli-session-windowed") cache = InMemoryCache() mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs([]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([])) await recover_key_metadata_from_spend_logs( mock_prisma, {digest}, (datetime(2026, 9, 1), datetime(2026, 9, 4)), cache=cache ) - mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(digest, "later-alias", None, None)]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, "later-alias", None, None)])) result = await recover_key_metadata_from_spend_logs( mock_prisma, {digest}, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=cache ) assert result[digest]["key_alias"] == "later-alias" - assert mock_prisma.db.query_raw.await_count == 1 + assert query_raw.await_count == 1 @pytest.mark.asyncio @@ -391,13 +421,13 @@ async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_ window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) mock_prisma = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("statement timeout")) + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("statement timeout"))) started = time.time() assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} - mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(digest, "back-online", None, None)]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, "back-online", None, None)])) assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {} - mock_prisma.db.query_raw.assert_not_awaited() + query_raw.assert_not_awaited() miss_key = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before")) assert cache.ttl_dict[miss_key] - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1 cache.ttl_dict[miss_key] = time.time() - 1 @@ -412,14 +442,11 @@ async def test_recover_key_metadata_from_spend_logs_drops_the_owner_of_a_digest_ shared_ui_digest = hash_token("ui-token") window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs( - [ - { - **_digest_row(shared_ui_digest, "ui-token", "litellm-dashboard", "bob"), - "first_owner": "alice", - "last_owner": "bob", - } - ] + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [{**_spend_log_row(shared_ui_digest, "ui-token", "litellm-dashboard", None), "first_owner": "alice", "last_owner": "bob"}] + ), ) result = await recover_key_metadata_from_spend_logs( @@ -434,8 +461,11 @@ async def test_recover_key_metadata_from_spend_logs_keeps_the_owner_when_every_n digest = hash_token("cli-session-one-owner") window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs( - [{**_digest_row(digest, None, None, "carol"), "first_owner": "carol", "last_owner": "carol"}] + query_raw = _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [_spend_log_row(digest, None, None, "carol")] + ), ) result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) @@ -450,7 +480,7 @@ async def test_recover_key_metadata_from_spend_logs_forgets_a_miss_long_before_a window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs([_digest_row(found, "found-alias", None, None)]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(found, "found-alias", None, None)])) started = time.time() await recover_key_metadata_from_spend_logs(mock_prisma, {found, unknown}, window, cache=cache) @@ -471,9 +501,9 @@ async def test_recover_key_metadata_from_spend_logs_runs_one_query_for_concurren async def slow_query_raw(sql: str, *params: object) -> list[dict[str, str | None]]: await asyncio.sleep(0.01) - return [_digest_row(digest, "shared-alias", None, None)] + return [_spend_log_row(digest, "shared-alias", None, None)] - mock_prisma.db.query_raw = AsyncMock(side_effect=slow_query_raw) + query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=slow_query_raw)) results = await asyncio.gather( *( @@ -483,7 +513,7 @@ async def test_recover_key_metadata_from_spend_logs_runs_one_query_for_concurren ) assert all(result[digest]["key_alias"] == "shared-alias" for result in results) - assert mock_prisma.db.query_raw.await_count == 1 + assert query_raw.await_count == 1 @pytest.mark.asyncio @@ -492,7 +522,7 @@ async def test_recover_key_metadata_from_spend_logs_keeps_a_repeated_miss_as_lon window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL) mock_prisma = MagicMock() - mock_prisma.db.query_raw = _query_raw_spend_logs([]) + query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([])) await recover_key_metadata_from_spend_logs(mock_prisma, {unknown}, window, cache=cache) first_miss_key = next(key for key in cache.ttl_dict if unknown in key and not key.endswith(":missed-before")) cache.ttl_dict[first_miss_key] = time.time() - 1 @@ -500,5 +530,61 @@ async def test_recover_key_metadata_from_spend_logs_keeps_a_repeated_miss_as_lon await recover_key_metadata_from_spend_logs(mock_prisma, {unknown}, window, cache=cache) - assert mock_prisma.db.query_raw.await_count == 2 + assert query_raw.await_count == 2 assert cache.ttl_dict[first_miss_key] - started >= SPEND_LOG_KEY_METADATA_CACHE_TTL - 1 + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_keeps_the_owner_older_rows_agree_on_when_the_newest_is_nameless(): + digest = hash_token("cli-session-owner-from-older-rows") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, None, "team-x", "alice")])) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert result[digest] == {"key_alias": None, "team_id": "team-x", "user_id": "alice"} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_names_nothing_for_a_field_whose_rows_disagree(): + digest = hash_token("cli-session-disagreeing-rows") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + _spend_log_transaction( + mock_prisma, + _query_raw_spend_logs( + [ + { + **_spend_log_row(digest, None, None, "carol"), + "first_alias": "old-alias", + "last_alias": "renamed-alias", + "first_team": "team-a", + "last_team": "team-b", + } + ] + ), + ) + + result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert result[digest] == {"key_alias": None, "team_id": None, "user_id": "carol"} + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_bounds_the_scan_with_a_statement_timeout(): + digest = hash_token("cli-session-bounded-scan") + window = (datetime(2026, 9, 7), datetime(2026, 9, 10)) + mock_prisma = MagicMock() + calls: list[str] = [] + transaction = MagicMock() + transaction.execute_raw = AsyncMock(side_effect=lambda sql: calls.append(sql) or 0) + transaction.query_raw = AsyncMock(side_effect=lambda sql, *args: calls.append("scan") or []) + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + + await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) + + assert calls == [f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}", "scan"] + assert mock_prisma.db.tx.call_args.kwargs["timeout"] == timedelta( + milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS + ) From 529b8706ee4a35ca479301242cb07f74e45d6775 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:51:13 -0700 Subject: [PATCH 9/9] test(spend-tracking): mock the spend-log scan through the transaction its caller now opens --- .../test_common_daily_activity.py | 54 +++++++++++-------- 1 file changed, 33 insertions(+), 21 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index d214ab6b5ac..6cd900cb041 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -518,6 +518,26 @@ async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_s assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups) +def _spend_log_transaction(mock_prisma: MagicMock, rows: list[dict[str, str | None]]) -> AsyncMock: + transaction = MagicMock() + transaction.execute_raw = AsyncMock(return_value=0) + transaction.query_raw = AsyncMock(return_value=rows) + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return transaction.query_raw + + +def _spend_log_row(digest: str, key_alias: str, user_id: str) -> dict[str, str | None]: + return { + "digest": digest, + "first_alias": key_alias, + "last_alias": key_alias, + "first_team": None, + "last_team": None, + "first_owner": user_id, + "last_owner": user_id, + } + + @pytest.mark.asyncio async def test_get_api_key_metadata_permanent_miss_with_a_window_reads_spend_logs_once_within_it(): from litellm.proxy.utils import hash_token @@ -529,14 +549,13 @@ async def test_get_api_key_metadata_permanent_miss_with_a_window_reads_spend_log mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma.db.query_raw = AsyncMock(return_value=[]) + spend_log_query_raw = _spend_log_transaction(mock_prisma, []) result = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={double_hashed}, spend_logs_window=window) assert double_hashed not in result - assert mock_prisma.db.query_raw.await_count == 3 - ((_, digests, start, end),) = [ - call.args for call in mock_prisma.db.query_raw.call_args_list if "LiteLLM_SpendLogs" in call.args[0] - ] + assert mock_prisma.db.query_raw.await_count == 2 + ((_, digests, start, end),) = [call.args for call in spend_log_query_raw.call_args_list] assert digests == [double_hashed] assert (start, end) == window @@ -559,12 +578,10 @@ async def test_get_daily_activity_recovers_a_session_key_alias_from_spend_logs_a return_value=[SimpleNamespace(user_id="session-user", user_email="session@example.com")] ) - async def query_raw(sql, *params): - if "LiteLLM_SpendLogs" in sql: - return [{"digest": session_digest, "key_alias": "cli-session-alias", "team_id": None, "user_id": "session-user"}] - return [] - - mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + spend_log_query_raw = _spend_log_transaction( + mock_prisma, [_spend_log_row(session_digest, "cli-session-alias", "session-user")] + ) result = await get_daily_activity( prisma_client=mock_prisma, @@ -583,9 +600,7 @@ async def test_get_daily_activity_recovers_a_session_key_alias_from_spend_logs_a key_metadata = result.results[0].breakdown.api_keys[session_digest].metadata assert key_metadata.key_alias == "cli-session-alias" assert key_metadata.user_email == "session@example.com" - ((_, digests, start, end),) = [ - call.args for call in mock_prisma.db.query_raw.call_args_list if "LiteLLM_SpendLogs" in call.args[0] - ] + ((_, digests, start, end),) = [call.args for call in spend_log_query_raw.call_args_list] assert digests == [session_digest] assert (start, end) == (datetime(2023, 12, 31), datetime(2024, 1, 3)) @@ -2191,12 +2206,10 @@ async def test_get_api_key_metadata_resolves_session_key_via_spend_log_window(): return_value=[SimpleNamespace(user_id="user-42", user_email="user42@example.com")] ) - async def query_raw(sql, *params): - if "LiteLLM_SpendLogs" in sql: - return [{"digest": session_digest, "key_alias": "cli-session-user-42", "team_id": None, "user_id": "user-42"}] - return [] - - mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + spend_log_query_raw = _spend_log_transaction( + mock_prisma, [_spend_log_row(session_digest, "cli-session-user-42", "user-42")] + ) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -2207,8 +2220,7 @@ async def test_get_api_key_metadata_resolves_session_key_via_spend_log_window(): assert result[session_digest]["key_alias"] == "cli-session-user-42" assert result[session_digest]["user_id"] == "user-42" assert result[session_digest]["user_email"] == "user42@example.com" - spend_log_calls = [call.args for call in mock_prisma.db.query_raw.call_args_list if "LiteLLM_SpendLogs" in call.args[0]] - ((_, digests, start, end),) = spend_log_calls + ((_, digests, start, end),) = [call.args for call in spend_log_query_raw.call_args_list] assert digests == [session_digest] assert (start, end) == (datetime(2026, 9, 7), datetime(2026, 9, 10))