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.

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(

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(
             {

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"

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"