mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(team_endpoints): don't block /team/update on unchanged team budget (#29525)
On /team/update for a standalone (no-org) team, _check_user_team_limits() compared the request max_budget against the caller's personal max_budget whenever max_budget was present in the payload. A team admin whose personal budget is lower than the team's budget could not edit any field (tpm_limit, team name, etc.) because the UI re-sends the unchanged max_budget on every update, tripping the personal-budget check. Pass the team's current max_budget into _check_user_team_limits() and skip the personal-budget comparison when the incoming value is unchanged or lower than the team's current budget. Only genuine increases above the team's current budget are still validated against the caller's personal limit, so no over-relaxation. Proxy admins and the org-scoped path are unaffected. Adds two regression tests for the standalone update path (unchanged budget + tpm_limit change, and lowering the budget), both for a caller whose personal budget is below the team budget. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
milan-berri
GitHub
parent
a7ecf6b5b1
commit
9ccda11919
2 changed files with 226 additions and 15 deletions
|
|
@ -779,6 +779,7 @@ async def _check_user_team_limits(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: Any,
|
||||
existing_team_max_budget: Optional[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Check user team limits for standalone teams (not org-scoped).
|
||||
|
|
@ -789,28 +790,44 @@ async def _check_user_team_limits(
|
|||
|
||||
Should only be called for standalone teams (when organization_id is None).
|
||||
For org-scoped teams, use _check_org_team_limits() instead.
|
||||
|
||||
`existing_team_max_budget` is the team's current `max_budget` on the
|
||||
/team/update path. When the incoming `max_budget` is unchanged or lower
|
||||
than the team's current budget, the personal-budget comparison is skipped
|
||||
so a team admin can edit other fields (e.g. tpm_limit, team name) without
|
||||
being blocked by a budget the team already has. The UI sends the full team
|
||||
object on every update, so the unchanged `max_budget` would otherwise fail.
|
||||
"""
|
||||
# Validate team budget against user's max_budget
|
||||
if data.max_budget is not None and user_api_key_dict.user_id is not None:
|
||||
user_obj = await get_user_object(
|
||||
user_id=user_api_key_dict.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
# On /team/update, allow unchanged or lower budgets without checking
|
||||
# the caller's personal max_budget. Only increases above the team's
|
||||
# current budget are validated against the user's personal limit.
|
||||
budget_unchanged_or_lower = (
|
||||
existing_team_max_budget is not None
|
||||
and data.max_budget <= existing_team_max_budget
|
||||
)
|
||||
|
||||
if (
|
||||
user_obj is not None
|
||||
and user_obj.max_budget is not None
|
||||
and data.max_budget > user_obj.max_budget
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}"
|
||||
},
|
||||
if not budget_unchanged_or_lower:
|
||||
user_obj = await get_user_object(
|
||||
user_id=user_api_key_dict.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
)
|
||||
|
||||
if (
|
||||
user_obj is not None
|
||||
and user_obj.max_budget is not None
|
||||
and data.max_budget > user_obj.max_budget
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}"
|
||||
},
|
||||
)
|
||||
|
||||
# Validate team models against user's allowed models
|
||||
if data.models is not None and len(user_api_key_dict.models) > 0:
|
||||
for m in data.models:
|
||||
|
|
@ -1824,6 +1841,7 @@ async def update_team( # noqa: PLR0915
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
existing_team_max_budget=existing_team_row.max_budget,
|
||||
)
|
||||
|
||||
updated_kv = data.json(exclude_unset=True)
|
||||
|
|
|
|||
|
|
@ -4531,6 +4531,199 @@ async def test_update_team_standalone_budget_exceeds_user_limit():
|
|||
assert "budget" in str(exc_info.value.message).lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_standalone_unchanged_budget_allowed():
|
||||
"""
|
||||
Test that /team/update for a standalone team does NOT compare against the
|
||||
caller's personal max_budget when the budget is unchanged.
|
||||
|
||||
This is the LiteLLM UI scenario: the UI sends the full team object on every
|
||||
update (including the unchanged max_budget). A team admin only changing
|
||||
tpm_limit should not be blocked by a budget the team already has.
|
||||
|
||||
Scenario:
|
||||
- User (team admin) has personal max_budget=$100
|
||||
- Standalone team exists with current budget=$500
|
||||
- User updates tpm_limit and re-sends the unchanged max_budget=$500
|
||||
- Expected: Should succeed (budget unchanged, not an increase)
|
||||
"""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_UserTable,
|
||||
UpdateTeamRequest,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import update_team
|
||||
|
||||
team_admin_user = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="standalone-unchanged-budget-admin",
|
||||
models=[],
|
||||
)
|
||||
|
||||
# UI re-sends the unchanged max_budget alongside the tpm_limit change.
|
||||
update_request = UpdateTeamRequest(
|
||||
team_id="standalone-unchanged-budget-123",
|
||||
max_budget=500.0, # Unchanged from the team's current budget
|
||||
tpm_limit=50000,
|
||||
)
|
||||
|
||||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()
|
||||
) as mock_audit,
|
||||
):
|
||||
# Mock existing standalone team (no organization_id) with budget=$500
|
||||
mock_existing_team = MagicMock()
|
||||
mock_existing_team.team_id = "standalone-unchanged-budget-123"
|
||||
mock_existing_team.organization_id = None
|
||||
mock_existing_team.max_budget = 500.0
|
||||
mock_existing_team.model_id = None
|
||||
mock_existing_team.model_dump.return_value = {
|
||||
"team_id": "standalone-unchanged-budget-123",
|
||||
"organization_id": None,
|
||||
"max_budget": 500.0,
|
||||
"members_with_roles": [
|
||||
{"user_id": "standalone-unchanged-budget-admin", "role": "admin"}
|
||||
],
|
||||
}
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_existing_team
|
||||
)
|
||||
mock_prisma.jsonify_team_object = lambda db_data: db_data
|
||||
|
||||
# User has a restrictive personal budget that is lower than the team's.
|
||||
mock_user_obj = LiteLLM_UserTable(
|
||||
user_id="standalone-unchanged-budget-admin",
|
||||
max_budget=100.0,
|
||||
)
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj)
|
||||
mock_cache.async_set_cache = AsyncMock()
|
||||
|
||||
mock_updated_team = MagicMock()
|
||||
mock_updated_team.team_id = "standalone-unchanged-budget-123"
|
||||
mock_updated_team.organization_id = None
|
||||
mock_updated_team.max_budget = 500.0
|
||||
mock_updated_team.litellm_model_table = None
|
||||
mock_updated_team.model_dump.return_value = {
|
||||
"team_id": "standalone-unchanged-budget-123",
|
||||
"organization_id": None,
|
||||
"max_budget": 500.0,
|
||||
"tpm_limit": 50000,
|
||||
}
|
||||
mock_prisma.db.litellm_teamtable.update = AsyncMock(
|
||||
return_value=mock_updated_team
|
||||
)
|
||||
|
||||
# Should NOT raise - unchanged budget skips the personal-budget check.
|
||||
result = await update_team(
|
||||
data=update_request,
|
||||
http_request=dummy_request,
|
||||
user_api_key_dict=team_admin_user,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result["data"].max_budget == 500.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_standalone_lower_budget_allowed():
|
||||
"""
|
||||
Test that /team/update for a standalone team allows lowering the budget
|
||||
below the team's current value even when the new value still exceeds the
|
||||
caller's personal max_budget.
|
||||
|
||||
Scenario:
|
||||
- User (team admin) has personal max_budget=$100
|
||||
- Standalone team exists with current budget=$500
|
||||
- User lowers team budget to $300 (a decrease, still above user's $100)
|
||||
- Expected: Should succeed (decrease is not an increase above team budget)
|
||||
"""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_UserTable,
|
||||
UpdateTeamRequest,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import update_team
|
||||
|
||||
team_admin_user = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="standalone-lower-budget-admin",
|
||||
models=[],
|
||||
)
|
||||
|
||||
update_request = UpdateTeamRequest(
|
||||
team_id="standalone-lower-budget-123",
|
||||
max_budget=300.0, # Lower than current $500, still above user's $100
|
||||
)
|
||||
|
||||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()
|
||||
) as mock_audit,
|
||||
):
|
||||
mock_existing_team = MagicMock()
|
||||
mock_existing_team.team_id = "standalone-lower-budget-123"
|
||||
mock_existing_team.organization_id = None
|
||||
mock_existing_team.max_budget = 500.0
|
||||
mock_existing_team.model_id = None
|
||||
mock_existing_team.model_dump.return_value = {
|
||||
"team_id": "standalone-lower-budget-123",
|
||||
"organization_id": None,
|
||||
"max_budget": 500.0,
|
||||
"members_with_roles": [
|
||||
{"user_id": "standalone-lower-budget-admin", "role": "admin"}
|
||||
],
|
||||
}
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=mock_existing_team
|
||||
)
|
||||
mock_prisma.jsonify_team_object = lambda db_data: db_data
|
||||
|
||||
mock_user_obj = LiteLLM_UserTable(
|
||||
user_id="standalone-lower-budget-admin",
|
||||
max_budget=100.0,
|
||||
)
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj)
|
||||
mock_cache.async_set_cache = AsyncMock()
|
||||
|
||||
mock_updated_team = MagicMock()
|
||||
mock_updated_team.team_id = "standalone-lower-budget-123"
|
||||
mock_updated_team.organization_id = None
|
||||
mock_updated_team.max_budget = 300.0
|
||||
mock_updated_team.litellm_model_table = None
|
||||
mock_updated_team.model_dump.return_value = {
|
||||
"team_id": "standalone-lower-budget-123",
|
||||
"organization_id": None,
|
||||
"max_budget": 300.0,
|
||||
}
|
||||
mock_prisma.db.litellm_teamtable.update = AsyncMock(
|
||||
return_value=mock_updated_team
|
||||
)
|
||||
|
||||
result = await update_team(
|
||||
data=update_request,
|
||||
http_request=dummy_request,
|
||||
user_api_key_dict=team_admin_user,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result["data"].max_budget == 300.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_org_scoped_budget_exceeds_org_limit():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue