From f5382ff09e795365d4cca0839d30a64c0a7b6cd4 Mon Sep 17 00:00:00 2001 From: daanhendrio <255322319+daanhendrio@users.noreply.github.com> Date: Tue, 12 May 2026 15:53:06 +0200 Subject: [PATCH] fix(proxy): allow clearing budget_duration and team_member fields by sending null on /key/update and /team/update Fixes #27734 Sending null for budget_duration, team_member_budget, team_member_budget_duration, team_member_rpm_limit, or team_member_tpm_limit via /key/update or /team/update returned 200 OK but silently ignored the null value. The fields remained unchanged in the database. Root causes: - /key/update: prepare_key_update_data() popped budget_duration from the update dict but never re-added it (or budget_reset_at) when the value was None. - /team/update: _set_budget_reset_at() only acted when budget_duration was non-None, leaving a stale budget_reset_at in the DB. - /team/update: team_member_* null values bypassed the budget table update entirely because should_create_budget() requires at least one non-None field. --- .../key_management_endpoints.py | 11 +- .../management_endpoints/team_endpoints.py | 69 ++++- .../test_key_management_endpoints.py | 78 +++++ .../test_team_endpoints.py | 286 ++++++++++++++++++ 4 files changed, 433 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4439a55c1c6..128bc840271 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1736,11 +1736,10 @@ async def prepare_key_update_data( if "budget_duration" in non_default_values: budget_duration = non_default_values.pop("budget_duration") - if ( - budget_duration - and (isinstance(budget_duration, str)) - and len(budget_duration) > 0 - ): + if budget_duration is None: + non_default_values["budget_duration"] = None + non_default_values["budget_reset_at"] = None + elif isinstance(budget_duration, str) and len(budget_duration) > 0: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time key_reset_at = get_budget_reset_time(budget_duration=budget_duration) @@ -2388,7 +2387,7 @@ async def update_key_fn( # noqa: PLR0915 }, ) - data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: dict = data.model_dump(exclude_unset=True) key = data_json.pop("key") # get the row from db diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 259624f1e18..4d66ef9d6c5 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -312,6 +312,42 @@ class TeamMemberBudgetHandler: data_dict.pop("team_member_rpm_limit", None) data_dict.pop("team_member_tpm_limit", None) + @staticmethod + async def clear_team_member_budget_fields( + team_table: LiteLLM_TeamTable, + user_api_key_dict: "UserAPIKeyAuth", + updated_kv: dict, + explicitly_set_fields: set, + ) -> dict: + """Clear explicitly-nulled fields on the team member budget row.""" + from litellm.proxy._types import BudgetNewRequest + from litellm.proxy.management_endpoints.budget_management_endpoints import ( + update_budget, + ) + + if team_table.metadata is None: + team_table.metadata = {} + + team_member_budget_id = team_table.metadata.get("team_member_budget_id") + if team_member_budget_id is not None and isinstance(team_member_budget_id, str): + budget_request = BudgetNewRequest(budget_id=team_member_budget_id) + if "team_member_budget" in explicitly_set_fields: + budget_request.max_budget = None + if "team_member_budget_duration" in explicitly_set_fields: + budget_request.budget_duration = None + budget_request.budget_reset_at = None + if "team_member_rpm_limit" in explicitly_set_fields: + budget_request.rpm_limit = None + if "team_member_tpm_limit" in explicitly_set_fields: + budget_request.tpm_limit = None + await update_budget( + budget_obj=budget_request, + user_api_key_dict=user_api_key_dict, + ) + + TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) + return updated_kv + @staticmethod async def backfill_team_member_budget_entries( team_id: str, @@ -1766,11 +1802,25 @@ async def update_team( # noqa: PLR0915 # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) - if TeamMemberBudgetHandler.should_create_budget( - team_member_budget=data.team_member_budget, - team_member_rpm_limit=data.team_member_rpm_limit, - team_member_tpm_limit=data.team_member_tpm_limit, - team_member_budget_duration=data.team_member_budget_duration, + _team_member_fields_in_request = { + field + for field in [ + "team_member_budget", + "team_member_rpm_limit", + "team_member_tpm_limit", + "team_member_budget_duration", + ] + if field in updated_kv + } + + if ( + _team_member_fields_in_request + and TeamMemberBudgetHandler.should_create_budget( + team_member_budget=data.team_member_budget, + team_member_rpm_limit=data.team_member_rpm_limit, + team_member_tpm_limit=data.team_member_tpm_limit, + team_member_budget_duration=data.team_member_budget_duration, + ) ): updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=existing_team_row, @@ -1793,6 +1843,13 @@ async def update_team( # noqa: PLR0915 team_member_budget_id=_backfill_budget_id, prisma_client=prisma_client, ) + elif _team_member_fields_in_request: + updated_kv = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=existing_team_row, + user_api_key_dict=user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields=_team_member_fields_in_request, + ) else: TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) @@ -1873,6 +1930,8 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: reset_at = get_budget_reset_time(budget_duration=data.budget_duration) updated_kv["budget_reset_at"] = reset_at + elif "budget_duration" in updated_kv and updated_kv["budget_duration"] is None: + updated_kv["budget_reset_at"] = None if data.budget_limits is not None and len(data.budget_limits) > 0: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 66716400c4f..7f529b7e483 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -10596,3 +10596,81 @@ async def test_bulk_update_team_keys_blocks_metadata_allowed_passthrough_routes( assert exc.value.status_code == 403 assert "allowed_passthrough_routes" in str(exc.value.detail) mock.update_data.assert_not_called() + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_null_clears_fields(): + """ + When budget_duration is explicitly set to null, prepare_key_update_data + should produce budget_duration=None and budget_reset_at=None so Prisma + clears them in the DB. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration=None) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" in result + assert result["budget_duration"] is None + assert "budget_reset_at" in result + assert result["budget_reset_at"] is None + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_not_sent_excluded(): + """ + When budget_duration is NOT sent in the request (unset), it should not + appear in the result dict at all — the existing DB value stays unchanged. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"]) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" not in result + assert "budget_reset_at" not in result + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_valid_sets_reset(): + """ + When budget_duration is set to a valid duration string, both + budget_duration and budget_reset_at should be populated. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration="30d") + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert result["budget_duration"] == "30d" + assert result["budget_reset_at"] is not None 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 e668672dd2a..c7968799ffe 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -7828,3 +7828,289 @@ async def test_team_member_me_returns_404_for_unknown_team(mock_db_client): user_api_key_dict=caller_auth, ) assert exc_info.value.status_code == 404 + + +def test_set_budget_reset_at_clears_when_budget_duration_null(): + """ + When budget_duration is explicitly set to null, _set_budget_reset_at + should set budget_reset_at=None in updated_kv so Prisma clears it in the DB. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team", budget_duration=None) + updated_kv = {"team_id": "test-team", "budget_duration": None} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" in updated_kv + assert updated_kv["budget_reset_at"] is None + + +def test_set_budget_reset_at_noop_when_budget_duration_not_sent(): + """ + When budget_duration is NOT sent (unset), _set_budget_reset_at should + not add budget_reset_at to updated_kv. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team") + updated_kv = {"team_id": "test-team"} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" not in updated_kv + + +def test_set_budget_reset_at_sets_value_when_budget_duration_provided(): + """ + When budget_duration is set to a valid string, _set_budget_reset_at + should compute and set budget_reset_at. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team", budget_duration="30d") + updated_kv = {"team_id": "test-team", "budget_duration": "30d"} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" in updated_kv + assert updated_kv["budget_reset_at"] is not None + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_duration_calls_update_budget(): + """ + When team_member_budget_duration is explicitly null and a budget row + exists, clear_team_member_budget_fields should call update_budget + with budget_duration=None and budget_reset_at=None. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-123"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget_duration": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget_duration"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-123" + assert "budget_duration" in budget_request.model_fields_set + assert budget_request.budget_duration is None + assert "budget_reset_at" in budget_request.model_fields_set + assert budget_request.budget_reset_at is None + assert "team_member_budget_duration" not in result + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_clears_max_budget(): + """ + When team_member_budget is explicitly null, clear_team_member_budget_fields + should call update_budget with max_budget=None. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-456"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-456" + assert "max_budget" in budget_request.model_fields_set + assert budget_request.max_budget is None + assert "team_member_budget" not in result + + +@pytest.mark.asyncio +async def test_clear_team_member_rpm_tpm_limits(): + """ + When team_member_rpm_limit and team_member_tpm_limit are explicitly null, + clear_team_member_budget_fields should clear both on the budget row. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-789"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_rpm_limit": None, + "team_member_tpm_limit": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_rpm_limit", "team_member_tpm_limit"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-789" + assert "rpm_limit" in budget_request.model_fields_set + assert budget_request.rpm_limit is None + assert "tpm_limit" in budget_request.model_fields_set + assert budget_request.tpm_limit is None + assert "team_member_rpm_limit" not in result + assert "team_member_tpm_limit" not in result + + +@pytest.mark.asyncio +async def test_clear_all_team_member_fields_at_once(): + """ + When all team_member fields are explicitly null, all corresponding + budget row fields should be cleared in a single update. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-all"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + "team_member_budget_duration": None, + "team_member_rpm_limit": None, + "team_member_tpm_limit": None, + } + + all_fields = { + "team_member_budget", + "team_member_budget_duration", + "team_member_rpm_limit", + "team_member_tpm_limit", + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields=all_fields, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-all" + assert budget_request.max_budget is None + assert budget_request.budget_duration is None + assert budget_request.budget_reset_at is None + assert budget_request.rpm_limit is None + assert budget_request.tpm_limit is None + for field in all_fields: + assert field not in result + + +@pytest.mark.asyncio +async def test_team_member_budget_duration_not_sent_does_not_update(): + """ + When team_member_budget_duration is NOT sent in the request, no budget + update should occur and the field should not appear in updated_kv. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + updated_kv = {"team_id": "test-team", "max_budget": 200} + + _team_member_fields_in_request = { + field + for field in [ + "team_member_budget", + "team_member_rpm_limit", + "team_member_tpm_limit", + "team_member_budget_duration", + ] + if field in updated_kv + } + + assert len(_team_member_fields_in_request) == 0 + + TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) + + assert "team_member_budget_duration" not in updated_kv + assert "team_member_budget" not in updated_kv