diff --git a/litellm/constants.py b/litellm/constants.py index 10c943656f7..0ed87b189a2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1904,6 +1904,7 @@ 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 +SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE: Final = 100 # 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 1686e055524..6c69572b12c 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -17,6 +17,7 @@ from litellm.constants import ( SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, ) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient @@ -48,10 +49,15 @@ def _named_spend_log_edge_row_sql(direction: Literal["ASC", "DESC"]) -> str: 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 = keys.digest - AND "startTime" >= $2::timestamp - AND "startTime" < $3::timestamp + FROM ( + SELECT "startTime", metadata, team_id, "user" + FROM "LiteLLM_SpendLogs" + WHERE api_key = keys.digest + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp + ORDER BY "startTime" {direction} + LIMIT {SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE} + ) edge ) named WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL ORDER BY "startTime" {direction} @@ -68,8 +74,8 @@ SELECT keys.digest, first_row.user_id AS first_owner, last_row.user_id AS last_owner FROM unnest($1::text[]) AS keys(digest) -CROSS JOIN LATERAL ({_named_spend_log_edge_row_sql("ASC")}) first_row -CROSS JOIN LATERAL ({_named_spend_log_edge_row_sql("DESC")}) last_row +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("ASC")}) first_row ON true +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("DESC")}) last_row ON true """ _SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" 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 bf09649a89f..b9b5b83df0b 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 @@ -21,6 +21,7 @@ from litellm.constants import ( SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, ) from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_details, @@ -741,6 +742,49 @@ async def test_recover_key_metadata_from_spend_logs_reads_two_rows_per_key_howev assert rows_read is not None and rows_read[0] <= 2 * len(owners) +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_walks_a_bounded_number_of_nameless_rows_per_key( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + nameless_rows: Final = 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + named_late: Final[Mapping[str, str]] = {hash_token(f"cli-session-late-{i}"): f"user-{i}" for i in range(3)} + never_named: Final = frozenset(hash_token(f"cli-session-never-{i}") for i in range(3)) + for digest in (*named_late, *never_named): + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime") + SELECT %(digest)s || '-' || g, %(digest)s, %(start)s + g * interval '1 minute' + FROM generate_series(1, %(rows)s) g + """, + {"digest": digest, "start": datetime(2026, 9, 7), "rows": nameless_rows}, + ) + for digest, owner in named_late.items(): + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", metadata) + VALUES (%(digest)s || '-newest', %(digest)s, %(logged_at)s, %(owner)s, + jsonb_build_object('user_api_key_alias', 'cli-session-' || %(owner)s)) + """, + {"digest": digest, "owner": owner, "logged_at": datetime(2026, 9, 9)}, + ) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + frozenset(named_late) | never_named, + (datetime(2026, 9, 7), datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert {digest: meta.get("user_id") for digest, meta in result.items()} == named_late + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None + assert rows_read[0] <= 2 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE * (len(named_late) + len(never_named)) + + @pytest.mark.asyncio async def test_recover_cli_session_key_metadata_names_the_owner_only_when_the_suffix_is_a_real_user(): mock_prisma = MagicMock()