diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 7a784ee4622..f2eafcbf839 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index d580f1f7703..400ed287802 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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(): """