diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a624234cf5f..c21113f8c29 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -865,6 +865,7 @@ class LiteLLMRoutes(enum.Enum): "/management/v1/teams/{team_id}/members/bulk_update", "/team/member_update", "/team/{team_id}/member/{user_id}/reset_spend", + "/team/{team_id}/member/{user_id}/reset_budget", "/team/permissions_list", "/team/permissions_update", "/team/daily/activity", diff --git a/tests/proxy_behavior/management/test_team_member_reset_budget.py b/tests/proxy_behavior/management/test_team_member_reset_budget.py new file mode 100644 index 00000000000..42f327c33ef --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_reset_budget.py @@ -0,0 +1,201 @@ +import uuid + +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_SEED_SPEND = 5.0 +_TEAM_DEFAULT_MAX_BUDGET = 100.0 +_CUSTOM_MAX_BUDGET = 50.0 + +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_budget(prisma, budget_id: str, max_budget: float) -> str: + await prisma.db.litellm_budgettable.create( + data={ + "budget_id": budget_id, + "max_budget": max_budget, + "created_by": "phase4-scratch", + "updated_by": "phase4-scratch", + } + ) + return budget_id + + +async def _seed_team_with_default_budget(prisma, world, shape: str, team_id: str, scratch) -> str: + default_budget_id = await _seed_budget(prisma, scratch.tag("team-default-budget"), _TEAM_DEFAULT_MAX_BUDGET) + metadata = {"team_member_budget_id": default_budget_id} + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + metadata=metadata, + ) + elif shape == "beta": + await create_scratch_team(prisma, team_id, organization_id=world.org_b_id, metadata=metadata) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + return default_budget_id + + +async def _seed_custom_member(prisma, team_id: str, member_id: str, scratch) -> str: + custom_budget_id = await _seed_budget(prisma, scratch.tag("custom-budget"), _CUSTOM_MAX_BUDGET) + await prisma.db.litellm_teammembership.create( + data={ + "user_id": member_id, + "team_id": team_id, + "spend": _SEED_SPEND, + "litellm_budget_table": {"connect": {"budget_id": custom_budget_id}}, + } + ) + return custom_budget_id + + +async def _membership(prisma, team_id: str, member_id: str): + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": member_id, "team_id": team_id}} + ) + assert row is not None + return row + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_reset_budget_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + member_id = scratch.tag("member") + default_budget_id = await _seed_team_with_default_budget(prisma, world, shape, scratch.prefix, scratch) + custom_budget_id = await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + caller = world.keys[actor] + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert resp.status_code == expected_status, f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.spend == _SEED_SPEND, "reset_budget must never touch spend" + if expected_status == 200: + assert row.budget_id == default_budget_id + body = resp.json() + assert body["budget_id"] == default_budget_id + assert body["previous_budget_id"] == custom_budget_id + assert body["budget_source"] == "team_default" + else: + assert row.budget_id == custom_budget_id, "denied but budget relinked" + + +async def test_team_member_reset_budget_leaves_shared_default_row_untouched(proxy_client, prisma, scratch, world): + """Relinking must point the member at the shared row, not copy or edit it, so a later + /team/update to team_member_budget reaches this member again.""" + member_id = scratch.tag("member") + default_budget_id = await _seed_team_with_default_budget(prisma, world, "alpha", scratch.prefix, scratch) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + + default_row = await prisma.db.litellm_budgettable.find_unique(where={"budget_id": default_budget_id}) + assert default_row is not None and default_row.max_budget == _TEAM_DEFAULT_MAX_BUDGET + + info = await proxy_client.get( + f"/team/info?team_id={scratch.prefix}", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert info.status_code == 200, info.text + memberships = {tm["user_id"]: tm for tm in info.json()["team_memberships"]} + assert memberships[member_id]["budget_source"] == "team_default" + assert memberships[member_id]["litellm_budget_table"]["max_budget"] == _TEAM_DEFAULT_MAX_BUDGET + + +async def test_team_member_reset_budget_without_team_default_detaches_member(proxy_client, prisma, scratch, world): + member_id = scratch.tag("member") + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["budget_id"] is None + assert resp.json()["budget_source"] == "none" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.budget_id is None + assert row.spend == _SEED_SPEND + + +async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(proxy_client, prisma, scratch, world): + """metadata.team_member_budget_id can outlive its budget row; a stale id must not be + relinked to (the FK would fail) and must read as no budget, not as the team default.""" + member_id = scratch.tag("member") + await create_scratch_team( + prisma, + scratch.prefix, + organization_id=world.org_a_id, + metadata={"team_member_budget_id": scratch.tag("deleted-budget")}, + ) + await _seed_custom_member(prisma, scratch.prefix, member_id, scratch) + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["budget_id"] is None + assert resp.json()["budget_source"] == "none" + + row = await _membership(prisma, scratch.prefix, member_id) + assert row.budget_id is None + + +async def test_team_member_reset_budget_missing_team_is_404(proxy_client, world): + resp = await proxy_client.post( + f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_budget_missing_membership_is_404(proxy_client, prisma, scratch, world): + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_budget", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 404, resp.text