From 9bb0ed2f0f6b7b64bda379fe16aa7503737caa5d Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 13 Aug 2026 13:40:30 -0700 Subject: [PATCH 1/6] fix(access groups): sync assigned_team_ids from the team write paths `assigned_team_ids` was only ever written from the access-group side, so a team that dropped an access group on the Teams page kept showing up under the group's Attached Teams forever. That column is not display-only. `get_authorized_resources_from_key_access_groups` reads it as an authorization input, so the stale entry also kept granting the group's models, MCP servers and agents to keys on that team. `/team/new`, `/team/update` (and `PATCH /team/{id}`, which delegates to it) and team delete now mirror the change back onto every affected access group and invalidate its cache entry. --- .../access_group_endpoints.py | 20 +- .../management_endpoints/team_endpoints.py | 10 + .../access_group_team_sync.py | 121 +++++++ .../test_access_group_team_sync.py | 209 ++++++++++++ .../test_team_endpoints.py | 312 ++++++++++++++++++ 5 files changed, 654 insertions(+), 18 deletions(-) create mode 100644 litellm/proxy/management_helpers/access_group_team_sync.py create mode 100644 tests/proxy_admin_ui_tests/test_access_group_team_sync.py diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index c00c2a5ba4c8..2271501d480a 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -14,11 +14,11 @@ _cache_access_object, _cache_key_object, _cache_team_object, - _delete_cache_access_object, _get_team_object_from_cache, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache from litellm.proxy.utils import get_prisma_client_or_throw from litellm.repositories.table_repositories import AccessGroupRepository from litellm.types.access_group import ( @@ -146,22 +146,6 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None: ) -async def _invalidate_cache_access_group(access_group_id: str) -> None: - """ - Invalidate (delete) an access group entry from both in-memory and Redis caches. - - Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server - to avoid circular imports, following the same pattern as key_management_endpoints. - """ - from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache - - await _delete_cache_access_object( - access_group_id=access_group_id, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - # --------------------------------------------------------------------------- # DB sync helpers (called inside a Prisma transaction) # --------------------------------------------------------------------------- @@ -595,7 +579,7 @@ async def delete_access_group( from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache - await _invalidate_cache_access_group(access_group_id) + await invalidate_access_group_cache(access_group_id) await _patch_team_caches_remove_access_group( affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae1bb52278af..d3134373a2ca 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -104,6 +104,9 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) +from litellm.proxy.management_helpers.access_group_team_sync import ( + sync_team_access_group_membership, +) from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, enforce_all_proxy_mcp_servers_grant_is_admin_only, @@ -1509,6 +1512,8 @@ async def new_team( include={"litellm_model_table": True}, ) + await sync_team_access_group_membership(prisma_client=prisma_client, team_id=team_row.team_id) + ## ADD TEAM ID TO USER TABLE ## team_member_add_request: Final = TeamMemberAddRequest( team_id=data.team_id, @@ -2210,6 +2215,7 @@ async def update_team( ) verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) + await sync_team_access_group_membership(prisma_client=prisma_client, team_id=team_row.team_id) await _refresh_cached_team( team_row=team_row, user_api_key_cache=user_api_key_cache, @@ -3790,6 +3796,10 @@ async def delete_team( ## DELETE TEAMS deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team") + + for deleted_team in team_rows: + await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id) + return deleted_teams diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py new file mode 100644 index 000000000000..f40e4d850ab9 --- /dev/null +++ b/litellm/proxy/management_helpers/access_group_team_sync.py @@ -0,0 +1,121 @@ +""" +Reverse sync for the team side of the team <-> access group relationship. + +`litellm_accessgrouptable.assigned_team_ids` and `litellm_teamtable.access_group_ids` +are two copies of the same relationship, and both are read: the access group's +attached-teams view reads the former, and so does the key-side grant check in +`auth_checks.get_authorized_resources_from_key_access_groups`. The access-group +endpoints maintain both copies already; this module is what the team write paths +call so an edit from that side is mirrored back. + +It deliberately lives outside `access_group_endpoints`, which is a lazily +registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that +module eagerly from `team_endpoints` would put it in `sys.modules` without its +router ever being included, which drops its routes from the OpenAPI schema. +""" + +from collections.abc import Mapping, Sequence +from typing import Final, Protocol + +from pydantic import BaseModel, TypeAdapter + +from litellm.proxy.auth.auth_checks import _delete_cache_access_object + +# hashtext collisions only cost two unrelated teams a little serialization, and the +# lock is never taken by the access-group endpoints, so it cannot join their +# access-group-then-team lock order to form a cycle. +_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" + +_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1' + +_ATTACH_SQL: Final = """ +UPDATE "LiteLLM_AccessGroupTable" +SET assigned_team_ids = array_append(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]), $1) +WHERE access_group_id = ANY($2::TEXT[]) + AND NOT ($1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))) +RETURNING access_group_id +""" + +_DETACH_SQL: Final = """ +UPDATE "LiteLLM_AccessGroupTable" +SET assigned_team_ids = array_remove(assigned_team_ids, $1) +WHERE $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])) + AND NOT (access_group_id = ANY($2::TEXT[])) +RETURNING access_group_id +""" + + +class _ChangedGroup(BaseModel): + access_group_id: str + + +class _TeamGroups(BaseModel): + access_group_ids: tuple[str, ...] | None = None + + +_ChangedGroups: Final = TypeAdapter(tuple[_ChangedGroup, ...]) +_TeamRows: Final = TypeAdapter(tuple[_TeamGroups, ...]) + + +class _RawQueryExecutor(Protocol): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + +class _Transaction(Protocol): + async def __aenter__(self) -> _RawQueryExecutor: ... + + async def __aexit__(self, *exc_info: object) -> None: ... + + +class _PrismaDb(Protocol): + def tx(self) -> _Transaction: ... + + +class _PrismaClient(Protocol): + @property + def db(self) -> _PrismaDb: ... + + +async def invalidate_access_group_cache(access_group_id: str) -> None: + """ + Drop an access group entry from both the in-memory and Redis caches. + + Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server + to avoid circular imports, following the same pattern as key_management_endpoints. + """ + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + await _delete_cache_access_object( + access_group_id=access_group_id, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None: + """ + Reconcile every access group's `assigned_team_ids` against the team's own committed + `access_group_ids`, so a team-side edit is visible on the access-group side. + + Call this after the team row is written, or after it is deleted: a team with no row + reconciles to an empty set, which detaches it from every group. + + The team row is read here rather than passed in, under an advisory lock held for the + rest of the transaction. That is what makes concurrent writes to the same team + converge, since each mirror reconciles against the row as committed instead of + against the snapshot its own caller happened to see. It also means a retry heals a + sync that failed partway, where a before/after delta would compute nothing to do. + + Both mirror statements are set-based and mutate the array inside the statement, so a + concurrent write for a different team cannot be lost the way a read-modify-write of + the whole array can, and the pair commits together or not at all. + """ + async with prisma_client.db.tx() as tx: + await tx.query_raw(_LOCK_TEAM_SQL, team_id) + team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id)) + desired: Final = list(team_rows[0].access_group_ids or ()) if team_rows else [] + attached: Final = _ChangedGroups.validate_python(await tx.query_raw(_ATTACH_SQL, team_id, desired)) + detached: Final = _ChangedGroups.validate_python(await tx.query_raw(_DETACH_SQL, team_id, desired)) + + for group in (*attached, *detached): + await invalidate_access_group_cache(group.access_group_id) diff --git a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py new file mode 100644 index 000000000000..d47960cf2996 --- /dev/null +++ b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py @@ -0,0 +1,209 @@ +""" +Real-Postgres coverage for the team -> access group mirror. + +`sync_team_access_group_membership` reconciles `assigned_team_ids` with two raw +statements, and a mocked prisma cannot tell whether that SQL is right: a fake has to +reimplement the array semantics in Python, so it passes no matter what the SQL says. +These tests run the statements against the same Postgres CI seeds for the admin UI +suite, which is the only place a `NOT (... = ANY(...))` guard going missing shows up. +""" + +import asyncio +import os +import sys +from contextlib import asynccontextmanager +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy.management_helpers.access_group_team_sync import ( + sync_team_access_group_membership, +) + +TEAM = "ags-team-a" +OTHER_TEAM = "ags-team-b" +GROUPS = ("ags-group-1", "ags-group-2", "ags-group-3") +_DELETE_SEEDED = 'DELETE FROM "LiteLLM_AccessGroupTable" WHERE access_group_id = ANY($1::TEXT[])' +_DELETE_TEAMS = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::TEXT[])' + + +@asynccontextmanager +async def _clean_db(): + """Connects inside the running test's loop. An async fixture would be torn up on a + different loop than the test body, which prisma's engine lock refuses outright.""" + from prisma import Prisma + + if not os.getenv("DATABASE_URL"): + pytest.fail("DATABASE_URL is required; these tests must not silently skip") + + db = Prisma() + await db.connect() + try: + await db.execute_raw(_DELETE_SEEDED, list(GROUPS)) + await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM]) + yield db + finally: + await db.execute_raw(_DELETE_SEEDED, list(GROUPS)) + await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM]) + await db.disconnect() + + +async def _seed(db, assignments): + for group_id, team_ids in assignments.items(): + await db.litellm_accessgrouptable.create( + data={ + "access_group_id": group_id, + "access_group_name": group_id, + "assigned_team_ids": team_ids, + } + ) + + +async def _read(db): + rows = await db.query_raw( + 'SELECT access_group_id, assigned_team_ids FROM "LiteLLM_AccessGroupTable" ' + "WHERE access_group_id = ANY($1::TEXT[])", + list(GROUPS), + ) + return {row["access_group_id"]: sorted(row["assigned_team_ids"] or []) for row in rows} + + +async def _set_team_groups(db, team_id, access_group_ids): + """The mirror reads the committed team row, so the desired state is written there.""" + if access_group_ids is None: + await db.execute_raw(_DELETE_TEAMS, [team_id]) + return + await db.litellm_teamtable.upsert( + where={"team_id": team_id}, + data={ + "create": {"team_id": team_id, "access_group_ids": list(access_group_ids)}, + "update": {"access_group_ids": list(access_group_ids)}, + }, + ) + + +async def _sync(db, team_id, access_group_ids): + await _set_team_groups(db, team_id, access_group_ids) + with patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ) as invalidate: + await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=team_id) + return {call.args[0] for call in invalidate.call_args_list} + + +@pytest.mark.asyncio +async def test_reconcile_attaches_and_detaches_without_touching_other_teams(): + """The detach must be scoped to groups the team dropped. Losing that scope would + strip the team from the very groups it just kept, silently revoking live grants.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]}) + + invalidated = await _sync(db, TEAM, [GROUPS[1], GROUPS[2]]) + + assert await _read(db) == { + GROUPS[0]: [OTHER_TEAM], + GROUPS[1]: [TEAM], + GROUPS[2]: sorted([TEAM, OTHER_TEAM]), + } + assert invalidated == {GROUPS[0], GROUPS[2]} + + +@pytest.mark.asyncio +async def test_reconcile_is_idempotent_so_a_retry_heals_rather_than_duplicates(): + """Reconciling to the same desired state twice must be a no-op. A delta-based mirror + would instead go quiet after the team row commits, leaving a half-applied sync stuck.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: [TEAM], GROUPS[2]: []}) + + first = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]]) + after_first = await _read(db) + second = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]]) + + assert after_first == {GROUPS[0]: [TEAM], GROUPS[1]: [TEAM], GROUPS[2]: []} + assert await _read(db) == after_first + assert first == {GROUPS[0]} + assert second == set() + + +@pytest.mark.asyncio +async def test_reconcile_handles_a_null_array_column(): + """`assigned_team_ids` is nullable in Postgres. Without COALESCE both statements + evaluate their guard to NULL, skip the row, and the grant silently never syncs.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: []}) + await db.execute_raw( + 'UPDATE "LiteLLM_AccessGroupTable" SET assigned_team_ids = NULL WHERE access_group_id = $1', + GROUPS[0], + ) + + await _sync(db, TEAM, [GROUPS[0]]) + + assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []} + + +@pytest.mark.asyncio +async def test_passing_none_detaches_the_team_from_every_group(): + """Team deletion. A group the deleted row never listed must still let the team go, + otherwise the id dangles under Attached Teams and grants again if it is reused.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]}) + + invalidated = await _sync(db, TEAM, None) + + assert await _read(db) == {GROUPS[0]: [OTHER_TEAM], GROUPS[1]: [], GROUPS[2]: [OTHER_TEAM]} + assert invalidated == {GROUPS[0], GROUPS[1]} + + +@pytest.mark.asyncio +async def test_a_concurrent_writer_cannot_replay_a_stale_team_row_over_a_newer_one(): + """ + Two writers edit one team at once. Whichever team row commits last is the admin's + final intent and the mirror must match it, so the mirror has to hold the team's + advisory lock across its read and its writes. + + A second connection holds that lock and changes the team underneath, which pins the + interleaving instead of hoping a sleep lands in the gap. With the lock the sync waits + and then reads the new row. Without it the sync reads the old row and writes a group + the admin already moved off, which keeps granting to that team. + """ + from prisma import Prisma + + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: []}) + await _sync(db, TEAM, [GROUPS[0]]) + assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []} + + blocker = Prisma() + await blocker.connect() + sync_started = asyncio.Event() + + async def competing_sync(): + sync_started.set() + with patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ): + await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=TEAM) + + try: + async with blocker.tx(timeout=timedelta(seconds=30)) as held: + await held.query_raw("SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked", TEAM) + task = asyncio.create_task(competing_sync()) + await sync_started.wait() + await asyncio.sleep(0.2) + assert not task.done(), "the mirror did not wait on the team's advisory lock" + await held.execute_raw( + 'UPDATE "LiteLLM_TeamTable" SET access_group_ids = $1 WHERE team_id = $2', + [GROUPS[1]], + TEAM, + ) + await asyncio.wait_for(task, timeout=30) + finally: + await blocker.disconnect() + + assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [TEAM]} 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 5190df7521ef..3fbc85df81bf 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 os import sys from datetime import datetime, timezone +from types import SimpleNamespace from typing import Optional, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -6409,6 +6410,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_created_team.rpm_limit = 1000 mock_created_team.metadata = None mock_created_team.members_with_roles = [] + mock_created_team.access_group_ids = None mock_created_team.model_dump.return_value = { "team_id": "new-bypass-team-id", "team_alias": "org-bypass-test-team", @@ -6698,6 +6700,7 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_updated_team.team_id = "org-team-update-bypass-123" mock_updated_team.tpm_limit = 10000 mock_updated_team.rpm_limit = 1000 + mock_updated_team.access_group_ids = None mock_updated_team.model_dump.return_value = { "team_id": "org-team-update-bypass-123", "tpm_limit": 10000, @@ -6851,6 +6854,7 @@ async def test_update_team_guardrails_with_org_id(): "guardrails": ["aporia-pre-call", "aporia-post-call"] } mock_updated_team.litellm_model_table = None + mock_updated_team.access_group_ids = None mock_updated_team.model_dump.return_value = { "team_id": "team-guardrails-123", "organization_id": "test-org-guardrails", @@ -11195,3 +11199,311 @@ async def test_new_team_explicit_null_max_budget_still_takes_configured_default( team_data = mock_team_create.call_args.kwargs["data"] assert team_data.get("max_budget") == 100.0 + + +class _FakeMirrorDb: + """Stands in for prisma inside the access-group mirror. + + Dispatches on the statement so a change to the SQL's shape is visible here, but it + cannot validate the SQL itself: it reimplements the array semantics in Python, so it + passes whatever the statement says. Correctness of the SQL is pinned against a real + Postgres in tests/proxy_admin_ui_tests/test_access_group_team_sync.py. + """ + + def __init__(self, access_groups, teams, plain_lists=False): + self._access_groups = access_groups + self._teams = teams + self._plain_lists = plain_lists + self.transactions = [] + + def _team_ids(self, group_id): + stored = self._access_groups[group_id] + return stored if self._plain_lists else stored["assigned_team_ids"] + + async def _query_raw(self, sql, *args): + assert self._open, "mirror statement ran outside a transaction" + if "pg_advisory_xact_lock" in sql: + self.transactions[-1].append("lock") + return [{"locked": False}] + if "LiteLLM_TeamTable" in sql: + self.transactions[-1].append("read") + team_id = args[0] + if team_id not in self._teams: + return [] + return [{"access_group_ids": list(self._teams[team_id])}] + + team_id, desired = args + if "array_append" in sql: + self.transactions[-1].append("attach") + changed = [ + g for g in desired if g in self._access_groups and team_id not in self._team_ids(g) + ] + for group_id in changed: + self._team_ids(group_id).append(team_id) + else: + self.transactions[-1].append("detach") + changed = [ + g for g in self._access_groups if team_id in self._team_ids(g) and g not in desired + ] + for group_id in changed: + self._team_ids(group_id).remove(team_id) + return [{"access_group_id": group_id} for group_id in changed] + + def tx(self, *_args, **_kwargs): + outer = self + + class _Tx: + async def __aenter__(self): + outer.transactions.append([]) + outer._open = True + return SimpleNamespace(query_raw=outer._query_raw) + + async def __aexit__(self, *_exc_info): + outer._open = False + return None + + return _Tx() + + _open = False + + +@pytest.mark.asyncio +async def test_update_team_syncs_access_group_assigned_team_ids_in_both_directions(): + """ + A team-side edit of `access_group_ids` must be mirrored onto every affected access + group's `assigned_team_ids`, in one transaction, in both directions. + + `assigned_team_ids` is not display-only. `get_authorized_resources_from_key_access_groups` + reads it as an authorization input, so a group the team dropped must stop granting its + resources to keys on that team, and a group the team added must start granting them. + A single-direction assertion would pass against a fix that only ever removes (or only + ever adds), so this covers add, remove, untouched, and the authorization consequence. + """ + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import ( + get_authorized_resources_from_key_access_groups, + ) + + access_groups = { + "ag-drop": {"assigned_team_ids": ["team-a"], "access_model_names": ["dropped-model"]}, + "ag-keep": {"assigned_team_ids": ["team-a"], "access_model_names": ["kept-model"]}, + "ag-add": {"assigned_team_ids": [], "access_model_names": ["added-model"]}, + "ag-other-team": {"assigned_team_ids": ["team-b"], "access_model_names": ["other-model"]}, + } + committed_team_groups = ["ag-keep", "ag-add"] + fake_db = _FakeMirrorDb(access_groups, {"team-a": committed_team_groups}) + + existing_team = MagicMock() + existing_team.access_group_ids = ["ag-drop", "ag-keep"] + existing_team.metadata = {} + existing_team.max_budget = None + existing_team.organization_id = None + existing_team.team_alias = "team-a" + existing_team.model_dump.return_value = {"team_id": "team-a", "team_alias": "team-a"} + + updated_team = MagicMock() + updated_team.team_id = "team-a" + updated_team.access_group_ids = committed_team_groups + updated_team.model_dump.return_value = {"team_id": "team-a"} + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as prisma, + patch("litellm.proxy.proxy_server.llm_router"), + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team"), + patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ) as invalidate_cache, + ): + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + prisma.db.tx = fake_db.tx + prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + await update_team( + data=UpdateTeamRequest(team_id="team-a", access_group_ids=committed_team_groups), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert access_groups["ag-drop"]["assigned_team_ids"] == [] + assert access_groups["ag-add"]["assigned_team_ids"] == ["team-a"] + assert access_groups["ag-keep"]["assigned_team_ids"] == ["team-a"] + assert access_groups["ag-other-team"]["assigned_team_ids"] == ["team-b"] + + assert fake_db.transactions == [["lock", "read", "attach", "detach"]] + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-drop", "ag-add"} + + async def _get_access_object(*, access_group_id, **_kwargs): + stored = access_groups[access_group_id] + return LiteLLM_AccessGroupTable( + access_group_id=access_group_id, + access_group_name=access_group_id, + access_model_names=list(stored["access_model_names"]), + assigned_team_ids=list(stored["assigned_team_ids"]), + assigned_key_ids=[], + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + side_effect=_get_access_object, + ), + ): + authorized_models = await get_authorized_resources_from_key_access_groups( + valid_token=UserAPIKeyAuth( + token="sk-hash", + models=[], + team_id="team-a", + access_group_ids=["ag-drop", "ag-keep", "ag-add"], + ), + team_object=LiteLLM_TeamTable(team_id="team-a", models=[]), + resource_field="access_model_names", + ) + + assert sorted(authorized_models) == ["added-model", "kept-model"] + + +@pytest.mark.asyncio +async def test_sync_reads_the_committed_team_row_rather_than_the_callers_snapshot(): + """ + The mirror takes no desired-state argument on purpose. It locks the team and reads + the row as committed, so two concurrent writers for one team converge on the row the + last one committed instead of each replaying its own stale snapshot. Reconciling also + means a retry heals a half-applied sync, where a before/after delta computes nothing. + + A team with no row at all is deletion, and must detach from every group. + """ + from litellm.proxy.management_helpers.access_group_team_sync import ( + sync_team_access_group_membership, + ) + + access_groups = {"ag-1": ["team-a", "team-b"], "ag-2": ["team-a"], "ag-3": []} + teams = {"team-a": ["ag-2", "ag-3"]} + fake_db = _FakeMirrorDb(access_groups, teams, plain_lists=True) + prisma_client = SimpleNamespace(db=SimpleNamespace(tx=fake_db.tx)) + + with patch( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + new_callable=AsyncMock, + ) as invalidate_cache: + await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") + assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]} + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1", "ag-3"} + + invalidate_cache.reset_mock() + await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") + assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]} + assert invalidate_cache.call_args_list == [] + + invalidate_cache.reset_mock() + del teams["team-a"] + await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") + assert access_groups == {"ag-1": ["team-b"], "ag-2": [], "ag-3": []} + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"} + + assert fake_db.transactions == [["lock", "read", "attach", "detach"]] * 3 + + +@pytest.mark.asyncio +async def test_new_team_and_delete_team_both_drive_the_mirror(): + """Every writer of `team.access_group_ids` has to reach the mirror, not just update. + These pin the wiring on the other two paths; the mirror's own behavior is covered above.""" + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import DeleteTeamRequest, NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import delete_team, new_team + + created = MagicMock() + created.team_id = "team-new" + created.access_group_ids = ["ag-1"] + created.model_dump.return_value = {"team_id": "team-new"} + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as prisma, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch("litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", new_callable=AsyncMock), + patch( + "litellm.proxy.management_endpoints.team_endpoints.sync_team_access_group_membership", + new_callable=AsyncMock, + ) as sync, + ): + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + prisma.db.litellm_teamtable.create = AsyncMock(return_value=created) + prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + prisma.get_data = AsyncMock(return_value=None) + + await new_team( + data=NewTeamRequest(team_id="team-new", team_alias="new"), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert sync.await_args_list[0].kwargs["team_id"] == "team-new" + + team_row = LiteLLM_TeamTable(team_id="team-gone", models=[], access_group_ids=["ag-1"]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as prisma, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.management_endpoints.team_endpoints._persist_deleted_team_records", new_callable=AsyncMock), + patch("litellm.proxy.management_endpoints.team_endpoints._verify_team_access", new_callable=AsyncMock), + patch( + "litellm.proxy.management_endpoints.team_endpoints.sync_team_access_group_membership", + new_callable=AsyncMock, + ) as sync, + ): + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma.delete_data = AsyncMock(return_value=[team_row]) + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-gone"]), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert sync.await_args_list[0].kwargs["team_id"] == "team-gone" + + +@pytest.mark.asyncio +async def test_invalidate_access_group_cache_deletes_the_cached_object(): + """The mirror's cache step is what stops a revoked group granting from cache until TTL, + so pin that it actually reaches the delete rather than only being called.""" + from litellm.proxy.management_helpers.access_group_team_sync import ( + invalidate_access_group_cache, + ) + + cache, logging_obj = MagicMock(), MagicMock() + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", logging_obj), + patch( + "litellm.proxy.management_helpers.access_group_team_sync._delete_cache_access_object", + new_callable=AsyncMock, + ) as delete_cached, + ): + await invalidate_access_group_cache("ag-1") + + assert delete_cached.await_args.kwargs == { + "access_group_id": "ag-1", + "user_api_key_cache": cache, + "proxy_logging_obj": logging_obj, + } From d0c4d2f326f223908a63f68e7eeac545c500b6e2 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 14 Aug 2026 00:56:25 +0000 Subject: [PATCH 2/6] fix(access groups): create the team and its mirror in one transaction A sync that failed after the insert committed left a team whose groups never learned about it, and the retry came back as a duplicate team id. Invalidate off the reconciled set so a retry after an unreachable cache still drops the stale grants. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 21 ++++- .../access_group_team_sync.py | 74 +++++++++++---- .../test_access_group_team_sync.py | 31 +++++- .../test_team_endpoints.py | 94 +++++++++++++++---- .../proxy/management_endpoints/test_ui_sso.py | 18 ++++ .../test_access_group_team_sync.py | 39 ++++++++ 6 files changed, 230 insertions(+), 47 deletions(-) create mode 100644 tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d3134373a2ca..8f69a9909c64 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -105,6 +105,9 @@ get_daily_activity, ) from litellm.proxy.management_helpers.access_group_team_sync import ( + AccessGroupSyncTx, + invalidate_access_group_caches, + reconcile_team_access_group_membership, sync_team_access_group_membership, ) from litellm.proxy.management_helpers.object_permission_utils import ( @@ -316,6 +319,11 @@ class _TeamIdInFilter(TypedDict, total=False): team_id: Mapping[str, Sequence[str]] +class _TeamCreateTx(AccessGroupSyncTx, Protocol): + @property + def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ... + + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) @@ -1507,12 +1515,15 @@ async def new_team( complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict - team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create( - data=team_creation_data, - include={"litellm_model_table": True}, - ) + tx: _TeamCreateTx + async with prisma_client.db.tx() as tx: + team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create( + data=team_creation_data, + include={"litellm_model_table": True}, + ) + affected_access_groups: Final = await reconcile_team_access_group_membership(tx, team_row.team_id) - await sync_team_access_group_membership(prisma_client=prisma_client, team_id=team_row.team_id) + await invalidate_access_group_caches(affected_access_groups) ## ADD TEAM ID TO USER TABLE ## team_member_add_request: Final = TeamMemberAddRequest( diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py index f40e4d850ab9..ce3c80db421b 100644 --- a/litellm/proxy/management_helpers/access_group_team_sync.py +++ b/litellm/proxy/management_helpers/access_group_team_sync.py @@ -14,6 +14,7 @@ router ever being included, which drops its routes from the OpenAPI schema. """ +import asyncio from collections.abc import Mapping, Sequence from typing import Final, Protocol @@ -28,6 +29,15 @@ _READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1' +# The groups the team is on either side of the reconcile, so the cache step is driven by +# desired state rather than by which rows this attempt happened to change. A retry after a +# failed invalidation finds the same set even though its statements are already no-ops. +_AFFECTED_SQL: Final = """ +SELECT access_group_id FROM "LiteLLM_AccessGroupTable" +WHERE access_group_id = ANY($2::TEXT[]) + OR $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])) +""" + _ATTACH_SQL: Final = """ UPDATE "LiteLLM_AccessGroupTable" SET assigned_team_ids = array_append(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]), $1) @@ -45,7 +55,7 @@ """ -class _ChangedGroup(BaseModel): +class _AffectedGroup(BaseModel): access_group_id: str @@ -53,16 +63,16 @@ class _TeamGroups(BaseModel): access_group_ids: tuple[str, ...] | None = None -_ChangedGroups: Final = TypeAdapter(tuple[_ChangedGroup, ...]) +_AffectedGroups: Final = TypeAdapter(tuple[_AffectedGroup, ...]) _TeamRows: Final = TypeAdapter(tuple[_TeamGroups, ...]) -class _RawQueryExecutor(Protocol): +class AccessGroupSyncTx(Protocol): async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... class _Transaction(Protocol): - async def __aenter__(self) -> _RawQueryExecutor: ... + async def __aenter__(self) -> AccessGroupSyncTx: ... async def __aexit__(self, *exc_info: object) -> None: ... @@ -92,30 +102,54 @@ async def invalidate_access_group_cache(access_group_id: str) -> None: ) -async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None: +async def invalidate_access_group_caches(access_group_ids: Sequence[str]) -> None: + """ + Drop every given access group from the caches, then raise if any drop failed. + + Every entry is attempted even when one raises, so a single unreachable cache cannot + leave the rest of the reconciled groups serving a grant the admin revoked. + """ + outcomes: Final = await asyncio.gather( + *(invalidate_access_group_cache(access_group_id) for access_group_id in access_group_ids), + return_exceptions=True, + ) + for outcome in outcomes: + if isinstance(outcome, BaseException): + raise outcome + + +async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: str) -> tuple[str, ...]: """ - Reconcile every access group's `assigned_team_ids` against the team's own committed - `access_group_ids`, so a team-side edit is visible on the access-group side. + Reconcile every access group's `assigned_team_ids` against the team's own + `access_group_ids`, and return the groups whose cache the caller has to drop once the + transaction commits. - Call this after the team row is written, or after it is deleted: a team with no row - reconciles to an empty set, which detaches it from every group. + Call this inside the transaction that writes the team row, or after that row is + written or deleted: a team with no row reconciles to an empty set, which detaches it + from every group. The team row is read here rather than passed in, under an advisory lock held for the rest of the transaction. That is what makes concurrent writes to the same team - converge, since each mirror reconciles against the row as committed instead of - against the snapshot its own caller happened to see. It also means a retry heals a - sync that failed partway, where a before/after delta would compute nothing to do. + converge, since each mirror reconciles against the row as the transaction sees it + instead of against the snapshot its own caller happened to see. It also means a retry + heals a sync that failed partway, where a before/after delta would compute nothing. Both mirror statements are set-based and mutate the array inside the statement, so a concurrent write for a different team cannot be lost the way a read-modify-write of the whole array can, and the pair commits together or not at all. """ + await tx.query_raw(_LOCK_TEAM_SQL, team_id) + team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id)) + desired: Final = list(team_rows[0].access_group_ids or ()) if team_rows else [] + affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired)) + await tx.query_raw(_ATTACH_SQL, team_id, desired) + await tx.query_raw(_DETACH_SQL, team_id, desired) + return tuple(group.access_group_id for group in affected) + + +async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None: + """Reconcile the mirror for an already committed team write, in its own transaction.""" async with prisma_client.db.tx() as tx: - await tx.query_raw(_LOCK_TEAM_SQL, team_id) - team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id)) - desired: Final = list(team_rows[0].access_group_ids or ()) if team_rows else [] - attached: Final = _ChangedGroups.validate_python(await tx.query_raw(_ATTACH_SQL, team_id, desired)) - detached: Final = _ChangedGroups.validate_python(await tx.query_raw(_DETACH_SQL, team_id, desired)) - - for group in (*attached, *detached): - await invalidate_access_group_cache(group.access_group_id) + affected: Final = await reconcile_team_access_group_membership(tx, team_id) + + await invalidate_access_group_caches(affected) diff --git a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py index d47960cf2996..629d77f20fc5 100644 --- a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py +++ b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py @@ -21,6 +21,7 @@ sys.path.insert(0, os.path.abspath("../..")) from litellm.proxy.management_helpers.access_group_team_sync import ( + reconcile_team_access_group_membership, sync_team_access_group_membership, ) @@ -110,13 +111,15 @@ async def test_reconcile_attaches_and_detaches_without_touching_other_teams(): GROUPS[1]: [TEAM], GROUPS[2]: sorted([TEAM, OTHER_TEAM]), } - assert invalidated == {GROUPS[0], GROUPS[2]} + assert invalidated == {GROUPS[0], GROUPS[1], GROUPS[2]} @pytest.mark.asyncio async def test_reconcile_is_idempotent_so_a_retry_heals_rather_than_duplicates(): - """Reconciling to the same desired state twice must be a no-op. A delta-based mirror - would instead go quiet after the team row commits, leaving a half-applied sync stuck.""" + """Reconciling to the same desired state twice must leave the rows alone and still name + the team's groups for the cache step, so a retry after a failed cache drop reaches them. + A delta-based mirror would instead go quiet once the rows match, leaving the caches + serving a grant the admin already revoked.""" async with _clean_db() as db: await _seed(db, {GROUPS[0]: [], GROUPS[1]: [TEAM], GROUPS[2]: []}) @@ -126,8 +129,8 @@ async def test_reconcile_is_idempotent_so_a_retry_heals_rather_than_duplicates() assert after_first == {GROUPS[0]: [TEAM], GROUPS[1]: [TEAM], GROUPS[2]: []} assert await _read(db) == after_first - assert first == {GROUPS[0]} - assert second == set() + assert first == {GROUPS[0], GROUPS[1]} + assert second == first @pytest.mark.asyncio @@ -159,6 +162,24 @@ async def test_passing_none_detaches_the_team_from_every_group(): assert invalidated == {GROUPS[0], GROUPS[1]} +@pytest.mark.asyncio +async def test_a_failed_mirror_takes_the_new_team_row_with_it(): + """`/team/new` inserts the team and mirrors it in one transaction. Mirroring in a + transaction of its own instead leaves a committed team whose groups never learned about + it, and the retry with that same team id comes back as a duplicate.""" + async with _clean_db() as db: + await _seed(db, {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]}) + + with pytest.raises(RuntimeError): + async with db.tx() as tx: + await tx.litellm_teamtable.create(data={"team_id": TEAM, "access_group_ids": [GROUPS[0]]}) + await reconcile_team_access_group_membership(tx, TEAM) + raise RuntimeError("the cache handoff blew up") + + assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]} + assert await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) is None + + @pytest.mark.asyncio async def test_a_concurrent_writer_cannot_replay_a_stale_team_row_over_a_newer_one(): """ 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 3fbc85df81bf..3c5d79691506 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2,6 +2,7 @@ import json import os import sys +from contextlib import asynccontextmanager from datetime import datetime, timezone from types import SimpleNamespace from typing import Optional, cast @@ -68,6 +69,21 @@ # Setup TestClient client = TestClient(app) + +def _wire_team_create_tx(prisma_client): + """`/team/new` inserts the team and mirrors it onto the access groups in one transaction, + so a mocked client has to hand its team table back out of `db.tx()`.""" + + @asynccontextmanager + async def _tx(): + yield SimpleNamespace( + litellm_teamtable=prisma_client.db.litellm_teamtable, + query_raw=AsyncMock(return_value=[]), + ) + + prisma_client.db.tx = lambda *_args, **_kwargs: _tx() + + # Mock prisma_client mock_prisma_client = MagicMock() # Set up async mock for db operations @@ -400,6 +416,7 @@ async def test_new_team_rejects_a_duration_that_never_advances( mock_team_create = AsyncMock() mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) with pytest.raises(ProxyException) as exc_info: await new_team( @@ -481,6 +498,7 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): mock_team_count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = mock_team_count mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -570,6 +588,7 @@ async def mock_obj_perm_create(**kwargs): mock_db_client.db.litellm_teamtable.create = AsyncMock( return_value=team_create_result ) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -663,6 +682,7 @@ async def test_new_team_disable_auto_add_proxy_admin_flag( mock_db_client.db.litellm_teamtable.create = AsyncMock( return_value=team_create_result ) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) @@ -4273,6 +4293,7 @@ async def test_new_team_max_budget_within_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -4416,6 +4437,7 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -4564,6 +4586,7 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -6422,6 +6445,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -7422,6 +7446,7 @@ async def test_new_team_soft_budget_validation( mock_prisma.db.litellm_teamtable.create = AsyncMock( return_value=mock_created_team ) + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_created_team ) @@ -7721,6 +7746,7 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): mock_team_count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = mock_team_count mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -9111,6 +9137,7 @@ async def test_new_team_encrypts_callback_vars( team_create_result.model_dump.return_value = {"team_id": "team-456"} mock_team_create = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result @@ -10271,6 +10298,7 @@ async def test_new_team_validator_runs_without_metadata_and_rejection_blocks_cre ): mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_license.is_team_count_over_limit.return_value = False with pytest.raises(ProxyException) as exc_info: @@ -10305,6 +10333,7 @@ async def test_new_team_validator_accept_proceeds_to_create(mock_db_client, mock team_create_result.model_dump.return_value = {"team_id": "team-accept-1"} mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) mock_db_client.db.litellm_usertable = MagicMock() @@ -10344,6 +10373,7 @@ async def test_new_team_rejection_precedes_model_alias_write(): ): mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) mock_license.is_team_count_over_limit.return_value = False @@ -11111,6 +11141,7 @@ def _wire_new_team_prisma(mock_db_client): mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=created_team) + _wire_team_create_tx(mock_db_client) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=created_team) mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) @@ -11233,6 +11264,11 @@ async def _query_raw(self, sql, *args): return [{"access_group_ids": list(self._teams[team_id])}] team_id, desired = args + if sql.lstrip().startswith("SELECT"): + self.transactions[-1].append("affected") + affected = [g for g in self._access_groups if g in desired or team_id in self._team_ids(g)] + return [{"access_group_id": group_id} for group_id in affected] + if "array_append" in sql: self.transactions[-1].append("attach") changed = [ @@ -11249,6 +11285,16 @@ async def _query_raw(self, sql, *args): self._team_ids(group_id).remove(team_id) return [{"access_group_id": group_id} for group_id in changed] + async def _create_team(self, data, include=None): + self.transactions[-1].append("create") + team_id = data["team_id"] + self._teams[team_id] = list(data.get("access_group_ids") or ()) + return SimpleNamespace( + team_id=team_id, + access_group_ids=list(self._teams[team_id]), + model_dump=lambda: {"team_id": team_id}, + ) + def tx(self, *_args, **_kwargs): outer = self @@ -11256,7 +11302,10 @@ class _Tx: async def __aenter__(self): outer.transactions.append([]) outer._open = True - return SimpleNamespace(query_raw=outer._query_raw) + return SimpleNamespace( + query_raw=outer._query_raw, + litellm_teamtable=SimpleNamespace(create=outer._create_team), + ) async def __aexit__(self, *_exc_info): outer._open = False @@ -11338,8 +11387,8 @@ async def test_update_team_syncs_access_group_assigned_team_ids_in_both_directio assert access_groups["ag-keep"]["assigned_team_ids"] == ["team-a"] assert access_groups["ag-other-team"]["assigned_team_ids"] == ["team-b"] - assert fake_db.transactions == [["lock", "read", "attach", "detach"]] - assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-drop", "ag-add"} + assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]] + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-drop", "ag-keep", "ag-add"} async def _get_access_object(*, access_group_id, **_kwargs): stored = access_groups[access_group_id] @@ -11383,6 +11432,10 @@ async def test_sync_reads_the_committed_team_row_rather_than_the_callers_snapsho last one committed instead of each replaying its own stale snapshot. Reconciling also means a retry heals a half-applied sync, where a before/after delta computes nothing. + The same holds for the cache step: the groups to drop come from the reconciled set, + not from the rows this attempt happened to change, so a retry after an unreachable + cache still drops the entries even though its statements are now no-ops. + A team with no row at all is deletion, and must detach from every group. """ from litellm.proxy.management_helpers.access_group_team_sync import ( @@ -11397,15 +11450,18 @@ async def test_sync_reads_the_committed_team_row_rather_than_the_callers_snapsho with patch( "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", new_callable=AsyncMock, + side_effect=[ConnectionError("redis unreachable"), None, None], ) as invalidate_cache: - await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") + with pytest.raises(ConnectionError): + await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]} - assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1", "ag-3"} + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1", "ag-2", "ag-3"} invalidate_cache.reset_mock() + invalidate_cache.side_effect = None await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a") assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]} - assert invalidate_cache.call_args_list == [] + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"} invalidate_cache.reset_mock() del teams["team-a"] @@ -11413,13 +11469,17 @@ async def test_sync_reads_the_committed_team_row_rather_than_the_callers_snapsho assert access_groups == {"ag-1": ["team-b"], "ag-2": [], "ag-3": []} assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"} - assert fake_db.transactions == [["lock", "read", "attach", "detach"]] * 3 + assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]] * 3 @pytest.mark.asyncio async def test_new_team_and_delete_team_both_drive_the_mirror(): """Every writer of `team.access_group_ids` has to reach the mirror, not just update. - These pin the wiring on the other two paths; the mirror's own behavior is covered above.""" + These pin the wiring on the other two paths; the mirror's own behavior is covered above. + + Creation has to insert the team row and mirror it in one transaction. With the mirror + in a transaction of its own, a sync that fails leaves a committed team whose groups + never learned about it, and the retry is rejected as a duplicate team id.""" from unittest.mock import Mock from fastapi import Request @@ -11427,10 +11487,8 @@ async def test_new_team_and_delete_team_both_drive_the_mirror(): from litellm.proxy._types import DeleteTeamRequest, NewTeamRequest from litellm.proxy.management_endpoints.team_endpoints import delete_team, new_team - created = MagicMock() - created.team_id = "team-new" - created.access_group_ids = ["ag-1"] - created.model_dump.return_value = {"team_id": "team-new"} + access_groups = {"ag-1": [], "ag-2": []} + fake_db = _FakeMirrorDb(access_groups, {}, plain_lists=True) with ( patch("litellm.proxy.proxy_server.prisma_client") as prisma, @@ -11439,23 +11497,25 @@ async def test_new_team_and_delete_team_both_drive_the_mirror(): patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch("litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", new_callable=AsyncMock), patch( - "litellm.proxy.management_endpoints.team_endpoints.sync_team_access_group_membership", + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", new_callable=AsyncMock, - ) as sync, + ) as invalidate_cache, ): prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) - prisma.db.litellm_teamtable.create = AsyncMock(return_value=created) + prisma.db.tx = fake_db.tx prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) prisma.get_data = AsyncMock(return_value=None) await new_team( - data=NewTeamRequest(team_id="team-new", team_alias="new"), + data=NewTeamRequest(team_id="team-new", team_alias="new", access_group_ids=["ag-1"]), http_request=Mock(spec=Request), user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), ) - assert sync.await_args_list[0].kwargs["team_id"] == "team-new" + assert access_groups == {"ag-1": ["team-new"], "ag-2": []} + assert fake_db.transactions == [["create", "lock", "read", "affected", "attach", "detach"]] + assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1"} team_row = LiteLLM_TeamTable(team_id="team-gone", models=[], access_group_ids=["ag-1"]) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 979eb09d7dbb..147aeda745d4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2,6 +2,7 @@ import json import os import sys +from contextlib import asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -37,6 +38,20 @@ ) +def _wire_team_create_tx(prisma_client): + """`/team/new` inserts the team and mirrors it onto the access groups in one transaction, + so a mocked client has to hand its team table back out of `db.tx()`.""" + + @asynccontextmanager + async def _tx(): + yield SimpleNamespace( + litellm_teamtable=prisma_client.db.litellm_teamtable, + query_raw=AsyncMock(return_value=[]), + ) + + prisma_client.db.tx = lambda *_args, **_kwargs: _tx() + + def test_microsoft_sso_handler_openid_from_response_user_principal_name(): # Arrange # Create a mock response similar to what Microsoft SSO would return @@ -577,6 +592,7 @@ def mock_jsonify_team_object(db_data): mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object) @@ -624,6 +640,7 @@ async def test_default_team_params_organization_id_reaches_sso_created_team(team mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) @@ -671,6 +688,7 @@ def mock_jsonify_team_object(db_data): mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.create = AsyncMock() + _wire_team_create_tx(mock_prisma) mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object) diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py new file mode 100644 index 000000000000..eb11292cf422 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py @@ -0,0 +1,39 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.management_helpers.access_group_team_sync import ( + invalidate_access_group_caches, +) + + +@pytest.mark.asyncio +async def test_one_unreachable_cache_does_not_skip_the_other_groups(monkeypatch): + """ + `assigned_team_ids` is an authorization input, so a group whose cache still holds the + revoked grant keeps serving it until the entry is dropped. + + A sequential loop would stop at the first failing group and leave the groups behind it + serving stale grants, and swallowing the failure would report success to the admin for + a revoke that never took effect. Every group has to be attempted, and the endpoint has + to fail so the caller can retry. + """ + attempted: list[str] = [] + + async def _invalidate(access_group_id: str) -> None: + attempted.append(access_group_id) + if access_group_id == "ag-redis-down": + raise ConnectionError("redis unreachable") + + monkeypatch.setattr( + "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache", + _invalidate, + ) + + with pytest.raises(ConnectionError): + await invalidate_access_group_caches(("ag-redis-down", "ag-2", "ag-3")) + + assert attempted == ["ag-redis-down", "ag-2", "ag-3"] From 5399f0df5fb0888044810f6b15c0c1aa886f02a0 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 14 Aug 2026 01:24:43 +0000 Subject: [PATCH 3/6] fix(access groups): keep the mirror wired after the delete-path rework Staging moved team delete's cache eviction and reference sweep after the row delete, so the mirror runs last where it still sees the team gone. Immutable desired set and include mapping to stay inside the LIT002 ceiling. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/team_endpoints.py | 5 ++++- litellm/proxy/management_helpers/access_group_team_sync.py | 2 +- .../proxy/management_endpoints/test_team_endpoints.py | 2 ++ 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 872e043e5d29..3d7f0808fb9d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -15,6 +15,7 @@ import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast import fastapi @@ -330,6 +331,8 @@ def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ... UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams) """ +_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True}) + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) @@ -1526,7 +1529,7 @@ async def new_team( async with prisma_client.db.tx() as tx: team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create( data=team_creation_data, - include={"litellm_model_table": True}, + include=_INCLUDE_MODEL_TABLE, ) affected_access_groups: Final = await reconcile_team_access_group_membership(tx, team_row.team_id) diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py index ce3c80db421b..55c0346e375e 100644 --- a/litellm/proxy/management_helpers/access_group_team_sync.py +++ b/litellm/proxy/management_helpers/access_group_team_sync.py @@ -140,7 +140,7 @@ async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: """ await tx.query_raw(_LOCK_TEAM_SQL, team_id) team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id)) - desired: Final = list(team_rows[0].access_group_ids or ()) if team_rows else [] + desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else () affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired)) await tx.query_raw(_ATTACH_SQL, team_id, desired) await tx.query_raw(_DETACH_SQL, team_id, desired) 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 3de275852160..c6960ecda5ad 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -12052,6 +12052,8 @@ async def test_new_team_and_delete_team_both_drive_the_mirror(): prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) prisma.delete_data = AsyncMock(return_value=[team_row]) + prisma.db.execute_raw = AsyncMock(return_value=0) + prisma.db.litellm_teammembership.delete_many = AsyncMock(return_value=0) await delete_team( data=DeleteTeamRequest(team_ids=["team-gone"]), From 81ac9d48d45614c7721c7870df93c2d1558edd8f Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 14 Aug 2026 02:51:49 +0000 Subject: [PATCH 4/6] chore: retrigger ci Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From 672be61b53356bcc946d41d11bc09619931c58af Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 14 Aug 2026 03:40:10 +0000 Subject: [PATCH 5/6] fix(tests): stop a monkeypatched prisma_client from leaking across the xdist worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/conftest.py | 60 ++++++++++++------- .../proxy/test_proxy_global_isolation.py | 49 +++++++++++++++ 2 files changed, 87 insertions(+), 22 deletions(-) create mode 100644 tests/test_litellm/proxy/test_proxy_global_isolation.py diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 230ccaf5fd4e..61e6c4c9a1ac 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -8,7 +8,9 @@ import asyncio import os import tempfile -from typing import Dict, Optional +from collections.abc import Mapping +from types import MappingProxyType +from typing import Dict, Final, Optional import pytest import yaml @@ -20,6 +22,9 @@ "prisma_client", ) +_MISSING: Final = object() +_PROXY_GLOBALS_SNAPSHOT: Final = pytest.StashKey[Mapping[str, object]]() + class StubClientNotConnectedError(ClientNotConnectedError): pass @@ -43,32 +48,43 @@ def disconnected_prisma() -> DisconnectedPrisma: return DisconnectedPrisma() -@pytest.fixture(autouse=True) -def _isolate_proxy_module_globals(): +@pytest.hookimpl(tryfirst=True) +def pytest_runtest_setup(item: pytest.Item) -> None: """ - Snapshot and restore module-level globals on litellm.proxy.proxy_server - that tests sometimes mutate via raw setattr (not monkeypatch). + Snapshot module-level globals on litellm.proxy.proxy_server that tests mutate. - Without this, a leaked value — e.g. master_key set by a sibling test — - flips the auth short-circuit in user_api_key_auth and causes unrelated - tests in the same xdist worker to return 401 instead of 200. + Without this, a leaked value, e.g. master_key set by a sibling test, flips the auth + short-circuit in user_api_key_auth and causes unrelated tests in the same xdist worker + to return 401 instead of 200. """ from litellm.proxy import proxy_server - sentinel = object() - snapshot = { - name: getattr(proxy_server, name, sentinel) - for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE - } - try: - yield - finally: - for name, value in snapshot.items(): - if value is sentinel: - if hasattr(proxy_server, name): - delattr(proxy_server, name) - else: - setattr(proxy_server, name, value) + item.stash[_PROXY_GLOBALS_SNAPSHOT] = MappingProxyType( + {name: getattr(proxy_server, name, _MISSING) for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE} + ) + + +@pytest.hookimpl(trylast=True) +def pytest_runtest_teardown(item: pytest.Item) -> None: + """ + Restore the snapshot after every fixture finalizer, `monkeypatch` undo included. + + A fixture cannot do this: `monkeypatch` tears down after the autouse fixtures, so its undo + reinstalls whatever the value was when the test called setattr, which for a test that also + holds an autouse `patch` of the same global is that patch's mock, leaked for the rest of the + worker's session. + """ + from litellm.proxy import proxy_server + + snapshot: Final = item.stash.get(_PROXY_GLOBALS_SNAPSHOT, None) + if snapshot is None: + return + for name, value in snapshot.items(): + if value is _MISSING: + if hasattr(proxy_server, name): + delattr(proxy_server, name) + else: + setattr(proxy_server, name, value) @pytest.fixture(autouse=True) diff --git a/tests/test_litellm/proxy/test_proxy_global_isolation.py b/tests/test_litellm/proxy/test_proxy_global_isolation.py new file mode 100644 index 000000000000..3d211be57b35 --- /dev/null +++ b/tests/test_litellm/proxy/test_proxy_global_isolation.py @@ -0,0 +1,49 @@ +"""Regression coverage for the proxy_server global isolation hooks in conftest.py.""" + +from pathlib import Path +from typing import Final + +import pytest + +pytest_plugins: Final = ("pytester",) + +_PARENT_CONFTEST_SHAPE: Final = ''' + +@pytest.fixture(autouse=True) +def _early_monkeypatch_user(monkeypatch): + """Mirrors tests/test_litellm/conftest.py, whose autouse env isolation pulls in monkeypatch + before anything else, so monkeypatch undo lands after every other finalizer.""" + yield +''' + +_CONFTEST_SOURCE: Final = (Path(__file__).parent / "conftest.py").read_text() + _PARENT_CONFTEST_SHAPE + +_LEAKY_MODULE: Final = ''' +from unittest.mock import MagicMock, patch + +import pytest + + +@pytest.fixture(autouse=True) +def _autouse_patched_prisma_client(): + with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()): + yield + + +def test_monkeypatches_the_already_patched_global(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) +''' + +_WITNESS_MODULE: Final = ''' +def test_the_real_global_is_back(): + from litellm.proxy import proxy_server + + assert proxy_server.prisma_client is None +''' + + +def test_a_monkeypatched_prisma_client_cannot_outlive_its_test(pytester: pytest.Pytester) -> None: + pytester.makeconftest(_CONFTEST_SOURCE) + pytester.makepyfile(test_a_leaks=_LEAKY_MODULE, test_b_witness=_WITNESS_MODULE) + + pytester.runpytest("-p", "no:randomly").assert_outcomes(passed=2) From 4b274230f51b6e61dcbea0acd71d9e50bf547c0b Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 14 Aug 2026 04:05:30 +0000 Subject: [PATCH 6/6] test: run the isolation regression in a clean subprocess so CI addopts cannot break it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/test_proxy_global_isolation.py | 42 ++++++++++++++----- 1 file changed, 32 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_global_isolation.py b/tests/test_litellm/proxy/test_proxy_global_isolation.py index 3d211be57b35..9ebda8edb8b2 100644 --- a/tests/test_litellm/proxy/test_proxy_global_isolation.py +++ b/tests/test_litellm/proxy/test_proxy_global_isolation.py @@ -1,11 +1,13 @@ """Regression coverage for the proxy_server global isolation hooks in conftest.py.""" +import os +import subprocess +import sys from pathlib import Path from typing import Final -import pytest - -pytest_plugins: Final = ("pytester",) +_PROXY_TESTS_DIR: Final = Path(__file__).parent +_REPO_ROOT: Final = _PROXY_TESTS_DIR.parents[2] _PARENT_CONFTEST_SHAPE: Final = ''' @@ -16,8 +18,6 @@ def _early_monkeypatch_user(monkeypatch): yield ''' -_CONFTEST_SOURCE: Final = (Path(__file__).parent / "conftest.py").read_text() + _PARENT_CONFTEST_SHAPE - _LEAKY_MODULE: Final = ''' from unittest.mock import MagicMock, patch @@ -42,8 +42,30 @@ def test_the_real_global_is_back(): ''' -def test_a_monkeypatched_prisma_client_cannot_outlive_its_test(pytester: pytest.Pytester) -> None: - pytester.makeconftest(_CONFTEST_SOURCE) - pytester.makepyfile(test_a_leaks=_LEAKY_MODULE, test_b_witness=_WITNESS_MODULE) - - pytester.runpytest("-p", "no:randomly").assert_outcomes(passed=2) +def test_a_monkeypatched_prisma_client_cannot_outlive_its_test(tmp_path: Path) -> None: + (tmp_path / "conftest.py").write_text((_PROXY_TESTS_DIR / "conftest.py").read_text() + _PARENT_CONFTEST_SHAPE) + (tmp_path / "test_a_leaks.py").write_text(_LEAKY_MODULE) + (tmp_path / "test_b_witness.py").write_text(_WITNESS_MODULE) + + completed: Final = subprocess.run( + ( + sys.executable, + "-m", + "pytest", + "test_a_leaks.py", + "test_b_witness.py", + "-q", + "-o", + "addopts=", + "-p", + "no:randomly", + "-p", + "no:cacheprovider", + ), + cwd=tmp_path, + capture_output=True, + text=True, + env={**os.environ, "PYTEST_ADDOPTS": "", "PYTHONPATH": str(_REPO_ROOT)}, + ) + + assert completed.returncode == 0, completed.stdout + completed.stderr