From d7198f48c0548666f3f50c48f0e773c7965c5bb0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:58:46 -0700 Subject: [PATCH 1/2] Revert "fix(spend-tracking): keep internal service-account key names readable in spend logs (#39572)" This reverts commit c2e18a4320b64442de6ce2325fda2c58d572120f. --- .../spend_tracking/spend_tracking_utils.py | 13 ++---- .../test_spend_tracking_utils.py | 43 ------------------- 2 files changed, 3 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 87c2aadba64..43709e4e6ff 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -14,8 +14,6 @@ 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 ( @@ -74,18 +72,13 @@ 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 in _NON_SECRET_KEY_ALIASES or is_valid_sha256_hash(value) or _HASHED_JWT_RE.fullmatch(value) is not None + value == LITELLM_PROXY_MASTER_KEY_ALIAS + or is_valid_sha256_hash(value) + or _HASHED_JWT_RE.fullmatch(value) is not None ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 74e45095bb3..59828932858 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -12,13 +12,9 @@ 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, @@ -3082,45 +3078,6 @@ 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(): From dec2c2a72a0f903e8c12ef36dcc305df43f85d0c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:58:47 -0700 Subject: [PATCH 2/2] Revert "fix(spend-tracking): keep batch spend keys joinable after v1.99 provenance gate (#39568)" This reverts commit 803e0f736e08905dae203458676a73e33e81b1eb. --- .../proxy/common_utils/check_batch_cost.py | 1 - litellm/integrations/cloudzero/database.py | 16 +- litellm/integrations/focus/database.py | 14 +- litellm/proxy/_lazy_openapi_snapshot.json | 11 - .../callback_logs_endpoints.py | 4 +- .../common_daily_activity.py | 84 +++---- litellm/proxy/proxy_server.py | 1 - .../spend_tracking/key_metadata_recovery.py | 236 ------------------ .../common_daily_activity.py | 1 - .../proxy_unit_tests/test_check_batch_cost.py | 43 ---- .../integrations/cloudzero/test_cloudzero.py | 29 --- .../integrations/focus/test_focus_database.py | 22 -- .../test_callback_logs_endpoints.py | 1 - .../test_common_daily_activity.py | 145 ----------- .../test_key_metadata_recovery.py | 220 ---------------- .../test_spend_tracking_utils.py | 35 --- .../EntityUsage/entityUsageAggregations.ts | 4 +- .../_components/components/UsagePageView.tsx | 4 +- .../src/components/EntityUsageExport/utils.ts | 3 +- .../UsagePage/keyActivityLabel.test.ts | 15 -- .../components/UsagePage/keyActivityLabel.ts | 8 - .../src/components/UsagePage/types.ts | 1 - .../src/components/activity_metrics.test.tsx | 13 +- .../src/components/activity_metrics.tsx | 5 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 25 files changed, 52 insertions(+), 866 deletions(-) delete mode 100644 litellm/proxy/spend_tracking/key_metadata_recovery.py delete mode 100644 tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py delete mode 100644 ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts delete mode 100644 ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index ea207294858..3b09dc9272e 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -155,7 +155,6 @@ 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)), } diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 2fb10ad8a96..b050ee8e1ed 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -94,18 +94,8 @@ class LiteLLMDatabase: try: db_response: Final = await client.db.query_raw(query, *params) - 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) + # 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) except Exception as e: raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 657c7e0d264..815c38b9e9c 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -96,19 +96,7 @@ class FocusLiteLLMDatabase: try: db_response: Final = await client.db.query_raw(query, *query_params) - 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) + return pl.DataFrame(db_response, infer_schema_length=None) except Exception as exc: raise RuntimeError(f"Error retrieving usage data: {exc}") from exc diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7c95f039f0f..5af45b29226 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3151,17 +3151,6 @@ } ], "title": "Team Id" - }, - "user_email": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "User Email" } }, "title": "KeyMetadata", diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py index cecadc03d71..66057f0dc16 100644 --- a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -88,10 +88,8 @@ 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": user_api_key_hash, - "user_api_key_hash": user_api_key_hash, + "user_api_key": metadata.get("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"), diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index ce6a97708ab..91cd80b3c81 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -2,19 +2,15 @@ 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 MappingProxyType, SimpleNamespace +from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Protocol from fastapi import HTTPException, status -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import 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 @@ -115,19 +111,8 @@ class DailySpendRecord(Protocol): 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] - - -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"), - ) + key_alias: str | None + team_id: str | None _WhereValue = str | dict[str, object] @@ -298,7 +283,10 @@ 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=_key_metadata(api_key_metadata, record.api_key), + 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), + ), ) breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.models[model_key].api_key_breakdown[record.api_key].metrics, @@ -322,7 +310,10 @@ 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=_key_metadata(api_key_metadata, record.api_key), + 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), + ), ) 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, @@ -344,7 +335,10 @@ def update_breakdown_metrics( breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = ( KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=_key_metadata(api_key_metadata, record.api_key), + 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), + ), ) ) @@ -369,7 +363,10 @@ 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=_key_metadata(api_key_metadata, record.api_key), + 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), + ), ) breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, @@ -391,7 +388,10 @@ 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=_key_metadata(api_key_metadata, record.api_key), + 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), + ), ) 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,7 +403,10 @@ def update_breakdown_metrics( if record.api_key not in breakdown.api_keys: breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), - metadata=_key_metadata(api_key_metadata, record.api_key), + 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 ) breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record) @@ -423,7 +426,10 @@ 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=_key_metadata(api_key_metadata, record.api_key), + 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), + ), ) breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics( breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, @@ -436,23 +442,17 @@ def update_breakdown_metrics( async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], -) -> Mapping[str, _KeyMetadataDict]: +) -> dict[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. Also recovers aliases for api_key - values that were double-hashed by the v1.99 spend-log provenance gate. + even after a key is deleted or regenerated. """ 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, - "user_id": getattr(k, "user_id", None), - } - for k in key_records + k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records } # For any keys not found in the active table, check the deleted keys table @@ -471,7 +471,6 @@ 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( @@ -480,13 +479,7 @@ async def get_api_key_metadata( e, ) - 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) + return result def _adjust_dates_for_timezone( @@ -958,6 +951,11 @@ 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], diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 44158631717..887716a383a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15066,7 +15066,6 @@ 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 diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py deleted file mode 100644 index 7de18521edd..00000000000 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ /dev/null @@ -1,236 +0,0 @@ -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 - ) diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 2b39c5dbb9b..101405abf50 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -43,7 +43,6 @@ class KeyMetadata(BaseModel): key_alias: str | None = None team_id: str | None = None - user_email: str | None = None class KeyMetricWithMetadata(MetricBase): diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9a6ab08e9b6..ff5e8f89d64 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2445,7 +2445,6 @@ 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" @@ -2554,48 +2553,6 @@ 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 diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py index 6ddb8cbaa7c..2d51eeb9944 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -1,4 +1,3 @@ -import hashlib from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -75,8 +74,6 @@ 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 @@ -149,9 +146,6 @@ 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 @@ -166,26 +160,3 @@ 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"] diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py index 06240eac387..5c13665f1f1 100644 --- a/tests/test_litellm/integrations/focus/test_focus_database.py +++ b/tests/test_litellm/integrations/focus/test_focus_database.py @@ -1,6 +1,5 @@ """Tests for FocusLiteLLMDatabase query construction.""" -import hashlib from datetime import datetime, timezone from types import SimpleNamespace from unittest.mock import AsyncMock @@ -88,24 +87,3 @@ 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"] diff --git a/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py index 590d63fd868..40e89329b8d 100644 --- a/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py +++ b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py @@ -61,7 +61,6 @@ 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" diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 37a54c4901a..a258127acff 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -454,151 +454,6 @@ 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. diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py deleted file mode 100644 index 7a80319239d..00000000000 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ /dev/null @@ -1,220 +0,0 @@ -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() diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 59828932858..5022dab32be 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2754,41 +2754,6 @@ 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") diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts index d482a5576ae..a53b1d2827b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts @@ -1,4 +1,3 @@ -import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; import { BreakdownMetrics, DailyData, KeyMetricWithMetadata, TagUsage } from "@/components/UsagePage/types"; export type ExtendedDailyData = DailyData & { @@ -119,7 +118,6 @@ 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] || [], }, }; @@ -139,7 +137,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number return Object.entries(keySpend) .map(([api_key, metrics]) => ({ api_key, - key_alias: keyActivityLabel(metrics.metadata), + key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias tags: metrics.metadata.tags || "-", spend: metrics.metrics.spend, })) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 29a81e1ae3f..cbdfc8f39e6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -44,7 +44,6 @@ 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 { @@ -427,7 +426,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { metadata: { key_alias: metrics.metadata.key_alias, team_id: null, - user_email: metrics.metadata.user_email, tags: metrics.metadata.tags || [], }, }; @@ -447,7 +445,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { return Object.entries(keySpend) .map(([api_key, metrics]) => ({ api_key, - key_alias: keyActivityLabel(metrics.metadata), + key_alias: metrics.metadata.key_alias || "-", tags: metrics.metadata.tags || [], spend: metrics.metrics.spend, })) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 8fd75134bcc..de637d5d627 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -1,7 +1,6 @@ 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 = ( @@ -187,7 +186,7 @@ export const generateDailyWithKeysData = ( // Iterate through each API key in the breakdown Object.entries(apiKeyBreakdown).forEach(([keyId, keyData]: [string, any]) => { - const keyAlias = keyActivityLabel(keyData?.metadata, "") || null; + const keyAlias = keyData?.metadata?.key_alias || null; // Create unique key for aggregation: Date_EntityID_KeyID const uniqueKey = `${day.date}_${entityId}_${keyId}`; diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts deleted file mode 100644 index eaf1985c5fa..00000000000 --- a/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.test.ts +++ /dev/null @@ -1,15 +0,0 @@ -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"); - }); -}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts deleted file mode 100644 index 8b3a7eec916..00000000000 --- a/ui/litellm-dashboard/src/components/UsagePage/keyActivityLabel.ts +++ /dev/null @@ -1,8 +0,0 @@ -import type { KeyMetadata } from "./types"; - -export function keyActivityLabel( - metadata: Pick | null | undefined, - fallback = "-", -): string { - return metadata?.key_alias || metadata?.user_email || fallback; -} diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index a10e9e68c4d..8e7c1869df2 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -46,7 +46,6 @@ export interface KeyMetricWithMetadata { export interface KeyMetadata { key_alias: string | null; team_id: string | null; - user_email?: string | null; tags?: { tag: string; usage: number }[]; } diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index b0fc8dc7866..74d258e2bd0 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx @@ -101,7 +101,7 @@ const createMockDailyData = ( }); const createMockKeyMetricWithMetadata = ( - metadata: { key_alias: string | null; team_id: string | null; user_email?: string | null }, + metadata: { key_alias: string | null; team_id: string | null }, metrics: typeof EMPTY_SPEND_METRICS = EMPTY_SPEND_METRICS, ): KeyMetricWithMetadata => ({ metrics, @@ -1450,17 +1450,6 @@ 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", diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index f4348fb65ae..a3fff08faae 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -7,7 +7,6 @@ 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"; @@ -434,7 +433,7 @@ export const ActivityMetrics: React.FC = ({ modelMetrics, // Helper function to format key label export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string, teams: Team[]): string => { - const keyAlias = keyActivityLabel(modelData.metadata, `key-hash-${model}`); + const keyAlias = modelData.metadata.key_alias || `key-hash-${model}`; const teamId = modelData.metadata.team_id; if (teamId) { const teamAlias = resolveTeamAliasFromTeamID(teamId, teams); @@ -517,7 +516,7 @@ export const processActivityData = ( if (!apiKeyBreakdown[apiKey]) { apiKeyBreakdown[apiKey] = { api_key: apiKey, - key_alias: keyActivityLabel(keyData.metadata, "") || null, + key_alias: keyData.metadata.key_alias, team_id: keyData.metadata.team_id, spend: 0, requests: 0, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9d31efab5df..137c67e837c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27617,8 +27617,6 @@ export interface components { key_alias?: string | null; /** Team Id */ team_id?: string | null; - /** User Email */ - user_email?: string | null; }; /** * KeyMetricWithMetadata