scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(spend): share double-hash key alias recovery with CloudZero and Focus

Browse source 
Extract the Usage reverse-hash / SpendLogs alias recovery into a shared
helper and apply it when CloudZero and Focus export DailyUserSpend rows,
so BI pulls get api_key_alias back for historical v1.99 double-hashed keys
instead of null.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-09-03 14:24:08 +00:00
parent e07d58a4ef
commit 9dae07175c
No known key found for this signature in database
5 changed files with 308 additions and 156 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/integrations/cloudzero/database.py
View file

@ -94,8 +94,13 @@ class LiteLLMDatabase:
try:
db_response: Final = await client.db.query_raw(query, *params)
# Convert the response to polars DataFrame with full schema inference
# This prevents schema mismatch errors when data types vary across rows
return pl.DataFrame(db_response, infer_schema_length=None)
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
)
# v1.99 double-hashed DailyUserSpend.api_key values miss the
# VerificationToken join above; recover alias/team for those rows.
recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response)
return pl.DataFrame(list(recovered_rows), infer_schema_length=None)
except Exception as e:
raise Exception(f"Error retrieving usage data: {e}")

9
litellm/integrations/focus/database.py
View file

@ -96,7 +96,14 @@ class FocusLiteLLMDatabase:
try:
db_response: Final = await client.db.query_raw(query, *query_params)
return pl.DataFrame(db_response, infer_schema_length=None)
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
)
# v1.99 double-hashed DailyUserSpend.api_key values miss the
# VerificationToken join above; recover alias/team for those rows.
recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response)
return pl.DataFrame(list(recovered_rows), infer_schema_length=None)
except Exception as exc:
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc

157
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -10,10 +10,12 @@ from typing_extensions import TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.spend_tracking.key_metadata_recovery import (
recover_double_hashed_key_metadata,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.utils import PrismaClient, hash_token
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
@ -116,25 +118,6 @@ class _KeyMetadataDict(TypedDict, total=False):
team_id: str | None
# Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key
# values cannot pull an unbounded VerificationToken table into memory.
_MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000
_SPEND_LOGS_KEY_METADATA_SQL: Final = """
SELECT DISTINCT ON (api_key)
api_key,
metadata->>'user_api_key_alias' AS key_alias,
metadata->>'user_api_key_team_id' AS team_id
FROM "LiteLLM_SpendLogs"
WHERE api_key = ANY($1::text[])
AND (
NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL
OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL
)
ORDER BY api_key, "startTime" DESC NULLS LAST
"""
_WhereValue = str | dict[str, object]
@ -459,136 +442,6 @@ def update_breakdown_metrics(
return breakdown
class _TokenAliasRecord(Protocol):
@property
def token(self) -> str: ...
@property
def key_alias(self) -> str | None: ...
@property
def team_id(self) -> str | None: ...
def _token_digest_metadata(
records: Sequence[_TokenAliasRecord],
wanted: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
return {
digested: {"key_alias": record.key_alias, "team_id": record.team_id}
for record in records
for digested in (hash_token(record.token),)
if digested in wanted
}
async def _reverse_hash_active_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
try:
active_records: Final[Sequence[_TokenAliasRecord]] = await VerificationTokenRepository(
prisma_client
).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN)
except Exception as e:
verbose_proxy_logger.warning(
"Failed reverse-hash recovery against active keys for %d missing keys: %s",
len(wanted),
e,
)
return {}
return _token_digest_metadata(active_records, wanted)
async def _reverse_hash_deleted_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
try:
deleted_records: Final[Sequence[_TokenAliasRecord]] = await DeletedVerificationTokenRepository(
prisma_client
).table.find_many(
take=_MAX_DOUBLE_HASH_TOKEN_SCAN,
order={"deleted_at": "desc"},
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed reverse-hash recovery against deleted keys for %d missing keys: %s",
len(wanted),
e,
)
return {}
return _token_digest_metadata(deleted_records, wanted)
async def _reverse_hash_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted)
still_wanted: Final = wanted - frozenset(from_active)
if not still_wanted:
return from_active
return {**from_active, **(await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted))}
async def _spend_logs_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
try:
spend_log_rows: Final = await prisma_client.db.query_raw(
_SPEND_LOGS_KEY_METADATA_SQL,
list(wanted),
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed SpendLogs metadata recovery for %d missing keys: %s",
len(wanted),
e,
)
return {}
if not isinstance(spend_log_rows, list):
return {}
return {
row["api_key"]: {
"key_alias": row.get("key_alias"),
"team_id": row.get("team_id"),
}
for row in spend_log_rows
if isinstance(row, dict)
and isinstance(row.get("api_key"), str)
and row["api_key"] in wanted
}
async def _recover_double_hashed_key_metadata(
prisma_client: PrismaClient,
missing_keys: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
"""
Recover key_alias/team_id for DailyUserSpend.api_key values that were
double-hashed by the v1.99 spend-log provenance gate.
Those rows store hash(VerificationToken.token) instead of the token, so the
exact join misses. Prefer a bounded reverse-hash against active/deleted
tokens; fall back to the alias/team stamped into SpendLogs metadata (which
stayed correct even when api_key did not).
"""
sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key))
if not sha_missing:
return {}
from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing)
still_missing: Final = sha_missing - frozenset(from_tokens)
if not still_missing:
return from_tokens
return {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))}
async def get_api_key_metadata(
prisma_client: PrismaClient,
api_keys: AbstractSet[str],
@ -635,7 +488,7 @@ async def get_api_key_metadata(
return result
return {
**result,
**(await _recover_double_hashed_key_metadata(prisma_client, still_missing)),
**(await recover_double_hashed_key_metadata(prisma_client, still_missing)),
}

223
litellm/proxy/spend_tracking/key_metadata_recovery.py Normal file
View file

@ -0,0 +1,223 @@
from collections.abc import Mapping, Sequence, Set as AbstractSet
from typing import Final, Protocol
from typing_extensions import TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
from litellm.proxy.utils import PrismaClient, hash_token
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
# Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key
# values cannot pull an unbounded VerificationToken table into memory.
_MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000
_SPEND_LOGS_KEY_METADATA_SQL: Final = """
SELECT DISTINCT ON (api_key)
api_key,
metadata->>'user_api_key_alias' AS key_alias,
metadata->>'user_api_key_team_id' AS team_id
FROM "LiteLLM_SpendLogs"
WHERE api_key = ANY($1::text[])
AND (
NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL
OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL
)
ORDER BY api_key, "startTime" DESC NULLS LAST
"""
class KeyMetadataDict(TypedDict, total=False):
key_alias: str | None
team_id: str | None
class _TokenAliasRecord(Protocol):
@property
def token(self) -> str: ...
@property
def key_alias(self) -> str | None: ...
@property
def team_id(self) -> str | None: ...
def _token_digest_metadata(
records: Sequence[_TokenAliasRecord],
wanted: AbstractSet[str],
) -> dict[str, KeyMetadataDict]:
return {
digested: {"key_alias": record.key_alias, "team_id": record.team_id}
for record in records
for digested in (hash_token(record.token),)
if digested in wanted
}
async def _reverse_hash_active_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, KeyMetadataDict]:
try:
active_records: Final[Sequence[_TokenAliasRecord]] = await VerificationTokenRepository(
prisma_client
).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN)
except Exception as e:
verbose_proxy_logger.warning(
"Failed reverse-hash recovery against active keys for %d missing keys: %s",
len(wanted),
e,
)
return {}
return _token_digest_metadata(active_records, wanted)
async def _reverse_hash_deleted_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, KeyMetadataDict]:
try:
deleted_records: Final[Sequence[_TokenAliasRecord]] = await DeletedVerificationTokenRepository(
prisma_client
).table.find_many(
take=_MAX_DOUBLE_HASH_TOKEN_SCAN,
order={"deleted_at": "desc"},
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed reverse-hash recovery against deleted keys for %d missing keys: %s",
len(wanted),
e,
)
return {}
return _token_digest_metadata(deleted_records, wanted)
async def _reverse_hash_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, KeyMetadataDict]:
from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted)
still_wanted: Final = wanted - frozenset(from_active)
if not still_wanted:
return from_active
return {**from_active, **(await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted))}
async def _spend_logs_key_metadata(
prisma_client: PrismaClient,
wanted: AbstractSet[str],
) -> dict[str, KeyMetadataDict]:
try:
spend_log_rows: Final = await prisma_client.db.query_raw(
_SPEND_LOGS_KEY_METADATA_SQL,
list(wanted),
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed SpendLogs metadata recovery for %d missing keys: %s",
len(wanted),
e,
)
return {}
if not isinstance(spend_log_rows, list):
return {}
return {
row["api_key"]: {
"key_alias": row.get("key_alias"),
"team_id": row.get("team_id"),
}
for row in spend_log_rows
if isinstance(row, dict)
and isinstance(row.get("api_key"), str)
and row["api_key"] in wanted
}
async def recover_double_hashed_key_metadata(
prisma_client: PrismaClient,
missing_keys: AbstractSet[str],
) -> dict[str, KeyMetadataDict]:
"""
Recover key_alias/team_id for DailyUserSpend.api_key values that were
double-hashed by the v1.99 spend-log provenance gate.
Those rows store hash(VerificationToken.token) instead of the token, so the
exact join misses. Prefer a bounded reverse-hash against active/deleted
tokens; fall back to the alias/team stamped into SpendLogs metadata (which
stayed correct even when api_key did not).
"""
sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key))
if not sha_missing:
return {}
from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing)
still_missing: Final = sha_missing - frozenset(from_tokens)
if not still_missing:
return from_tokens
return {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))}
def _row_with_recovered_alias(
row: Mapping[str, object],
recovered: Mapping[str, KeyMetadataDict],
*,
api_key_field: str,
alias_field: str,
team_id_field: str,
) -> Mapping[str, object]:
api_key: Final = row.get(api_key_field)
if not isinstance(api_key, str) or api_key not in recovered:
return row
meta: Final = recovered[api_key]
return {
**row,
alias_field: meta.get("key_alias") or row.get(alias_field),
team_id_field: meta.get("team_id") or row.get(team_id_field),
}
async def fill_missing_api_key_aliases(
prisma_client: PrismaClient,
rows: Sequence[Mapping[str, object]],
*,
api_key_field: str = "api_key",
alias_field: str = "api_key_alias",
team_id_field: str = "team_id",
) -> tuple[Mapping[str, object], ...]:
"""
Fill null api_key_alias / team_id on export rows whose api_key was double-hashed.
Used by CloudZero and Focus, which join DailyUserSpend.api_key to
VerificationToken.token and otherwise export null aliases for those rows.
"""
missing_keys: Final = frozenset(
key
for row in rows
for key in (row.get(api_key_field),)
if isinstance(key, str) and key and row.get(alias_field) in (None, "")
)
if not missing_keys:
return tuple(rows)
recovered: Final = await recover_double_hashed_key_metadata(prisma_client, missing_keys)
if not recovered:
return tuple(rows)
return tuple(
_row_with_recovered_alias(
row,
recovered,
api_key_field=api_key_field,
alias_field=alias_field,
team_id_field=team_id_field,
)
for row in rows
)

64
tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py Normal file
View file

@ -0,0 +1,64 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
recover_double_hashed_key_metadata,
)
from litellm.proxy.utils import hash_token
@pytest.mark.asyncio
async def test_recover_double_hashed_key_metadata_via_reverse_hash():
token = "a" * 64
double_hashed = hash_token(token)
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")]
)
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.query_raw = AsyncMock(return_value=[])
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result[double_hashed]["key_alias"] == "batch-worker"
assert result[double_hashed]["team_id"] == "team-1"
mock_prisma.db.query_raw.assert_not_called()
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_updates_null_alias_rows():
token = "c" * 64
double_hashed = hash_token(token)
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[SimpleNamespace(token=token, key_alias="recovered-alias", team_id="team-9")]
)
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.query_raw = AsyncMock(return_value=[])
rows = (
{
"api_key": double_hashed,
"api_key_alias": None,
"team_id": None,
"user_email": "owner@example.com",
"spend": 12.5,
},
{
"api_key": "already-joined-token",
"api_key_alias": "named-key",
"team_id": "team-ok",
"user_email": "other@example.com",
"spend": 1.0,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled[0]["api_key_alias"] == "recovered-alias"
assert filled[0]["team_id"] == "team-9"
assert filled[0]["user_email"] == "owner@example.com"
assert filled[1]["api_key_alias"] == "named-key"