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"