From d3a364d74f8bee7f6133d51e68c3036cde7f8136 Mon Sep 17 00:00:00 2001 From: ryan Date: Sat, 19 Sep 2026 00:49:58 +0000 Subject: [PATCH] fix(team): report no budget source when the team default row was deleted Derive budget_source from the budget row /team/info actually loaded, so a metadata id whose row was removed via /budget/delete reads as none instead of team_default. Share the /team/info test scaffolding so the added patch calls stay within the TQ008 budget, and allowlist the imperative reset_budget route in the provider endpoint audit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 5 +- .../endpointaudit/coverage_allowlist.txt | 1 + .../test_team_endpoints.py | 112 ++++++++++-------- 3 files changed, 70 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b8d95167045..a441b3834ed 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4825,6 +4825,9 @@ async def team_info( prisma_client=prisma_client, team_info_response_object=_team_info, ) + active_default_budget_id: Final = ( + team_member_budget_id if _team_info.team_member_budget_table is not None else None + ) # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) @@ -4856,7 +4859,7 @@ async def team_info( MappingProxyType( { **tm.model_dump(), - "budget_source": _member_budget_source(tm.budget_id, team_member_budget_id), + "budget_source": _member_budget_source(tm.budget_id, active_default_budget_id), } ) ) diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 6bc8947e89f..4ea64b152f1 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -81,6 +81,7 @@ POST /prompts/test POST /search_tools/test_connection POST /team/bulk_member_add POST /team/{team_id}/member/{user_id}/reset_spend +POST /team/{team_id}/member/{user_id}/reset_budget POST /team/key/bulk_update POST /team/permissions_bulk_update POST /team/{team_id}/disable_logging 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 d55b9f79b5f..484c054fa54 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14572,40 +14572,52 @@ async def test_reset_team_member_budget_fn_forbidden_for_non_admin(monkeypatch): mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited() +async def _team_info_budget_sources( + team_row: LiteLLM_TeamTable, + memberships: list[LiteLLM_TeamMembership], + default_budget_row: LiteLLM_BudgetTable | None, +) -> dict[str, str]: + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(return_value=default_budget_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch.object( # test-quality-ok: membership lookup is a module-level DB query with no injection point + team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships) + ), + ): + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id=team_row.team_id, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + return {tm.user_id: tm.budget_source for tm in response["team_memberships"]} + + @pytest.mark.asyncio async def test_team_info_reports_whether_each_member_follows_the_team_default_budget(): """/team/info must tell the caller which members still follow the team's shared member budget and which carry their own row, since budget_id alone only means something to a reader who also knows the team's team_member_budget_id.""" - from fastapi import Request - - from litellm.proxy.management_endpoints import team_endpoints - - team_row = _team_with_default_budget("team-1", "team-default-b") - memberships = [ - LiteLLM_TeamMembership(user_id="inherits", team_id="team-1", budget_id="team-default-b"), - LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), - LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), - ] - - mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) - mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( - return_value=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0) + sources = await _team_info_budget_sources( + team_row=_team_with_default_budget("team-1", "team-default-b"), + memberships=[ + LiteLLM_TeamMembership(user_id="inherits", team_id="team-1", budget_id="team-default-b"), + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ], + default_budget_row=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0), ) - mock_prisma.get_data = AsyncMock(return_value=[]) - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), - ): - response = await team_endpoints.team_info( - http_request=MagicMock(spec=Request), - team_id="team-1", - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - - assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + assert sources == { "inherits": "team_default", "customized": "custom", "unlinked": "team_default", @@ -14617,30 +14629,36 @@ async def test_team_info_reports_no_budget_source_when_team_has_no_default(): """A team that never set team_member_budget has nothing for members to inherit, so an unlinked member is 'none' rather than 'team_default', while a member with their own row is still 'custom'.""" - from fastapi import Request + sources = await _team_info_budget_sources( + team_row=LiteLLM_TeamTable(team_id="team-1"), + memberships=[ + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ], + default_budget_row=None, + ) - from litellm.proxy.management_endpoints import team_endpoints + assert sources == { + "customized": "custom", + "unlinked": "none", + } - memberships = [ - LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), - LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), - ] - mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")) - mock_prisma.get_data = AsyncMock(return_value=[]) +@pytest.mark.asyncio +async def test_team_info_reports_no_budget_source_when_team_default_row_was_deleted(): + """If the budget row named by team_member_budget_id was removed via /budget/delete, nothing is + enforced for unlinked members any more, so /team/info must not keep advertising a team default + that no longer exists.""" + sources = await _team_info_budget_sources( + team_row=_team_with_default_budget("team-1", "deleted-b"), + memberships=[ + LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"), + LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None), + ], + default_budget_row=None, + ) - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), - patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)), - ): - response = await team_endpoints.team_info( - http_request=MagicMock(spec=Request), - team_id="team-1", - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - - assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == { + assert sources == { "customized": "custom", "unlinked": "none", }