Merge bf16693b5e into 6fa0477d84

litellm/proxy/management_endpoints/common_utils.py

@@ -1,5 +1,6 @@
 import json
 import math
+from collections import Counter
 from collections.abc import Mapping
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Optional, Union
@@ -69,8 +70,12 @@ from litellm.repositories.team_repository import TeamRepository
 from litellm.types.utils import BudgetConfig
 
 if TYPE_CHECKING:
+    from prisma import Prisma
+    from prisma import models as prisma_models
+
     from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest
     from litellm.proxy.utils import PrismaClient, ProxyLogging
+    from litellm.repositories.prisma_protocols import TableActions
 
 _user_has_admin_view: Final = user_api_key_has_admin_view
 
@@ -625,6 +630,20 @@ def _has_meaningful_budget_limit(budget_values: Mapping[str, object]) -> bool:
     return any(_is_set_budget_value(budget_values.get(field)) for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS)
 
 
+def _membership_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamMembership]":
+    return tx.litellm_teammembership  # pyright: ignore[reportReturnType]  # TableActions widens the generated inputs to Mapping, as the repositories do
+
+
+async def budget_ids_shared_by_memberships(tx: "Prisma", budget_ids: frozenset[str]) -> frozenset[str]:
+    """The rows in ``budget_ids`` more than one team membership points at, counted across every
+    team so a row shared with another team is protected too."""
+    if not budget_ids:
+        return frozenset()
+    rows: Final = await _membership_tx_db(tx).find_many(where={"budget_id": {"in": sorted(budget_ids)}})
+    references: Final = Counter(row.budget_id for row in rows)
+    return frozenset(budget_id for budget_id in budget_ids if references[budget_id] > 1)
+
+
 async def upsert_budget_and_membership(
     tx,
     *,
@@ -651,8 +670,8 @@ async def upsert_budget_and_membership(
     mutate the shared default that every other member points at.
 
     ``shared_budget_ids`` extends that protection to any other row more than one
-    membership points at, which a caller patching several members at once has
-    already counted; a row listed there is cloned rather than written in place.
+    membership points at (see ``budget_ids_shared_by_memberships``); a row listed
+    there is cloned rather than written in place.
     A patch whose every field already matches the shared row changes nothing and
     is a no-op rather than a clone.
     """

litellm/proxy/management_endpoints/team_endpoints.py

@@ -164,6 +164,7 @@ from litellm.proxy.management_endpoints.common_utils import (  # noqa: F401  # l
     _update_metadata_fields,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     _upsert_budget_and_membership,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     _user_has_admin_view,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
+    budget_ids_shared_by_memberships,
     check_disable_global_guardrails_caller_permission,
     check_passthrough_routes_caller_permission,
     check_require_trace_id_caller_permission,
@@ -4210,9 +4211,6 @@ async def team_member_update(
             identified_budget_id = tm.budget_id
             break
 
-    # If this membership still points at the team's shared default member
-    # budget, _upsert_budget_and_membership will clone-on-write so that the
-    # update only touches this user (not every member sharing the default).
     team_default_budget_id: str | None = None
     if team_table.metadata is not None:
         raw_default_budget_id: Final = team_table.metadata.get("team_member_budget_id")
@@ -4242,6 +4240,9 @@ async def team_member_update(
             user_api_key_dict=user_api_key_dict,
             budget_patch=budget_patch,
             team_default_budget_id=team_default_budget_id,
+            shared_budget_ids=await budget_ids_shared_by_memberships(
+                tx, frozenset({identified_budget_id} if identified_budget_id is not None else ())
+            ),
         )
     if budget_patch:
         await invalidate_team_member_spend_state(

litellm/proxy/management_helpers/bulk_team_member_budgets.py

@@ -27,6 +27,7 @@ from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN
 from litellm.proxy.management.teams.dependencies import get_team_access
 from litellm.proxy.management_endpoints.common_utils import (  # noqa: F401  # legacy module exports
     _upsert_budget_and_membership,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
+    budget_ids_shared_by_memberships,
     member_budget_patch,
     upsert_budget_and_membership,  # pyright: ignore[reportPrivateUsage]  # the single-member write, shared so the two surfaces cannot drift
 )
@@ -35,7 +36,6 @@ from litellm.proxy.management_helpers.bulk_user_deletion import (
     _duplicate_member_indexes,  # pyright: ignore[reportPrivateUsage]  # same duplicate rule as members/bulk_delete
     _eq_filter,  # pyright: ignore[reportPrivateUsage]  # same prisma filter shape as members/bulk_delete
     _forbidden,  # pyright: ignore[reportPrivateUsage]  # same problem shape as members/bulk_delete
-    _in_filter,  # pyright: ignore[reportPrivateUsage]  # same prisma filter shape as members/bulk_delete
     _team_not_found,  # pyright: ignore[reportPrivateUsage]  # same problem shape as members/bulk_delete
     _team_users_filter,  # pyright: ignore[reportPrivateUsage]  # same prisma filter shape as members/bulk_delete
 )
@@ -79,15 +79,6 @@ def _team_default_budget_id(team: LiteLLM_TeamTable) -> str | None:
     return raw if isinstance(raw, str) else None
 
 
-async def _shared_budget_ids(tx: "Prisma", budget_ids: frozenset[str]) -> frozenset[str]:
-    """The rows in ``budget_ids`` more than one membership points at, counted across every
-    team so a row shared with another team is protected too."""
-    if not budget_ids:
-        return frozenset()
-    rows: Final = await _membership_tx_db(tx).find_many(where=_in_filter("budget_id", budget_ids))
-    return frozenset(budget_id for budget_id in budget_ids if sum(1 for row in rows if row.budget_id == budget_id) > 1)
-
-
 class _AuditedMemberBudget(LiteLLMBaseModel):
     """One member's limits as the audit log's before/after values record them."""
 
@@ -210,7 +201,7 @@ async def bulk_update_team_member_budgets(
     async with prisma_client.tx(timeout=_BATCH_TX_TIMEOUT) as tx:
         memberships: Final = await _membership_tx_db(tx).find_many(where=team_members_filter, include=_WITH_BUDGET)
         budget_id_of: Final = MappingProxyType({m.user_id: m.budget_id for m in memberships})
-        shared: Final = await _shared_budget_ids(
+        shared: Final = await budget_ids_shared_by_memberships(
             tx, frozenset(budget_id for budget_id in budget_id_of.values() if budget_id is not None)
         )
         for index, user_id in applied:

tests/unit/proxy/management_endpoints/test_team_endpoints.py

@@ -14,7 +14,7 @@ from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, stat
 import pytest
 from fastapi import HTTPException, Request
 from fastapi.testclient import TestClient
-from pydantic import ValidationError
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from litellm._uuid import uuid
 from litellm.integrations.custom_logger import CustomLogger
@@ -17196,6 +17196,147 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(
     assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5
 
 
+class _StoredMemberBudget(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="allow")
+
+    budget_id: str
+    max_budget: float | None = None
+    tpm_limit: int | None = None
+    budget_duration: str | None = None
+    budget_reset_at: datetime | None = None
+    temp_budget_increase: float | None = None
+    temp_budget_expiry: datetime | None = None
+    allowed_models: list[str] = Field(default_factory=list)
+
+
+class _MemberBudgetTable:
+    def __init__(self, rows: dict[str, _StoredMemberBudget]) -> None:
+        self.rows: Final = rows
+
+    async def find_unique(self, where: Mapping[str, str]) -> _StoredMemberBudget | None:
+        return self.rows.get(where["budget_id"])
+
+    async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> _StoredMemberBudget:
+        updated: Final = self.rows[where["budget_id"]].model_copy(update=dict(data))
+        self.rows[where["budget_id"]] = updated
+        return updated
+
+    async def create(self, data: Mapping[str, object], include: object = None) -> _StoredMemberBudget:
+        created: Final = _StoredMemberBudget.model_validate({**data, "budget_id": f"created-{len(self.rows)}"})
+        self.rows[created.budget_id] = created
+        return created
+
+
+class _MembershipBudgetLinks:
+    def __init__(self, team_id: str, budget_id_of: dict[str, str | None]) -> None:
+        self.team_id: Final = team_id
+        self.budget_id_of: Final = budget_id_of
+
+    def rows(self) -> list[LiteLLM_TeamMembership]:
+        return [
+            LiteLLM_TeamMembership(user_id=user_id, team_id=self.team_id, budget_id=budget_id)
+            for user_id, budget_id in self.budget_id_of.items()
+        ]
+
+    async def find_many(self, where: Mapping[str, Mapping[str, Sequence[str]]]) -> list[LiteLLM_TeamMembership]:
+        wanted: Final = frozenset(where["budget_id"]["in"])
+        return [row for row in self.rows() if row.budget_id in wanted]
+
+    async def update(self, where: Mapping[str, Mapping[str, str]], data: Mapping[str, object]) -> None:
+        self.budget_id_of[where["user_id_team_id"]["user_id"]] = None
+
+    async def upsert(self, where: Mapping[str, Mapping[str, str]], data: Mapping[str, Mapping[str, object]]) -> None:
+        link: Final = data["update"]["litellm_budget_table"]
+        assert isinstance(link, Mapping)
+        self.budget_id_of[where["user_id_team_id"]["user_id"]] = link["connect"]["budget_id"]
+
+
+async def _update_member_budget(
+    monkeypatch: pytest.MonkeyPatch,
+    budgets: _MemberBudgetTable,
+    memberships: _MembershipBudgetLinks,
+    request: TeamMemberUpdateRequest,
+) -> None:
+    team_row: Final = LiteLLM_TeamTable(team_id=memberships.team_id, metadata={}, members_with_roles=[])
+    tx: Final = SimpleNamespace(litellm_budgettable=budgets, litellm_teammembership=memberships)
+
+    @asynccontextmanager
+    async def _tx():
+        yield tx
+
+    prisma_client: Final = SimpleNamespace(
+        db=SimpleNamespace(litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team_row))),
+        tx=_tx,
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
+    with patch(  # test-quality-ok: team_info is a separate endpoint; the membership rows it returns come from the same in-memory tables
+        "litellm.proxy.management_endpoints.team_endpoints.team_info",
+        AsyncMock(return_value={"team_info": team_row, "team_memberships": memberships.rows()}),
+    ):
+        await team_member_update(
+            data=request,
+            http_request=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user"),
+        )
+
+
+_SHARED_RESET_AT: Final = datetime(2030, 1, 1, tzinfo=timezone.utc)
+_TEMP_EXPIRY: Final = datetime(2030, 2, 1, tzinfo=timezone.utc)
+
+
+@pytest.mark.asyncio
+async def test_team_member_update_on_a_shared_non_default_budget_row_gives_only_that_member_new_limits(monkeypatch):
+    shared: Final = _StoredMemberBudget(
+        budget_id="shared-b", max_budget=100.0, tpm_limit=900, budget_duration="30d", budget_reset_at=_SHARED_RESET_AT
+    )
+    budgets: Final = _MemberBudgetTable({"shared-b": shared})
+    memberships: Final = _MembershipBudgetLinks("team-1", {"member-1": "shared-b", "member-2": "shared-b"})
+
+    await _update_member_budget(
+        monkeypatch,
+        budgets,
+        memberships,
+        TeamMemberUpdateRequest(
+            team_id="team-1",
+            user_id="member-1",
+            max_budget_in_team=5.0,
+            temp_budget_increase=20.0,
+            temp_budget_expiry=_TEMP_EXPIRY,
+        ),
+    )
+
+    assert memberships.budget_id_of["member-2"] == "shared-b"
+    assert budgets.rows["shared-b"] == shared
+    edited_budget_id: Final = memberships.budget_id_of["member-1"]
+    assert edited_budget_id not in (None, "shared-b")
+    edited: Final = budgets.rows[edited_budget_id]
+    assert (edited.max_budget, edited.temp_budget_increase, edited.temp_budget_expiry) == (5.0, 20.0, _TEMP_EXPIRY)
+    assert (edited.tpm_limit, edited.budget_duration, edited.budget_reset_at) == (900, "30d", _SHARED_RESET_AT)
+
+
+@pytest.mark.asyncio
+async def test_team_member_update_on_a_budget_row_only_that_member_points_at_updates_it_in_place(monkeypatch):
+    budgets: Final = _MemberBudgetTable(
+        {
+            "own-b": _StoredMemberBudget(budget_id="own-b", max_budget=100.0, tpm_limit=900),
+            "other-b": _StoredMemberBudget(budget_id="other-b", max_budget=100.0),
+        }
+    )
+    memberships: Final = _MembershipBudgetLinks("team-1", {"member-1": "own-b", "member-2": "other-b"})
+
+    await _update_member_budget(
+        monkeypatch,
+        budgets,
+        memberships,
+        TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=5.0),
+    )
+
+    assert memberships.budget_id_of == {"member-1": "own-b", "member-2": "other-b"}
+    assert set(budgets.rows) == {"own-b", "other-b"}
+    assert (budgets.rows["own-b"].max_budget, budgets.rows["own-b"].tpm_limit) == (5.0, 900)
+    assert budgets.rows["other-b"].max_budget == 100.0
+
+
 @pytest.mark.asyncio
 async def test_evict_created_membership_caches_drops_the_negative_sentinel():
     """

tests/unit/proxy/test_team_member_update.py

@@ -69,6 +69,7 @@ def happy_path_upsert(monkeypatch):
 
     class _FakeTx:
         litellm_teamtable = prisma_client.db.litellm_teamtable
+        litellm_teammembership = types.SimpleNamespace(find_many=AsyncMock(return_value=[]))
 
         async def __aenter__(self):
             return self