diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 644d54f0139..3fd6b886811 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -574,12 +574,16 @@ def _check_budget_limits_delegation_ceiling( delegation_ceiling: Optional[float], user_api_key_dict: UserAPIKeyAuth, is_ui_session_team_key: bool, + team_table: Optional[LiteLLM_TeamTableCachedObj], ) -> None: """ - Enforce two invariants on `budget_limits`: + Enforce three invariants on `budget_limits`: - Every `budget_limits[*].max_budget` must be a finite number; applies to every caller including proxy admin. + - A CLI session token caller may not set `budget_limits` on a personal + key (one with no `team_id`); mirrors the scalar `max_budget` guard in + `_common_key_generation_helper`. - Non-admin callers may not set a window above their delegation ceiling. """ if not budget_limits: @@ -594,6 +598,13 @@ def _check_budget_limits_delegation_ceiling( return if is_ui_session_team_key: return + if user_api_key_dict.is_session_token and team_table is None: + raise HTTPException( + status_code=400, + detail={ + "error": ("budget_limits cannot be set without specifying team_id when using a CLI session token.") + }, + ) if delegation_ceiling is None: return over_ceiling = next((w for w in budget_limits if w.max_budget > delegation_ceiling), None) @@ -811,6 +822,7 @@ async def _common_key_generation_helper( delegation_ceiling=delegation_ceiling, user_api_key_dict=user_api_key_dict, is_ui_session_team_key=is_ui_session_team_key, + team_table=team_table, ) _check_permissions_caller_permission( permissions=data.permissions, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index f2a4ffae747..8e4f6148dfe 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -12912,6 +12912,129 @@ async def test_budget_limits_window_non_finite_rejected_for_admin(monkeypatch, n assert "finite" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_budget_limits_session_token_personal_key_rejected(monkeypatch): + """A CLI session token caller may not set `budget_limits` on a + personal key (no `team_id`). Mirrors the scalar `max_budget` guard.""" + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params", + None, + raising=False, + ) + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-1", + is_session_token=True, + ) + request = GenerateKeyRequest( + budget_limits=[{"budget_duration": "1d", "max_budget": 1_000_000.0}], + ) + with pytest.raises(HTTPException) as exc_info: + await _common_key_generation_helper( + data=request, + user_api_key_dict=caller, + litellm_changed_by=None, + team_table=None, + ) + assert exc_info.value.status_code == 400 + assert "session token" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_budget_limits_session_token_team_key_uses_team_ceiling(monkeypatch): + """A CLI session token acting on a team key uses the team's + `max_budget` as the ceiling; values within it are permitted.""" + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params", + None, + raising=False, + ) + team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0) + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-1", + team_id="team-1", + is_session_token=True, + ) + request = GenerateKeyRequest( + budget_limits=[{"budget_duration": "1d", "max_budget": 25.0}], + team_id="team-1", + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"key": "sk-test", "expires": None, "user_id": "user-1"}, + ): + result = await _common_key_generation_helper( + data=request, + user_api_key_dict=caller, + litellm_changed_by=None, + team_table=team, + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_budget_limits_session_token_team_key_over_team_budget_rejected(monkeypatch): + """Same shape, but window exceeds the team's `max_budget`.""" + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params", + None, + raising=False, + ) + team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0) + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-1", + team_id="team-1", + is_session_token=True, + ) + request = GenerateKeyRequest( + budget_limits=[{"budget_duration": "1d", "max_budget": 1_000_000.0}], + team_id="team-1", + ) + with pytest.raises(HTTPException) as exc_info: + await _common_key_generation_helper( + data=request, + user_api_key_dict=caller, + litellm_changed_by=None, + team_table=team, + ) + assert exc_info.value.status_code == 400 + assert "cannot exceed" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_budget_limits_session_token_personal_key_admin_unaffected(monkeypatch): + """A proxy admin using a session token is exempt from the personal-key + reject; the role short-circuit runs first.""" + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params", + None, + raising=False, + ) + admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-1", + is_session_token=True, + ) + request = GenerateKeyRequest( + budget_limits=[{"budget_duration": "1d", "max_budget": 1_000_000.0}], + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"key": "sk-test", "expires": None, "user_id": "admin-1"}, + ): + result = await _common_key_generation_helper( + data=request, + user_api_key_dict=admin, + litellm_changed_by=None, + team_table=None, + ) + assert result is not None + + @pytest.mark.asyncio async def test_permissions_field_rejected_for_non_admin(monkeypatch): """A non-admin caller may not set the `permissions` field on a key