Merge pull request #23484 from michelligabriele/fix/team-member-budget-duration-on-create

fix(proxy): add team_member_budget_duration to NewTeamRequest

litellm/proxy/_types.py

@@ -1672,6 +1672,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

litellm/proxy/management_endpoints/team_endpoints.py

@@ -934,6 +934,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,
@@ -942,6 +943,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

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -6441,3 +6441,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"