From e76a18ca3251489af9954e7a8ea30b3b763d6c01 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 3 Sep 2026 14:56:26 +0000 Subject: [PATCH] fix(usage): return user_email on key activity metadata Recovery already resolved the owner email for double-hashed spend keys, then Key Activity dropped it. The Usage payload now carries user_email and the key label falls back to that email before key-hash-... Co-authored-by: Mateo Wang --- .../common_daily_activity.py | 29 ++++++++++++++----- .../spend_tracking/key_metadata_recovery.py | 4 +-- .../common_daily_activity.py | 1 + .../test_common_daily_activity.py | 19 ++++++++++++ .../src/components/UsagePage/types.ts | 1 + .../src/components/activity_metrics.test.tsx | 13 ++++++++- .../src/components/activity_metrics.tsx | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 ++ 8 files changed, 60 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 9853f05a068..833f463621a 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -6,12 +6,13 @@ 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 @@ -116,6 +117,8 @@ class DailySpendRecord(Protocol): class _KeyMetadataDict(TypedDict, total=False): key_alias: str | None team_id: str | None + user_id: ReadOnly[str | None] + user_email: ReadOnly[str | None] _WhereValue = str | dict[str, object] @@ -456,7 +459,12 @@ async def get_api_key_metadata( 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 @@ -475,6 +483,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( @@ -484,10 +493,12 @@ async def get_api_key_metadata( ) still_missing: Final = api_keys - frozenset(result) - if not still_missing: - return result - recovered: Final = await recover_double_hashed_key_metadata(prisma_client, still_missing) - return MappingProxyType({**result, **recovered}) + 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( @@ -961,7 +972,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")) + return KeyMetadata( + key_alias=meta.get("key_alias"), + team_id=meta.get("team_id"), + user_email=meta.get("user_email"), + ) def _aggregate_grouping_sets_records_sync( diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 10ce4548214..fae091d4948 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -209,7 +209,7 @@ def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMet return updated -async def _with_user_emails( +async def attach_user_emails( prisma_client: PrismaClient, recovered: Mapping[str, KeyMetadataDict], ) -> Mapping[str, KeyMetadataDict]: @@ -249,7 +249,7 @@ async def recover_double_hashed_key_metadata( if not still_missing else MappingProxyType({**from_tokens, **(await _spend_logs_key_metadata(prisma_client, still_missing))}) ) - return await _with_user_emails(prisma_client, recovered) + return await attach_user_emails(prisma_client, recovered) def _row_with_recovered_fields( diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 101405abf50..2b39c5dbb9b 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -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): 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 2418a1cf245..f7ad01c3bba 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 @@ -516,6 +516,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_spend_logs(): } ] ) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -527,6 +528,24 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_spend_logs(): mock_prisma.db.query_raw.assert_called_once() +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" + + @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/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index 8e7c1869df2..a10e9e68c4d 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -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 }[]; } diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index 74d258e2bd0..b0fc8dc7866 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 }, + 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", diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index a3fff08faae..e17263ab078 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -433,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 = modelData.metadata.key_alias || `key-hash-${model}`; + const keyAlias = modelData.metadata.key_alias || modelData.metadata.user_email || `key-hash-${model}`; const teamId = modelData.metadata.team_id; if (teamId) { const teamAlias = resolveTeamAliasFromTeamID(teamId, teams); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b0f8e618645..8085967cf6a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27994,6 +27994,8 @@ export interface components { key_alias?: string | null; /** Team Id */ team_id?: string | null; + /** User Email */ + user_email?: string | null; }; /** * KeyMetricWithMetadata