diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 0357bc7dbc6..1f91eeedf64 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -1,4 +1,5 @@ from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final, Protocol from fastapi import APIRouter, Depends, HTTPException, status @@ -20,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, @@ -74,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): @@ -119,8 +120,56 @@ def _require_admin_view(user_api_key_dict: UserAPIKeyAuth) -> None: ) -def _record_to_response(record: _AccessGroupRecord) -> AccessGroupResponse: - return AccessGroupResponse.model_validate(record.dict()) +def _record_to_response( + record: _AccessGroupRecord, *, assigned_team_ids: Sequence[str] | None = None +) -> AccessGroupResponse: + stored: Final = record.dict() + payload: Final = ( + stored if assigned_team_ids is None else MappingProxyType({**stored, "assigned_team_ids": assigned_team_ids}) + ) + return AccessGroupResponse.model_validate(payload) + + +def _attached_team_ids_by_group( + records: Sequence[_AccessGroupRecord], teams: Sequence[_TeamRecord] +) -> Mapping[str, tuple[str, ...]]: + """Teams really attached to each group: the stored column minus ghosts, plus teams the mirror missed.""" + real_team_ids: Final = frozenset(team.team_id for team in teams) + + def attached(record: _AccessGroupRecord) -> tuple[str, ...]: + stored: Final = (team_id for team_id in (record.assigned_team_ids or ()) if team_id in real_team_ids) + carrying: Final = (team.team_id for team in teams if record.access_group_id in (team.access_group_ids or ())) + return tuple(dict.fromkeys((*stored, *carrying))) + + return MappingProxyType({record.access_group_id: attached(record) for record in records}) + + +async def _attached_team_ids_for( + team_table: _TeamTable, records: Sequence[_AccessGroupRecord] +) -> Mapping[str, tuple[str, ...]]: + if not records: + return MappingProxyType({}) + group_ids: Final = tuple(record.access_group_id for record in records) + 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 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 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: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Unknown team ids: {', '.join(sorted(missing))}", + ) def _record_to_access_group_table(record: _AccessGroupRecord) -> LiteLLM_AccessGroupTable: @@ -330,6 +379,7 @@ async def create_access_group( status_code=status.HTTP_409_CONFLICT, detail=f"Access group '{data.access_group_name}' already exists", ) + await _require_teams_exist(tx, data.assigned_team_ids or ()) record: Final = await tx.litellm_accessgrouptable.create( data={ @@ -390,7 +440,8 @@ async def list_access_groups( table: Final = AccessGroupRepository(prisma_client).table records: Final = await table.find_many(order={"created_at": "desc"}) - return [_record_to_response(r) for r in 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] @router.get( @@ -411,7 +462,8 @@ async def get_access_group( status_code=status.HTTP_404_NOT_FOUND, detail=f"Access group '{access_group_id}' not found", ) - return _record_to_response(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]) @router.put( @@ -461,8 +513,10 @@ async def update_access_group( status_code=status.HTTP_404_NOT_FOUND, detail=f"Access group '{access_group_id}' not found", ) + 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 diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 18cf884f267..739d6ada71c 100644 --- a/litellm/repositories/table_repositories.py +++ b/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" diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index e8f768c14ef..81816e21c10 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -12,13 +12,12 @@ from fastapi.testclient import TestClient from prisma.errors import PrismaError import litellm.proxy.proxy_server as ps -from litellm.proxy.proxy_server import app from litellm.proxy._types import ( CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth, ) - +from litellm.proxy.proxy_server import app def _make_access_group_record( @@ -58,6 +57,10 @@ def _make_access_group_record( return record +def _make_team_record(team_id: str, access_group_ids: list[str] | None = None): + return types.SimpleNamespace(team_id=team_id, access_group_ids=access_group_ids or []) + + @pytest.fixture def client_and_mocks(monkeypatch): """Setup mock prisma and admin auth for access group endpoints.""" @@ -185,7 +188,8 @@ ACCESS_GROUP_PATHS = ["/v1/access_group", "/v1/unified_access_group"] ) def test_create_access_group_success(client_and_mocks, base_path, payload): """Create access group with various payloads returns 201.""" - client, _, mock_table, *_ = client_and_mocks + client, mock_prisma, mock_table, *_ = client_and_mocks + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[_make_team_record("team-1")]) resp = client.post(base_path, json=payload) assert resp.status_code == 201 @@ -277,13 +281,45 @@ def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) def test_list_access_groups_success_empty(client_and_mocks, base_path): - """List access groups returns empty list when none exist.""" - client, _, mock_table, *_ = client_and_mocks + """List access groups returns empty list when none exist, without querying teams.""" + client, mock_prisma, mock_table, *_ = client_and_mocks resp = client.get(base_path) assert resp.status_code == 200 assert resp.json() == [] mock_table.find_many.assert_awaited_once() + mock_prisma.db.litellm_teamtable.find_many.assert_not_awaited() + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_list_access_groups_attributes_teams_per_group_with_one_query(client_and_mocks, base_path): + """List derives each group's teams from the team table in a single query, attributed per group.""" + client, mock_prisma, mock_table, *_ = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + records = [ + _make_access_group_record(access_group_id="ag-1", access_group_name="group-1"), + _make_access_group_record(access_group_id="ag-2", access_group_name="group-2"), + ] + mock_table.find_many = AsyncMock(return_value=records) + mock_team_table.find_many = AsyncMock( + return_value=[ + _make_team_record("team-x", ["ag-1"]), + _make_team_record("team-y", ["ag-2"]), + _make_team_record("team-z", ["ag-1", "ag-2"]), + ] + ) + + resp = client.get(base_path) + assert resp.status_code == 200 + body = resp.json() + assert body[0]["assigned_team_ids"] == ["team-x", "team-z"] + assert body[1]["assigned_team_ids"] == ["team-y", "team-z"] + + mock_team_table.find_many.assert_awaited_once() + 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) @@ -373,6 +409,43 @@ def test_get_access_group_success(client_and_mocks, base_path, access_group_id): assert resp.json()["access_group_id"] == access_group_id +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_get_access_group_derives_assigned_teams_from_team_table(client_and_mocks, base_path): + """Get drops ghost ids from the stored column and adds teams that carry the group but were never mirrored.""" + client, mock_prisma, mock_table, *_ = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + record = _make_access_group_record(access_group_id="ag-123", assigned_team_ids=["team-a", "ghost-team"]) + mock_table.find_unique = AsyncMock(return_value=record) + mock_team_table.find_many = AsyncMock( + return_value=[ + _make_team_record("team-a", ["ag-123"]), + _make_team_record("team-b", ["ag-123"]), + _make_team_record("team-c", ["ag-123"]), + ] + ) + + resp = client.get(f"{base_path}/ag-123") + assert resp.status_code == 200 + assert resp.json()["assigned_team_ids"] == ["team-a", "team-b", "team-c"] + + carrying, listed = mock_team_table.find_many.call_args.kwargs["where"]["OR"] + assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-123"] + assert list(listed["team_id"]["in"]) == ["team-a", "ghost-team"] + + +def test_get_access_group_empty_column_and_no_teams_returns_empty(client_and_mocks): + """Get returns [] when the column is empty and no team carries the group.""" + client, mock_prisma, mock_table, *_ = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=_make_access_group_record(access_group_id="ag-123")) + mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + + resp = client.get("/v1/access_group/ag-123") + assert resp.status_code == 200 + assert resp.json()["assigned_team_ids"] == [] + + def test_get_access_group_not_found(client_and_mocks): """Get access group returns 404 when not found.""" client, _, mock_table, *_ = client_and_mocks @@ -985,6 +1058,28 @@ def test_record_to_access_group_table(): assert result.access_agent_ids == ["agent-1"] +def test_attached_team_ids_by_group_keeps_column_order_then_appends_unmirrored_teams(): + """Stored ids that resolve keep their order, ghosts drop, carriers the mirror missed append once, per group.""" + from litellm.proxy.management_endpoints.access_group_endpoints import ( + _attached_team_ids_by_group, + ) + + records = [ + _make_access_group_record(access_group_id="ag-1", assigned_team_ids=["team-b", "ghost", "team-a"]), + _make_access_group_record(access_group_id="ag-2", assigned_team_ids=[]), + ] + teams = [ + _make_team_record("team-a", ["ag-1"]), + _make_team_record("team-b", []), + _make_team_record("team-c", ["ag-1"]), + _make_team_record("team-d", ["ag-2"]), + ] + + result = _attached_team_ids_by_group(records, teams) + + assert dict(result) == {"ag-1": ("team-b", "team-a", "team-c"), "ag-2": ("team-d",)} + + # --------------------------------------------------------------------------- # Sync tests: CREATE # --------------------------------------------------------------------------- @@ -997,9 +1092,8 @@ def test_create_access_group_syncs_assigned_teams(client_and_mocks): ) mock_team_table = mock_prisma.db.litellm_teamtable - team_record = MagicMock() - team_record.team_id = "team-1" - team_record.access_group_ids = [] + team_record = _make_team_record("team-1") + mock_team_table.find_many = AsyncMock(return_value=[team_record]) mock_team_table.find_unique = AsyncMock(return_value=team_record) resp = client.post( @@ -1043,20 +1137,22 @@ def test_create_access_group_syncs_assigned_keys(client_and_mocks): assert "ag-new" in call_kwargs["data"]["access_group_ids"] -def test_create_access_group_skips_sync_for_nonexistent_team(client_and_mocks): - """Create skips updating a team that doesn't exist in DB.""" - client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks +def test_create_access_group_rejects_nonexistent_team(client_and_mocks): + """Create refuses to store a team id that does not resolve to a team row.""" + client, mock_prisma, mock_access_group_table, *_ = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - mock_team_table.find_unique = AsyncMock(return_value=None) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-real")]) resp = client.post( "/v1/access_group", json={ "access_group_name": "new-group", - "assigned_team_ids": ["nonexistent-team"], + "assigned_team_ids": ["team-real", "nonexistent-team", "also-missing"], }, ) - assert resp.status_code == 201 + assert resp.status_code == 400 + assert resp.json()["detail"] == "Unknown team ids: also-missing, nonexistent-team" + mock_access_group_table.create.assert_not_awaited() mock_team_table.update.assert_not_awaited() @@ -1065,9 +1161,8 @@ def test_create_access_group_idempotent_team_sync(client_and_mocks): client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - team_record = MagicMock() - team_record.team_id = "team-1" - team_record.access_group_ids = ["ag-new"] # already synced + team_record = _make_team_record("team-1", ["ag-new"]) + mock_team_table.find_many = AsyncMock(return_value=[team_record]) mock_team_table.find_unique = AsyncMock(return_value=team_record) resp = client.post( @@ -1095,9 +1190,8 @@ def test_update_access_group_syncs_added_teams(client_and_mocks): ) mock_access_group_table.find_unique = AsyncMock(return_value=existing) - team_record = MagicMock() - team_record.team_id = "team-new" - team_record.access_group_ids = [] + team_record = _make_team_record("team-new") + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-existing", ["ag-update"]), team_record]) mock_team_table.find_unique = AsyncMock(return_value=team_record) resp = client.put( @@ -1113,6 +1207,25 @@ def test_update_access_group_syncs_added_teams(client_and_mocks): assert "ag-update" in call_kwargs["data"]["access_group_ids"] +def test_update_access_group_rejects_nonexistent_team(client_and_mocks): + """Update refuses to store a team id that does not resolve to a team row and leaves the group untouched.""" + 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-existing"]) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-existing", ["ag-update"])]) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_team_ids": ["team-existing", "team-ghost"]}, + ) + assert resp.status_code == 400 + assert resp.json()["detail"] == "Unknown team ids: team-ghost" + mock_access_group_table.update.assert_not_awaited() + mock_team_table.update.assert_not_awaited() + + def test_update_access_group_syncs_removed_teams(client_and_mocks): """Update removes access_group_id from de-assigned teams.""" client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( @@ -1125,9 +1238,8 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks): ) mock_access_group_table.find_unique = AsyncMock(return_value=existing) - team_to_remove = MagicMock() - team_to_remove.team_id = "team-remove" - team_to_remove.access_group_ids = ["ag-update"] + 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"]), team_to_remove]) mock_team_table.find_unique = AsyncMock(return_value=team_to_remove) resp = client.put( @@ -1145,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 = (