From d037d863c12077deb6476c0beb95a933956ecc4b Mon Sep 17 00:00:00 2001 From: RoyVivat Date: Thu, 2 Apr 2026 15:27:41 -0700 Subject: [PATCH] fix: update existing budget row instead of orphaning it in _upsert_budget_and_membership When existing_budget_id is provided (member already has a budget), update the existing LiteLLM_BudgetTable row in-place rather than always creating a new one. The previous code accepted existing_budget_id but never used it, silently orphaning one budget row per team_member_update call that set any budget field. Adds two unit tests covering the update-existing and create-new paths. Co-Authored-By: Claude Sonnet 4.6 --- .../management_endpoints/common_utils.py | 25 ++++-- .../test_team_endpoints.py | 78 +++++++++++++++++++ 2 files changed, 97 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 98b59adbeb5..ca3120a5cad 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -399,10 +399,23 @@ async def _upsert_budget_and_membership( budget_duration=budget_duration ) - new_budget = await tx.litellm_budgettable.create( - data=create_data, - include={"team_membership": True}, - ) + if existing_budget_id is not None: + # Update in-place: patch only the fields that were supplied so we don't + # overwrite fields the caller didn't touch (e.g. keep max_budget when + # only budget_duration changes). Exclude created_by — that's set once. + update_data = {k: v for k, v in create_data.items() if k != "created_by"} + await tx.litellm_budgettable.update( + where={"budget_id": existing_budget_id}, + data=update_data, + ) + budget_id_to_connect = existing_budget_id + else: + new_budget = await tx.litellm_budgettable.create( + data=create_data, + include={"team_membership": True}, + ) + budget_id_to_connect = new_budget.budget_id + # upsert the team membership with the new/updated budget await tx.litellm_teammembership.upsert( where={ @@ -416,12 +429,12 @@ async def _upsert_budget_and_membership( "user_id": user_id, "team_id": team_id, "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, + "connect": {"budget_id": budget_id_to_connect}, }, }, "update": { "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, + "connect": {"budget_id": budget_id_to_connect}, }, }, }, 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 5c97495f203..54a49bab219 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6627,3 +6627,81 @@ async def test_create_team_member_budget_table_with_duration(): assert budget_request.budget_duration == "30d" assert budget_request.max_budget == 20.0 assert result["metadata"]["team_member_budget_id"] == "budget-abc" + + +@pytest.mark.asyncio +async def test_upsert_budget_and_membership_updates_existing_budget(): + """ + When existing_budget_id is provided, _upsert_budget_and_membership must UPDATE + the existing budget row rather than creating a new one — otherwise each call + silently orphans the previous LiteLLM_BudgetTable row. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership + + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + + mock_tx = MagicMock() + mock_tx.litellm_budgettable.update = AsyncMock() + mock_tx.litellm_budgettable.create = AsyncMock() + mock_tx.litellm_teammembership.upsert = AsyncMock() + + await _upsert_budget_and_membership( + tx=mock_tx, + team_id="team-1", + user_id="user-1", + max_budget=50.0, + existing_budget_id="budget-existing-123", + user_api_key_dict=mock_user, + ) + + # Must update, not create + mock_tx.litellm_budgettable.update.assert_awaited_once() + mock_tx.litellm_budgettable.create.assert_not_awaited() + + # Membership must be connected to the existing budget id + upsert_call = mock_tx.litellm_teammembership.upsert.call_args + connect_id = upsert_call.kwargs["data"]["update"]["litellm_budget_table"]["connect"]["budget_id"] + assert connect_id == "budget-existing-123" + + +@pytest.mark.asyncio +async def test_upsert_budget_and_membership_creates_new_budget_when_none(): + """ + When existing_budget_id is None, _upsert_budget_and_membership must CREATE + a new budget row and connect it to the membership. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership + + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + + mock_new_budget = MagicMock() + mock_new_budget.budget_id = "budget-new-456" + + mock_tx = MagicMock() + mock_tx.litellm_budgettable.update = AsyncMock() + mock_tx.litellm_budgettable.create = AsyncMock(return_value=mock_new_budget) + mock_tx.litellm_teammembership.upsert = AsyncMock() + + await _upsert_budget_and_membership( + tx=mock_tx, + team_id="team-1", + user_id="user-1", + max_budget=100.0, + existing_budget_id=None, + user_api_key_dict=mock_user, + ) + + # Must create, not update + mock_tx.litellm_budgettable.create.assert_awaited_once() + mock_tx.litellm_budgettable.update.assert_not_awaited() + + # Membership must be connected to the newly created budget id + upsert_call = mock_tx.litellm_teammembership.upsert.call_args + connect_id = upsert_call.kwargs["data"]["update"]["litellm_budget_table"]["connect"]["budget_id"] + assert connect_id == "budget-new-456"