fix(access_groups): reconcile update deltas against the derived team set

Team membership deltas on PUT now start from the teams that really carry
the group, so a team the mirror column missed can be detached. Read
endpoints go through a typed TeamRepository instead of the untyped db
handle, and the where clause always carries both OR arms.

Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL

litellm/proxy/management_endpoints/access_group_endpoints.py

@@ -21,7 +21,7 @@ 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.repositories.table_repositories import AccessGroupRepository, TeamRepository
 from litellm.types.access_group import (
     AccessGroupCreateRequest,
     AccessGroupResponse,
@@ -75,11 +75,11 @@ class _AccessGroupTable(Protocol):
 
 
 class _TeamTable(Protocol):
-    async def find_unique(self, where: Mapping[str, object]) -> _TeamRecord | None: ...
+    async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ...
 
-    async def find_many(self, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ...
+    async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ...
 
-    async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
+    async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
 
 
 class _KeyTable(Protocol):
@@ -153,17 +153,16 @@ async def _attached_team_ids_for(
     stored_team_ids: Final = tuple(
         dict.fromkeys(team_id for record in records for team_id in (record.assigned_team_ids or ()))
     )
-    carrying: Final = {"access_group_ids": {"hasSome": group_ids}}  # mutable-ok: prisma where must be a dict
-    listed: Final = {"team_id": {"in": stored_team_ids}}  # mutable-ok: prisma where must be a dict
-    clauses: Final = (carrying, listed) if stored_team_ids else (carrying,)
-    where: Final = {"OR": clauses}  # mutable-ok: prisma where must be a dict
-    return _attached_team_ids_by_group(records, await team_table.find_many(where=where))
+    carrying: Final = {"access_group_ids": {"hasSome": group_ids}}  # mutable-ok: prisma where is a dict
+    listed: Final = {"team_id": {"in": stored_team_ids}}  # mutable-ok: prisma where is a dict
+    teams: Final = await team_table.find_many(where={"OR": (carrying, listed)})  # mutable-ok: prisma where is a dict
+    return _attached_team_ids_by_group(records, teams)
 
 
 async def _require_teams_exist(tx: _AccessGroupTx, team_ids: Sequence[str]) -> None:
     if not team_ids:
         return
-    where: Final = {"team_id": {"in": team_ids}}  # mutable-ok: prisma where must be a dict
+    where: Final = {"team_id": {"in": team_ids}}  # mutable-ok: prisma where is a dict
     found: Final = await tx.litellm_teamtable.find_many(where=where)
     missing: Final = frozenset(team_ids) - frozenset(team.team_id for team in found)
     if missing:
@@ -441,7 +440,7 @@ async def list_access_groups(
 
     table: Final = AccessGroupRepository(prisma_client).table
     records: Final = await table.find_many(order={"created_at": "desc"})
-    attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, records)
+    attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, records)
     return [_record_to_response(r, assigned_team_ids=attached[r.access_group_id]) for r in records]
 
 
@@ -463,7 +462,7 @@ async def get_access_group(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=f"Access group '{access_group_id}' not found",
         )
-    attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, (record,))
+    attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, (record,))
     return _record_to_response(record, assigned_team_ids=attached[record.access_group_id])
 
 
@@ -516,7 +515,8 @@ async def update_access_group(
                 )
             await _require_teams_exist(tx, data.assigned_team_ids or ())
 
-            old_team_ids: Final[set[str]] = set(existing.assigned_team_ids or [])
+            attached: Final = await _attached_team_ids_for(tx.litellm_teamtable, (existing,))
+            old_team_ids: Final[set[str]] = set(attached[access_group_id])
             old_key_ids: Final[set[str]] = set(existing.assigned_key_ids or [])
             new_team_ids: Final[set[str]] = (
                 set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids

litellm/repositories/table_repositories.py

@@ -176,6 +176,10 @@ class PolicyAttachmentRepository(PrismaTableRepository["prisma_models.LiteLLM_Po
     table_name = "litellm_policyattachmenttable"
 
 
+class TeamRepository(PrismaTableRepository["prisma_models.LiteLLM_TeamTable"]):
+    table_name = "litellm_teamtable"
+
+
 class DeletedTeamRepository(PrismaTableRepository["prisma_models.LiteLLM_DeletedTeamTable"]):
     table_name = "litellm_deletedteamtable"
 

tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py

@@ -317,8 +317,9 @@ def test_list_access_groups_attributes_teams_per_group_with_one_query(client_and
     assert body[1]["assigned_team_ids"] == ["team-y", "team-z"]
 
     mock_team_table.find_many.assert_awaited_once()
-    (carrying,) = mock_team_table.find_many.call_args.kwargs["where"]["OR"]
+    carrying, listed = mock_team_table.find_many.call_args.kwargs["where"]["OR"]
     assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-1", "ag-2"]
+    assert list(listed["team_id"]["in"]) == []
 
 
 @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
@@ -1238,7 +1239,7 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks):
     mock_access_group_table.find_unique = AsyncMock(return_value=existing)
 
     team_to_remove = _make_team_record("team-remove", ["ag-update"])
-    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"])])
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), team_to_remove])
     mock_team_table.find_unique = AsyncMock(return_value=team_to_remove)
 
     resp = client.put(
@@ -1256,6 +1257,28 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks):
     assert "ag-update" not in call_kwargs["data"]["access_group_ids"]
 
 
+def test_update_access_group_detaches_team_the_mirror_missed(client_and_mocks):
+    """Update removes the group from a team that carries it but was never written to the stored column."""
+    client, mock_prisma, mock_access_group_table, *_ = client_and_mocks
+    mock_team_table = mock_prisma.db.litellm_teamtable
+
+    existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-keep"])
+    mock_access_group_table.find_unique = AsyncMock(return_value=existing)
+
+    unmirrored = _make_team_record("team-unmirrored", ["ag-update", "ag-other"])
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), unmirrored])
+    mock_team_table.find_unique = AsyncMock(return_value=unmirrored)
+
+    resp = client.put("/v1/access_group/ag-update", json={"assigned_team_ids": ["team-keep"]})
+    assert resp.status_code == 200
+
+    mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-unmirrored"})
+    mock_team_table.update.assert_awaited_once()
+    call_kwargs = mock_team_table.update.call_args.kwargs
+    assert call_kwargs["where"] == {"team_id": "team-unmirrored"}
+    assert call_kwargs["data"]["access_group_ids"] == ["ag-other"]
+
+
 def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks):
     """Update does not sync teams when assigned_team_ids is absent from the payload."""
     client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
@@ -1271,7 +1294,6 @@ def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_moc
     resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"})
     assert resp.status_code == 200
 
-    mock_team_table.find_many.assert_not_awaited()
     mock_team_table.find_unique.assert_not_awaited()
     mock_team_table.update.assert_not_awaited()
 
@@ -1405,7 +1427,7 @@ def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks)
 
 def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks):
     """Update with explicit null for assigned_*_ids clears the list and writes [] to DB."""
-    client, mock_prisma, mock_table, *_ = client_and_mocks
+    client, _, mock_table, *_ = client_and_mocks
 
     existing = _make_access_group_record(
         access_group_id="ag-update",
@@ -1425,4 +1447,3 @@ def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks
     update_call_kwargs = mock_table.update.call_args.kwargs
     assert update_call_kwargs["data"]["assigned_team_ids"] == []
     assert update_call_kwargs["data"]["assigned_key_ids"] == []
-    mock_prisma.db.litellm_teamtable.find_many.assert_not_awaited()