From bf16693b5ec360a0482f83c357491f668ef8201b Mon Sep 17 00:00:00 2001 From: Jesus Velarde Date: Tue, 3 Nov 2026 16:30:43 -0800 Subject: [PATCH] fix(proxy): clone shared team member budget rows on single member update --- .../management_endpoints/common_utils.py | 23 ++- .../management_endpoints/team_endpoints.py | 7 +- .../bulk_team_member_budgets.py | 13 +- .../test_team_endpoints.py | 143 +++++++++++++++++- tests/unit/proxy/test_team_member_update.py | 1 + 5 files changed, 170 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 983c3d24513..308990d89f1 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/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. """ diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 54b9b26123e..106d157f9ff 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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, @@ -4209,9 +4210,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") @@ -4241,6 +4239,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( diff --git a/litellm/proxy/management_helpers/bulk_team_member_budgets.py b/litellm/proxy/management_helpers/bulk_team_member_budgets.py index 08169eab761..b27e74627a9 100644 --- a/litellm/proxy/management_helpers/bulk_team_member_budgets.py +++ b/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: diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 697d9810284..725a59e164d 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/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(): """ diff --git a/tests/unit/proxy/test_team_member_update.py b/tests/unit/proxy/test_team_member_update.py index d77904fc05d..0e0b33cf470 100644 --- a/tests/unit/proxy/test_team_member_update.py +++ b/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