Merge pull request #40176 from BerriAI/litellm_backport_1_100_x_spend_key_hash

chore(release): backport #39568 and #39572 to stable/1.100.x and cut 1.100.1
This commit is contained in:
Mateo Wang 2026-09-07 17:18:00 -07:00 committed by GitHub
commit ecc04bf811
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
28 changed files with 922 additions and 58 deletions

View file

@ -155,6 +155,7 @@ class CheckBatchCost:
metadata: Dict[str, Any] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}

View file

@ -94,8 +94,18 @@ 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,
)
usage_rows: Final = (
db_response.to_dicts()
if isinstance(db_response, pl.DataFrame)
else db_response
if isinstance(db_response, list)
else []
)
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
except Exception as e:
raise Exception(f"Error retrieving usage data: {e}")

View file

@ -96,7 +96,19 @@ 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,
)
usage_rows: Final = (
db_response.to_dicts()
if isinstance(db_response, pl.DataFrame)
else db_response
if isinstance(db_response, list)
else []
)
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
except Exception as exc:
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc

View file

@ -3151,6 +3151,17 @@
}
],
"title": "Team Id"
},
"user_email": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "User Email"
}
},
"title": "KeyMetadata",

View file

@ -88,8 +88,10 @@ class CallbackLogsReplayer:
)
metadata: Final[dict[str, Any]] = payload.get("metadata") or {}
user_api_key_hash: Final = metadata.get("user_api_key_hash")
litellm_metadata: Final[dict[str, Any]] = {
"user_api_key": metadata.get("user_api_key_hash"),
"user_api_key": user_api_key_hash,
"user_api_key_hash": user_api_key_hash,
"user_api_key_alias": metadata.get("user_api_key_alias"),
"user_api_key_user_id": metadata.get("user_api_key_user_id"),
"user_api_key_team_id": metadata.get("user_api_key_team_id"),

View file

@ -2,15 +2,19 @@ import asyncio
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from types import MappingProxyType, SimpleNamespace
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import HTTPException, status
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.spend_tracking.key_metadata_recovery import (
attach_user_emails,
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
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
@ -111,8 +115,19 @@ class DailySpendRecord(Protocol):
class _KeyMetadataDict(TypedDict, total=False):
key_alias: str | None
team_id: str | None
key_alias: ReadOnly[str | None]
team_id: ReadOnly[str | None]
user_id: ReadOnly[str | None]
user_email: ReadOnly[str | None]
def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata:
meta: Final = api_key_metadata.get(api_key, {})
return KeyMetadata(
key_alias=meta.get("key_alias"),
team_id=meta.get("team_id"),
user_email=meta.get("user_email"),
)
_WhereValue = str | dict[str, object]
@ -283,10 +298,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.models[model_key].api_key_breakdown:
breakdown.models[model_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics,
@ -310,10 +322,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown:
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics,
@ -335,10 +344,7 @@ def update_breakdown_metrics(
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = (
KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
)
@ -363,10 +369,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.providers[provider].api_key_breakdown:
breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics,
@ -388,10 +391,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics,
@ -403,10 +403,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.api_keys:
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
), # Add any api_key-specific metadata here
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record)
@ -426,10 +423,7 @@ def update_breakdown_metrics(
if record.api_key not in breakdown.entities[entity_value].api_key_breakdown:
breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
),
metadata=_key_metadata(api_key_metadata, record.api_key),
)
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics(
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics,
@ -442,17 +436,23 @@ def update_breakdown_metrics(
async def get_api_key_metadata(
prisma_client: PrismaClient,
api_keys: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
) -> Mapping[str, _KeyMetadataDict]:
"""Get api key metadata, falling back to deleted keys table for keys not found in active table.
This ensures that key_alias and team_id are preserved in historical activity logs
even after a key is deleted or regenerated.
even after a key is deleted or regenerated. Also recovers aliases for api_key
values that were double-hashed by the v1.99 spend-log provenance gate.
"""
key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
where={"token": {"in": list(api_keys)}}
)
result: Final[dict[str, _KeyMetadataDict]] = {
k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records
k.token: {
"key_alias": k.key_alias,
"team_id": k.team_id,
"user_id": getattr(k, "user_id", None),
}
for k in key_records
}
# For any keys not found in the active table, check the deleted keys table
@ -471,6 +471,7 @@ async def get_api_key_metadata(
result[k.token] = {
"key_alias": k.key_alias,
"team_id": k.team_id,
"user_id": getattr(k, "user_id", None),
}
except Exception as e:
verbose_proxy_logger.warning(
@ -479,7 +480,13 @@ async def get_api_key_metadata(
e,
)
return result
still_missing: Final = api_keys - frozenset(result)
combined: Final = (
result
if not still_missing
else MappingProxyType({**result, **(await recover_double_hashed_key_metadata(prisma_client, still_missing))})
)
return await attach_user_emails(prisma_client, combined)
def _adjust_dates_for_timezone(
@ -951,11 +958,6 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
)
def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata:
meta: Final = api_key_metadata.get(api_key, {})
return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id"))
def _aggregate_grouping_sets_records_sync(
*,
records: Sequence[_GroupingSetsRow],

View file

@ -15066,6 +15066,7 @@ async def async_queue_request(
# extra_body); see above for the same guard upstream.
data["metadata"] = {}
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
data["metadata"]["user_api_key_hash"] = user_api_key_dict.api_key
data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)
_headers: Final = _safe_get_request_headers(request).copy()
_headers.pop("authorization", None) # do not store the original `sk-..` api key in the db

View file

@ -0,0 +1,236 @@
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from types import MappingProxyType
from typing import Final, TypeVar
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, 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
from litellm.repositories.user_repository import UserRepository
_T = TypeVar("_T")
_ACTIVE_TOKEN_DIGEST_SQL: Final = """
SELECT encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id
FROM "LiteLLM_VerificationToken"
WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[])
"""
_DELETED_TOKEN_DIGEST_SQL: Final = """
SELECT DISTINCT ON (token)
encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id
FROM "LiteLLM_DeletedVerificationToken"
WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[])
ORDER BY token, deleted_at DESC
"""
class KeyMetadataDict(TypedDict, total=False):
key_alias: ReadOnly[str | None]
team_id: ReadOnly[str | None]
user_id: ReadOnly[str | None]
user_email: ReadOnly[str | None]
class _TokenDigestRow(BaseModel):
digest: str
key_alias: str | None = None
team_id: str | None = None
user_id: str | None = None
_TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...])
_EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({})
_EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({})
async def _db_or_empty(
load: Callable[[], Awaitable[_T]],
warning: str,
count: int,
) -> _T | None:
from prisma.errors import PrismaError
try:
return await load()
except PrismaError as e:
verbose_proxy_logger.warning(warning, count, e)
return None
async def _reverse_hash_key_metadata(
prisma_client: PrismaClient,
sql: str,
wanted: AbstractSet[str],
*,
warning: str,
) -> Mapping[str, KeyMetadataDict]:
rows: Final = await _db_or_empty(
lambda: prisma_client.db.query_raw(sql, sorted(wanted)),
warning,
len(wanted),
)
if rows is None:
return _EMPTY_KEY_METADATA
return MappingProxyType(
{
row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=row.user_id)
for row in _TOKEN_DIGEST_ROWS.validate_python(rows)
if row.digest in wanted
}
)
async def _emails_for_user_ids(
prisma_client: PrismaClient,
user_ids: AbstractSet[str],
) -> Mapping[str, str]:
if not user_ids:
return _EMPTY_EMAILS
users: Final = await _db_or_empty(
lambda: UserRepository(prisma_client).table.find_many(
where={"user_id": {"in": list(user_ids)}}, # mutable-ok: Prisma find_many where= is a dict
),
"Failed user_email recovery for %d user ids: %s",
len(user_ids),
)
if users is None:
return _EMPTY_EMAILS
return MappingProxyType(
{
user.user_id: user.user_email
for user in users
if getattr(user, "user_id", None) and getattr(user, "user_email", None)
}
)
def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMetadataDict:
if meta.get("user_email"):
return meta
user_id: Final = meta.get("user_id")
if not isinstance(user_id, str) or user_id not in emails:
return meta
updated: Final[KeyMetadataDict] = {**meta, "user_email": emails[user_id]}
return updated
async def attach_user_emails(
prisma_client: PrismaClient,
recovered: Mapping[str, KeyMetadataDict],
) -> Mapping[str, KeyMetadataDict]:
needing_email: Final = frozenset(
user_id
for meta in recovered.values()
for user_id in (meta.get("user_id"),)
if isinstance(user_id, str) and user_id and not meta.get("user_email")
)
emails: Final = await _emails_for_user_ids(prisma_client, needing_email)
if not emails:
return recovered
return MappingProxyType({api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()})
async def recover_double_hashed_key_metadata(
prisma_client: PrismaClient,
missing_keys: AbstractSet[str],
) -> Mapping[str, KeyMetadataDict]:
"""
Recover key_alias/team_id/user_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. Postgres hashes the token column itself, one pass over
active keys and one over deleted keys, so no key row crosses the wire.
"""
sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key))
if not sha_missing:
return _EMPTY_KEY_METADATA
from_active: Final = await _reverse_hash_key_metadata(
prisma_client,
_ACTIVE_TOKEN_DIGEST_SQL,
sha_missing,
warning="Failed reverse-hash recovery against active keys for %d missing keys: %s",
)
still_missing: Final = sha_missing - frozenset(from_active)
if not still_missing:
return from_active
from_deleted: Final = await _reverse_hash_key_metadata(
prisma_client,
_DELETED_TOKEN_DIGEST_SQL,
still_missing,
warning="Failed reverse-hash recovery against deleted keys for %d missing keys: %s",
)
return MappingProxyType({**from_active, **from_deleted})
def _row_with_recovered_fields(
row: Mapping[str, object],
recovered: Mapping[str, KeyMetadataDict],
*,
api_key_field: str,
alias_field: str,
team_id_field: str,
user_email_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 MappingProxyType(
{
**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),
user_email_field: row.get(user_email_field) or meta.get("user_email"),
}
)
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",
user_email_field: str = "user_email",
) -> tuple[Mapping[str, object], ...]:
"""
Fill null api_key_alias / team_id / user_email 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 attach_user_emails(
prisma_client,
await recover_double_hashed_key_metadata(prisma_client, missing_keys),
)
if not recovered:
return tuple(rows)
return tuple(
_row_with_recovered_fields(
row,
recovered,
api_key_field=api_key_field,
alias_field=alias_field,
team_id_field=team_id_field,
user_email_field=user_email_field,
)
for row in rows
)

View file

@ -14,6 +14,8 @@ from litellm.constants import (
LITELLM_PROXY_MASTER_KEY_ALIAS,
LITELLM_TRUNCATED_PAYLOAD_FIELD,
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
LITTELM_CLI_SERVICE_ACCOUNT_NAME,
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
REDACTED_BY_LITELM_STRING,
)
from litellm.constants import (
@ -72,13 +74,18 @@ def _is_master_key(api_key: str | None, _master_key: str | None) -> bool:
_HASHED_JWT_RE = re.compile(r"hashed-jwt-[a-fA-F0-9]{64}")
_NON_SECRET_KEY_ALIASES: Final = frozenset(
{
LITELLM_PROXY_MASTER_KEY_ALIAS,
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
LITTELM_CLI_SERVICE_ACCOUNT_NAME,
}
)
def _is_non_secret_key_value(value: str) -> bool:
return (
value == LITELLM_PROXY_MASTER_KEY_ALIAS
or is_valid_sha256_hash(value)
or _HASHED_JWT_RE.fullmatch(value) is not None
value in _NON_SECRET_KEY_ALIASES or is_valid_sha256_hash(value) or _HASHED_JWT_RE.fullmatch(value) is not None
)

View file

@ -43,6 +43,7 @@ class KeyMetadata(BaseModel):
key_alias: str | None = None
team_id: str | None = None
user_email: str | None = None
class KeyMetricWithMetadata(MetricBase):

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.100.0"
version = "1.100.1"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -311,7 +311,7 @@ members = ["enterprise", "litellm-proxy-extras"]
profile = "black"
[tool.commitizen]
version = "1.100.0"
version = "1.100.1"
version_files = [
"pyproject.toml:^version",
]

View file

@ -2445,6 +2445,7 @@ class TestBatchCostAttribution:
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert metadata["user_api_key"] == "hash-alice"
assert metadata["user_api_key_hash"] == "hash-alice"
assert metadata["user_api_key_user_id"] == "alice"
assert metadata["user_api_key_team_id"] == "team-alpha"
assert metadata["user_api_key_alias"] == "prod-key"
@ -2553,6 +2554,48 @@ class TestBatchCostAttribution:
assert metadata["user_api_key_alias"] == "prod-key"
@pytest.mark.asyncio
async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self):
"""
CheckBatchCost stores the VerificationToken hash on the managed object. The
spend-log writer must receive matching user_api_key_hash provenance so it
does not re-hash that value; otherwise DailyUserSpend.api_key no longer joins
VerificationToken and Usage shows key-hash-... with a null alias/email.
"""
from datetime import datetime, timezone
from types import SimpleNamespace
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import hash_token
token_hash = hash_token("sk-batch-creator-key")
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key"),
user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None),
)
metadata = await instance._build_creator_attribution_metadata(
self._job(api_key=token_hash), "batch-1"
)
assert metadata["user_api_key"] == token_hash
assert metadata["user_api_key_hash"] == token_hash
payload = get_logging_payload(
kwargs={
"model": "gpt-4o",
"call_type": "aretrieve_batch",
"litellm_params": {"metadata": metadata},
},
response_obj={
"id": "batch_123",
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
start_time=datetime.now(timezone.utc),
end_time=datetime.now(timezone.utc),
)
assert payload["api_key"] == token_hash
assert payload["api_key"] != hash_token(token_hash)
class TestPollPageStarvation:
"""LIT-5462 regression: a row that can never be costed used to keep its slot in the

View file

@ -1,3 +1,4 @@
import hashlib
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
@ -74,6 +75,8 @@ class TestCloudZeroHourlyExport:
fake_db = MagicMock()
async def query_raw_mock(query: str, *params):
if "sha256(" in query:
return []
start_time_utc = params[0] if len(params) > 0 else None
end_time_utc = params[1] if len(params) > 1 else None
limit = params[2] if len(params) > 2 else None
@ -146,6 +149,9 @@ class TestCloudZeroHourlyExport:
return joined
fake_db.query_raw = AsyncMock(side_effect=query_raw_mock)
fake_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
fake_db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
fake_db.litellm_usertable.find_many = AsyncMock(return_value=[])
fake_client.db = fake_db
mock_prisma_client_getter.return_value = fake_client
@ -160,3 +166,26 @@ class TestCloudZeroHourlyExport:
logger = CloudZeroLogger(api_key="test", connection_id="test")
await logger._hourly_usage_data_export()
class TestLiteLLMDatabaseUsageData:
@pytest.mark.asyncio
async def test_builds_frame_from_rows_recovered_for_double_hashed_keys(self, monkeypatch: pytest.MonkeyPatch):
double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
async def query_raw(query: str, *params):
if "sha256(" in query:
return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
return [joined_row, dirty_row]
fake_client = MagicMock()
fake_client.db.query_raw = AsyncMock(side_effect=query_raw)
db = LiteLLMDatabase()
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: fake_client)
result = await db.get_usage_data()
assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
assert result["team_id"].to_list() == ["team-0", "team-1"]

View file

@ -1,5 +1,6 @@
"""Tests for FocusLiteLLMDatabase query construction."""
import hashlib
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock
@ -87,3 +88,24 @@ async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch):
)
assert "ot.organization_alias as organization_alias" in query_text
assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text
@pytest.mark.asyncio
async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(monkeypatch: pytest.MonkeyPatch):
double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
async def query_raw(query: str, *params):
if "sha256(" in query:
return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
return [joined_row, dirty_row]
mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(side_effect=query_raw)))
db = FocusLiteLLMDatabase()
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
result = await db.get_usage_data()
assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
assert result["team_id"].to_list() == ["team-0", "team-1"]

View file

@ -61,6 +61,7 @@ def test_build_logging_obj_seeds_model_call_details():
# Metadata is mapped to the keys the cost-tracking callback reads.
md = details["litellm_params"]["metadata"]
assert md["user_api_key"] == "rust-gateway-test-key"
assert md["user_api_key_hash"] == "rust-gateway-test-key"
assert md["user_api_key_user_id"] == "user-cb-logs-test"
assert md["user_api_key_team_id"] == "team-cb-logs-test"

View file

@ -454,6 +454,151 @@ async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_rec
assert result["old-key-hash"]["team_id"] == "latest-team"
@pytest.mark.asyncio
async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash():
"""
v1.99 spend logging re-hashed already-hashed api_key values when provenance was
missing. Usage joins DailyUserSpend.api_key to VerificationToken.token, so those
rows looked like key-hash-... with a null alias. Recovery asks Postgres for the
key whose hashed token matches the dirty value and maps it back to its alias.
"""
from litellm.proxy.utils import hash_token
double_hashed = hash_token("a" * 64)
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")]
)
mock_prisma.db.query_raw = AsyncMock(
return_value=[
{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}
]
)
result = await get_api_key_metadata(
prisma_client=mock_prisma,
api_keys={double_hashed},
)
assert result[double_hashed]["key_alias"] == "batch-worker"
assert result[double_hashed]["team_id"] == "team-1"
assert result[double_hashed]["user_email"] == "alice@example.com"
((digest_sql, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list]
assert '"LiteLLM_VerificationToken"' in digest_sql
assert digests == [double_hashed]
@pytest.mark.asyncio
async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_spend_logs():
"""A dirty key no table can explain costs two digest lookups, never a token page walk or a SpendLogs scan."""
from litellm.proxy.utils import hash_token
double_hashed = hash_token("b" * 64)
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma.db.query_raw = AsyncMock(return_value=[])
result = await get_api_key_metadata(
prisma_client=mock_prisma,
api_keys={double_hashed},
)
assert double_hashed not in result
issued_sql = [call.args[0] for call in mock_prisma.db.query_raw.call_args_list]
assert len(issued_sql) == 2
assert not any("LiteLLM_SpendLogs" in sql for sql in issued_sql)
token_lookups = (
mock_prisma.db.litellm_verificationtoken.find_many.call_args_list
+ mock_prisma.db.litellm_deletedverificationtoken.find_many.call_args_list
)
assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups)
def test_key_metadata_includes_recovered_user_email():
from litellm.proxy.management_endpoints.common_daily_activity import _key_metadata
meta = _key_metadata(
{
"dirty-key": {
"key_alias": "batch-worker",
"team_id": "team-1",
"user_email": "alice@example.com",
}
},
"dirty-key",
)
assert meta.key_alias == "batch-worker"
assert meta.user_email == "alice@example.com"
def test_update_breakdown_metrics_includes_user_email():
from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
breakdown = BreakdownMetrics()
record = SimpleNamespace(
api_key="dirty-key",
model="gpt-4o-mini",
model_group="grp",
mcp_namespaced_tool_name="srv/tool",
custom_llm_provider="openai",
endpoint="/v1/chat/completions",
spend=1.23,
prompt_tokens=1,
completion_tokens=1,
cache_read_input_tokens=0,
cache_creation_input_tokens=0,
compression_saved_tokens=0,
compression_savings_spend=0,
prompt_caching_savings_spend=0,
gateway_injected_caching_savings_spend=0,
autorouter_savings_spend=0,
total_tokens=2,
api_requests=1,
successful_requests=1,
failed_requests=0,
ptu_flat_cost=0.0,
user_id="alice",
)
api_key_metadata = {
"dirty-key": {
"key_alias": "batch-worker",
"team_id": "team-1",
"user_email": "alice@example.com",
}
}
update_breakdown_metrics(
breakdown,
record,
{},
{},
api_key_metadata,
entity_id_field="user_id",
)
expected = ("batch-worker", "alice@example.com")
top = breakdown.api_keys["dirty-key"].metadata
assert (top.key_alias, top.user_email) == expected
assert (
breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.key_alias,
breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.user_email,
) == expected
assert (
breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.key_alias,
breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.user_email,
) == expected
assert (
breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.key_alias,
breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.user_email,
) == expected
@pytest.mark.asyncio
async def test_tag_daily_activity_metadata_totals_not_zero():
"""Test that tag daily activity returns correct metadata totals.

View file

@ -0,0 +1,220 @@
from collections.abc import Sequence
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from prisma.errors import PrismaError
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
def _digest_row(digest: str, key_alias: str, team_id: str | None, user_id: str | None) -> dict[str, str | None]:
return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id}
def _query_raw_by_table(
active_rows: Sequence[dict[str, str | None]],
deleted_rows: Sequence[dict[str, str | None]],
) -> AsyncMock:
async def query_raw(sql: str, *params: object) -> list[dict[str, str | None]]:
if '"LiteLLM_VerificationToken"' in sql:
return list(active_rows)
if '"LiteLLM_DeletedVerificationToken"' in sql:
return list(deleted_rows)
raise AssertionError(f"unexpected query: {sql}")
return AsyncMock(side_effect=query_raw)
@pytest.mark.asyncio
async def test_recover_double_hashed_key_metadata_via_active_token_digest():
double_hashed = hash_token("a" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(double_hashed, "batch-worker", "team-1", "alice")],
deleted_rows=[],
)
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"
assert result[double_hashed]["user_id"] == "alice"
((_, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list]
assert digests == [double_hashed]
@pytest.mark.asyncio
async def test_recover_double_hashed_key_metadata_falls_back_to_deleted_tokens():
double_hashed = hash_token("y" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[],
deleted_rows=[_digest_row(double_hashed, "deleted-key", "team-del", "erin")],
)
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result[double_hashed]["key_alias"] == "deleted-key"
assert result[double_hashed]["team_id"] == "team-del"
assert result[double_hashed]["user_id"] == "erin"
assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [[double_hashed], [double_hashed]]
@pytest.mark.asyncio
async def test_recover_only_asks_deleted_tokens_for_digests_active_keys_missed():
found_active = hash_token("1" * 64)
found_deleted = hash_token("2" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(found_active, "active-key", None, None)],
deleted_rows=[_digest_row(found_deleted, "deleted-key", None, None)],
)
result = await recover_double_hashed_key_metadata(mock_prisma, {found_active, found_deleted})
assert result[found_active]["key_alias"] == "active-key"
assert result[found_deleted]["key_alias"] == "deleted-key"
assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [
sorted((found_active, found_deleted)),
[found_deleted],
]
@pytest.mark.asyncio
async def test_recover_permanent_miss_costs_two_digest_lookups_and_no_table_walk():
double_hashed = hash_token("b" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(active_rows=[], deleted_rows=[])
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result == {}
assert len(mock_prisma.db.query_raw.call_args_list) == 2
mock_prisma.db.litellm_verificationtoken.find_many.assert_not_called()
mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called()
@pytest.mark.asyncio
async def test_recover_skips_keys_that_are_not_sha256_digests():
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(return_value=[])
result = await recover_double_hashed_key_metadata(mock_prisma, {"sk-plain-key", "key-hash-short"})
assert result == {}
mock_prisma.db.query_raw.assert_not_called()
@pytest.mark.asyncio
async def test_recover_returns_empty_when_digest_lookup_raises_prisma_error():
double_hashed = hash_token("c" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down"))
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
assert result == {}
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows():
double_hashed = hash_token("d" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(double_hashed, "recovered-alias", "team-9", "bob")],
deleted_rows=[],
)
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="bob", user_email="bob@example.com")]
)
rows = (
{
"api_key": double_hashed,
"api_key_alias": None,
"team_id": None,
"user_email": None,
"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"] == "bob@example.com"
assert filled[0]["spend"] == 12.5
assert filled[1]["api_key_alias"] == "named-key"
assert mock_prisma.db.litellm_usertable.find_many.call_args.kwargs["where"] == {"user_id": {"in": ["bob"]}}
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_leaves_rows_untouched_when_nothing_is_missing():
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(return_value=[])
rows = ({"api_key": hash_token("e" * 64), "api_key_alias": "named", "user_email": "x@example.com"},)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled == rows
mock_prisma.db.query_raw.assert_not_called()
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_keeps_spend_user_email_when_alias_is_missing():
double_hashed = hash_token("f" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(double_hashed, "team-key", "team-9", "key-owner")],
deleted_rows=[],
)
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="key-owner", user_email="owner@example.com")]
)
rows = (
{
"api_key": double_hashed,
"api_key_alias": None,
"team_id": None,
"user_email": "spender@example.com",
"spend": 4.0,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled[0]["api_key_alias"] == "team-key"
assert filled[0]["team_id"] == "team-9"
assert filled[0]["user_email"] == "spender@example.com"
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_skips_named_keys_that_have_no_email():
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(return_value=[])
rows = (
{
"api_key": hash_token("g" * 64),
"api_key_alias": "service-key",
"team_id": "team-svc",
"user_email": None,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled == rows
mock_prisma.db.query_raw.assert_not_called()

View file

@ -12,9 +12,13 @@ import litellm
from litellm.constants import (
LITELLM_TRUNCATED_PAYLOAD_FIELD,
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
LITTELM_CLI_SERVICE_ACCOUNT_NAME,
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
REDACTED_BY_LITELM_STRING,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.spend_tracking.spend_tracking_utils import (
_get_messages_for_spend_logs_payload,
_get_proxy_server_request_for_spend_logs_payload,
@ -2754,6 +2758,41 @@ def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed():
assert meta["user_api_key"] == hash_token(already_hashed)
def test_get_logging_payload_batch_attribution_keeps_verification_token_hash():
"""
Batch cost rebuilds metadata with the managed object's already-hashed api_key.
That hash must land in SpendLogs.api_key unchanged so Usage/CloudZero can join
LiteLLM_VerificationToken for api_key_alias and user_email. Regression: without
user_api_key_hash provenance, v1.99+ re-hashed the token and broke the join.
"""
token_hash = hash_token("sk-batch-creator-key")
kwargs = {
"model": "gpt-4o",
"call_type": "aretrieve_batch",
"litellm_params": {
"metadata": {
"user_api_key": token_hash,
"user_api_key_hash": token_hash,
"user_api_key_alias": "batch-creator",
"user_api_key_user_id": "alice",
"user_api_key_team_id": "team-1",
}
},
}
payload = get_logging_payload(
kwargs=kwargs,
response_obj={"id": "batch_123", "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}},
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
assert payload["api_key"] == token_hash
assert payload["api_key"] != hash_token(token_hash)
parsed_meta = json.loads(payload["metadata"])
assert parsed_meta["user_api_key"] == token_hash
assert parsed_meta["user_api_key_alias"] == "batch-creator"
def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match():
already_hashed = hash_token("sk-some-key")
different_hash = hash_token("sk-other-key")
@ -3043,6 +3082,45 @@ def test_get_logging_payload_keeps_master_key_alias_readable():
assert parsed_meta["user_api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS
@pytest.mark.parametrize(
"service_account",
[LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, LITTELM_CLI_SERVICE_ACCOUNT_NAME],
)
def test_get_logging_payload_keeps_internal_service_account_key_readable(service_account: str):
data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data={"metadata": {}},
user_api_key_dict=UserAPIKeyAuth(
api_key=service_account,
team_id=service_account,
key_alias=service_account,
team_alias=service_account,
),
_metadata_variable_name="metadata",
)
kwargs = {
"model": "openai/gpt-4.1",
"messages": [{"role": "user", "content": "Hello"}],
"call_type": "acompletion",
"litellm_params": {"metadata": data["metadata"]},
}
payload = get_logging_payload(
kwargs=kwargs,
response_obj=Exception("error"),
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
assert payload["api_key"] == service_account
parsed_meta = json.loads(payload["metadata"])
assert parsed_meta["user_api_key"] == service_account
assert parsed_meta["user_api_key_alias"] == service_account
def test_redact_logged_api_key_service_account_name_without_provenance_is_hashed():
result = _redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME)
assert result == hash_token(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME)
@patch("litellm.proxy.proxy_server.master_key", None)
@patch("litellm.proxy.proxy_server.general_settings", {})
def test_get_logging_payload_hashes_bearer_prefixed_api_key():

View file

@ -1,3 +1,4 @@
import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel";
import { BreakdownMetrics, DailyData, KeyMetricWithMetadata, TagUsage } from "@/components/UsagePage/types";
export type ExtendedDailyData = DailyData & {
@ -118,6 +119,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number
metadata: {
key_alias: metrics.metadata.key_alias,
team_id: metrics.metadata.team_id || null,
user_email: metrics.metadata.user_email,
tags: tagDictionary[key] || [],
},
};
@ -137,7 +139,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number
return Object.entries(keySpend)
.map(([api_key, metrics]) => ({
api_key,
key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias
key_alias: keyActivityLabel(metrics.metadata),
tags: metrics.metadata.tags || "-",
spend: metrics.metrics.spend,
}))

View file

@ -44,6 +44,7 @@ import { Tag } from "@/components/tag_management/types";
import UserAgentActivity from "@/components/user_agent_activity";
import ViewUserSpend from "@/components/view_user_spend";
import { usePaginatedDailyActivity } from "../hooks/usePaginatedDailyActivity";
import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel";
import { DailyData, KeyMetricWithMetadata, MetricWithMetadata } from "@/components/UsagePage/types";
import { valueFormatterSpend } from "@/components/UsagePage/utils/value_formatters";
import {
@ -426,6 +427,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
metadata: {
key_alias: metrics.metadata.key_alias,
team_id: null,
user_email: metrics.metadata.user_email,
tags: metrics.metadata.tags || [],
},
};
@ -445,7 +447,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
return Object.entries(keySpend)
.map(([api_key, metrics]) => ({
api_key,
key_alias: metrics.metadata.key_alias || "-",
key_alias: keyActivityLabel(metrics.metadata),
tags: metrics.metadata.tags || [],
spend: metrics.metrics.spend,
}))

View file

@ -1,6 +1,7 @@
import { formatNumberWithCommas } from "@/utils/dataUtils";
import type { DateRangePickerValue } from "@/components/shared/date_picker_types";
import Papa from "papaparse";
import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel";
import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types";
const resolveEntityDisplay = (
@ -186,7 +187,7 @@ export const generateDailyWithKeysData = (
// Iterate through each API key in the breakdown
Object.entries(apiKeyBreakdown).forEach(([keyId, keyData]: [string, any]) => {
const keyAlias = keyData?.metadata?.key_alias || null;
const keyAlias = keyActivityLabel(keyData?.metadata, "") || null;
// Create unique key for aggregation: Date_EntityID_KeyID
const uniqueKey = `${day.date}_${entityId}_${keyId}`;

View file

@ -0,0 +1,15 @@
import { keyActivityLabel } from "./keyActivityLabel";
describe("keyActivityLabel", () => {
it("prefers key_alias", () => {
expect(keyActivityLabel({ key_alias: "batch-worker", user_email: "alice@example.com" })).toBe("batch-worker");
});
it("falls back to user_email when alias is missing", () => {
expect(keyActivityLabel({ key_alias: null, user_email: "alice@example.com" })).toBe("alice@example.com");
});
it("uses the fallback when both alias and email are missing", () => {
expect(keyActivityLabel({ key_alias: null, user_email: null }, "key-hash-abc")).toBe("key-hash-abc");
});
});

View file

@ -0,0 +1,8 @@
import type { KeyMetadata } from "./types";
export function keyActivityLabel(
metadata: Pick<KeyMetadata, "key_alias" | "user_email"> | null | undefined,
fallback = "-",
): string {
return metadata?.key_alias || metadata?.user_email || fallback;
}

View file

@ -46,6 +46,7 @@ export interface KeyMetricWithMetadata {
export interface KeyMetadata {
key_alias: string | null;
team_id: string | null;
user_email?: string | null;
tags?: { tag: string; usage: number }[];
}

View file

@ -101,7 +101,7 @@ const createMockDailyData = (
});
const createMockKeyMetricWithMetadata = (
metadata: { key_alias: string | null; team_id: string | null },
metadata: { key_alias: string | null; team_id: string | null; user_email?: string | null },
metrics: typeof EMPTY_SPEND_METRICS = EMPTY_SPEND_METRICS,
): KeyMetricWithMetadata => ({
metrics,
@ -1450,6 +1450,17 @@ describe("formatKeyLabel", () => {
expect(result).toBe("key-hash-actual-key (team: Test Team 1)");
});
it("should use user_email when key_alias is null", () => {
const modelData = createMockKeyMetricWithMetadata({
key_alias: null,
team_id: "team1",
user_email: "alice@example.com",
});
const result = formatKeyLabel(modelData, "actual-key", MOCK_TEAMS);
expect(result).toBe("alice@example.com (team: Test Team 1)");
});
it("should return key_alias with team_id when teams array is empty", () => {
const modelData = createMockKeyMetricWithMetadata({
key_alias: "my-key",

View file

@ -7,6 +7,7 @@ import { ChevronDown } from "lucide-react";
import React, { useState } from "react";
import { Team } from "./key_team_helpers/key_list";
import KeyModelUsageView from "./UsagePage/components/KeyModelUsageView";
import { keyActivityLabel } from "./UsagePage/keyActivityLabel";
import { DailyData, KeyMetricWithMetadata, ModelActivityData, TopApiKeyData, TopModelData } from "./UsagePage/types";
import { valueFormatter } from "./UsagePage/utils/value_formatters";
@ -433,7 +434,7 @@ export const ActivityMetrics: React.FC<ActivityMetricsProps> = ({ modelMetrics,
// Helper function to format key label
export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string, teams: Team[]): string => {
const keyAlias = modelData.metadata.key_alias || `key-hash-${model}`;
const keyAlias = keyActivityLabel(modelData.metadata, `key-hash-${model}`);
const teamId = modelData.metadata.team_id;
if (teamId) {
const teamAlias = resolveTeamAliasFromTeamID(teamId, teams);
@ -516,7 +517,7 @@ export const processActivityData = (
if (!apiKeyBreakdown[apiKey]) {
apiKeyBreakdown[apiKey] = {
api_key: apiKey,
key_alias: keyData.metadata.key_alias,
key_alias: keyActivityLabel(keyData.metadata, "") || null,
team_id: keyData.metadata.team_id,
spend: 0,
requests: 0,

View file

@ -27617,6 +27617,8 @@ export interface components {
key_alias?: string | null;
/** Team Id */
team_id?: string | null;
/** User Email */
user_email?: string | null;
};
/**
* KeyMetricWithMetadata

2
uv.lock generated
View file

@ -4266,7 +4266,7 @@ wheels = [
[[package]]
name = "litellm"
version = "1.100.0"
version = "1.100.1"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },