diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py index 66057f0dc16..cecadc03d71 100644 --- a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -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"), diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 91cd80b3c81..c1fec88d43d 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -10,9 +10,10 @@ from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY +from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy._types import CommonProxyErrors from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import PrismaClient, hash_token from litellm.repositories.table_repositories import DeletedVerificationTokenRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, @@ -115,6 +116,25 @@ class _KeyMetadataDict(TypedDict, total=False): team_id: str | None +# Cap reverse-hash scans so a Usage page with orphaned double-hashed api_key +# values cannot pull an unbounded VerificationToken table into memory. +_MAX_DOUBLE_HASH_TOKEN_SCAN: Final = 10_000 + +_SPEND_LOGS_KEY_METADATA_SQL: Final = """ +SELECT DISTINCT ON (api_key) + api_key, + metadata->>'user_api_key_alias' AS key_alias, + metadata->>'user_api_key_team_id' AS team_id +FROM "LiteLLM_SpendLogs" +WHERE api_key = ANY($1::text[]) + AND ( + NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL + OR NULLIF(metadata->>'user_api_key_team_id', '') IS NOT NULL + ) +ORDER BY api_key, "startTime" DESC NULLS LAST +""" + + _WhereValue = str | dict[str, object] @@ -439,6 +459,136 @@ def update_breakdown_metrics( return breakdown +class _TokenAliasRecord(Protocol): + @property + def token(self) -> str: ... + + @property + def key_alias(self) -> str | None: ... + + @property + def team_id(self) -> str | None: ... + + +def _token_digest_metadata( + records: Sequence[_TokenAliasRecord], + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + return { + digested: {"key_alias": record.key_alias, "team_id": record.team_id} + for record in records + for digested in (hash_token(record.token),) + if digested in wanted + } + + +async def _reverse_hash_active_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + try: + active_records: Final[Sequence[_TokenAliasRecord]] = await VerificationTokenRepository( + prisma_client + ).table.find_many(take=_MAX_DOUBLE_HASH_TOKEN_SCAN) + except Exception as e: + verbose_proxy_logger.warning( + "Failed reverse-hash recovery against active keys for %d missing keys: %s", + len(wanted), + e, + ) + return {} + return _token_digest_metadata(active_records, wanted) + + +async def _reverse_hash_deleted_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + try: + deleted_records: Final[Sequence[_TokenAliasRecord]] = await DeletedVerificationTokenRepository( + prisma_client + ).table.find_many( + take=_MAX_DOUBLE_HASH_TOKEN_SCAN, + order={"deleted_at": "desc"}, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed reverse-hash recovery against deleted keys for %d missing keys: %s", + len(wanted), + e, + ) + return {} + return _token_digest_metadata(deleted_records, wanted) + + +async def _reverse_hash_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + from_active: Final = await _reverse_hash_active_key_metadata(prisma_client, wanted) + still_wanted: Final = wanted - frozenset(from_active) + if not still_wanted: + return from_active + return {**from_active, **(await _reverse_hash_deleted_key_metadata(prisma_client, still_wanted))} + + +async def _spend_logs_key_metadata( + prisma_client: PrismaClient, + wanted: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + try: + spend_log_rows: Final = await prisma_client.db.query_raw( + _SPEND_LOGS_KEY_METADATA_SQL, + list(wanted), + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed SpendLogs metadata recovery for %d missing keys: %s", + len(wanted), + e, + ) + return {} + + if not isinstance(spend_log_rows, list): + return {} + + return { + row["api_key"]: { + "key_alias": row.get("key_alias"), + "team_id": row.get("team_id"), + } + for row in spend_log_rows + if isinstance(row, dict) + and isinstance(row.get("api_key"), str) + and row["api_key"] in wanted + } + + +async def _recover_double_hashed_key_metadata( + prisma_client: PrismaClient, + missing_keys: AbstractSet[str], +) -> dict[str, _KeyMetadataDict]: + """ + Recover key_alias/team_id for DailyUserSpend.api_key values that were + double-hashed by the v1.99 spend-log provenance gate. + + Those rows store hash(VerificationToken.token) instead of the token, so the + exact join misses. Prefer a bounded reverse-hash against active/deleted + tokens; fall back to the alias/team stamped into SpendLogs metadata (which + stayed correct even when api_key did not). + """ + sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key)) + if not sha_missing: + return {} + + from_tokens: Final = await _reverse_hash_key_metadata(prisma_client, sha_missing) + still_missing: Final = sha_missing - frozenset(from_tokens) + if not still_missing: + return from_tokens + + return {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))} + + async def get_api_key_metadata( prisma_client: PrismaClient, api_keys: AbstractSet[str], @@ -446,7 +596,8 @@ async def get_api_key_metadata( """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)}} @@ -479,7 +630,13 @@ async def get_api_key_metadata( e, ) - return result + still_missing: Final = api_keys - set(result.keys()) + if not still_missing: + return result + return { + **result, + **(await _recover_double_hashed_key_metadata(prisma_client, still_missing)), + } def _adjust_dates_for_timezone( 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 40e89329b8d..590d63fd868 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,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" 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 a258127acff..78752495f25 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,6 +454,68 @@ 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. Reverse-hash recovery must map + hash(token) back to the key's alias for historical dirty spend. + """ + from litellm.proxy.utils import hash_token + + token = "a" * 64 + double_hashed = hash_token(token) + mock_prisma = MagicMock() + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + side_effect=[ + [], # exact join miss + [SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")], + ] + ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + + result = await 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" + mock_prisma.db.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_api_key_metadata_recovers_double_hashed_key_via_spend_logs(): + """When the token tables cannot reverse-hash the dirty key, use SpendLogs metadata.""" + 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.query_raw = AsyncMock( + return_value=[ + { + "api_key": double_hashed, + "key_alias": "from-spend-log", + "team_id": "team-spend", + } + ] + ) + + result = await get_api_key_metadata( + prisma_client=mock_prisma, + api_keys={double_hashed}, + ) + + assert result[double_hashed]["key_alias"] == "from-spend-log" + assert result[double_hashed]["team_id"] == "team-spend" + mock_prisma.db.query_raw.assert_called_once() + + @pytest.mark.asyncio async def test_tag_daily_activity_metadata_totals_not_zero(): """Test that tag daily activity returns correct metadata totals.