fix(proxy): add team_member_budget_duration to NewTeamRequest

NewTeamRequest was missing the team_member_budget_duration field,
causing Pydantic to silently drop the value when creating a team via
POST /team/new. The template budget row was created without
budget_duration or budget_reset_at, so the ResetBudgetJob never
found it and team member spend was never reset.

Add the field to NewTeamRequest and pass it through to
should_create_budget and create_team_member_budget_table in the
new_team handler (matching the existing update_team path which
already works correctly).

Fixes #16057
This commit is contained in:
michelligabriele 2026-03-12 22:30:52 +01:00
parent 89d8401d72
commit 78927138d2
3 changed files with 53 additions and 0 deletions

View file

@ -1669,6 +1669,7 @@ class NewTeamRequest(TeamBase):
int
] = None # allow user to set TPM limit for all team members
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo"
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
enforced_batch_output_expires_after: Optional[dict] = None
enforced_file_expires_after: Optional[dict] = None

View file

@ -898,6 +898,7 @@ async def new_team( # noqa: PLR0915
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,
):
data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
data=data,
@ -906,6 +907,7 @@ async def new_team( # noqa: PLR0915
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,
)
## ADD TO TEAM TABLE

View file

@ -6157,3 +6157,53 @@ async def test_list_team_v1_batches_key_queries():
assert result[0].keys == [key1, key2]
assert result[1].team_id == "team-2"
assert result[1].keys == [key3]
def test_new_team_request_accepts_team_member_budget_duration():
"""Test that NewTeamRequest does not silently drop team_member_budget_duration."""
from litellm.proxy._types import NewTeamRequest
request = NewTeamRequest(
team_member_budget=20.0,
team_member_budget_duration="30d",
)
assert request.team_member_budget == 20.0
assert request.team_member_budget_duration == "30d"
@pytest.mark.asyncio
async def test_create_team_member_budget_table_with_duration():
"""Verify that create_team_member_budget_table passes budget_duration
through to the new_budget call when team_member_budget_duration is provided."""
from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
mock_budget_response = MagicMock(budget_id="budget-abc")
mock_admin = UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
data = NewTeamRequest(
team_alias="test-team",
team_member_budget=20.0,
team_member_budget_duration="30d",
)
with patch(
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
new_callable=AsyncMock,
return_value=mock_budget_response,
) as mock_new_budget:
result = await TeamMemberBudgetHandler.create_team_member_budget_table(
data=data,
new_team_data_json={"metadata": None},
user_api_key_dict=mock_admin,
team_member_budget=20.0,
team_member_budget_duration="30d",
)
mock_new_budget.assert_awaited_once()
budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
assert budget_request.budget_duration == "30d"
assert budget_request.max_budget == 20.0
assert result["metadata"]["team_member_budget_id"] == "budget-abc"