scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

fix(auth): skip fallback budget cap when team lookup fails for cli session token

Browse source
This commit is contained in:
Sameer Kankute 2026-06-25 12:01:29 +05:30
parent e3a6666e25
commit 25a64d125f
No known key found for this signature in database
2 changed files with 7 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -2339,6 +2339,7 @@ async def cli_poll_key(
user_budget = user_db_obj.max_budget if user_db_obj is not None else None
team_budget: Optional[float] = None
team_budget_resolved = False
if team_id is not None:
try:
team_obj = await get_team_object(
@ -2347,12 +2348,14 @@ async def cli_poll_key(
user_api_key_cache=user_api_key_cache,
)
team_budget = team_obj.max_budget
team_budget_resolved = True
except Exception:
pass
session_max_budget = (
litellm.max_ui_session_budget
if user_budget is None and team_budget is None
if user_budget is None
and (team_id is None or (team_budget_resolved and team_budget is None))
else None
)

4
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
View file

@ -2898,11 +2898,13 @@ class TestCLIKeyRegenerationFlow:
assert result["team_id"] == selected_team
assert result["teams"] == ["team-a", "team-b", "team-c"]
# Verify JWT was generated with correct team
# Verify JWT was generated with correct team and no budget cap
# (team lookup failed, but team_id is set, so fallback cap must not apply)
mock_get_jwt.assert_called_once()
jwt_call_args = mock_get_jwt.call_args
assert jwt_call_args.kwargs["team_id"] == selected_team
assert jwt_call_args.kwargs["team_alias"] == "Team B"
assert jwt_call_args.kwargs["max_budget"] is None
# Verify session was deleted after JWT generation
mock_cache.delete_cache.assert_called_once()