diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 54b9b26123e..cf2ae392677 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1827,6 +1827,7 @@ async def new_team( "tpm_limit", "rpm_limit", "team_member_permissions", + "access_group_ids", ): if getattr(data, field, None) is None and ( field != "budget_duration" or field not in data.model_fields_set diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index e68008135aa..5b961f1320b 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -47,10 +47,12 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, ) from litellm.proxy.utils import CONFIG_PARAMS_TARGET, invalidate_config_param +from litellm.repositories.chunked_in import find_many_in from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( + AccessGroupRepository, SSOConfigRepository, UISettingsRepository, ) @@ -1172,6 +1174,33 @@ async def update_internal_user_settings( ) +async def _validate_default_access_groups_exist(access_group_ids: Sequence[str]) -> None: + if not access_group_ids: + return + + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected. Please connect a database."}, + ) + + existing_groups: Final = await find_many_in( + AccessGroupRepository(prisma_client, use_writer=True).table, "access_group_id", dict.fromkeys(access_group_ids) + ) + existing_ids: Final = {group.access_group_id for group in existing_groups} + missing_ids: Final = tuple(dict.fromkeys(ag_id for ag_id in access_group_ids if ag_id not in existing_ids)) + if missing_ids: + raise HTTPException( + status_code=400, + detail={ + "error": f"Access group(s) not found: {', '.join(missing_ids)}. " + "Create the access group first or remove it from the default team settings." + }, + ) + + @router.patch( "/update/default_team_settings", tags=["SSO Settings"], @@ -1188,6 +1217,8 @@ async def update_default_team_settings( """ if settings.organization_id is not None: await _validate_default_organization_exists(settings.organization_id) + if settings.access_group_ids is not None: + await _validate_default_access_groups_exist(settings.access_group_ids) return await _update_litellm_setting( settings=settings, diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 38ed799f38b..aa1ca63d258 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from typing import Literal from pydantic import ConfigDict, Field @@ -241,3 +242,7 @@ class DefaultTeamSSOParams(LiteLLMPydanticObjectBase): default=None, description="Default organization for new teams created without an explicit organization", ) + access_group_ids: Sequence[str] | None = Field( + default=None, + description="Default access groups attached to new teams created without explicit access_group_ids", + ) diff --git a/tests/unit/proxy/management_endpoints/test_team_default_params.py b/tests/unit/proxy/management_endpoints/test_team_default_params.py index cf4961dc82b..386584655d1 100644 --- a/tests/unit/proxy/management_endpoints/test_team_default_params.py +++ b/tests/unit/proxy/management_endpoints/test_team_default_params.py @@ -3,6 +3,7 @@ Tests for applying default team params during team creation and loading default_team_params from DB on startup. """ +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -11,6 +12,7 @@ from fastapi import HTTPException import litellm from litellm.proxy._types import ( + LiteLLM_TeamTable, LiteLLM_BudgetTable, LiteLLM_OrganizationTable, NewTeamRequest, @@ -214,6 +216,59 @@ class TestNewTeamDefaultParamsApplied: assert data.rpm_limit == 500 assert data.team_member_permissions == ["/key/generate", "/key/update"] + async def _create_team( + self, monkeypatch: pytest.MonkeyPatch, data: NewTeamRequest + ) -> tuple[dict[str, object], AsyncMock, AsyncMock]: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints import team_endpoints + from litellm.proxy.management_endpoints.team_endpoints import new_team + + prisma_client: Final = proxy_server.prisma_client + prisma_client.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + async def _create(data: dict[str, object], include: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(**data) + + create: Final = AsyncMock(side_effect=_create) + tx: Final = MagicMock() + tx.litellm_teamtable.create = create + prisma_client.db.tx.return_value.__aenter__.return_value = tx + reconcile: Final = AsyncMock(return_value=()) + monkeypatch.setattr(team_endpoints, "reconcile_team_access_group_membership", reconcile) + monkeypatch.setattr(team_endpoints, "invalidate_access_group_caches", AsyncMock()) + monkeypatch.setattr(team_endpoints, "_add_team_members_to_team", AsyncMock()) + + response: Final = await new_team(data=data, user_api_key_dict=self._make_admin_auth(), http_request=MagicMock()) + return response, create, reconcile + + @pytest.mark.asyncio + async def test_default_access_group_ids_attach_to_a_team_created_without_any(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "default_team_params", {"access_group_ids": ["central-models"]}) + + response, create, reconcile = await self._create_team(monkeypatch, NewTeamRequest(team_alias="my-team")) + + assert create.await_args is not None + assert create.await_args.kwargs["data"]["access_group_ids"] == ["central-models"] + assert response["access_group_ids"] == ["central-models"] + reconcile.assert_awaited_once() + assert reconcile.await_args is not None + assert reconcile.await_args.args[1] == response["team_id"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("explicit_access_group_ids", [["research"], []]) + async def test_explicit_access_group_ids_win_over_the_default( + self, monkeypatch: pytest.MonkeyPatch, explicit_access_group_ids: list[str] + ): + monkeypatch.setattr(litellm, "default_team_params", {"access_group_ids": ["central-models"]}) + + response, create, _ = await self._create_team( + monkeypatch, NewTeamRequest(team_alias="my-team", access_group_ids=explicit_access_group_ids) + ) + + assert create.await_args is not None + assert create.await_args.kwargs["data"]["access_group_ids"] == explicit_access_group_ids + assert response["access_group_ids"] == explicit_access_group_ids + @pytest.mark.asyncio async def test_explicit_values_not_overridden(self, monkeypatch): """When request provides explicit values, defaults do not override them.""" diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 6ccb019b937..72416908c8d 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -784,10 +784,15 @@ def test_get_group_ids_from_graph_api_response(): [ # Test case 1: Using DefaultTeamSSOParams DefaultTeamSSOParams( - max_budget=10, budget_duration="1d", models=["special-gpt-5"] + max_budget=10, budget_duration="1d", models=["special-gpt-5"], access_group_ids=["central-models"] ), # Test case 2: Using Dict - {"max_budget": 10, "budget_duration": "1d", "models": ["special-gpt-5"]}, + { + "max_budget": 10, + "budget_duration": "1d", + "models": ["special-gpt-5"], + "access_group_ids": ["central-models"], + }, ], ) async def test_default_team_params(team_params): @@ -832,6 +837,7 @@ async def test_default_team_params(team_params): assert create_call_args["max_budget"] == 10 assert create_call_args["budget_duration"] == "1d" assert create_call_args["models"] == ["special-gpt-5"] + assert create_call_args["access_group_ids"] == ["central-models"] @pytest.mark.asyncio diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 77d0132ec72..7a67d4c5024 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1,8 +1,8 @@ import json import os -from collections.abc import Iterator +from collections.abc import Callable, Iterator from typing import Final -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient @@ -3498,6 +3498,89 @@ def test_update_default_team_settings_without_organization_skips_lookup( assert mock_proxy_config["save_call_count"]() == 1 +@pytest.fixture +def mock_access_group_lookup(monkeypatch: pytest.MonkeyPatch) -> Callable[[frozenset[str]], AsyncMock]: + import litellm + import litellm.proxy.proxy_server as proxy_server_module + + def _install(existing_access_group_ids: frozenset[str]) -> AsyncMock: + async def _find_many(where: dict[str, dict[str, list[str]]]) -> list[MagicMock]: + return [ + MagicMock(access_group_id=access_group_id) + for access_group_id in where["access_group_id"]["in"] + if access_group_id in existing_access_group_ids + ] + + find_many: Final = AsyncMock(side_effect=_find_many) + fake_prisma: Final = MagicMock() + fake_prisma.writer_db.litellm_accessgrouptable.find_many = find_many + + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.setattr(litellm, "default_team_params", {}) + return find_many + + return _install + + +def test_update_default_team_settings_rejects_unknown_access_group( + mock_proxy_config: dict[str, object], + mock_auth: None, + mock_access_group_lookup: Callable[[frozenset[str]], AsyncMock], +) -> None: + import litellm + + mock_access_group_lookup(frozenset({"real-group"})) + + resp: Final = client.patch( + "/update/default_team_settings", + json={"access_group_ids": ["real-group", "ghost-group"]}, + ) + + assert resp.status_code == 400, resp.text + assert "ghost-group" in resp.json()["detail"]["error"] + assert "real-group" not in resp.json()["detail"]["error"] + assert mock_proxy_config["save_call_count"]() == 0 + assert litellm.default_team_params == {} + + +def test_update_default_team_settings_saves_existing_access_groups( + mock_proxy_config: dict[str, object], + mock_auth: None, + mock_access_group_lookup: Callable[[frozenset[str]], AsyncMock], +) -> None: + import litellm + + mock_access_group_lookup(frozenset({"real-group", "other-group"})) + + resp: Final = client.patch( + "/update/default_team_settings", + json={"access_group_ids": ["real-group", "other-group"]}, + ) + + assert resp.status_code == 200, resp.text + assert resp.json()["settings"]["access_group_ids"] == ["real-group", "other-group"] + assert mock_proxy_config["save_call_count"]() == 1 + assert litellm.default_team_params["access_group_ids"] == ["real-group", "other-group"] + + +def test_update_default_team_settings_with_empty_access_groups_skips_lookup( + mock_proxy_config: dict[str, object], + mock_auth: None, + mock_access_group_lookup: Callable[[frozenset[str]], AsyncMock], +) -> None: + find_many: Final = mock_access_group_lookup(frozenset()) + + resp: Final = client.patch( + "/update/default_team_settings", + json={"access_group_ids": []}, + ) + + assert resp.status_code == 200, resp.text + find_many.assert_not_awaited() + assert mock_proxy_config["save_call_count"]() == 1 + + def test_update_mcp_semantic_filter_settings_requires_proxy_admin(monkeypatch): """Non-admin callers must not mutate global MCP semantic filter settings.""" from litellm.proxy._types import UserAPIKeyAuth diff --git a/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx b/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx index 4a042c32cda..5a4ce705b82 100644 --- a/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamSSOSettings.test.tsx @@ -89,6 +89,32 @@ vi.mock("./common_components/OrganizationDropdown", () => ({ ), })); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ + data: [ + { access_group_id: "ag-1", access_group_name: "Central Models" }, + { access_group_id: "ag-2", access_group_name: "Research" }, + ], + isLoading: false, + isError: false, + }), +})); + +vi.mock("./common_components/AccessGroupSelector", () => ({ + default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( + + ), +})); + vi.mock("./ModelSelect/ModelSelect", () => { const ModelSelect = ({ value, onChange }: { value: string[]; onChange: (value: string[]) => void }) => (