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 <noreply@anthropic.com>

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},
                 },
             },
         },

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"