mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-01 02:02:20 +00:00
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.
This commit is contained in:
parent
31c9355183
commit
660203a99b
3 changed files with 29 additions and 6 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue