From dbeb4058e1e782f3e77f1bdf7758db9b8bf390ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 28 Sep 2026 19:01:12 -0700 Subject: [PATCH 1/4] fix(proxy): look up hashed key names with two spend log rows per key The spend-log fallback for keys missing from the key table read every row per key to check that all named rows agreed, which passed the 5s statement timeout on busy keys even with the (api_key, startTime) index. Probe only the oldest and newest named row per key, so the lookup stays two index reads per key however much the key logged. --- .../spend_tracking/key_metadata_recovery.py | 57 ++++--- .../test_key_metadata_recovery.py | 151 +++++++++++++++++- 2 files changed, 185 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index ce96dc627801..1686e055524f 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from datetime import datetime, timedelta from types import MappingProxyType -from typing import Final, TypeVar +from typing import Final, Literal, TypeVar from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -39,26 +39,37 @@ ORDER BY token, deleted_at DESC """ -_SPEND_LOG_ALIAS_SQL: Final = """ -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, - 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 -GROUP BY api_key + +def _named_spend_log_edge_row_sql(direction: Literal["ASC", "DESC"]) -> str: + return f""" + SELECT key_alias, team_id, user_id + FROM ( + SELECT "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 = keys.digest + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp + ) named + WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL + ORDER BY "startTime" {direction} + LIMIT 1 + """ + + +_SPEND_LOG_ALIAS_SQL: Final = f""" +SELECT keys.digest, + first_row.key_alias AS first_alias, + last_row.key_alias AS last_alias, + first_row.team_id AS first_team, + last_row.team_id AS last_team, + 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 """ _SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" @@ -84,7 +95,9 @@ class _TokenDigestRow(BaseModel): def _unanimous(first: str | None, last: str | None) -> str | None: - return first if first == last else None + if first is None: + return last + return first if last is None or first == last else None class _SpendLogDigestRow(BaseModel): 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 acd03964bf34..bf09649a89f2 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,12 +1,20 @@ import asyncio +import re import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta +from pathlib import Path from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock +import litellm_proxy_extras +import psycopg import pytest from prisma.errors import PrismaError +from psycopg.rows import dict_row +from psycopg.types.json import Jsonb +from pytest_postgresql import factories from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( @@ -592,6 +600,147 @@ async def test_recover_key_metadata_from_spend_logs_bounds_the_scan_with_a_state ) +_spend_logs_postgresql_proc: Final = factories.postgresql_proc() +_spend_logs_postgresql: Final = factories.postgresql("_spend_logs_postgresql_proc") + +_SPEND_LOGS_DDL: Final = """ + CREATE TABLE "LiteLLM_SpendLogs" ( + request_id TEXT PRIMARY KEY, + api_key TEXT NOT NULL DEFAULT '', + "startTime" TIMESTAMP(3) NOT NULL, + "user" TEXT DEFAULT '', + team_id TEXT, + metadata JSONB DEFAULT '{}' + ) +""" + +_API_KEY_START_TIME_INDEX_MIGRATION: Final = ( + Path(litellm_proxy_extras.__file__).parent + / "migrations" + / "20260823000000_add_spend_logs_api_key_starttime_index" + / "migration.sql" +) + +_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL: Final = """ + SELECT COALESCE(seq_tup_read, 0) + COALESCE(idx_tup_fetch, 0) AS rows_read + FROM pg_stat_xact_user_tables + WHERE relname = 'LiteLLM_SpendLogs' +""" + + +def _create_spend_logs_table(conn: psycopg.Connection) -> None: + conn.execute(_SPEND_LOGS_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.execute(_API_KEY_START_TIME_INDEX_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # migration file + + +def _psycopg_prisma(conn: psycopg.Connection) -> MagicMock: + async def query_raw(sql: str, *params: object) -> list[dict[str, object]]: + with conn.cursor(row_factory=dict_row) as cur: + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", sql), # pyright: ignore[reportArgumentType] # proxy SQL is not a literal + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + return cur.fetchall() + + async def execute_raw(sql: str) -> int: + conn.execute(sql) # pyright: ignore[reportArgumentType] # proxy SQL is not a literal + return 0 + + mock_prisma: Final = MagicMock() + transaction: Final = MagicMock() + transaction.query_raw = AsyncMock(side_effect=query_raw) + transaction.execute_raw = AsyncMock(side_effect=execute_raw) + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return mock_prisma + + +def _named_spend_log( + digest: str, logged_at: datetime, alias: str | None, user: str | None, team: str | None = None +) -> tuple[str, str, datetime, str, str | None, Jsonb]: + return ( + f"{digest}-{logged_at.isoformat()}", + digest, + logged_at, + user or "", + team, + Jsonb({"user_api_key_alias": alias} if alias else {}), + ) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_names_a_key_by_its_oldest_and_newest_named_rows_in_the_window( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + unnamed_edges, owner_logged_late, reowned, outside_window, never_named = ( + hash_token(f"cli-session-{name}") for name in ("edges", "late", "reowned", "window", "never") + ) + with conn.cursor() as cur: + cur.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", team_id, metadata)' + " VALUES (%s, %s, %s, %s, %s, %s)", + ( + _named_spend_log(unnamed_edges, datetime(2026, 9, 7, 1), None, None), + _named_spend_log(unnamed_edges, datetime(2026, 9, 8), "cli-a", "alice", "team-a"), + _named_spend_log(unnamed_edges, datetime(2026, 9, 9), "cli-a", "alice", "team-a"), + _named_spend_log(unnamed_edges, datetime(2026, 9, 9, 23), None, None), + _named_spend_log(owner_logged_late, datetime(2026, 9, 7, 1), "cli-b", None), + _named_spend_log(owner_logged_late, datetime(2026, 9, 9), "cli-b", "bob"), + _named_spend_log(reowned, datetime(2026, 9, 7, 1), "cli-c", "carol"), + _named_spend_log(reowned, datetime(2026, 9, 9), "cli-c", "dave"), + _named_spend_log(outside_window, datetime(2026, 9, 6), "stale-alias", "erin"), + _named_spend_log(outside_window, datetime(2026, 9, 8), "cli-d", "erin"), + _named_spend_log(outside_window, datetime(2026, 9, 10), "later-alias", "erin"), + _named_spend_log(never_named, datetime(2026, 9, 8), None, None), + ), + ) + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + {unnamed_edges, owner_logged_late, reowned, outside_window, never_named}, + (datetime(2026, 9, 7), datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert dict(result) == { + unnamed_edges: {"key_alias": "cli-a", "team_id": "team-a", "user_id": "alice"}, + owner_logged_late: {"key_alias": "cli-b", "team_id": None, "user_id": "bob"}, + reowned: {"key_alias": "cli-c", "team_id": None, "user_id": None}, + outside_window: {"key_alias": "cli-d", "team_id": None, "user_id": "erin"}, + } + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_reads_two_rows_per_key_however_many_the_key_logged( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + owners: Final[Mapping[str, str]] = {hash_token(f"cli-session-busy-{i}"): f"user-{i}" for i in range(5)} + for digest, owner in owners.items(): + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", metadata) + SELECT %(digest)s || '-' || g, %(digest)s, %(start)s + g * interval '1 minute', %(owner)s, + jsonb_build_object('user_api_key_alias', 'cli-session-' || %(owner)s) + FROM generate_series(1, 2000) g + """, + {"digest": digest, "owner": owner, "start": datetime(2026, 9, 7)}, + ) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), frozenset(owners), (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert {digest: meta.get("user_id") for digest, meta in result.items()} == owners + 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 and rows_read[0] <= 2 * len(owners) + + @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() From bea6f49c9c5331555dbafffc61fd8c0a3453f49b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 29 Sep 2026 09:45:11 -0700 Subject: [PATCH 2/4] fix(proxy): cap each spend log name probe at 100 rows per key --- litellm/constants.py | 1 + .../spend_tracking/key_metadata_recovery.py | 18 +++++--- .../test_key_metadata_recovery.py | 44 +++++++++++++++++++ 3 files changed, 57 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 10c943656f7d..0ed87b189a2c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1904,6 +1904,7 @@ 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 1686e055524f..6c69572b12cf 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -17,6 +17,7 @@ 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 @@ def _named_spend_log_edge_row_sql(direction: Literal["ASC", "DESC"]) -> str: 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 bf09649a89f2..b9b5b83df0b6 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 @@ 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() From f2d7062482b35e4dc3385e48d17dcd9fc81f6d79 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 29 Sep 2026 14:44:01 -0700 Subject: [PATCH 3/4] fix(proxy): bound the newest-row probe at where the oldest probe stopped The newest-row probe now starts at the row where the oldest-row probe gave up, so a key with under 200 rows in the window is read once instead of twice, and the lookup transaction turns bitmap scans off so the planner walks the (api_key, startTime) index instead of every row of a busy key when statistics or the visibility map are stale. --- .../spend_tracking/key_metadata_recovery.py | 28 +++- .../test_key_metadata_recovery.py | 140 ++++++++++++++++-- 2 files changed, 152 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 6c69572b12cf..b1b2f7d732a2 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -41,9 +41,11 @@ """ -def _named_spend_log_edge_row_sql(direction: Literal["ASC", "DESC"]) -> str: +def _named_spend_log_edge_row_sql( + direction: Literal["ASC", "DESC"], since: Literal["$2::timestamp", "oldest_probe.stopped_at"] +) -> str: return f""" - SELECT key_alias, team_id, user_id + SELECT "startTime", key_alias, team_id, user_id FROM ( SELECT "startTime", NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, @@ -53,7 +55,7 @@ def _named_spend_log_edge_row_sql(direction: Literal["ASC", "DESC"]) -> str: SELECT "startTime", metadata, team_id, "user" FROM "LiteLLM_SpendLogs" WHERE api_key = keys.digest - AND "startTime" >= $2::timestamp + AND "startTime" >= {since} AND "startTime" < $3::timestamp ORDER BY "startTime" {direction} LIMIT {SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE} @@ -65,6 +67,19 @@ def _named_spend_log_edge_row_sql(direction: Literal["ASC", "DESC"]) -> str: """ +_OLDEST_PROBE_STOPPED_AT_SQL: Final = f""" + SELECT COALESCE(first_row."startTime", ( + SELECT "startTime" + FROM "LiteLLM_SpendLogs" + WHERE api_key = keys.digest + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp + ORDER BY "startTime" ASC + OFFSET {SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE - 1} + LIMIT 1 + )) AS stopped_at +""" + _SPEND_LOG_ALIAS_SQL: Final = f""" SELECT keys.digest, first_row.key_alias AS first_alias, @@ -74,11 +89,13 @@ def _named_spend_log_edge_row_sql(direction: Literal["ASC", "DESC"]) -> str: first_row.user_id AS first_owner, last_row.user_id AS last_owner FROM unnest($1::text[]) AS keys(digest) -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 +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("ASC", "$2::timestamp")}) first_row ON true +LEFT JOIN LATERAL ({_OLDEST_PROBE_STOPPED_AT_SQL}) oldest_probe ON true +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("DESC", "oldest_probe.stopped_at")}) last_row ON true """ _SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" +_SPEND_LOG_NO_BITMAP_SCAN_SQL: Final = "SET LOCAL enable_bitmapscan = off" _SPEND_LOG_TRANSACTION_TIMEOUT: Final = timedelta(milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS) _HASHED_JWT_PREFIX: Final = "hashed-jwt-" @@ -336,6 +353,7 @@ async def _spend_log_rows_within_the_statement_timeout( 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) + await transaction.execute_raw(_SPEND_LOG_NO_BITMAP_SCAN_SQL) return await transaction.query_raw(_SPEND_LOG_ALIAS_SQL, sorted(digests), start, end) 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 b9b5b83df0b6..845a3ed86c4b 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 @@ -595,7 +595,11 @@ async def test_recover_key_metadata_from_spend_logs_bounds_the_scan_with_a_state 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 calls == [ + f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}", + "SET LOCAL enable_bitmapscan = off", + "scan", + ] assert mock_prisma.db.tx.call_args.kwargs["timeout"] == timedelta( milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS ) @@ -655,6 +659,24 @@ async def execute_raw(sql: str) -> int: return mock_prisma +def _commit_and_vacuum(conn: psycopg.Connection) -> None: + conn.commit() + conn.set_autocommit(True) + conn.execute('VACUUM (ANALYZE) "LiteLLM_SpendLogs"') + conn.set_autocommit(False) + + +def _insert_nameless_spend_logs(conn: psycopg.Connection, digest: str, rows: int) -> None: + 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": rows}, + ) + + def _named_spend_log( digest: str, logged_at: datetime, alias: str | None, user: str | None, team: str | None = None ) -> tuple[str, str, datetime, str, str | None, Jsonb]: @@ -748,18 +770,10 @@ async def test_recover_key_metadata_from_spend_logs_walks_a_bounded_number_of_na ): 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}, - ) + _insert_nameless_spend_logs(conn, digest, 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) for digest, owner in named_late.items(): conn.execute( """ @@ -782,7 +796,111 @@ async def test_recover_key_metadata_from_spend_logs_walks_a_bounded_number_of_na 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)) + assert rows_read[0] <= 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE * (len(named_late) + len(never_named)) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_reads_a_short_nameless_key_once( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + rows_per_key: Final = SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE // 2 + never_named: Final = frozenset(hash_token(f"cli-session-short-{i}") for i in range(20)) + for digest in never_named: + _insert_nameless_spend_logs(conn, digest, rows_per_key) + _commit_and_vacuum(conn) + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), never_named, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert dict(result) == {} + 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] <= rows_per_key * len(never_named) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_bounds_a_busy_nameless_key_among_short_keys_before_any_vacuum( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + busy: Final = frozenset(hash_token(f"cli-session-busy-nameless-{i}") for i in range(3)) + for short_key in range(200): + _insert_nameless_spend_logs( + conn, hash_token(f"cli-session-short-{short_key}"), SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE // 5 + ) + for digest in busy: + _insert_nameless_spend_logs(conn, digest, 30 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), busy, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert dict(result) == {} + 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] <= 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE * len(busy) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_finds_a_name_logged_where_the_oldest_probe_stopped( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + start: Final = datetime(2026, 9, 7) + past_the_stop, tied_with_the_stop = (hash_token(f"cli-session-{name}") for name in ("past", "tied")) + same_millisecond: Final = tuple( + start + timedelta(minutes=SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, microseconds=n) for n in (100, 200, 300) + ) + with conn.cursor() as cur: + cur.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", team_id, metadata)' + " VALUES (%s, %s, %s, %s, %s, %s)", + ( + *( + _named_spend_log(past_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range(1, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 20) + ), + _named_spend_log( + past_the_stop, + start + timedelta(minutes=SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 20), + "cli-p", + "pat", + ), + *( + _named_spend_log(past_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range( + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 21, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 51 + ) + ), + *( + _named_spend_log(tied_with_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range(1, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + ), + _named_spend_log(tied_with_the_stop, same_millisecond[0], None, None), + _named_spend_log(tied_with_the_stop, same_millisecond[1], None, None), + _named_spend_log(tied_with_the_stop, same_millisecond[2], "cli-t", "tess"), + ), + ) + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + {past_the_stop, tied_with_the_stop}, + (start, datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert dict(result) == { + past_the_stop: {"key_alias": "cli-p", "team_id": None, "user_id": "pat"}, + tied_with_the_stop: {"key_alias": "cli-t", "team_id": None, "user_id": "tess"}, + } @pytest.mark.asyncio From 54da4e523bdc7fb66d1d11cb0ecb856414fa0b68 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:05:50 -0700 Subject: [PATCH 4/4] test(integration): add spend log alias probe cells for the daily activity routes --- tests/integration/_support/daily_activity.py | 65 +++ .../test_daily_activity_key_alias_probes.py | 490 ++++++++++++++++++ .../test_daily_activity_key_owner_traffic.py | 48 +- 3 files changed, 602 insertions(+), 1 deletion(-) create mode 100644 tests/integration/spend/test_daily_activity_key_alias_probes.py diff --git a/tests/integration/_support/daily_activity.py b/tests/integration/_support/daily_activity.py index debb8c4cdb43..346fc156ea9b 100644 --- a/tests/integration/_support/daily_activity.py +++ b/tests/integration/_support/daily_activity.py @@ -3,6 +3,8 @@ from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass +from datetime import datetime, timedelta +from hashlib import sha256 from itertools import chain from typing import Final @@ -34,7 +36,16 @@ " VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s)" ) DELETE_SPEND_LOG: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +INSERT_SPEND_LOG_ROW: Final = ( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, "startTime", "endTime", metadata, team_id, "user")' + " VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s, %s, %s)" +) +DELETE_SPEND_LOG_ROWS: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)' +DELETE_KEY_ROW: Final = 'DELETE FROM "LiteLLM_VerificationToken" WHERE token = %s' +DELETE_ARCHIVED_KEY_ROW: Final = 'DELETE FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s' LOCK_TABLE: Final = sql.SQL("LOCK TABLE {table} IN ACCESS EXCLUSIVE MODE") +SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs" +FIRST_SPEND_LOG_AT: Final = datetime(2026, 2, 3, 12, 0, 0) @dataclass(frozen=True, slots=True) @@ -67,6 +78,10 @@ def key_no_key_table_holds() -> str: return f"integration-ownerless-{uuid.uuid4().hex}" +def digest_no_key_table_holds() -> str: + return sha256(uuid.uuid4().bytes).hexdigest() + + def activity_of_key( gateway: Gateway, path: str, api_key: str, *, reader: str | None = None, **filters: str ) -> httpx.Response: @@ -148,6 +163,56 @@ def spend_log_naming_only_an_alias(request_id: str, api_key: str, started: str, connection.execute(DELETE_SPEND_LOG, (request_id,)) +@dataclass(frozen=True, slots=True) +class SpendLogRow: + started: str + metadata: JsonValue = None + team_id: str | None = None + user: str | None = None + + +def started_at(index: int) -> str: + return (FIRST_SPEND_LOG_AT + timedelta(seconds=index)).strftime("%Y-%m-%d %H:%M:%S") + + +def nameless_rows(count: int, first_index: int = 0) -> tuple[SpendLogRow, ...]: + return tuple(SpendLogRow(started_at(first_index + offset), {}) for offset in range(count)) + + +def named_row(index: int, alias: str) -> SpendLogRow: + return SpendLogRow(started_at(index), {"user_api_key_alias": alias}) + + +@contextmanager +def spend_logs_of_key( + api_key: str, rows: Sequence[SpendLogRow], *, database_url: str | None = None +) -> Iterator[tuple[str, ...]]: + request_ids: Final = tuple(f"integration-{uuid.uuid4().hex}" for _ in rows) + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.cursor().executemany( + INSERT_SPEND_LOG_ROW, + tuple( + (request_id, api_key, row.started, row.started, Jsonb(row.metadata), row.team_id, row.user) + for request_id, row in zip(request_ids, rows, strict=True) + ), + ) + try: + yield request_ids + finally: + delete_spend_logs(request_ids, database_url=database_url) + + +def delete_spend_logs(request_ids: Sequence[str], *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.execute(DELETE_SPEND_LOG_ROWS, (list(request_ids),)) + + +def purge_key_from_the_key_tables(digest: str, *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.execute(DELETE_KEY_ROW, (digest,)) + connection.execute(DELETE_ARCHIVED_KEY_ROW, (digest,)) + + @contextmanager def locked_table(table: str, *, database_url: str | None = None) -> Iterator[None]: with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: diff --git a/tests/integration/spend/test_daily_activity_key_alias_probes.py b/tests/integration/spend/test_daily_activity_key_alias_probes.py new file mode 100644 index 000000000000..8d9b86164358 --- /dev/null +++ b/tests/integration/spend/test_daily_activity_key_alias_probes.py @@ -0,0 +1,490 @@ +import time +import uuid +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.daily_activity import ( + AGGREGATED_USER_ACTIVITY, + DAY, + ROUTES, + SPEND_LOGS_TABLE, + USER_SPEND, + Route, + SpendLogRow, + activity_of_key, + assert_key_reported, + daily_rows, + digest_no_key_table_holds, + key_metadata, + locked_table, + named_row, + nameless_rows, + records_of_key, + seeded_metrics, + seeded_row, + spend_logs_of_key, + started_at, + user_row, + user_with_an_email, +) +from integration._support.database import read_rows, scratch_database +from integration._support.process import OwnedProxy, owned_proxy_process +from pydantic import JsonValue + +DAY_OUTSIDE_THE_WINDOW: Final = "2026-02-10" +GIVES_UP_WITHIN_SECONDS: Final = 10 +CONCURRENT_READS: Final = 20 +CACHED_MISS_CLEARS_WITHIN_SECONDS: Final = 45 +ALIAS_OF_ONE_SPEND_LOG: Final = ( + "SELECT metadata->>'user_api_key_alias' AS alias FROM \"LiteLLM_SpendLogs\" WHERE request_id = %s" +) + + +def _alias() -> str: + return f"integration-alias-{uuid.uuid4().hex}" + + +def _named_between_fifty_and_fifty(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(50), named_row(50, alias), *nameless_rows(50, 51)) + + +def _oldest_named(alias: str) -> tuple[SpendLogRow, ...]: + return (named_row(0, alias), *nameless_rows(150, 1)) + + +def _newest_named(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(150), named_row(150, alias)) + + +def _both_edges_named(alias: str) -> tuple[SpendLogRow, ...]: + return (named_row(0, alias), *nameless_rows(150, 1), named_row(151, alias)) + + +def _named_after_one_hundred(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(100), named_row(100, alias), *nameless_rows(99, 101)) + + +def _named_after_ninety_nine(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(99), named_row(99, alias), *nameless_rows(100, 100)) + + +def _named_only_in_the_middle(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(100), named_row(100, alias), *nameless_rows(100, 101)) + + +def _renamed_and_renamed_back(alias: str, other: str) -> tuple[SpendLogRow, ...]: + return ( + named_row(0, alias), + *nameless_rows(100, 1), + named_row(101, other), + *nameless_rows(100, 102), + named_row(202, alias), + ) + + +def _team_in_the_column(team: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {}, team_id=team) + + +def _team_in_the_metadata(team: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {"user_api_key_team_id": team}) + + +def _user_in_the_column(user: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {}, user=user) + + +def _user_in_the_metadata(user: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {"user_api_key_user_id": user}) + + +def _activity_on_route(gateway: Gateway, route: Route, api_key: str, entity: str) -> httpx.Response: + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + return activity_of_key(gateway, route.path, api_key, **filters) + + +def _reported_aliases(response: httpx.Response, api_key: str) -> tuple[JsonValue, ...]: + if response.status_code != 200: + return () + return tuple( + object_value(object_value(record)["metadata"])["key_alias"] + for record in records_of_key(object_value(response.json()), api_key) + ) + + +def _names_the_key(api_key: str, alias: str) -> Callable[[httpx.Response], bool]: + def names(response: httpx.Response) -> bool: + reported: Final = _reported_aliases(response, api_key) + return bool(reported) and frozenset(reported) == frozenset((alias,)) + + return names + + +@contextmanager +def _proxy_on(gateway: Gateway, directory: Path, database_url: str, *, workers: int = 1) -> Iterator[OwnedProxy]: + with owned_proxy_process( + gateway, + directory, + {"DATABASE_URL": database_url}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=workers, + ) as owned: + yield owned + + +def _owner_on(candidate: Gateway) -> tuple[str, str]: + owner: Final = f"integration-{uuid.uuid4().hex}" + email: Final = f"{owner}@example.com" + candidate.post("/user/new", {"user_id": owner, "user_email": email, "auto_create_key": False}) + return owner, email + + +@pytest.mark.parametrize("route", ROUTES, ids=lambda route: route.path.strip("/").replace("/", "_")) +def test_alias_named_only_by_a_spend_log_is_reported_on_every_daily_activity_route( + gateway: Gateway, route: Route +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + entity_rows: Final = ( + () if route.table == USER_SPEND else (seeded_row(route.table, route.entity_column, entity, api_key, DAY),) + ) + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + with ( + daily_rows((user_row(owner, api_key, DAY), *entity_rows)), + spend_logs_of_key(api_key, (named_row(0, alias),)), + ): + assert_key_reported( + activity_of_key(gateway, route.path, api_key, **filters), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "layout", + ( + pytest.param(_named_between_fifty_and_fifty, id="named_between_50_and_50_nameless"), + pytest.param(_oldest_named, id="oldest_named_150_nameless_newer"), + pytest.param(_newest_named, id="newest_named_150_nameless_older"), + pytest.param(_both_edges_named, id="both_edges_named_150_nameless_between"), + pytest.param(_named_after_one_hundred, id="100_nameless_named_99_nameless"), + pytest.param(_named_after_ninety_nine, id="99_nameless_named_100_nameless"), + ), +) +def test_alias_on_an_edge_of_the_window_is_reported_whatever_surrounds_it( + gateway: Gateway, layout: Callable[[str], tuple[SpendLogRow, ...]] +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, layout(alias)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_alias_named_only_in_the_middle_of_two_hundred_nameless_rows_is_not_picked_up(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with ( + daily_rows((user_row(owner, api_key, DAY),)), + spend_logs_of_key(api_key, _named_only_in_the_middle(_alias())), + ): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +def test_key_renamed_and_renamed_back_is_reported_with_the_alias_on_both_edges(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + rows: Final = _renamed_and_renamed_back(alias, _alias()) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "spend_log_of_team", + ( + pytest.param(_team_in_the_column, id="team_id_column"), + pytest.param(_team_in_the_metadata, id="team_id_in_metadata"), + ), +) +def test_team_named_only_by_a_spend_log_is_reported_next_to_the_daily_owner( + gateway: Gateway, spend_log_of_team: Callable[[str], SpendLogRow] +) -> None: + api_key: Final = digest_no_key_table_holds() + team: Final = f"integration-team-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (spend_log_of_team(team),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(team=team, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "spend_log_of_user", + ( + pytest.param(_user_in_the_column, id="user_column"), + pytest.param(_user_in_the_metadata, id="user_id_in_metadata"), + ), +) +def test_user_named_by_a_spend_log_beats_the_owner_the_daily_rows_name( + gateway: Gateway, spend_log_of_user: Callable[[str], SpendLogRow] +) -> None: + api_key: Final = digest_no_key_table_holds() + with gateway.scenario() as scenario: + daily_owner, _ = user_with_an_email(scenario) + log_user, log_email = user_with_an_email(scenario) + with ( + daily_rows((user_row(daily_owner, api_key, DAY),)), + spend_logs_of_key(api_key, (spend_log_of_user(log_user),)), + ): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=log_user, email=log_email), + seeded_metrics(1), + ) + + +def test_hashed_jwt_digest_is_named_by_its_spend_log(gateway: Gateway) -> None: + api_key: Final = f"hashed-jwt-{sha256(uuid.uuid4().bytes).hexdigest()}" + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (named_row(0, alias),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + ("started", "inside_the_window"), + ( + pytest.param("2026-02-01 23:59:59", False, id="second_before_the_window"), + pytest.param("2026-02-02 00:00:00", True, id="first_second_of_the_window"), + pytest.param("2026-02-04 23:59:59", True, id="last_second_of_the_window"), + pytest.param("2026-02-05 00:00:00", False, id="first_second_after_the_window"), + ), +) +def test_spend_log_names_the_key_only_from_one_day_before_to_two_days_after_the_read( + gateway: Gateway, started: str, inside_the_window: bool +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + row: Final = SpendLogRow(started, {"user_api_key_alias": alias}) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias if inside_the_window else None, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_two_aliases_on_the_two_edges_leave_the_key_unnamed(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + rows: Final = (named_row(0, _alias()), *nameless_rows(150, 1), named_row(151, _alias())) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "unnamed_rows", + ( + pytest.param((SpendLogRow(started_at(0), {"user_api_key_alias": ""}),), id="empty_string_alias"), + pytest.param( + (SpendLogRow(started_at(0), ["x"]), SpendLogRow(started_at(1), "x")), id="array_then_string_metadata" + ), + ), +) +def test_rows_without_a_usable_alias_do_not_hide_the_named_row_after_them( + gateway: Gateway, unnamed_rows: tuple[SpendLogRow, ...] +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + rows: Final = (*unnamed_rows, named_row(len(unnamed_rows), alias)) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "stored_alias", + ( + pytest.param(123, id="json_int"), + pytest.param(["a"], id="json_list"), + pytest.param("a" * 5000, id="five_kb_string"), + ), +) +def test_alias_of_an_unexpected_shape_is_reported_as_postgres_renders_it( + gateway: Gateway, stored_alias: JsonValue +) -> None: + api_key: Final = digest_no_key_table_holds() + row: Final = SpendLogRow(started_at(0), {"user_api_key_alias": stored_alias}) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)) as request_ids: + rendered: Final = read_rows(ALIAS_OF_ONE_SPEND_LOG, (request_ids[0],))[0]["alias"] + assert isinstance(rendered, str) and rendered, rendered + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=rendered, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.timeout(300) +def test_alias_found_once_is_served_from_the_cache_for_the_same_window_only(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned: + owner, email = _owner_on(owned.gateway) + rows: Final = (user_row(owner, api_key, DAY), user_row(owner, api_key, DAY_OUTSIDE_THE_WINDOW)) + with daily_rows(rows, database_url=database_url): + with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url): + first: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + cached: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + other_window: Final = owned.gateway.request( + "GET", + AGGREGATED_USER_ACTIVITY, + params={"start_date": DAY_OUTSIDE_THE_WINDOW, "end_date": DAY_OUTSIDE_THE_WINDOW, "api_key": api_key}, + ) + named: Final = key_metadata(alias=alias, user=owner, email=email) + assert_key_reported(first, api_key, DAY, named, seeded_metrics(1)) + assert_key_reported(cached, api_key, DAY, named, seeded_metrics(1)) + assert_key_reported( + other_window, api_key, DAY_OUTSIDE_THE_WINDOW, key_metadata(user=owner, email=email), seeded_metrics(1) + ) + + +@pytest.mark.timeout(300) +def test_alias_logged_after_a_cached_miss_shows_once_the_miss_expires(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned: + owner, email = _owner_on(owned.gateway) + with daily_rows((user_row(owner, api_key, DAY),), database_url=database_url): + missed: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url): + named: Final = eventually( + lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), + _names_the_key(api_key, alias), + seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, + ) + assert_key_reported(missed, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_key_reported(named, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) + + +@pytest.mark.timeout(300) +def test_alias_lookup_gives_up_while_spend_logs_are_locked_and_answers_once_they_are_not( + gateway: Gateway, tmp_path: Path +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url, workers=2) as owned: + owner, email = _owner_on(owned.gateway) + with ( + daily_rows((user_row(owner, api_key, DAY),), database_url=database_url), + spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url), + ): + with locked_table(SPEND_LOGS_TABLE, database_url=database_url): + started: Final = time.monotonic() + locked: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + waited: Final = time.monotonic() - started + unlocked: Final = eventually( + lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), + _names_the_key(api_key, alias), + seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, + ) + assert waited < GIVES_UP_WITHIN_SECONDS, waited + assert_key_reported(locked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_key_reported(unlocked, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) + + +def test_concurrent_reads_over_every_route_all_name_a_fresh_key(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + entity_columns: Final = {route.table: route.entity_column for route in ROUTES if route.table != USER_SPEND} + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + rows: Final = ( + user_row(owner, api_key, DAY), + *(seeded_row(table, column, entity, api_key, DAY) for table, column in entity_columns.items()), + ) + with ( + daily_rows(rows), + spend_logs_of_key(api_key, (named_row(0, alias),)), + ThreadPoolExecutor(CONCURRENT_READS) as pool, + ): + reads: Final = tuple( + pool.submit(_activity_on_route, gateway, ROUTES[index % len(ROUTES)], api_key, entity) + for index in range(CONCURRENT_READS) + ) + responses: Final = tuple(read.result() for read in reads) + for response in responses: + assert_key_reported( + response, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1) + ) diff --git a/tests/integration/spend/test_daily_activity_key_owner_traffic.py b/tests/integration/spend/test_daily_activity_key_owner_traffic.py index e0ec13104851..b8f113aca494 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_traffic.py +++ b/tests/integration/spend/test_daily_activity_key_owner_traffic.py @@ -12,7 +12,7 @@ import httpx import pytest -from integration._support.client import Gateway, Scenario, eventually +from integration._support.client import Gateway, Scenario, eventually, string_value from integration._support.daily_activity import ( AGGREGATED_USER_ACTIVITY, DAY, @@ -25,6 +25,7 @@ daily_rows, key_metadata, key_no_key_table_holds, + purge_key_from_the_key_tables, seeded_metrics, seeded_row, user_row, @@ -42,6 +43,10 @@ 'SELECT COALESCE(SUM(api_requests), 0)::int AS requests FROM "LiteLLM_DailyUserSpend" ' "WHERE api_key=%s AND user_id=%s" ) +NAMED_SPEND_LOGS_OF_KEY: Final = ( + 'SELECT COUNT(*)::int AS named FROM "LiteLLM_SpendLogs" ' + "WHERE api_key=%s AND NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL" +) UNIFIED_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses") REQUESTS_OF_A_BURST: Final = 21 READS_DURING_A_BURST: Final = 30 @@ -210,6 +215,14 @@ def _wait_for_requests(api_key: str, user: str, requests: int) -> None: ) +def _wait_for_named_spend_logs(api_key: str, requests: int) -> None: + eventually( + lambda: read_rows(NAMED_SPEND_LOGS_OF_KEY, (api_key,)), + lambda rows: rows[0]["named"] == requests, + seconds=70, + ) + + def _cli_session_token(user: str, team: str) -> str: cli_user: Final = LiteLLM_UserTable(user_id=user, user_role="internal_user", teams=[team], models=[]) return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=cli_user, team_id=team, team_alias="cli-team") @@ -244,6 +257,39 @@ def test_key_used_on_every_unified_endpoint_is_reported_with_its_own_alias_and_u ) +def test_key_purged_from_the_key_tables_is_reported_with_the_alias_its_spend_logs_name(gateway: Gateway) -> None: + prompts: Final = (_prompt(), _prompt(), _prompt()) + with wire_server(_provider) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire.url) + owner, email = user_with_an_email(scenario) + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + generated: Final = gateway.post("/key/generate", {"user_id": owner, "key_alias": alias, "models": [model]}) + key: Final = string_value(generated["key"]) + stored: Final = sha256(key.encode()).hexdigest() + try: + answers: Final = tuple( + gateway.request("POST", endpoint, _request_body(endpoint, model, prompt), key=key) + for endpoint, prompt in zip(UNIFIED_ENDPOINTS, prompts, strict=True) + ) + assert [answer.status_code for answer in answers] == [200, 200, 200], [answer.text for answer in answers] + received: Final = _sent_for_callers(wire.drain()) + assert [request.target for request in received] == [ + "/v1/chat/completions", + "/v1/responses", + "/v1/responses", + ] + _wait_for_requests(stored, owner, 3) + _wait_for_named_spend_logs(stored, 3) + finally: + purge_key_from_the_key_tables(stored) + assert_key_owner_and_totals( + _activity_around_today(gateway, stored), + stored, + key_metadata(alias=alias, user=owner, email=email, exists=False), + _totals_of_requests(3), + ) + + def test_cli_session_spend_is_reported_with_the_user_and_team_of_the_session( gateway: Gateway, monkeypatch: pytest.MonkeyPatch ) -> None: