scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(proxy): look up hashed key names with two spend log rows per key (#43656)

Browse source 
* 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.

* fix(proxy): cap each spend log name probe at 100 rows per key

* 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.

* test(integration): add spend log alias probe cells for the daily activity routes

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-29 20:21:55 -07:00 • committed by GitHub
parent 8afabe81f1
commit 61a73c59b0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 984 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -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

92
litellm/proxy/spend_tracking/key_metadata_recovery.py
View file

@ -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),
)

65
tests/integration/_support/daily_activity.py
View file

@ -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:

490
tests/integration/spend/test_daily_activity_key_alias_probes.py Normal file
View file

@ -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)
)

48
tests/integration/spend/test_daily_activity_key_owner_traffic.py
View file

@ -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:

314
tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py
View file

@ -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()