From 31c9355183a70dfee6f69b16b1089eaadc22dfd0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:52:26 -0700 Subject: [PATCH] fix(spend): resolve the CLI session team from the user's first team in usage metadata A cli-session key carries no team of its own in the DB, so the usage breakdown showed team_id None for it and the export grouped it as Unassigned. The login attaches the user's first team to the session, so the recovery mirrors that rule for cli-session keys only. --- .../common_daily_activity.py | 4 +- .../spend_tracking/key_metadata_recovery.py | 75 +++++++++++++------ .../test_common_daily_activity.py | 3 +- .../test_key_metadata_recovery.py | 26 ++++++- 4 files changed, 82 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 3eda9669eea..5afdf297b3f 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -12,7 +12,7 @@ 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, + attach_user_details, recover_cli_session_key_metadata, recover_double_hashed_key_metadata, recover_key_metadata_from_spend_logs, @@ -515,7 +515,7 @@ async def get_api_key_metadata( else _EMPTY_KEY_METADATA ) combined: Final = MappingProxyType({**after_token_recovery, **from_spend_logs}) - return await attach_user_emails(prisma_client, combined) + return await attach_user_details(prisma_client, combined) def _adjust_dates_for_timezone( diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index bc28685c0ad..202c43b2a21 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -1,6 +1,7 @@ import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet +from dataclasses import dataclass from datetime import datetime, timedelta from types import MappingProxyType from typing import Final, TypeVar @@ -110,7 +111,6 @@ _SPEND_LOG_METADATA_CACHE: Final = InMemoryCache( ) _SPEND_LOG_QUERY_LOCK: Final = asyncio.Lock() _EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({}) -_EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({}) async def _db_or_empty( @@ -150,54 +150,85 @@ async def _reverse_hash_key_metadata( ) -async def _emails_for_user_ids( +@dataclass(frozen=True, slots=True) +class _UserDetails: + email: str | None + first_team: str | None + + +_EMPTY_USER_DETAILS: Final[Mapping[str, _UserDetails]] = MappingProxyType({}) + + +async def _details_for_user_ids( prisma_client: PrismaClient, user_ids: AbstractSet[str], -) -> Mapping[str, str]: +) -> Mapping[str, _UserDetails]: if not user_ids: - return _EMPTY_EMAILS + return _EMPTY_USER_DETAILS 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", + "Failed user detail recovery for %d user ids: %s", len(user_ids), ) if users is None: - return _EMPTY_EMAILS + return _EMPTY_USER_DETAILS return MappingProxyType( { - user.user_id: user.user_email + user.user_id: _UserDetails( + email=getattr(user, "user_email", None) or None, + first_team=next(iter(getattr(user, "teams", None) or ()), None), + ) for user in users - if getattr(user, "user_id", None) and getattr(user, "user_email", None) + if getattr(user, "user_id", None) } ) -def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMetadataDict: - if meta.get("user_email"): - return meta +def _is_cli_session_key(api_key: str) -> bool: + return api_key.startswith(_CLI_SESSION_KEY_PREFIX) and len(api_key) > len(_CLI_SESSION_KEY_PREFIX) + + +def _meta_with_user_details( + api_key: str, meta: KeyMetadataDict, details: Mapping[str, _UserDetails] +) -> KeyMetadataDict: user_id: Final = meta.get("user_id") - if not isinstance(user_id, str) or user_id not in emails: + if not isinstance(user_id, str) or user_id not in details: return meta - updated: Final[KeyMetadataDict] = {**meta, "user_email": emails[user_id]} + user: Final = details[user_id] + email: Final = meta.get("user_email") or user.email + team_id: Final = meta.get("team_id") or (user.first_team if _is_cli_session_key(api_key) else None) + updated: Final[KeyMetadataDict] = { + **meta, + **({"user_email": email} if email else {}), + **({"team_id": team_id} if team_id else {}), + } return updated -async def attach_user_emails( +async def attach_user_details( prisma_client: PrismaClient, recovered: Mapping[str, KeyMetadataDict], ) -> Mapping[str, KeyMetadataDict]: - needing_email: Final = frozenset( + """ + Fill user_email from the owner's user row, and for a cli-session key also + the team the CLI login attaches to that user (its first team). + """ + needing_details: Final = frozenset( user_id - for meta in recovered.values() + for api_key, meta in recovered.items() for user_id in (meta.get("user_id"),) - if isinstance(user_id, str) and user_id and not meta.get("user_email") + if isinstance(user_id, str) + and user_id + and (not meta.get("user_email") or (_is_cli_session_key(api_key) and not meta.get("team_id"))) ) - emails: Final = await _emails_for_user_ids(prisma_client, needing_email) - if not emails: + details: Final = await _details_for_user_ids(prisma_client, needing_details) + if not details: return recovered - return MappingProxyType({api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()}) + return MappingProxyType( + {api_key: _meta_with_user_details(api_key, meta, details) for api_key, meta in recovered.items()} + ) def recover_cli_session_key_metadata(missing_keys: AbstractSet[str]) -> Mapping[str, KeyMetadataDict]: @@ -205,7 +236,7 @@ def recover_cli_session_key_metadata(missing_keys: AbstractSet[str]) -> Mapping[ { key: KeyMetadataDict(key_alias=key, user_id=key.removeprefix(_CLI_SESSION_KEY_PREFIX)) for key in missing_keys - if key.startswith(_CLI_SESSION_KEY_PREFIX) and len(key) > len(_CLI_SESSION_KEY_PREFIX) + if _is_cli_session_key(key) } ) @@ -396,7 +427,7 @@ async def fill_missing_api_key_aliases( return tuple(rows) from_session_keys: Final = recover_cli_session_key_metadata(missing_keys) - recovered: Final = await attach_user_emails( + recovered: Final = await attach_user_details( prisma_client, MappingProxyType( { 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 32a91bac506..47db278352c 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 @@ -2246,10 +2246,11 @@ async def test_get_api_key_metadata_resolves_cli_session_keys_from_the_key_itsel mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) mock_prisma.db.query_raw = AsyncMock(side_effect=AssertionError("no reverse-hash or spend-log scan expected")) mock_prisma.db.litellm_usertable.find_many = AsyncMock( - return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com", teams=["team-a", "team-b"])] ) result = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={"cli-session-alice"}) assert result["cli-session-alice"]["key_alias"] == "cli-session-alice" assert result["cli-session-alice"]["user_email"] == "alice@example.com" + assert result["cli-session-alice"]["team_id"] == "team-a" 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 index ad709a299c7..bd4ce840a2d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -15,6 +15,7 @@ from litellm.constants import ( SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, ) from litellm.proxy.spend_tracking.key_metadata_recovery import ( + attach_user_details, fill_missing_api_key_aliases, recover_cli_session_key_metadata, recover_double_hashed_key_metadata, @@ -602,7 +603,7 @@ async def test_fill_missing_api_key_aliases_resolves_cli_session_keys_without_a_ mock_prisma = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=AssertionError("no reverse-hash lookup expected")) mock_prisma.db.litellm_usertable.find_many = AsyncMock( - return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com", teams=["team-a", "team-b"])] ) rows = ({"api_key": "cli-session-alice", "api_key_alias": None, "team_id": None, "user_email": None, "spend": 2.0},) @@ -610,5 +611,28 @@ async def test_fill_missing_api_key_aliases_resolves_cli_session_keys_without_a_ assert filled[0]["api_key_alias"] == "cli-session-alice" assert filled[0]["user_email"] == "alice@example.com" + assert filled[0]["team_id"] == "team-a" assert filled[0]["spend"] == 2.0 assert mock_prisma.db.litellm_usertable.find_many.call_args.kwargs["where"] == {"user_id": {"in": ["alice"]}} + + +@pytest.mark.asyncio +async def test_attach_user_details_gives_the_login_team_only_to_cli_session_keys(): + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com", teams=["team-a"])] + ) + personal_key = hash_token("sk-personal") + + attached = await attach_user_details( + mock_prisma, + { + "cli-session-alice": {"key_alias": "cli-session-alice", "user_id": "alice"}, + personal_key: {"key_alias": "personal", "user_id": "alice"}, + }, + ) + + assert attached["cli-session-alice"]["team_id"] == "team-a" + assert attached["cli-session-alice"]["user_email"] == "alice@example.com" + assert "team_id" not in attached[personal_key] + assert attached[personal_key]["user_email"] == "alice@example.com"