diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fcc224e848..138469312e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -112,6 +112,14 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( router = APIRouter() +def _sanitize_for_log(value: Any) -> str: + """Strip CR/LF from user-controlled values to prevent log injection.""" + try: + text = str(value) + except Exception: + text = repr(value) + return text.replace("\r", "").replace("\n", "") + async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, @@ -285,6 +293,61 @@ class TeamMemberBudgetHandler: data_dict.pop("team_member_rpm_limit", None) data_dict.pop("team_member_tpm_limit", None) + @staticmethod + async def backfill_team_member_budget_entries( + team_id: str, + members_with_roles: List[Union[Member, dict]], + team_member_budget_id: str, + prisma_client: PrismaClient, + ) -> None: + """ + Create team_memberships entries for existing members that don't have one. + + Called after team_member_budget is set/updated on a team to ensure + members who joined before the budget was configured also get budget + enforcement. + + Only creates missing entries — does not touch existing memberships + (which may carry individual per-member budgets). + """ + if not members_with_roles: + return + + # Batch-fetch existing memberships for this team (avoids N+1 queries) + existing_memberships = ( + await prisma_client.db.litellm_teammembership.find_many( + where={"team_id": team_id} + ) + ) + existing_user_ids = {m.user_id for m in existing_memberships} + + # Identify members with no existing membership row. + # members_with_roles may contain Member instances or raw dicts depending + # on how the team was fetched/deserialized. + missing = [] + for m in members_with_roles: + user_id = m.get("user_id") if isinstance(m, dict) else m.user_id + if user_id is not None and user_id not in existing_user_ids: + missing.append( + { + "team_id": team_id, + "user_id": user_id, + "budget_id": team_member_budget_id, + } + ) + + if missing: + await prisma_client.db.litellm_teammembership.create_many( + data=missing, + skip_duplicates=True, # safety net against concurrent races + ) + verbose_proxy_logger.info( + "Backfilled %d team_memberships for team %s with budget %s", + len(missing), + _sanitize_for_log(team_id), + _sanitize_for_log(team_member_budget_id), + ) + def _get_default_team_param(field: str) -> Any: """ @@ -1551,6 +1614,18 @@ async def update_team( # noqa: PLR0915 team_member_tpm_limit=data.team_member_tpm_limit, team_member_budget_duration=data.team_member_budget_duration, ) + # Backfill team_memberships for members who joined before the + # budget was configured — they won't have a membership row yet. + _backfill_budget_id = (updated_kv.get("metadata") or {}).get( + "team_member_budget_id" + ) + if _backfill_budget_id and existing_team_row.members_with_roles: + await TeamMemberBudgetHandler.backfill_team_member_budget_entries( + team_id=data.team_id, + members_with_roles=existing_team_row.members_with_roles, + team_member_budget_id=_backfill_budget_id, + prisma_client=prisma_client, + ) else: TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) 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 20c3e3c0b5..bee6642dec 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1766,6 +1766,143 @@ async def test_update_team_with_team_member_budget_duration(): assert "team_member_budget_duration" not in update_data +@pytest.mark.asyncio +async def test_backfill_team_member_budget_entries_creates_missing_memberships(): + """ + When backfill_team_member_budget_entries is called, it should create + team_memberships rows only for members that don't already have one. + + Regression test for: https://github.com/BerriAI/litellm/issues/25506 + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import Member + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + + team_id = "team-abc" + budget_id = "budget-xyz" + + # user-A already has a membership; user-B does not + existing_membership = MagicMock() + existing_membership.user_id = "user-A" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teammembership.find_many = AsyncMock( + return_value=[existing_membership] + ) + mock_prisma.db.litellm_teammembership.create_many = AsyncMock(return_value=None) + + # Test with Member instances + members = [ + Member(user_id="user-A", role="user"), + Member(user_id="user-B", role="user"), + ] + + await TeamMemberBudgetHandler.backfill_team_member_budget_entries( + team_id=team_id, + members_with_roles=members, + team_member_budget_id=budget_id, + prisma_client=mock_prisma, + ) + + # find_many should have been called to fetch existing memberships + mock_prisma.db.litellm_teammembership.find_many.assert_awaited_once_with( + where={"team_id": team_id} + ) + + # create_many should only create an entry for user-B (user-A already has one) + mock_prisma.db.litellm_teammembership.create_many.assert_awaited_once_with( + data=[{"team_id": team_id, "user_id": "user-B", "budget_id": budget_id}], + skip_duplicates=True, + ) + + # Also test with raw dicts (members_with_roles may be dicts when deserialized from DB) + mock_prisma.db.litellm_teammembership.find_many.reset_mock() + mock_prisma.db.litellm_teammembership.create_many.reset_mock() + + members_as_dicts = [ + {"user_id": "user-A", "role": "user"}, + {"user_id": "user-B", "role": "user"}, + ] + + await TeamMemberBudgetHandler.backfill_team_member_budget_entries( + team_id=team_id, + members_with_roles=members_as_dicts, + team_member_budget_id=budget_id, + prisma_client=mock_prisma, + ) + + mock_prisma.db.litellm_teammembership.create_many.assert_awaited_once_with( + data=[{"team_id": team_id, "user_id": "user-B", "budget_id": budget_id}], + skip_duplicates=True, + ) + + +@pytest.mark.asyncio +async def test_backfill_team_member_budget_entries_no_op_when_all_exist(): + """ + backfill_team_member_budget_entries should not call create_many when all + members already have a team_memberships entry. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import Member + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + + team_id = "team-abc" + budget_id = "budget-xyz" + + existing_a = MagicMock() + existing_a.user_id = "user-A" + existing_b = MagicMock() + existing_b.user_id = "user-B" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teammembership.find_many = AsyncMock( + return_value=[existing_a, existing_b] + ) + mock_prisma.db.litellm_teammembership.create_many = AsyncMock(return_value=None) + + members = [ + Member(user_id="user-A", role="user"), + Member(user_id="user-B", role="user"), + ] + + await TeamMemberBudgetHandler.backfill_team_member_budget_entries( + team_id=team_id, + members_with_roles=members, + team_member_budget_id=budget_id, + prisma_client=mock_prisma, + ) + + mock_prisma.db.litellm_teammembership.create_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_backfill_team_member_budget_entries_empty_members(): + """ + backfill_team_member_budget_entries should be a no-op when the member list + is empty (no DB queries at all). + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_teammembership.create_many = AsyncMock(return_value=None) + + await TeamMemberBudgetHandler.backfill_team_member_budget_entries( + team_id="team-abc", + members_with_roles=[], + team_member_budget_id="budget-xyz", + prisma_client=mock_prisma, + ) + + mock_prisma.db.litellm_teammembership.find_many.assert_not_awaited() + mock_prisma.db.litellm_teammembership.create_many.assert_not_awaited() + + @pytest.mark.asyncio async def test_bulk_team_member_add_success(): """