fix(spend): restore user_email for double-hashed keys and persist it in spend logs

Recovery now resolves the key owner's email from UserTable via the
recovered token user_id, and SpendLogsMetadata keeps user_api_key_user_email
so new batch/export consumers see email without a separate user join.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

litellm/proxy/_types.py

@@ -3645,6 +3645,7 @@ class SpendLogsMetadata(TypedDict):
     user_api_key_project_alias: str | None
     user_api_key_org_id: str | None
     user_api_key_user_id: str | None
+    user_api_key_user_email: str | None
     user_api_key_team_alias: str | None
     spend_logs_metadata: dict | None  # special param to log k,v pairs to spendlogs for a call
     requester_ip_address: str | None

litellm/proxy/spend_tracking/key_metadata_recovery.py

@@ -7,6 +7,7 @@ from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
 from litellm.proxy.utils import PrismaClient, hash_token
 from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
+from litellm.repositories.user_repository import UserRepository
 from litellm.repositories.verification_token_repository import (
     VerificationTokenRepository,
 )
@@ -19,12 +20,15 @@ _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
+    metadata->>'user_api_key_team_id' AS team_id,
+    metadata->>'user_api_key_user_id' AS user_id,
+    metadata->>'user_api_key_user_email' AS user_email
 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
+    OR NULLIF(metadata->>'user_api_key_user_email', '') IS NOT NULL
   )
 ORDER BY api_key, "startTime" DESC NULLS LAST
 """
@@ -33,6 +37,8 @@ ORDER BY api_key, "startTime" DESC NULLS LAST
 class KeyMetadataDict(TypedDict, total=False):
     key_alias: str | None
     team_id: str | None
+    user_id: str | None
+    user_email: str | None
 
 
 class _TokenAliasRecord(Protocol):
@@ -45,13 +51,20 @@ class _TokenAliasRecord(Protocol):
     @property
     def team_id(self) -> str | None: ...
 
+    @property
+    def user_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}
+        digested: {
+            "key_alias": record.key_alias,
+            "team_id": record.team_id,
+            "user_id": getattr(record, "user_id", None),
+        }
         for record in records
         for digested in (hash_token(record.token),)
         if digested in wanted
@@ -132,6 +145,8 @@ async def _spend_logs_key_metadata(
         row["api_key"]: {
             "key_alias": row.get("key_alias"),
             "team_id": row.get("team_id"),
+            "user_id": row.get("user_id"),
+            "user_email": row.get("user_email"),
         }
         for row in spend_log_rows
         if isinstance(row, dict)
@@ -140,18 +155,67 @@ async def _spend_logs_key_metadata(
     }
 
 
+async def _emails_for_user_ids(
+    prisma_client: PrismaClient,
+    user_ids: AbstractSet[str],
+) -> Mapping[str, str]:
+    if not user_ids:
+        return {}
+    try:
+        users: Final = await UserRepository(prisma_client).table.find_many(
+            where={"user_id": {"in": list(user_ids)}}
+        )
+    except Exception as e:
+        verbose_proxy_logger.warning(
+            "Failed user_email recovery for %d user ids: %s",
+            len(user_ids),
+            e,
+        )
+        return {}
+    return {
+        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
+    return {**meta, "user_email": emails[user_id]}
+
+
+async def _with_user_emails(
+    prisma_client: PrismaClient,
+    recovered: Mapping[str, KeyMetadataDict],
+) -> dict[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 dict(recovered)
+    return {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],
 ) -> 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.
+    Recover key_alias/team_id/user_email 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).
+    tokens; fall back to SpendLogs metadata. Emails come from SpendLogs when
+    present, otherwise from UserTable via the recovered key's user_id.
     """
     sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key))
     if not sha_missing:
@@ -159,19 +223,22 @@ async def recover_double_hashed_key_metadata(
 
     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))}
+    recovered: Final = (
+        from_tokens
+        if not still_missing
+        else {**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))}
+    )
+    return await _with_user_emails(prisma_client, recovered)
 
 
-def _row_with_recovered_alias(
+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:
@@ -181,6 +248,7 @@ def _row_with_recovered_alias(
         **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: meta.get("user_email") or row.get(user_email_field),
     }
 
 
@@ -191,9 +259,11 @@ async def fill_missing_api_key_aliases(
     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 on export rows whose api_key was double-hashed.
+    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.
@@ -202,7 +272,9 @@ async def fill_missing_api_key_aliases(
         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 isinstance(key, str)
+        and key
+        and (row.get(alias_field) in (None, "") or row.get(user_email_field) in (None, ""))
     )
     if not missing_keys:
         return tuple(rows)
@@ -212,12 +284,13 @@ async def fill_missing_api_key_aliases(
         return tuple(rows)
 
     return tuple(
-        _row_with_recovered_alias(
+        _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
     )

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -138,6 +138,7 @@ def _get_spend_logs_metadata(
             user_api_key_project_alias=None,
             user_api_key_org_id=None,
             user_api_key_user_id=None,
+            user_api_key_user_email=None,
             user_api_key_team_alias=None,
             spend_logs_metadata=None,
             requester_ip_address=None,

tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py

@@ -471,10 +471,20 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash(
     mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
         side_effect=[
             [],  # exact join miss
-            [SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")],
+            [
+                SimpleNamespace(
+                    token=token,
+                    key_alias="batch-worker",
+                    team_id="team-1",
+                    user_id="alice",
+                )
+            ],
         ]
     )
     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=[])
 
     result = await get_api_key_metadata(
@@ -484,6 +494,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash(
 
     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"
     mock_prisma.db.query_raw.assert_not_called()
 
 

tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py

@@ -16,27 +16,48 @@ async def test_recover_double_hashed_key_metadata_via_reverse_hash():
     double_hashed = hash_token(token)
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
-        return_value=[SimpleNamespace(token=token, key_alias="batch-worker", team_id="team-1")]
+        return_value=[
+            SimpleNamespace(
+                token=token,
+                key_alias="batch-worker",
+                team_id="team-1",
+                user_id="alice",
+            )
+        ]
     )
     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=[])
 
     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_email"] == "alice@example.com"
     mock_prisma.db.query_raw.assert_not_called()
 
 
 @pytest.mark.asyncio
-async def test_fill_missing_api_key_aliases_updates_null_alias_rows():
+async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows():
     token = "c" * 64
     double_hashed = hash_token(token)
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
-        return_value=[SimpleNamespace(token=token, key_alias="recovered-alias", team_id="team-9")]
+        return_value=[
+            SimpleNamespace(
+                token=token,
+                key_alias="recovered-alias",
+                team_id="team-9",
+                user_id="bob",
+            )
+        ]
     )
     mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
+    mock_prisma.db.litellm_usertable.find_many = AsyncMock(
+        return_value=[SimpleNamespace(user_id="bob", user_email="bob@example.com")]
+    )
     mock_prisma.db.query_raw = AsyncMock(return_value=[])
 
     rows = (
@@ -44,7 +65,7 @@ async def test_fill_missing_api_key_aliases_updates_null_alias_rows():
             "api_key": double_hashed,
             "api_key_alias": None,
             "team_id": None,
-            "user_email": "owner@example.com",
+            "user_email": None,
             "spend": 12.5,
         },
         {
@@ -60,5 +81,5 @@ async def test_fill_missing_api_key_aliases_updates_null_alias_rows():
 
     assert filled[0]["api_key_alias"] == "recovered-alias"
     assert filled[0]["team_id"] == "team-9"
-    assert filled[0]["user_email"] == "owner@example.com"
+    assert filled[0]["user_email"] == "bob@example.com"
     assert filled[1]["api_key_alias"] == "named-key"