fix #25506

litellm/proxy/management_endpoints/team_endpoints.py

@@ -254,6 +254,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[Member],
+        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),
+                team_id,
+                team_member_budget_id,
+            )
+
 
 def _get_default_team_param(field: str) -> Any:
     """
@@ -1514,6 +1569,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)
 

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -1749,6 +1749,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():
     """