fix(proxy): stop /tag/update 500 when the body carries the tag's budget_id (#50324)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-06 08:55:44 -08:00 • committed by GitHub
parent 5d69250cd0
commit ddd162439f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 50 additions and 1 deletions

View file

@ -217,7 +217,7 @@ async def handle_budget_for_entity(
)
# Get all budget field names
budget_params: Final = LiteLLM_BudgetTable.model_fields.keys()
budget_params: Final = LiteLLM_BudgetTable.model_fields.keys() - {"budget_id"}
# Extract budget fields from data
_json_data: Final = _BUDGET_DATA_MAPPING.validate_python(

View file

@ -0,0 +1,20 @@
from typing import Final
from uuid import uuid4
from tests.integration._support.client import Gateway, object_value, string_value
def test_tag_budget_can_be_updated_with_its_existing_budget_id(gateway: Gateway) -> None:
name: Final = f"integration-tag-{uuid4().hex}"
gateway.post("/tag/new", {"name": name, "max_budget": 10})
with gateway.scenario() as scenario:
scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": name})
before: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name])
budget_id: Final = string_value(object_value(before["litellm_budget_table"])["budget_id"])
scenario.cleanups.callback(gateway.post, "/budget/delete", {"id": budget_id})
response: Final = gateway.request(
"POST", "/tag/update", {"name": name, "budget_id": budget_id, "max_budget": 20}
)
assert response.status_code == 200, response.text
after: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name])
assert object_value(after["litellm_budget_table"])["max_budget"] == 20, after

View file

@ -249,6 +249,35 @@ async def test_add_new_member_links_default_team_budget_id():
assert create_data["budget_id"] == test_default_budget_id
@pytest.mark.asyncio
async def test_handle_budget_for_entity_updates_with_existing_budget_id_in_request() -> None:
from litellm.proxy.management_helpers.utils import handle_budget_for_entity
from litellm.types.tag_management import TagUpdateRequest
budget_update: Final = AsyncMock()
prisma_client: Final = cast(
PrismaClient,
SimpleNamespace(db=SimpleNamespace(litellm_budgettable=SimpleNamespace(update=budget_update))),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
):
budget_id: Final = await handle_budget_for_entity(
data=TagUpdateRequest(name="budget-tag", budget_id="budget-1", max_budget=20),
existing_budget_id="budget-1",
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
prisma_client=prisma_client,
litellm_proxy_admin_name="admin",
)
assert budget_id == "budget-1"
budget_update.assert_awaited_once_with(
where={"budget_id": "budget-1"},
data={"budget_id": "budget-1", "max_budget": 20.0, "updated_by": "admin"},
)
@pytest.mark.parametrize(
("request_fields", "cleared_budget_fields", "should_update", "expected_max_budget"),
cast(