From 5e8247a1c004458e5ffd3db49dbc4a0130ee3768 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 17:19:06 +0000 Subject: [PATCH 01/10] fix(team): emit audit events for member_delete and role changes and carry the final roster on team create Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 90 +++++- .../test_team_endpoints.py | 304 +++++++++++++++++- 2 files changed, 379 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 28c12173ea7..8e5d62976fa 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1784,7 +1784,10 @@ async def new_team( ) if is_audit_logging_enabled(): - _updated_values = complete_team_data.json(exclude_none=True) + created_team_snapshot: Final = complete_team_data.model_copy( + update={"members_with_roles": list(team_row.members_with_roles)} + ) + _updated_values = created_team_snapshot.json(exclude_none=True) _updated_values = json.dumps(_updated_values, default=str) @@ -3160,6 +3163,27 @@ def _members_audit_value(members: Sequence[Member]) -> str: ) +async def _create_team_membership_audit_log( + team_id: str, + before_members: Sequence[Member], + after_members: Sequence[Member], + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, +) -> None: + from litellm.proxy.management_helpers.audit_logs import create_object_audit_log + + await create_object_audit_log( + object_id=team_id, + action="updated", + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + before_value=_members_audit_value(before_members), + after_value=_members_audit_value(after_members), + ) + + async def _create_team_member_add_audit_logs( team_id: str, updated_users: Sequence[LiteLLM_UserTable], @@ -3191,15 +3215,12 @@ async def _create_team_member_add_audit_logs( if user.user_id is not None and user.user_id not in existing_user_ids ) - membership_entry: Final = create_object_audit_log( - object_id=team_id, - action="updated", - litellm_changed_by=None, + membership_entry: Final = _create_team_membership_audit_log( + team_id=team_id, + before_members=before_members, + after_members=after_members, user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, - table_name=LitellmTableNames.TEAM_TABLE_NAME, - before_value=_members_audit_value(before_members), - after_value=_members_audit_value(after_members), ) await asyncio.gather(*created_user_entries, membership_entry) @@ -3508,7 +3529,33 @@ async def team_member_delete( }' ``` """ - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + existing_team_row, before_members, after_members = await _team_member_delete( + data=data, user_api_key_dict=user_api_key_dict + ) + + if before_members != after_members: + await _create_team_membership_audit_log( + team_id=existing_team_row.team_id, + before_members=before_members, + after_members=after_members, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + + return existing_team_row + + +async def _team_member_delete( + data: TeamMemberDeleteRequest, + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[LiteLLM_TeamTable, tuple[Member, ...], tuple[Member, ...]]: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3672,7 +3719,7 @@ async def team_member_delete( _emit_team_members_metric(existing_team_row) - return existing_team_row + return existing_team_row, tuple(fresh_members), tuple(new_team_members) @router.post( @@ -3692,7 +3739,12 @@ async def team_member_update( Update team member budgets and team member role """ - from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + premium_user, + prisma_client, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3800,8 +3852,12 @@ async def team_member_update( ### update team member role if data.role is not None: + members_before_role_update: Final = tuple( + Member(user_id=member.user_id, user_email=member.user_email, role=member.role) + for member in team_table.members_with_roles + ) team_members: Final[list[Member]] = [] - for member in team_table.members_with_roles: + for member in members_before_role_update: if member.user_id == received_user_id: team_members.append( Member( @@ -3820,6 +3876,14 @@ async def team_member_update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, ) + if members_before_role_update != tuple(team_members): + await _create_team_membership_audit_log( + team_id=data.team_id, + before_members=members_before_role_update, + after_members=team_members, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) return TeamMemberUpdateResponse( team_id=data.team_id, @@ -4298,7 +4362,7 @@ async def delete_team( tasks = [] for team_member in team_members: tasks.append( - team_member_delete( + _team_member_delete( data=TeamMemberDeleteRequest( team_id=team_row.team_id, user_id=team_member.user_id, 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 690b5ae80b6..72c854f1d6c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -12,6 +12,7 @@ from fastapi.testclient import TestClient from pydantic import ValidationError from litellm._uuid import uuid +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, @@ -23,11 +24,14 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, + LitellmTableNames, LitellmUserRoles, Member, ProxyErrorTypes, ProxyException, ResetSpendRequest, + TeamInfoMember, + TeamInfoResponseObjectTeamTable, TeamMemberAddRequest, TeamMemberUpdateRequest, UpdateTeamRequest, @@ -68,6 +72,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkTeamMemberAddResponse, TeamMemberAddResult, ) +from litellm.types.utils import StandardAuditLogPayload from tests.test_litellm.proxy.management_endpoints.jwt_key_mapping_doubles import ( CascadingJWTMappingTable, JWTMappingRow, @@ -8696,8 +8701,8 @@ async def test_delete_team_persists_deleted_teams( "admin", ) monkeypatch.setattr( - "litellm.proxy.management_endpoints.team_endpoints.team_member_delete", - AsyncMock(return_value=team1), + "litellm.proxy.management_endpoints.team_endpoints._team_member_delete", + AsyncMock(return_value=(team1, (), ())), ) data = DeleteTeamRequest(team_ids=["team-1"]) @@ -13316,6 +13321,301 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"] +class _RecordingAuditLogger(CustomLogger): + """An audit_log_callbacks sink that keeps every payload it is handed.""" + + def __init__(self) -> None: + super().__init__() + self.payloads: list[StandardAuditLogPayload] = [] + + async def async_log_audit_log_event(self, audit_log_payload: StandardAuditLogPayload) -> None: + self.payloads.append(audit_log_payload) + + +def _wire_audit_log_callback(monkeypatch: pytest.MonkeyPatch) -> _RecordingAuditLogger: + """Turn audit logging on and register one recording callback, the way an operator's + `litellm_settings.audit_log_callbacks` entry would be.""" + audit_logger = _RecordingAuditLogger() + monkeypatch.setattr("litellm.store_audit_logs", True) + monkeypatch.setattr("litellm.audit_log_callbacks", [audit_logger]) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + return audit_logger + + +async def _settle_audit_log_tasks() -> None: + """Audit callbacks run on `asyncio.create_task`, so give the loop a few turns.""" + for _ in range(5): + await asyncio.sleep(0) + + +def _team_roster_events(audit_logger: _RecordingAuditLogger, action: str) -> list[StandardAuditLogPayload]: + return [ + p for p in audit_logger.payloads if p["table_name"] == LitellmTableNames.TEAM_TABLE_NAME and p["action"] == action + ] + + +def _roster_user_roles(members_json: str | None) -> dict[str, str]: + assert members_json is not None + return {m["user_id"]: m["role"] for m in json.loads(members_json)["members_with_roles"]} + + +@pytest.mark.asyncio +async def test_new_team_created_audit_event_carries_the_final_roster(monkeypatch): + """The `created` event a `/team/new` hands to audit_log_callbacks must list the members + the team was created with. The team row is inserted empty and the members attached + afterwards, so a snapshot taken from the pre-insert object reports no members and a + downstream consumer syncing membership from the event has nothing to sync.""" + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + audit_logger = _wire_audit_log_callback(monkeypatch) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_prisma.get_data = AsyncMock(return_value=None) + mock_prisma.update_data = AsyncMock() + created_team = MagicMock() + created_team.team_id = "team-audit-roster" + created_team.members_with_roles = [] + created_team.metadata = None + created_team.default_team_member_models = None + created_team.model_dump.return_value = {"team_id": "team-audit-roster", "members_with_roles": []} + mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=created_team) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=created_team) + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) + user_row = MagicMock() + user_row.user_id = "alice" + user_row.model_dump.return_value = {"user_id": "alice", "teams": ["team-audit-roster"]} + mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=user_row) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=user_row) + membership_row = MagicMock() + membership_row.model_dump.return_value = {"team_id": "team-audit-roster", "user_id": "alice", "budget_id": None} + mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=membership_row) + mock_prisma.db.litellm_auditlog.create = AsyncMock() + _wire_team_create_tx(mock_prisma) + + mock_license = MagicMock() + mock_license.is_team_count_over_limit.return_value = False + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check", mock_license) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await new_team( + data=NewTeamRequest( + team_id="team-audit-roster", + team_alias="audit-roster", + members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")], + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1", api_key="sk-a"), + ) + await _settle_audit_log_tasks() + + created_events = _team_roster_events(audit_logger, "created") + assert [e["object_id"] for e in created_events] == ["team-audit-roster"] + assert _roster_user_roles(created_events[0]["updated_values"]) == { + "admin-1": "admin", + "alice": "admin", + "bob": "user", + } + + +@pytest.mark.asyncio +async def test_team_member_delete_emits_a_roster_audit_event(monkeypatch, mock_db_client, mock_admin_auth): + """Removing a member must reach audit_log_callbacks as a TEAM_TABLE `updated` event whose + before and after rosters differ by exactly the removed user, the same shape + `/team/member_add` already emits, so one consumer can diff both directions.""" + from litellm.proxy._types import TeamMemberDeleteRequest + + audit_logger = _wire_audit_log_callback(monkeypatch) + + team_row = MagicMock() + team_row.model_dump.return_value = { + "team_id": "team-del-audit", + "members_with_roles": [ + {"user_id": "alice", "user_email": None, "role": "admin"}, + {"user_id": "bob", "user_email": None, "role": "user"}, + ], + "team_member_permissions": [], + "metadata": {}, + "models": [], + "spend": 0.0, + } + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + user_row = MagicMock() + user_row.user_id = "bob" + user_row.teams = ["team-del-audit"] + mock_db_client.db.litellm_usertable.find_many = AsyncMock(return_value=[user_row]) + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + mock_db_client.db.litellm_teammembership = MagicMock() + mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock()) + mock_db_client.db.litellm_verificationtoken = MagicMock() + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) + _wire_member_delete_tx(mock_db_client) + + await team_member_delete( + data=TeamMemberDeleteRequest(team_id="team-del-audit", user_id="bob"), + user_api_key_dict=mock_admin_auth, + ) + await _settle_audit_log_tasks() + + updated_events = _team_roster_events(audit_logger, "updated") + assert [e["object_id"] for e in updated_events] == ["team-del-audit"] + assert _roster_user_roles(updated_events[0]["before_value"]) == {"alice": "admin", "bob": "user"} + assert _roster_user_roles(updated_events[0]["updated_values"]) == {"alice": "admin"} + + stale_user_row = MagicMock() + stale_user_row.user_id = "carol" + stale_user_row.teams = ["team-del-audit"] + mock_db_client.db.litellm_usertable.find_many = AsyncMock(return_value=[stale_user_row]) + + await team_member_delete( + data=TeamMemberDeleteRequest(team_id="team-del-audit", user_id="carol"), + user_api_key_dict=mock_admin_auth, + ) + await _settle_audit_log_tasks() + + assert len(_team_roster_events(audit_logger, "updated")) == 1, ( + "scrubbing a stale team reference off a user row leaves the roster as it was, so no roster event" + ) + + +@pytest.mark.asyncio +async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeypatch): + """Changing a member's role must reach audit_log_callbacks as a TEAM_TABLE `updated` + event whose before roster carries the old role and whose after roster carries the new one.""" + audit_logger = _wire_audit_log_callback(monkeypatch) + + mock_prisma_client = MagicMock() + team_row = LiteLLM_TeamTable( + team_id="team-role-audit", + metadata={}, + members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")], + ) + + def _team_info_as_read_from_db(bob_role: str): + return { + "team_info": TeamInfoResponseObjectTeamTable( + team_id="team-role-audit", + metadata={}, + members_with_roles=( + TeamInfoMember(user_id="alice", role="admin", user_alias="Alice"), + TeamInfoMember(user_id="bob", role=bob_role, user_alias="Bob"), + ), + ), + "team_memberships": [LiteLLM_TeamMembership(user_id="bob", team_id="team-role-audit", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + mock_prisma_client.db.litellm_auditlog.create = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(side_effect=[_team_info_as_read_from_db("user"), _team_info_as_read_from_db("admin")]), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-role-audit", user_id="bob", role="admin"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + await _settle_audit_log_tasks() + + updated_events = _team_roster_events(audit_logger, "updated") + assert [e["object_id"] for e in updated_events] == ["team-role-audit"] + assert _roster_user_roles(updated_events[0]["before_value"]) == {"alice": "admin", "bob": "user"} + assert _roster_user_roles(updated_events[0]["updated_values"]) == {"alice": "admin", "bob": "admin"} + + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-role-audit", user_id="bob", role="admin"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + await _settle_audit_log_tasks() + + assert len(_team_roster_events(audit_logger, "updated")) == 1, ( + "re-sending the role a member already holds leaves the roster as it was, so no roster event" + ) + + +@pytest.mark.asyncio +async def test_delete_team_emits_only_the_deleted_audit_event(monkeypatch): + """`/team/delete` removes every member on its way out through the same code path + `/team/member_delete` uses. Those removals must not surface as TEAM_TABLE `updated` + roster events trailing the `deleted` one: the team is gone, and the `deleted` event + already carries the roster it went out with.""" + from litellm.proxy._types import DeleteTeamRequest + + audit_logger = _wire_audit_log_callback(monkeypatch) + + members = (Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")) + team = LiteLLM_TeamTable( + team_id="team-gone", + team_alias="gone", + members_with_roles=list(members), + metadata={}, + model_max_budget={}, + model_spend={}, + ) + mock_prisma = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma.get_data = AsyncMock( + return_value=SimpleNamespace(json=lambda **_kwargs: team.model_dump_json(exclude_none=True)) + ) + mock_prisma.delete_data = AsyncMock(return_value={"deleted_keys": 0}) + mock_prisma.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_auditlog.create = AsyncMock() + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma.db.tx = MagicMock(return_value=mock_tx_cm) + _wire_team_delete_tx(mock_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + removals = [(team, members, members[1:]), (team, members[1:], ())] + monkeypatch.setattr( + "litellm.proxy.management_endpoints.team_endpoints._team_member_delete", + AsyncMock(side_effect=lambda **_kwargs: removals.pop(0)), + ) + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-gone"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1", api_key="sk-a"), + litellm_changed_by=None, + ) + await _settle_audit_log_tasks() + + team_events = [(p["object_id"], p["action"]) for p in audit_logger.payloads if p["table_name"] == "LiteLLM_TeamTable"] + assert team_events == [("team-gone", "deleted")] + + def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back(): """A large member list must not echo every id back in the error body.""" from litellm.proxy.management_endpoints.team_endpoints import ( From 0e74dd2811f42a5e9a8f4d79b808ea22c704ce71 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 21:18:37 +0000 Subject: [PATCH 02/10] test(team): drop the docstrings from the roster audit event tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_team_endpoints.py | 18 ------------------ 1 file changed, 18 deletions(-) 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 72c854f1d6c..5a245c96a09 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13322,8 +13322,6 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp class _RecordingAuditLogger(CustomLogger): - """An audit_log_callbacks sink that keeps every payload it is handed.""" - def __init__(self) -> None: super().__init__() self.payloads: list[StandardAuditLogPayload] = [] @@ -13333,8 +13331,6 @@ class _RecordingAuditLogger(CustomLogger): def _wire_audit_log_callback(monkeypatch: pytest.MonkeyPatch) -> _RecordingAuditLogger: - """Turn audit logging on and register one recording callback, the way an operator's - `litellm_settings.audit_log_callbacks` entry would be.""" audit_logger = _RecordingAuditLogger() monkeypatch.setattr("litellm.store_audit_logs", True) monkeypatch.setattr("litellm.audit_log_callbacks", [audit_logger]) @@ -13343,7 +13339,6 @@ def _wire_audit_log_callback(monkeypatch: pytest.MonkeyPatch) -> _RecordingAudit async def _settle_audit_log_tasks() -> None: - """Audit callbacks run on `asyncio.create_task`, so give the loop a few turns.""" for _ in range(5): await asyncio.sleep(0) @@ -13361,10 +13356,6 @@ def _roster_user_roles(members_json: str | None) -> dict[str, str]: @pytest.mark.asyncio async def test_new_team_created_audit_event_carries_the_final_roster(monkeypatch): - """The `created` event a `/team/new` hands to audit_log_callbacks must list the members - the team was created with. The team row is inserted empty and the members attached - afterwards, so a snapshot taken from the pre-insert object reports no members and a - downstream consumer syncing membership from the event has nothing to sync.""" from fastapi import Request from litellm.proxy._types import NewTeamRequest @@ -13428,9 +13419,6 @@ async def test_new_team_created_audit_event_carries_the_final_roster(monkeypatch @pytest.mark.asyncio async def test_team_member_delete_emits_a_roster_audit_event(monkeypatch, mock_db_client, mock_admin_auth): - """Removing a member must reach audit_log_callbacks as a TEAM_TABLE `updated` event whose - before and after rosters differ by exactly the removed user, the same shape - `/team/member_add` already emits, so one consumer can diff both directions.""" from litellm.proxy._types import TeamMemberDeleteRequest audit_logger = _wire_audit_log_callback(monkeypatch) @@ -13490,8 +13478,6 @@ async def test_team_member_delete_emits_a_roster_audit_event(monkeypatch, mock_d @pytest.mark.asyncio async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeypatch): - """Changing a member's role must reach audit_log_callbacks as a TEAM_TABLE `updated` - event whose before roster carries the old role and whose after roster carries the new one.""" audit_logger = _wire_audit_log_callback(monkeypatch) mock_prisma_client = MagicMock() @@ -13562,10 +13548,6 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp @pytest.mark.asyncio async def test_delete_team_emits_only_the_deleted_audit_event(monkeypatch): - """`/team/delete` removes every member on its way out through the same code path - `/team/member_delete` uses. Those removals must not surface as TEAM_TABLE `updated` - roster events trailing the `deleted` one: the team is gone, and the `deleted` event - already carries the roster it went out with.""" from litellm.proxy._types import DeleteTeamRequest audit_logger = _wire_audit_log_callback(monkeypatch) From 2a7dcc77b25ddc77fc842aafad633aa44d42f1ef Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 21:47:32 +0000 Subject: [PATCH 03/10] test(team): mock the membership upsert the member add now issues on team create Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_team_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 5a245c96a09..67b8f3bc5f3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13387,7 +13387,7 @@ async def test_new_team_created_audit_event_carries_the_final_roster(monkeypatch mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=user_row) membership_row = MagicMock() membership_row.model_dump.return_value = {"team_id": "team-audit-roster", "user_id": "alice", "budget_id": None} - mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=membership_row) + mock_prisma.db.litellm_teammembership.upsert = AsyncMock(return_value=membership_row) mock_prisma.db.litellm_auditlog.create = AsyncMock() _wire_team_create_tx(mock_prisma) From 209eba6718b09689e25c3caf3cc358a926a67ff0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 23:28:39 +0000 Subject: [PATCH 04/10] feat(team): carry team_alias on member add, delete and role-change audit payloads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 15 ++++++++---- .../test_team_endpoints.py | 23 ++++++++++++++++--- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8e5d62976fa..81fd7f44538 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3150,7 +3150,7 @@ def _validate_member_user_id_provisioning( ) -def _members_audit_value(members: Sequence[Member]) -> str: +def _members_audit_value(team_alias: str | None, members: Sequence[Member]) -> str: """Serialize a team's member list for an audit-log value. The audit-log columns hold a JSON object, so the member list is nested @@ -3158,13 +3158,15 @@ def _members_audit_value(members: Sequence[Member]) -> str: """ return safe_dumps( { # mutable-ok: the audit-log JSON column rejects a top-level array, so this value must be an object - "members_with_roles": tuple(member.model_dump() for member in members) + "team_alias": team_alias, + "members_with_roles": tuple(member.model_dump() for member in members), } ) async def _create_team_membership_audit_log( team_id: str, + team_alias: str | None, before_members: Sequence[Member], after_members: Sequence[Member], user_api_key_dict: UserAPIKeyAuth, @@ -3179,13 +3181,14 @@ async def _create_team_membership_audit_log( user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=litellm_proxy_admin_name, table_name=LitellmTableNames.TEAM_TABLE_NAME, - before_value=_members_audit_value(before_members), - after_value=_members_audit_value(after_members), + before_value=_members_audit_value(team_alias, before_members), + after_value=_members_audit_value(team_alias, after_members), ) async def _create_team_member_add_audit_logs( team_id: str, + team_alias: str | None, updated_users: Sequence[LiteLLM_UserTable], existing_user_ids: frozenset[str], before_members: Sequence[Member], @@ -3217,6 +3220,7 @@ async def _create_team_member_add_audit_logs( membership_entry: Final = _create_team_membership_audit_log( team_id=team_id, + team_alias=team_alias, before_members=before_members, after_members=after_members, user_api_key_dict=user_api_key_dict, @@ -3460,6 +3464,7 @@ async def team_member_add( await _create_team_member_add_audit_logs( team_id=data.team_id, + team_alias=complete_team_data.team_alias, updated_users=updated_users, existing_user_ids=pre_existing_user_ids, before_members=members_before_add, @@ -3538,6 +3543,7 @@ async def team_member_delete( if before_members != after_members: await _create_team_membership_audit_log( team_id=existing_team_row.team_id, + team_alias=existing_team_row.team_alias, before_members=before_members, after_members=after_members, user_api_key_dict=user_api_key_dict, @@ -3879,6 +3885,7 @@ async def team_member_update( if members_before_role_update != tuple(team_members): await _create_team_membership_audit_log( team_id=data.team_id, + team_alias=team_table.team_alias, before_members=members_before_role_update, after_members=team_members, user_api_key_dict=user_api_key_dict, 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 67b8f3bc5f3..0077be9bfbb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13245,9 +13245,12 @@ def test_members_audit_value_serializes_to_a_json_object(): """The audit-log columns hold a JSON object; a top-level array is rejected by the DB.""" from litellm.proxy.management_endpoints.team_endpoints import _members_audit_value - payload = json.loads(_members_audit_value([Member(user_id="u1", role="admin"), Member(user_id="u2", role="user")])) + payload = json.loads( + _members_audit_value("my-team", [Member(user_id="u1", role="admin"), Member(user_id="u2", role="user")]) + ) assert isinstance(payload, dict) + assert payload["team_alias"] == "my-team" assert [m["user_id"] for m in payload["members_with_roles"]] == ["u1", "u2"] @@ -13272,7 +13275,7 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id") - team_row = LiteLLM_TeamTable(team_id=team_id, members_with_roles=[]) + team_row = LiteLLM_TeamTable(team_id=team_id, team_alias="list-audit", members_with_roles=[]) created_user = LiteLLM_UserTable( user_id=created_user_id, user_email="invitee@example.com", max_budget=None, spend=0.0, models=[] ) @@ -13319,6 +13322,7 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp mock_audit.assert_called_once() assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"] + assert mock_audit.call_args.kwargs["team_alias"] == "list-audit" class _RecordingAuditLogger(CustomLogger): @@ -13345,7 +13349,9 @@ async def _settle_audit_log_tasks() -> None: def _team_roster_events(audit_logger: _RecordingAuditLogger, action: str) -> list[StandardAuditLogPayload]: return [ - p for p in audit_logger.payloads if p["table_name"] == LitellmTableNames.TEAM_TABLE_NAME and p["action"] == action + p + for p in audit_logger.payloads + if p["table_name"] == LitellmTableNames.TEAM_TABLE_NAME and p["action"] == action ] @@ -13354,6 +13360,11 @@ def _roster_user_roles(members_json: str | None) -> dict[str, str]: return {m["user_id"]: m["role"] for m in json.loads(members_json)["members_with_roles"]} +def _roster_team_alias(members_json: str | None) -> str | None: + assert members_json is not None + return json.loads(members_json)["team_alias"] + + @pytest.mark.asyncio async def test_new_team_created_audit_event_carries_the_final_roster(monkeypatch): from fastapi import Request @@ -13426,6 +13437,7 @@ async def test_team_member_delete_emits_a_roster_audit_event(monkeypatch, mock_d team_row = MagicMock() team_row.model_dump.return_value = { "team_id": "team-del-audit", + "team_alias": "del-audit", "members_with_roles": [ {"user_id": "alice", "user_email": None, "role": "admin"}, {"user_id": "bob", "user_email": None, "role": "user"}, @@ -13459,6 +13471,8 @@ async def test_team_member_delete_emits_a_roster_audit_event(monkeypatch, mock_d assert [e["object_id"] for e in updated_events] == ["team-del-audit"] assert _roster_user_roles(updated_events[0]["before_value"]) == {"alice": "admin", "bob": "user"} assert _roster_user_roles(updated_events[0]["updated_values"]) == {"alice": "admin"} + assert _roster_team_alias(updated_events[0]["before_value"]) == "del-audit" + assert _roster_team_alias(updated_events[0]["updated_values"]) == "del-audit" stale_user_row = MagicMock() stale_user_row.user_id = "carol" @@ -13483,6 +13497,7 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp mock_prisma_client = MagicMock() team_row = LiteLLM_TeamTable( team_id="team-role-audit", + team_alias="role-audit", metadata={}, members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")], ) @@ -13491,6 +13506,7 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp return { "team_info": TeamInfoResponseObjectTeamTable( team_id="team-role-audit", + team_alias="role-audit", metadata={}, members_with_roles=( TeamInfoMember(user_id="alice", role="admin", user_alias="Alice"), @@ -13531,6 +13547,7 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp assert [e["object_id"] for e in updated_events] == ["team-role-audit"] assert _roster_user_roles(updated_events[0]["before_value"]) == {"alice": "admin", "bob": "user"} assert _roster_user_roles(updated_events[0]["updated_values"]) == {"alice": "admin", "bob": "admin"} + assert _roster_team_alias(updated_events[0]["updated_values"]) == "role-audit" await team_member_update( data=TeamMemberUpdateRequest(team_id="team-role-audit", user_id="bob", role="admin"), From c4272d894f867e1cae416727cf0f57b435d71d94 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 23:29:11 +0000 Subject: [PATCH 05/10] style(team): wrap the roster audit helper comprehensions at 120 columns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_team_endpoints.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 0077be9bfbb..d391a3ebf9c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13611,7 +13611,9 @@ async def test_delete_team_emits_only_the_deleted_audit_event(monkeypatch): ) await _settle_audit_log_tasks() - team_events = [(p["object_id"], p["action"]) for p in audit_logger.payloads if p["table_name"] == "LiteLLM_TeamTable"] + team_events = [ + (p["object_id"], p["action"]) for p in audit_logger.payloads if p["table_name"] == "LiteLLM_TeamTable" + ] assert team_events == [("team-gone", "deleted")] From 181406e05fd3e0c8788a9053d901e478469bb9e4 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 22:33:22 +0000 Subject: [PATCH 06/10] fix(team): schedule membership audit writes after commit and lock the roster on role updates The member add, delete and role-change audit rows were awaited on the request path, so a slow audit sink held the response, and the roster was serialized before checking whether audit logging is enabled at all. Membership audit work is now scheduled after the transaction commits and skipped outright when auditing is off. member_update read the roster outside the team advisory lock and wrote it back, so a concurrent add or delete could be lost. It now takes the lock, rereads the roster, and builds the before and after snapshots from that read. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 172 +++++++------- .../test_team_endpoints.py | 213 +++++++++++++++++- .../proxy/test_team_member_update.py | 14 +- 3 files changed, 310 insertions(+), 89 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 81fd7f44538..982e0926b26 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3053,6 +3053,35 @@ async def _add_team_members_to_team( return updated_team, updated_users, updated_team_memberships +async def _update_team_member_role( + prisma_client: PrismaClient, + team_id: str, + user_id: str, + role: Literal["admin", "user"], + user_email: str | None, +) -> tuple[tuple[Member, ...], tuple[Member, ...]]: + """Rewrite one member's role from the roster read under the team lock; returns (before, after).""" + async with prisma_client.tx() as tx: + await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) + + locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team_id) + if locked_members is None: + raise HTTPException(status_code=404, detail={"error": f"Team id={team_id} does not exist in db"}) + + before: Final = tuple(locked_members) + after: Final = tuple( + Member(user_id=member.user_id, role=role, user_email=user_email or member.user_email) + if member.user_id == user_id + else member + for member in before + ) + await _team_tx_db(tx).update( + where={"team_id": team_id}, + data={"members_with_roles": json.dumps([m.model_dump() for m in after])}, + ) + return before, after + + def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: """Update the Prometheus team members gauge after a membership change. @@ -3164,7 +3193,7 @@ def _members_audit_value(team_alias: str | None, members: Sequence[Member]) -> s ) -async def _create_team_membership_audit_log( +def _schedule_team_membership_audit_log( team_id: str, team_alias: str | None, before_members: Sequence[Member], @@ -3172,21 +3201,29 @@ async def _create_team_membership_audit_log( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, ) -> None: - from litellm.proxy.management_helpers.audit_logs import create_object_audit_log + from litellm.proxy.management_helpers.audit_logs import ( + create_object_audit_log, + is_audit_logging_enabled, + ) - await create_object_audit_log( - object_id=team_id, - action="updated", - litellm_changed_by=None, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - table_name=LitellmTableNames.TEAM_TABLE_NAME, - before_value=_members_audit_value(team_alias, before_members), - after_value=_members_audit_value(team_alias, after_members), + if not is_audit_logging_enabled() or tuple(before_members) == tuple(after_members): + return + + asyncio.create_task( + create_object_audit_log( + object_id=team_id, + action="updated", + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + before_value=_members_audit_value(team_alias, before_members), + after_value=_members_audit_value(team_alias, after_members), + ) ) -async def _create_team_member_add_audit_logs( +def _schedule_team_member_add_audit_logs( team_id: str, team_alias: str | None, updated_users: Sequence[LiteLLM_UserTable], @@ -3196,29 +3233,32 @@ async def _create_team_member_add_audit_logs( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, ) -> None: - """Record the membership change, and any user row it created, in the audit log. - - The entries are written concurrently so a request adding many members does - not pay for them one after another. - """ - from litellm.proxy.management_helpers.audit_logs import create_object_audit_log - - created_user_entries: Final = tuple( - create_object_audit_log( - object_id=user.user_id, - action="created", - litellm_changed_by=None, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - table_name=LitellmTableNames.USER_TABLE_NAME, - before_value=None, - after_value=safe_dumps(user.model_dump(exclude_none=True)), - ) - for user in updated_users - if user.user_id is not None and user.user_id not in existing_user_ids + """Record the membership change, and any user row it created, in the audit log.""" + from litellm.proxy.management_helpers.audit_logs import ( + create_object_audit_log, + is_audit_logging_enabled, ) - membership_entry: Final = _create_team_membership_audit_log( + if not is_audit_logging_enabled(): + return + + for user in updated_users: + if user.user_id is None or user.user_id in existing_user_ids: + continue + asyncio.create_task( + create_object_audit_log( + object_id=user.user_id, + action="created", + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + table_name=LitellmTableNames.USER_TABLE_NAME, + before_value=None, + after_value=safe_dumps(user.model_dump(exclude_none=True)), + ) + ) + + _schedule_team_membership_audit_log( team_id=team_id, team_alias=team_alias, before_members=before_members, @@ -3227,8 +3267,6 @@ async def _create_team_member_add_audit_logs( litellm_proxy_admin_name=litellm_proxy_admin_name, ) - await asyncio.gather(*created_user_entries, membership_entry) - async def _validate_and_populate_member_user_info( member: Member, @@ -3462,7 +3500,7 @@ async def team_member_add( _emit_team_members_metric(complete_team_data) - await _create_team_member_add_audit_logs( + _schedule_team_member_add_audit_logs( team_id=data.team_id, team_alias=complete_team_data.team_alias, updated_users=updated_users, @@ -3540,15 +3578,14 @@ async def team_member_delete( data=data, user_api_key_dict=user_api_key_dict ) - if before_members != after_members: - await _create_team_membership_audit_log( - team_id=existing_team_row.team_id, - team_alias=existing_team_row.team_alias, - before_members=before_members, - after_members=after_members, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) + _schedule_team_membership_audit_log( + team_id=existing_team_row.team_id, + team_alias=existing_team_row.team_alias, + before_members=before_members, + after_members=after_members, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) return existing_team_row @@ -3858,39 +3895,22 @@ async def team_member_update( ### update team member role if data.role is not None: - members_before_role_update: Final = tuple( - Member(user_id=member.user_id, user_email=member.user_email, role=member.role) - for member in team_table.members_with_roles + members_before_role_update, team_members = await _update_team_member_role( + prisma_client=prisma_client, + team_id=data.team_id, + user_id=received_user_id, + role=data.role, + user_email=data.user_email, ) - team_members: Final[list[Member]] = [] - for member in members_before_role_update: - if member.user_id == received_user_id: - team_members.append( - Member( - user_id=member.user_id, - role=data.role, - user_email=data.user_email or member.user_email, - ) - ) - else: - team_members.append(member) - - team_table.members_with_roles = team_members - - _db_team_members: Final[list[dict]] = [m.model_dump() for m in team_members] - await _team_db(prisma_client).update( - where={"team_id": data.team_id}, - data={"members_with_roles": json.dumps(_db_team_members)}, + team_table.members_with_roles = list(team_members) + _schedule_team_membership_audit_log( + team_id=data.team_id, + team_alias=team_table.team_alias, + before_members=members_before_role_update, + after_members=team_members, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, ) - if members_before_role_update != tuple(team_members): - await _create_team_membership_audit_log( - team_id=data.team_id, - team_alias=team_table.team_alias, - before_members=members_before_role_update, - after_members=team_members, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) return TeamMemberUpdateResponse( team_id=data.team_id, 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 d391a3ebf9c..87530cc8526 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -3,6 +3,7 @@ import json from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timezone from types import SimpleNamespace +from collections.abc import Sequence from typing import Final, Optional, cast from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch @@ -148,11 +149,11 @@ def _wire_member_add_tx(prisma_client): def _wire_member_delete_tx(prisma_client): - """/team/member_delete's four cleanups, plus the advisory-lock re-read that now guards - them, run inside one transaction, so a mocked client has to hand back its own table - mocks (and a `query_raw` that answers the locked re-read from the same team row the - test already configured on `find_unique`) out of `tx()` for the existing per-table - assertions to keep seeing the calls.""" + """/team/member_delete's four cleanups and /team/member_update's role rewrite, plus the + advisory-lock re-read that guards them, run inside one transaction, so a mocked client + has to hand back its own table mocks (and a `query_raw` that answers the locked re-read + from the same team row the test already configured on `find_unique`) out of `tx()` for + the existing per-table assertions to keep seeing the calls.""" async def _query_raw(sql, team_id): if sql != TEAM_ADVISORY_LOCK_SQL: @@ -168,10 +169,12 @@ def _wire_member_delete_tx(prisma_client): return getattr(prisma_client.db, table_name) tx = _Tx() + tx.query_raw = AsyncMock(side_effect=_query_raw) tx_cm = MagicMock() tx_cm.__aenter__ = AsyncMock(return_value=tx) tx_cm.__aexit__ = AsyncMock(return_value=None) prisma_client.tx = MagicMock(return_value=tx_cm) + return tx def _wire_team_delete_tx(prisma_client): @@ -13311,8 +13314,7 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp side_effect=fake_add_team_members_to_team, ), patch( - "litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs", - new_callable=AsyncMock, + "litellm.proxy.management_endpoints.team_endpoints._schedule_team_member_add_audit_logs", ) as mock_audit, ): await team_member_add( @@ -13517,12 +13519,10 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp } mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) - mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=_roster_writer(team_row)) mock_prisma_client.db.litellm_auditlog.create = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mock_tx = AsyncMock() - mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) - mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + _wire_member_delete_tx(mock_prisma_client) with ( patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests @@ -13563,6 +13563,197 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp ) +def _roster_writer(team_row: LiteLLM_TeamTable): + """An `update` side effect that lands `members_with_roles` on the team row later reads see.""" + + async def _update(where, data): + team_row.members_with_roles = [Member(**m) for m in json.loads(data["members_with_roles"])] + return team_row + + return _update + + +def _member_update_patches(team_snapshot: LiteLLM_TeamTable): + return ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock( + return_value={ + "team_info": TeamInfoResponseObjectTeamTable(**team_snapshot.model_dump()), + "team_memberships": [ + LiteLLM_TeamMembership(user_id="bob", team_id=team_snapshot.team_id, budget_id=None) + ], + } + ), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ) + + +@pytest.mark.asyncio +async def test_team_member_update_role_change_rewrites_the_roster_it_read_under_the_lock(monkeypatch): + """Regression: a member added between /team/member_update's permission checks and its write + was dropped, because the new roster was built from the pre-check snapshot.""" + audit_logger = _wire_audit_log_callback(monkeypatch) + + stale_snapshot = LiteLLM_TeamTable( + team_id="team-race", + team_alias="race", + metadata={}, + members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")], + ) + team_row = LiteLLM_TeamTable( + **{ + **stale_snapshot.model_dump(), + "members_with_roles": [*stale_snapshot.members_with_roles, Member(user_id="carol", role="user")], + } + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=_roster_writer(team_row)) + mock_prisma_client.db.litellm_auditlog.create = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + tx = _wire_member_delete_tx(mock_prisma_client) + + team_info_patch, upsert_patch = _member_update_patches(stale_snapshot) + with team_info_patch, upsert_patch: + response = await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-race", user_id="bob", role="admin"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + await _settle_audit_log_tasks() + + assert {m.user_id: m.role for m in team_row.members_with_roles} == { + "alice": "admin", + "bob": "admin", + "carol": "user", + } + assert response.team_id == "team-race" and response.user_id == "bob" + assert tx.query_raw.await_args_list[0].args == (TEAM_ADVISORY_LOCK_SQL, "team-race"), ( + "the roster must be read only after the team advisory lock is held" + ) + + [event] = _team_roster_events(audit_logger, "updated") + assert _roster_user_roles(event["before_value"]) == {"alice": "admin", "bob": "user", "carol": "user"} + assert _roster_user_roles(event["updated_values"]) == {"alice": "admin", "bob": "admin", "carol": "user"} + assert _roster_team_alias(event["updated_values"]) == "race" + + +@pytest.mark.asyncio +async def test_team_member_update_role_change_404s_when_the_team_is_gone_under_the_lock(monkeypatch): + _wire_audit_log_callback(monkeypatch) + snapshot = LiteLLM_TeamTable( + team_id="team-gone-race", + team_alias="gone-race", + metadata={}, + members_with_roles=[Member(user_id="bob", role="user")], + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=[snapshot, None]) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + _wire_member_delete_tx(mock_prisma_client) + + team_info_patch, upsert_patch = _member_update_patches(snapshot) + with team_info_patch, upsert_patch, pytest.raises(HTTPException) as exc_info: + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-gone-race", user_id="bob", role="admin"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert exc_info.value.status_code == 404 + mock_prisma_client.db.litellm_teamtable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_member_delete_response_does_not_wait_for_the_audit_insert( + monkeypatch, mock_db_client, mock_admin_auth +): + """Regression: the roster audit row was awaited on the request path, so a slow audit table + held every /team/member_delete response.""" + from litellm.proxy._types import TeamMemberDeleteRequest + + audit_logger = _wire_audit_log_callback(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + team_row = LiteLLM_TeamTable( + team_id="team-slow-audit", + team_alias="slow-audit", + metadata={}, + members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")], + ) + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + mock_db_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock()) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) + _wire_member_delete_tx(mock_db_client) + + audit_table_answers = asyncio.Event() + audit_rows = [] + + async def _blocked_create(data): + await audit_table_answers.wait() + audit_rows.append(data) + + mock_db_client.db.litellm_auditlog.create = AsyncMock(side_effect=_blocked_create) + + await asyncio.wait_for( + team_member_delete( + data=TeamMemberDeleteRequest(team_id="team-slow-audit", user_id="bob"), + user_api_key_dict=mock_admin_auth, + ), + timeout=1, + ) + assert audit_rows == [], "the response returned while the audit table was still blocked" + + audit_table_answers.set() + await _settle_audit_log_tasks() + + assert [row["object_id"] for row in audit_rows] == ["team-slow-audit"] + [event] = _team_roster_events(audit_logger, "updated") + assert _roster_user_roles(event["before_value"]) == {"alice": "admin", "bob": "user"} + assert _roster_user_roles(event["updated_values"]) == {"alice": "admin"} + + +class _UntouchableRoster(Sequence[Member]): + """A roster that fails the test the moment anything reads it.""" + + def __getitem__(self, index): + raise AssertionError("the roster was read while audit logging is off") + + def __len__(self) -> int: + raise AssertionError("the roster was read while audit logging is off") + + +def test_membership_audit_scheduling_skips_the_roster_entirely_when_audit_logging_is_off(monkeypatch): + """Regression: the before/after rosters were serialized on every membership change, even + when audit logs are not stored.""" + from litellm.proxy.management_endpoints.team_endpoints import _schedule_team_membership_audit_log + + monkeypatch.setattr("litellm.store_audit_logs", False) + + _schedule_team_membership_audit_log( + team_id="team-quiet", + team_alias="quiet", + before_members=_UntouchableRoster(), + after_members=_UntouchableRoster(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user"), + litellm_proxy_admin_name="admin", + ) + + @pytest.mark.asyncio async def test_delete_team_emits_only_the_deleted_audit_event(monkeypatch): from litellm.proxy._types import DeleteTeamRequest diff --git a/tests/test_litellm/proxy/test_team_member_update.py b/tests/test_litellm/proxy/test_team_member_update.py index 352c68d491c..ace4c4e65af 100644 --- a/tests/test_litellm/proxy/test_team_member_update.py +++ b/tests/test_litellm/proxy/test_team_member_update.py @@ -14,7 +14,10 @@ from litellm.proxy._types import ( TeamMemberUpdateRequest, UserAPIKeyAuth, ) -from litellm.proxy.management_endpoints.team_endpoints import team_member_update +from litellm.proxy.management_endpoints.team_endpoints import ( + TEAM_ADVISORY_LOCK_SQL, + team_member_update, +) @pytest.mark.asyncio @@ -65,13 +68,20 @@ def happy_path_upsert(monkeypatch): prisma_client.db.litellm_teamtable.update = AsyncMock() class _FakeTx: + litellm_teamtable = prisma_client.db.litellm_teamtable + async def __aenter__(self): return self async def __aexit__(self, *args): return False - prisma_client.db.tx = MagicMock(return_value=_FakeTx()) + async def query_raw(self, sql, team_id): + if sql == TEAM_ADVISORY_LOCK_SQL: + return [] + return [{"members_with_roles": team_row.model_dump()["members_with_roles"]}] + + prisma_client.tx = MagicMock(return_value=_FakeTx()) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "premium_user", False) From cc41b80770827afbe1a336953fe9008268b07dc5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 22:46:37 +0000 Subject: [PATCH 07/10] test(team): patch the scheduled member-add audit helper in the cache eviction test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_team_endpoints.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 3da900d4d9a..fe461eca4a1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13866,8 +13866,7 @@ async def test_team_member_add_evicts_the_new_members_cached_user_row_on_every_w side_effect=fake_add_team_members_to_team, ), patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers - "litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs", - new_callable=AsyncMock, + "litellm.proxy.management_endpoints.team_endpoints._schedule_team_member_add_audit_logs", ), ): await team_member_add( From 431ddbdd2200cf9cd3f99fb5c2839ea944b4c34a Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 23:17:18 +0000 Subject: [PATCH 08/10] test(team): exercise the member-add audit helper directly and drop its dead user_id None guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 2 +- .../test_team_endpoints.py | 83 +++++++++++++++++-- 2 files changed, 79 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 222cf8e299f..6eaed62013e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3243,7 +3243,7 @@ def _schedule_team_member_add_audit_logs( return for user in updated_users: - if user.user_id is None or user.user_id in existing_user_ids: + if user.user_id in existing_user_ids: continue asyncio.create_task( create_object_audit_log( 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 fe461eca4a1..dc72a6a2256 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13737,12 +13737,16 @@ class _UntouchableRoster(Sequence[Member]): raise AssertionError("the roster was read while audit logging is off") -def test_membership_audit_scheduling_skips_the_roster_entirely_when_audit_logging_is_off(monkeypatch): - """Regression: the before/after rosters were serialized on every membership change, even - when audit logs are not stored.""" - from litellm.proxy.management_endpoints.team_endpoints import _schedule_team_membership_audit_log +class _UntouchableUsers(Sequence[LiteLLM_UserTable]): + def __getitem__(self, index): + raise AssertionError("the created users were read while audit logging is off") - monkeypatch.setattr("litellm.store_audit_logs", False) + def __len__(self) -> int: + raise AssertionError("the created users were read while audit logging is off") + + +def _schedule_membership_audit_with_untouchable_roster() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _schedule_team_membership_audit_log _schedule_team_membership_audit_log( team_id="team-quiet", @@ -13754,6 +13758,75 @@ def test_membership_audit_scheduling_skips_the_roster_entirely_when_audit_loggin ) +def _schedule_member_add_audit_with_untouchable_roster() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _schedule_team_member_add_audit_logs + + _schedule_team_member_add_audit_logs( + team_id="team-quiet", + team_alias="quiet", + updated_users=_UntouchableUsers(), + existing_user_ids=frozenset(), + before_members=_UntouchableRoster(), + after_members=_UntouchableRoster(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user"), + litellm_proxy_admin_name="admin", + ) + + +@pytest.mark.parametrize( + "schedule", + [_schedule_membership_audit_with_untouchable_roster, _schedule_member_add_audit_with_untouchable_roster], +) +def test_membership_audit_scheduling_skips_the_roster_entirely_when_audit_logging_is_off(monkeypatch, schedule): + """Regression: the before/after rosters were serialized on every membership change, even + when audit logs are not stored.""" + monkeypatch.setattr("litellm.store_audit_logs", False) + + schedule() + + +@pytest.mark.asyncio +async def test_member_add_audit_reports_only_the_users_it_created_plus_the_roster_change(monkeypatch): + from litellm.proxy.management_endpoints.team_endpoints import _schedule_team_member_add_audit_logs + + audit_logger = _wire_audit_log_callback(monkeypatch) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_auditlog.create = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + before = (Member(user_id="alice", role="admin"),) + after = (*before, Member(user_id="bob", role="user"), Member(user_id="carol", role="user")) + + _schedule_team_member_add_audit_logs( + team_id="team-add-audit", + team_alias="add-audit", + updated_users=[ + LiteLLM_UserTable(user_id="bob", user_email="bob@example.com", teams=["team-add-audit"]), + LiteLLM_UserTable(user_id="carol", teams=["team-add-audit"]), + ], + existing_user_ids=frozenset({"bob"}), + before_members=before, + after_members=after, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-1"), + litellm_proxy_admin_name="admin", + ) + await _settle_audit_log_tasks() + + created_users = [ + p + for p in audit_logger.payloads + if p["table_name"] == LitellmTableNames.USER_TABLE_NAME and p["action"] == "created" + ] + assert [p["object_id"] for p in created_users] == ["carol"], "only the user this request created is audited" + assert json.loads(created_users[0]["updated_values"])["teams"] == ["team-add-audit"] + + [roster_event] = _team_roster_events(audit_logger, "updated") + assert roster_event["object_id"] == "team-add-audit" + assert _roster_user_roles(roster_event["before_value"]) == {"alice": "admin"} + assert _roster_user_roles(roster_event["updated_values"]) == {"alice": "admin", "bob": "user", "carol": "user"} + assert _roster_team_alias(roster_event["updated_values"]) == "add-audit" + + @pytest.mark.asyncio async def test_delete_team_emits_only_the_deleted_audit_event(monkeypatch): from litellm.proxy._types import DeleteTeamRequest From b7db48c7c1f162efda73ccd3bde1386c912c0786 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 19 Sep 2026 23:57:38 +0000 Subject: [PATCH 09/10] fix(team): 404 a role update whose target left the roster before the locked read Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 3 ++ .../test_team_endpoints.py | 39 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6eaed62013e..5917f219fde 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3069,6 +3069,9 @@ async def _update_team_member_role( raise HTTPException(status_code=404, detail={"error": f"Team id={team_id} does not exist in db"}) before: Final = tuple(locked_members) + if all(member.user_id != user_id for member in before): + raise HTTPException(status_code=404, detail={"error": f"User {user_id} is not a member of team {team_id}"}) + after: Final = tuple( Member(user_id=member.user_id, role=role, user_email=user_email or member.user_email) if member.user_id == user_id 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 dc72a6a2256..155f03e6d3e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13675,6 +13675,45 @@ async def test_team_member_update_role_change_404s_when_the_team_is_gone_under_t mock_prisma_client.db.litellm_teamtable.update.assert_not_awaited() +@pytest.mark.asyncio +async def test_team_member_update_role_change_404s_when_the_member_left_before_the_locked_read(monkeypatch): + """Regression: a member removed between the pre-lock read and the locked read was reported as updated.""" + audit_logger = _wire_audit_log_callback(monkeypatch) + snapshot = LiteLLM_TeamTable( + team_id="team-member-gone-race", + team_alias="member-gone-race", + metadata={}, + members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")], + ) + locked_row = LiteLLM_TeamTable( + team_id="team-member-gone-race", + team_alias="member-gone-race", + metadata={}, + members_with_roles=[Member(user_id="alice", role="admin")], + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=[snapshot, locked_row]) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + _wire_member_delete_tx(mock_prisma_client) + + team_info_patch, upsert_patch = _member_update_patches(snapshot) + with team_info_patch, upsert_patch, pytest.raises(HTTPException) as exc_info: + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-member-gone-race", user_id="bob", role="admin"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert exc_info.value.status_code == 404 + assert "bob" in str(exc_info.value.detail) + mock_prisma_client.db.litellm_teamtable.update.assert_not_awaited() + await _settle_audit_log_tasks() + assert audit_logger.payloads == [] + + @pytest.mark.asyncio async def test_team_member_delete_response_does_not_wait_for_the_audit_insert( monkeypatch, mock_db_client, mock_admin_auth From de70cf842a17674c9be9351926d1532e6fa50d45 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 20 Sep 2026 00:04:32 +0000 Subject: [PATCH 10/10] fix(team): run the role update and budget upsert in one transaction under the team lock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 59 ++++++++++--------- .../test_team_endpoints.py | 7 ++- 2 files changed, 37 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 5917f219fde..dbc709a1742 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3054,6 +3054,7 @@ async def _add_team_members_to_team( async def _update_team_member_role( + tx: "Prisma", prisma_client: PrismaClient, team_id: str, user_id: str, @@ -3061,27 +3062,26 @@ async def _update_team_member_role( user_email: str | None, ) -> tuple[tuple[Member, ...], tuple[Member, ...]]: """Rewrite one member's role from the roster read under the team lock; returns (before, after).""" - async with prisma_client.tx() as tx: - await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) + await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) - locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team_id) - if locked_members is None: - raise HTTPException(status_code=404, detail={"error": f"Team id={team_id} does not exist in db"}) + locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team_id) + if locked_members is None: + raise HTTPException(status_code=404, detail={"error": f"Team id={team_id} does not exist in db"}) - before: Final = tuple(locked_members) - if all(member.user_id != user_id for member in before): - raise HTTPException(status_code=404, detail={"error": f"User {user_id} is not a member of team {team_id}"}) + before: Final = tuple(locked_members) + if all(member.user_id != user_id for member in before): + raise HTTPException(status_code=404, detail={"error": f"User {user_id} is not a member of team {team_id}"}) - after: Final = tuple( - Member(user_id=member.user_id, role=role, user_email=user_email or member.user_email) - if member.user_id == user_id - else member - for member in before - ) - await _team_tx_db(tx).update( - where={"team_id": team_id}, - data={"members_with_roles": json.dumps([m.model_dump() for m in after])}, - ) + after: Final = tuple( + Member(user_id=member.user_id, role=role, user_email=user_email or member.user_email) + if member.user_id == user_id + else member + for member in before + ) + await _team_tx_db(tx).update( + where={"team_id": team_id}, + data={"members_with_roles": json.dumps([m.model_dump() for m in after])}, + ) return before, after @@ -3885,6 +3885,18 @@ async def team_member_update( ### upsert new budget budget_patch: Final = member_budget_patch(data) async with prisma_client.tx() as tx: + role_change: Final = ( + await _update_team_member_role( + tx=tx, + prisma_client=prisma_client, + team_id=data.team_id, + user_id=received_user_id, + role=data.role, + user_email=data.user_email, + ) + if data.role is not None + else None + ) await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, @@ -3901,15 +3913,8 @@ async def team_member_update( user_api_key_cache=user_api_key_cache, ) - ### update team member role - if data.role is not None: - members_before_role_update, team_members = await _update_team_member_role( - prisma_client=prisma_client, - team_id=data.team_id, - user_id=received_user_id, - role=data.role, - user_email=data.user_email, - ) + if role_change is not None: + members_before_role_update, team_members = role_change team_table.members_with_roles = list(team_members) _schedule_team_membership_audit_log( team_id=data.team_id, 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 155f03e6d3e..7cb62a8da11 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13698,9 +13698,11 @@ async def test_team_member_update_role_change_404s_when_the_member_left_before_t _wire_member_delete_tx(mock_prisma_client) team_info_patch, upsert_patch = _member_update_patches(snapshot) - with team_info_patch, upsert_patch, pytest.raises(HTTPException) as exc_info: + with team_info_patch, upsert_patch as upsert_budget, pytest.raises(HTTPException) as exc_info: await team_member_update( - data=TeamMemberUpdateRequest(team_id="team-member-gone-race", user_id="bob", role="admin"), + data=TeamMemberUpdateRequest( + team_id="team-member-gone-race", user_id="bob", role="admin", max_budget_in_team=5.0 + ), http_request=MagicMock(), user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" @@ -13710,6 +13712,7 @@ async def test_team_member_update_role_change_404s_when_the_member_left_before_t assert exc_info.value.status_code == 404 assert "bob" in str(exc_info.value.detail) mock_prisma_client.db.litellm_teamtable.update.assert_not_awaited() + upsert_budget.assert_not_awaited() await _settle_audit_log_tasks() assert audit_logger.payloads == []