From 660203a99bb7f2294e7751a83c6cf547ed0f418b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:10:16 -0700 Subject: [PATCH] fix(spend): claim the session team only for a single-team user The CLI login attaches a team on its own only when the user has exactly one; a user in several teams picks one per login, so usage metadata for the alias would otherwise name a team the login may not have used. --- .../spend_tracking/key_metadata_recovery.py | 16 ++++++++++++---- .../test_common_daily_activity.py | 2 +- .../test_key_metadata_recovery.py | 17 ++++++++++++++++- 3 files changed, 29 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 202c43b2a21..0d233ef8e4e 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -153,7 +153,7 @@ async def _reverse_hash_key_metadata( @dataclass(frozen=True, slots=True) class _UserDetails: email: str | None - first_team: str | None + only_team: str | None _EMPTY_USER_DETAILS: Final[Mapping[str, _UserDetails]] = MappingProxyType({}) @@ -178,7 +178,7 @@ async def _details_for_user_ids( { user.user_id: _UserDetails( email=getattr(user, "user_email", None) or None, - first_team=next(iter(getattr(user, "teams", None) or ()), None), + only_team=_only_team(getattr(user, "teams", None)), ) for user in users if getattr(user, "user_id", None) @@ -186,6 +186,13 @@ async def _details_for_user_ids( ) +def _only_team(teams: object) -> str | None: + if not isinstance(teams, list) or len(teams) != 1: + return None + team: Final = teams[0] + return team if isinstance(team, str) and team else None + + 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) @@ -198,7 +205,7 @@ def _meta_with_user_details( return meta 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) + team_id: Final = meta.get("team_id") or (user.only_team if _is_cli_session_key(api_key) else None) updated: Final[KeyMetadataDict] = { **meta, **({"user_email": email} if email else {}), @@ -213,7 +220,8 @@ async def attach_user_details( ) -> Mapping[str, KeyMetadataDict]: """ 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). + the team the CLI login attaches on its own: the user's only team. A user in + several teams picks one per login, so the alias claims none for them. """ needing_details: Final = frozenset( user_id 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 47db278352c..d3ede9d535e 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,7 +2246,7 @@ 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", teams=["team-a", "team-b"])] + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com", teams=["team-a"])] ) result = await get_api_key_metadata(prisma_client=mock_prisma, api_keys={"cli-session-alice"}) 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 bd4ce840a2d..657371802be 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 @@ -603,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", teams=["team-a", "team-b"])] + return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com", teams=["team-a"])] ) rows = ({"api_key": "cli-session-alice", "api_key_alias": None, "team_id": None, "user_email": None, "spend": 2.0},) @@ -636,3 +636,18 @@ async def test_attach_user_details_gives_the_login_team_only_to_cli_session_keys 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" + + +@pytest.mark.asyncio +async def test_attach_user_details_claims_no_team_for_a_multi_team_user_session_key(): + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="bob", user_email="bob@example.com", teams=["team-a", "team-b"])] + ) + + attached = await attach_user_details( + mock_prisma, {"cli-session-bob": {"key_alias": "cli-session-bob", "user_id": "bob"}} + ) + + assert "team_id" not in attached["cli-session-bob"] + assert attached["cli-session-bob"]["user_email"] == "bob@example.com"