From e3a82f2f66dc9ca3294cbb2da41c1b017a046a0b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 17:35:18 -0700 Subject: [PATCH] fix(proxy): stop team admins raising an org team's max_budget under the org cap The keep-or-lower budget rule only ran for standalone teams, so once max_budget is enabled a team admin on an org team could grow its own budget up to the organization's. It now applies to team admins on every team; org admins keep editing within the org cap. --- .../management_endpoints/team_endpoints.py | 18 ++--- tests/e2e/coverage_registry/mgmt.yaml | 2 +- .../management/test_team_management_e2e.py | 35 +++++++++- .../test_team_endpoints.py | 67 +++++++++++++++++-- 4 files changed, 106 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b2dc3551ced..40913784c8b 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1206,13 +1206,13 @@ def _check_team_budget_update_authority( existing_team_max_budget: float | None, ) -> None: """ - Restrict who can grow a standalone team's spend ceiling on /team/update. + Restrict who can grow a team's spend ceiling on /team/update. - A team admin (already authorized via _verify_team_access) may keep or lower - the team budget, but only a proxy admin may grow it - by raising max_budget - above the team's current value or by removing the cap (setting it to None). - Setting a finite budget on a team that has no cap is a restriction and is - allowed. Org-scoped teams are governed by _check_org_team_limits(). + A team admin may keep or lower the team budget, but only a proxy admin may + grow it - by raising max_budget above the team's current value or by + removing the cap (setting it to None). Setting a finite budget on a team + that has no cap is a restriction and is allowed. Org admins editing + org-scoped teams are governed by _check_org_team_limits() instead. """ if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return @@ -2339,9 +2339,9 @@ async def update_team( prisma_client=prisma_client, ) - # Only a proxy admin may grow a standalone team's spend ceiling. - # Org-scoped teams are validated by _check_org_team_limits() above. - if org_id_to_check is None: + # A team admin never grows its own team's spend ceiling. Org admins grow org-scoped teams + # within the org limits _check_org_team_limits() enforced above. + if org_id_to_check is None or access_role == "team_admin": _check_team_budget_update_authority( data=data, user_api_key_dict=user_api_key_dict, diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 0b7987c6ffb..85fbd0acd91 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -32,7 +32,7 @@ - {id: mgmt.team.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1582", rationale: "Metadata/budget updates persist"} - {id: mgmt.team.update.team_admin_forbidden_until_enabled, module: mgmt, tier: P0, surface: api, assertions: [team_admin_forbidden_until_enabled], source: "team_admin_field_permissions.py:156", rationale: "With no team admin editable fields enabled, a team admin's /team/update is 403 and /team/info reports editing disabled"} - {id: mgmt.team.update.team_admin_limited_to_enabled_fields, module: mgmt, tier: P0, surface: api, assertions: [team_admin_limited_to_enabled_fields], source: "team_admin_field_permissions.py:156", rationale: "A team admin may change only the enabled fields; a request that also changes any other field is 403 and writes nothing"} -- {id: mgmt.team.update.team_admin_cannot_grow_budget, module: mgmt, tier: P0, surface: api, assertions: [team_admin_cannot_grow_budget], source: "team_endpoints.py:1203", rationale: "With max_budget enabled, a team admin may keep or lower a standalone team's budget; raising or removing it is 403 and writes nothing"} +- {id: mgmt.team.update.team_admin_cannot_grow_budget, module: mgmt, tier: P0, surface: api, assertions: [team_admin_cannot_grow_budget], source: "team_endpoints.py:1203", fail_before_fix: proven, rationale: "With max_budget enabled, a team admin may keep or lower its team's budget; raising or removing it is 403 and writes nothing, also under an organization's larger cap"} - {id: mgmt.team.update.team_admin_resend_keeps_budget_reset, module: mgmt, tier: P1, surface: api, assertions: [team_admin_resend_keeps_budget_reset], source: "team_admin_field_permissions.py:147", fail_before_fix: proven, rationale: "A team admin resending unchanged budget settings with an enabled field must not push the team's budget reset times back"} - {id: mgmt.team.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1750", rationale: "Deletion prevents key access"} - {id: mgmt.team.block.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py", rationale: "Block suspends all members"} diff --git a/tests/e2e/management/test_team_management_e2e.py b/tests/e2e/management/test_team_management_e2e.py index 60e0047015c..3bbf2474718 100644 --- a/tests/e2e/management/test_team_management_e2e.py +++ b/tests/e2e/management/test_team_management_e2e.py @@ -32,6 +32,7 @@ from lifecycle import ResourceManager from management_client import ManagementClient from models import ( KeyGenerateBody, + OrgNewBody, TeamInfoParams, TeamMemberAddBody, TeamMemberDeleteBody, @@ -46,6 +47,7 @@ TeamRole = Literal["admin", "user"] _TEAM_TPM_LIMIT: Final = 1000 _TEAM_MAX_BUDGET: Final = 10.0 +_ORG_MAX_BUDGET: Final = 100.0 class TeamBlockBody(BaseModel): @@ -119,6 +121,10 @@ class TeamWithAdminNewBody(TeamNewBody): members_with_roles: list[TeamMemberEntry] +class OrgWithBudgetNewBody(OrgNewBody): + max_budget: float + + class TeamSettingsChange(PartialBody, TeamSettings): pass @@ -423,7 +429,10 @@ def rpm_limit_and_max_budget_editable_by_team_admins(client: ManagementClient) - def _team_with_admin( - client: ManagementClient, resources: ResourceManager, max_budget: float | None = None + client: ManagementClient, + resources: ResourceManager, + max_budget: float | None = None, + organization_id: str | None = None, ) -> tuple[str, str]: """A team with a tpm_limit, and the key of a user who is an admin of that team.""" admin_id = _create_user(client, resources, f"e2e-team-admin-{unique_marker()}@example.com") @@ -432,6 +441,7 @@ def _team_with_admin( team_alias=f"e2e-team-admin-{unique_marker()}", tpm_limit=_TEAM_TPM_LIMIT, max_budget=max_budget, + organization_id=organization_id, members_with_roles=[TeamMemberEntry(role="admin", user_id=admin_id)], ) ) @@ -654,3 +664,26 @@ class TestTeamAdminWithRpmLimitAndMaxBudgetEnabled: assert after == before, ( f"the refused update still wrote to the team, the rpm_limit included: before {before}, after {after}" ) + + @pytest.mark.covers("mgmt.team.update.team_admin_cannot_grow_budget") + def test_team_admin_cannot_raise_an_org_team_budget_under_the_org_cap( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + org_id = client.create_org( + OrgWithBudgetNewBody(organization_alias=f"e2e-team-admin-org-{unique_marker()}", max_budget=_ORG_MAX_BUDGET) + ) + resources.defer(lambda: client.delete_org(org_id)) + team_id, admin_key = _team_with_admin(client, resources, max_budget=_TEAM_MAX_BUDGET, organization_id=org_id) + before = _read_team(client, team_id).team_info + + outcome = _update_team_as( + client, admin_key, TeamSettingsUpdate(team_id=team_id, max_budget=_ORG_MAX_BUDGET / 2) + ) + + assert outcome.status_code == 403, ( + f"a team admin raising an org team's max_budget from {_TEAM_MAX_BUDGET} to {_ORG_MAX_BUDGET / 2}, " + f"under the org's {_ORG_MAX_BUDGET}, must be 403, got {outcome.status_code}: {outcome.body[:300]}" + ) + assert "Only a proxy admin can raise" in outcome.body, f"403 body should say why, got: {outcome.body[:300]}" + after = _read_team(client, team_id).team_info + assert after == before, f"the refused update still wrote to the team: before {before}, after {after}" 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 3dfd994bcee..41d0563bd6b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -7124,8 +7124,10 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit( mock_org.litellm_budget_table = mock_budget_table with ( - _team_admin_may_edit("max_budget"), - _not_org_admin(), + patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + AsyncMock(return_value=True), + ), 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"), @@ -7147,9 +7149,7 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit( "team_id": "org-team-update-budget-123", "organization_id": "test-org-update-budget", "max_budget": 30.0, - "members_with_roles": [ - {"user_id": "org-admin-update-budget-test", "role": "admin"} - ], + "members_with_roles": [], } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team @@ -15177,6 +15177,63 @@ async def test_update_team_holds_a_team_admin_to_the_org_tpm_limit(disable_audit assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 8000 +@pytest.mark.asyncio +async def test_update_team_stops_a_team_admin_raising_an_org_team_budget_under_the_org_cap( + disable_audit_logging_for_mocked_team, +): + """The org cap alone would let a team admin with max_budget enabled grow its own team's budget up to the org's.""" + import contextlib + + budgeted_org = LiteLLM_OrganizationTable( + organization_id="budgeted-org", + budget_id="budgeted-org-budget", + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + org_team = MagicMock() + org_team.metadata = {} + org_team.organization_id = "budgeted-org" + org_team.max_budget = 10.0 + org_team.model_max_budget = None + org_team.model_dump.return_value = { + "team_id": "test_team_id", + "team_alias": "test_team", + "organization_id": "budgeted-org", + "max_budget": 10.0, + "metadata": {}, + "members_with_roles": [{"user_id": "team-admin", "role": "admin"}], + } + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=org_team) + stack.enter_context(_team_admin_may_edit("max_budget")) + stack.enter_context(_not_org_admin()) + stack.enter_context( + patch( # test-quality-ok: update_team reads orgs through this module-level import; no seam to inject + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + AsyncMock(return_value=budgeted_org), + ) + ) + with pytest.raises(ProxyException) as raised: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", max_budget=50.0), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", max_budget=5.0), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(raised.value.code) == "403" + assert "Only a proxy admin can raise a team's max_budget" in str(raised.value.message) + assert prisma.db.litellm_teamtable.update.await_count == 1 + assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["max_budget"] == 5.0 + + @pytest.mark.asyncio async def test_update_team_org_admin_is_not_filtered_by_the_team_admin_field_list( disable_audit_logging_for_mocked_team,