From 11662ef18eebcb8cc60864f50afed9fa4bd58790 Mon Sep 17 00:00:00 2001 From: Kiyeon Jeon Date: Mon, 13 Apr 2026 04:32:17 +0900 Subject: [PATCH] fix(proxy): update existing budget row instead of creating new one in /team/member_update When updating a team member's budget via /team/member_update, the existing budget row was being orphaned because _upsert_budget_and_membership always created a new row even when existing_budget_id was provided. Now the function updates the existing budget row in-place when existing_budget_id is present, and only creates a new row when no prior budget exists. Closes #25508 --- .../management_endpoints/common_utils.py | 87 ++++++----- .../test_upsert_budget_membership.py | 141 ++++++------------ 2 files changed, 99 insertions(+), 129 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index efc42d3355c..1b083eece3c 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -378,45 +378,60 @@ async def _upsert_budget_and_membership( ) return - # create a new budget - create_data: Dict[str, Any] = { - "created_by": user_api_key_dict.user_id or "", - "updated_by": user_api_key_dict.user_id or "", - } - if max_budget is not None: - create_data["max_budget"] = max_budget - if tpm_limit is not None: - create_data["tpm_limit"] = tpm_limit - if rpm_limit is not None: - create_data["rpm_limit"] = rpm_limit + if existing_budget_id is not None: + # update the existing budget row in-place to avoid orphaning it + update_data: Dict[str, Any] = { + "updated_by": user_api_key_dict.user_id or "", + } + if max_budget is not None: + update_data["max_budget"] = max_budget + if tpm_limit is not None: + update_data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + update_data["rpm_limit"] = rpm_limit + await tx.litellm_budgettable.update( + where={"budget_id": existing_budget_id}, + data=update_data, + ) + else: + # create a new budget and link it to the team membership + create_data: Dict[str, Any] = { + "created_by": user_api_key_dict.user_id or "", + "updated_by": user_api_key_dict.user_id or "", + } + if max_budget is not None: + create_data["max_budget"] = max_budget + if tpm_limit is not None: + create_data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + create_data["rpm_limit"] = rpm_limit - new_budget = await tx.litellm_budgettable.create( - data=create_data, - include={"team_membership": True}, - ) - # upsert the team membership with the new/updated budget - await tx.litellm_teammembership.upsert( - where={ - "user_id_team_id": { - "user_id": user_id, - "team_id": team_id, - } - }, - data={ - "create": { - "user_id": user_id, - "team_id": team_id, - "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, + new_budget = await tx.litellm_budgettable.create( + data=create_data, + include={"team_membership": True}, + ) + await tx.litellm_teammembership.upsert( + where={ + "user_id_team_id": { + "user_id": user_id, + "team_id": team_id, + } + }, + data={ + "create": { + "user_id": user_id, + "team_id": team_id, + "litellm_budget_table": { + "connect": {"budget_id": new_budget.budget_id}, + }, + }, + "update": { + "litellm_budget_table": { + "connect": {"budget_id": new_budget.budget_id}, + }, }, }, - "update": { - "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, - }, - }, - }, - ) + ) def _update_metadata_field(updated_kv: dict, field_name: str) -> None: diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index a36dc7ff2e3..3e8674d28cb 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -12,18 +12,19 @@ from litellm.proxy.management_endpoints.common_utils import ( # Fixtures: a fake Prisma transaction and a fake UserAPIKeyAuth object # --------------------------------------------------------------------------- + @pytest.fixture def mock_tx(): """ Builds an object that looks just enough like the Prisma tx you use inside _upsert_budget_and_membership. """ - # membership “table” + # membership "table" membership = MagicMock() membership.update = AsyncMock() membership.upsert = AsyncMock() - # budget “table” + # budget "table" budget = MagicMock() budget.update = AsyncMock() # budget.create returns a fake row that has .budget_id @@ -42,6 +43,7 @@ def fake_user(): """Cheap stand-in for UserAPIKeyAuth.""" return types.SimpleNamespace(user_id="tester@example.com") + # TEST: max_budget is None, disconnect only @pytest.mark.asyncio async def test_upsert_disconnect(mock_tx, fake_user): @@ -63,54 +65,38 @@ async def test_upsert_disconnect(mock_tx, fake_user): mock_tx.litellm_teammembership.upsert.assert_not_called() -# TEST: existing budget id, creates new budget (current behavior) +# TEST: existing budget id -> update existing budget in-place @pytest.mark.asyncio -async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user): +async def test_upsert_with_existing_budget_id_updates_in_place(mock_tx, fake_user): """ - Test that even when existing_budget_id is provided, the function creates a new budget. - This reflects the current implementation behavior. + When existing_budget_id is provided, the function should update the + existing budget row instead of creating a new one. """ await _upsert_budget_and_membership( mock_tx, team_id="team-2", user_id="user-2", max_budget=42.0, - existing_budget_id="bud-999", # This parameter is currently unused + existing_budget_id="bud-999", user_api_key_dict=fake_user, ) - # Should create a new budget, not update existing - mock_tx.litellm_budgettable.create.assert_awaited_once_with( + # Should update the existing budget row + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "bud-999"}, data={ "max_budget": 42.0, - "created_by": fake_user.user_id, "updated_by": fake_user.user_id, }, - include={"team_membership": True}, ) - # Should upsert team membership with the new budget ID - new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id - mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "user-2", "team_id": "team-2"}}, - data={ - "create": { - "user_id": "user-2", - "team_id": "team-2", - "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, - }, - "update": { - "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, - }, - }, - ) - - # Should NOT update existing budget - mock_tx.litellm_budgettable.update.assert_not_called() + # Should NOT create a new budget or upsert membership + mock_tx.litellm_budgettable.create.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() mock_tx.litellm_teammembership.update.assert_not_called() -# TEST: create new budget and link membership +# TEST: create new budget and link membership (no existing budget) @pytest.mark.asyncio async def test_upsert_create_and_link(mock_tx, fake_user): await _upsert_budget_and_membership( @@ -152,12 +138,12 @@ async def test_upsert_create_and_link(mock_tx, fake_user): mock_tx.litellm_budgettable.update.assert_not_called() -# TEST: create new budget and link membership, then create another new budget +# TEST: create budget then update it via existing_budget_id @pytest.mark.asyncio -async def test_upsert_create_then_create_another(mock_tx, fake_user): +async def test_upsert_create_then_update_existing(mock_tx, fake_user): """ - Test that multiple calls to _upsert_budget_and_membership create separate budgets, - reflecting the current implementation behavior. + First call creates a new budget, second call updates it in-place + using the existing_budget_id. """ # FIRST CALL – create new budget and link membership await _upsert_budget_and_membership( @@ -176,62 +162,42 @@ async def test_upsert_create_then_create_another(mock_tx, fake_user): mock_tx.litellm_budgettable.create.assert_awaited_once() mock_tx.litellm_teammembership.upsert.assert_awaited_once() - # SECOND CALL – reset call history and create another budget + # SECOND CALL – reset call history and update the existing budget mock_tx.litellm_budgettable.create.reset_mock() mock_tx.litellm_teammembership.upsert.reset_mock() mock_tx.litellm_budgettable.update.reset_mock() - # Set up a new budget ID for the second create call - mock_tx.litellm_budgettable.create.return_value = types.SimpleNamespace(budget_id="new-budget-456") - await _upsert_budget_and_membership( mock_tx, team_id="team-42", user_id="user-42", - max_budget=25.0, # new limit - existing_budget_id=created_bid, # this is ignored in current implementation + max_budget=25.0, + existing_budget_id=created_bid, user_api_key_dict=fake_user, ) - # Should create another new budget (not update existing) - mock_tx.litellm_budgettable.create.assert_awaited_once_with( + # Should update the existing budget, not create a new one + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": created_bid}, data={ "max_budget": 25.0, - "created_by": fake_user.user_id, "updated_by": fake_user.user_id, }, - include={"team_membership": True}, ) - # Should upsert team membership with the new budget ID - new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id - mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "user-42", "team_id": "team-42"}}, - data={ - "create": { - "user_id": "user-42", - "team_id": "team-42", - "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, - }, - "update": { - "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, - }, - }, - ) - - # Should NOT call update - mock_tx.litellm_budgettable.update.assert_not_called() + mock_tx.litellm_budgettable.create.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() # TEST: update rpm_limit for member with existing budget_id @pytest.mark.asyncio -async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user): +async def test_upsert_rpm_limit_updates_existing_budget(mock_tx, fake_user): """ - Test that updating rpm_limit for a member with an existing budget_id - creates a new budget with the new rpm/tpm limits and assigns it to the user. + Updating rpm_limit for a member with an existing budget_id should + update the existing budget row in-place with all specified limits. """ existing_budget_id = "existing-budget-456" - + await _upsert_budget_and_membership( mock_tx, team_id="team-rpm-test", @@ -240,51 +206,35 @@ async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user): existing_budget_id=existing_budget_id, user_api_key_dict=fake_user, tpm_limit=1000, - rpm_limit=100, # updating rpm_limit + rpm_limit=100, ) - # Should create a new budget with all the specified limits - mock_tx.litellm_budgettable.create.assert_awaited_once_with( + # Should update the existing budget with all specified limits + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": existing_budget_id}, data={ "max_budget": 50.0, "tpm_limit": 1000, "rpm_limit": 100, - "created_by": fake_user.user_id, "updated_by": fake_user.user_id, }, - include={"team_membership": True}, ) - # Should NOT update the existing budget - mock_tx.litellm_budgettable.update.assert_not_called() - - # Should upsert team membership with the new budget ID - new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id - mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "user-rpm-test", "team_id": "team-rpm-test"}}, - data={ - "create": { - "user_id": "user-rpm-test", - "team_id": "team-rpm-test", - "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, - }, - "update": { - "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, - }, - }, - ) + # Should NOT create a new budget or touch membership + mock_tx.litellm_budgettable.create.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() -# TEST: create new budget with only rpm_limit (no max_budget) +# TEST: create new budget with only rpm_limit (no max_budget, no existing budget) @pytest.mark.asyncio async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): """ - Test that setting only rpm_limit creates a new budget with just the rpm_limit. + Setting only rpm_limit with no existing budget creates a new budget. """ await _upsert_budget_and_membership( mock_tx, team_id="team-rpm-only", - user_id="user-rpm-only", + user_id="user-rpm-only", max_budget=None, existing_budget_id=None, user_api_key_dict=fake_user, @@ -304,7 +254,12 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): # Should upsert team membership with the new budget ID new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "user-rpm-only", "team_id": "team-rpm-only"}}, + where={ + "user_id_team_id": { + "user_id": "user-rpm-only", + "team_id": "team-rpm-only", + } + }, data={ "create": { "user_id": "user-rpm-only",