scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(proxy): enforce team member budget alongside project budget

Browse source
This commit is contained in:
Shivam Rawat 2026-08-10 18:12:28 -07:00
parent 4daa7ae704
commit a9f01c7295
6 changed files with 36 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/auth/auth_checks.py
View file

@ -3900,7 +3900,6 @@ async def _check_team_member_budget(
and team_object.team_id is not None
and valid_token is not None
and valid_token.user_id is not None
and valid_token.project_id is None
):
team_membership: Final = await get_team_membership(
user_id=valid_token.user_id,

2
litellm/proxy/auth/user_api_key_auth.py
View file

@ -1049,7 +1049,7 @@ async def _team_member_budget_check_for_key(
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
) -> None:
if valid_token.team_member_spend is None or valid_token.project_id is not None:
if valid_token.team_member_spend is None:
return
if prisma_client is None:
return

3
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -540,9 +540,6 @@ async def _get_team_member_budget_counter(
if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None:
return None
if valid_token.project_id is not None:
return None
membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}"
cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key)
team_membership: LiteLLM_TeamMembership | None = None

39
tests/test_litellm/proxy/auth/test_team_member_budget.py
View file

@ -439,11 +439,11 @@ async def test_team_member_budget_check_personal_key_not_team():
@pytest.mark.asyncio
async def test_team_member_budget_check_skipped_for_project_scoped_key():
async def test_team_member_budget_check_enforced_for_project_scoped_key():
"""
Regression for project-scoped keys being blocked by the team member budget:
when a key carries a project_id, the project budget governs and the team
member budget check must not raise, even if the member is over budget.
When a key carries a project_id, the project budget is enforced in addition
to, not instead of, the team member budget: an over-budget member must
still be blocked on a project-scoped key.
"""
request_body = {
"model": "gpt-3.5-turbo",
@ -494,18 +494,21 @@ async def test_team_member_budget_check_skipped_for_project_scoped_key():
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache),
):
result = await common_checks(
request_body=request_body,
team_object=team_object,
user_object=user_object,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=mock_proxy_logging_obj,
valid_token=valid_token,
request=mock_request,
)
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await common_checks(
request_body=request_body,
team_object=team_object,
user_object=user_object,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=mock_proxy_logging_obj,
valid_token=valid_token,
request=mock_request,
)
assert result is True
assert "Budget has been exceeded" in str(exc_info.value)
assert "test-user-1" in str(exc_info.value)
assert "test-team-1" in str(exc_info.value)

20
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -5762,8 +5762,8 @@ async def test_temp_budget_increase_applied_for_cached_key():
class TestTeamMemberBudgetCheckForKey:
"""Tests for _team_member_budget_check_for_key: virtual keys with a
team member budget must be blocked once over budget, except project-scoped
keys, which are governed by the project budget instead."""
team member budget must be blocked once over budget, including
project-scoped keys, which enforce the project budget in addition."""
def _valid_token(self, project_id=None) -> UserAPIKeyAuth:
return UserAPIKeyAuth(
@ -5806,24 +5806,22 @@ class TestTeamMemberBudgetCheckForKey:
)
@pytest.mark.asyncio
async def test_skips_for_project_scoped_key(self):
async def test_raises_for_project_scoped_key(self):
from litellm.proxy.auth.user_api_key_auth import (
_team_member_budget_check_for_key,
)
cache = self._cache_with_membership()
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new_callable=AsyncMock,
return_value=0.03,
):
await _team_member_budget_check_for_key(
valid_token=self._valid_token(project_id="project-1"),
prisma_client=MagicMock(),
user_api_key_cache=cache,
)
cache.async_get_cache.assert_not_called()
with pytest.raises(litellm.BudgetExceededError):
await _team_member_budget_check_for_key(
valid_token=self._valid_token(project_id="project-1"),
prisma_client=MagicMock(),
user_api_key_cache=self._cache_with_membership(),
)
async def _proxy_exception_for_key(

9
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -2586,9 +2586,9 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat
@pytest.mark.asyncio
async def test_team_member_budget_counter_skipped_for_project_scoped_key():
"""Project-scoped keys are governed by the project budget, so no team
member budget counter should be reserved for them."""
async def test_team_member_budget_counter_reserved_for_project_scoped_key():
"""Project-scoped keys enforce the project budget in addition to the team
member budget, so both counters must be reserved for them."""
from litellm.proxy.spend_tracking.budget_reservation import (
_get_team_member_budget_counter,
)
@ -2625,7 +2625,8 @@ async def test_team_member_budget_counter_skipped_for_project_scoped_key():
user_object=user_object,
user_api_key_cache=cache,
)
assert counter is None
assert counter is not None
assert counter.counter_key == "spend:team_member:member-user:member-team"
@pytest.mark.asyncio