diff --git a/litellm/constants.py b/litellm/constants.py index 39c10d71709..530d678457d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1908,6 +1908,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 0e3e0598d17..560363ca7d7 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -4,7 +4,7 @@ from collections.abc import Set as AbstractSet 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 @@ -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 @@ -39,26 +40,58 @@ WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[]) 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"], since: Literal["$2::timestamp", "oldest_probe.stopped_at"] +) -> str: + return f""" + SELECT "startTime", 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 ( + SELECT "startTime", metadata, team_id, "user" + FROM "LiteLLM_SpendLogs" + WHERE api_key = keys.digest + AND "startTime" >= {since} + 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} + LIMIT 1 + """ + + +_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, + 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) +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 """ _DAILY_USER_SPEND_OWNER_SQL: Final = """ @@ -69,6 +102,7 @@ GROUP BY api_key """ _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-" @@ -91,7 +125,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): @@ -148,9 +184,12 @@ async def _rows_within_the_statement_timeout( prisma_client: PrismaClient, sql: str, *params: object, + planner_settings: tuple[str, ...] = (), ) -> Sequence[Mapping[str, object]]: async with prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction: await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) + for setting in planner_settings: + await transaction.execute_raw(setting) return await transaction.query_raw(sql, *params) @@ -364,7 +403,14 @@ async def _query_spend_log_metadata( ) -> Mapping[str, KeyMetadataDict] | None: start, end = window rows: Final = await _db_or_empty( - lambda: _rows_within_the_statement_timeout(prisma_client, _SPEND_LOG_ALIAS_SQL, sorted(digests), start, end), + lambda: _rows_within_the_statement_timeout( + prisma_client, + _SPEND_LOG_ALIAS_SQL, + sorted(digests), + start, + end, + planner_settings=(_SPEND_LOG_NO_BITMAP_SCAN_SQL,), + ), "Failed spend-log alias recovery for %d missing keys: %s", len(digests), ) diff --git a/tests/integration/_support/daily_activity.py b/tests/integration/_support/daily_activity.py index debb8c4cdb4..346fc156ea9 100644 --- a/tests/integration/_support/daily_activity.py +++ b/tests/integration/_support/daily_activity.py @@ -3,6 +3,8 @@ import uuid 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 @@ INSERT_SPEND_LOG: Final = ( " 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 00000000000..8d9b8616435 --- /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 e0ec1310485..b8f113aca49 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 @@ from typing import Final 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 @@ from integration._support.daily_activity import ( 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 @@ REQUESTS_OF_KEY: Final = ( '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: 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 74f3a2248c7..bc0bbd4dd38 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,19 +1,27 @@ 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 ( 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, @@ -588,12 +596,314 @@ 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 ) +_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 _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]: + 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_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) + 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): + _insert_nameless_spend_logs(conn, digest, 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + 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] <= 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 async def test_recover_cli_session_key_metadata_names_the_owner_only_when_the_suffix_is_a_real_user(): mock_prisma = MagicMock()