diff --git a/litellm/proxy/common_utils/rbac_utils.py b/litellm/proxy/common_utils/rbac_utils.py index b54b5e8f45..177ac2b62e 100644 --- a/litellm/proxy/common_utils/rbac_utils.py +++ b/litellm/proxy/common_utils/rbac_utils.py @@ -13,6 +13,14 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth FeatureName = Literal["agents", "vector_stores"] +OrgAdminFeatureName = Literal["key_generate", "team_create", "model_add"] + +_ORG_ADMIN_FEATURE_LABELS: dict = { + "key_generate": "key generation", + "team_create": "team creation", + "model_add": "model creation", +} + async def check_feature_access_for_user( user_api_key_dict: UserAPIKeyAuth, @@ -68,3 +76,36 @@ async def check_feature_access_for_user( "error": f"Access to {feature_name} is disabled for your role. Contact your proxy admin." }, ) + + +async def check_org_admin_feature_access( + user_api_key_dict: UserAPIKeyAuth, + feature_name: OrgAdminFeatureName, +) -> None: + """ + Raise HTTP 403 if the user is an org admin and the given feature is + disabled for org admins via UI settings. + + Only blocks the ORG_ADMIN role — proxy admins and all other roles are + unaffected, so those paths continue to be gated by their existing auth + checks. + """ + if user_api_key_dict.user_role not in ( + LitellmUserRoles.ORG_ADMIN, + LitellmUserRoles.ORG_ADMIN.value, + ): + return + + from litellm.proxy.proxy_server import general_settings + + disable_flag = f"disable_{feature_name}_for_org_admin" + if not general_settings.get(disable_flag, False): + return + + label = _ORG_ADMIN_FEATURE_LABELS.get(feature_name, feature_name) + raise HTTPException( + status_code=403, + detail={ + "error": f"{label} is disabled for org admins. Contact your proxy admin." + }, + ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b3cbf95401..c6701990dc 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.auth_utils import abbreviate_api_key from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.rbac_utils import check_org_admin_feature_access from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.common_utils import ( @@ -1256,6 +1257,10 @@ async def generate_key_fn( verbose_proxy_logger.debug("entered /key/generate") + await check_org_admin_feature_access( + user_api_key_dict=user_api_key_dict, feature_name="key_generate" + ) + # Validate budget values are not negative if data.max_budget is not None and data.max_budget < 0: raise HTTPException( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 754727d471..fbba94ffa4 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -37,6 +37,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.common_utils.rbac_utils import check_org_admin_feature_access from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( team_model_add, @@ -974,6 +975,10 @@ async def add_new_model( }, ) + await check_org_admin_feature_access( + user_api_key_dict=user_api_key_dict, feature_name="model_add" + ) + ## Auth check await ModelManagementAuthChecks.can_user_make_model_call( model_params=model_params, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 1ef40f0685..a6b034720f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -70,6 +70,7 @@ from litellm.proxy.auth.auth_checks import ( get_user_object, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.rbac_utils import check_org_admin_feature_access from litellm.proxy.management_endpoints.common_utils import ( _is_user_org_admin_for_team, _is_user_team_admin, @@ -896,6 +897,10 @@ async def new_team( # noqa: PLR0915 if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) + await check_org_admin_feature_access( + user_api_key_dict=user_api_key_dict, feature_name="team_create" + ) + # Validate budget values are not negative if data.max_budget is not None and data.max_budget < 0: raise HTTPException( diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 1523e1cb9b..6b175f940c 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -152,6 +152,21 @@ class UISettings(BaseModel): description="If true, users cannot specify custom key values. All keys must be auto-generated.", ) + disable_key_generate_for_org_admin: bool = Field( + default=False, + description="If true, org admins cannot generate API keys via /key/generate.", + ) + + disable_team_create_for_org_admin: bool = Field( + default=False, + description="If true, org admins cannot create teams via /team/new.", + ) + + disable_model_add_for_org_admin: bool = Field( + default=False, + description="If true, org admins cannot add models via /model/new.", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -174,6 +189,9 @@ ALLOWED_UI_SETTINGS_FIELDS = { "allow_vector_stores_for_team_admins", "scope_user_search_to_org", "disable_custom_api_keys", + "disable_key_generate_for_org_admin", + "disable_team_create_for_org_admin", + "disable_model_add_for_org_admin", } # Flags that must be synced from the persisted UISettings into @@ -185,6 +203,9 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS = [ "allow_agents_for_team_admins", "disable_vector_stores_for_internal_users", "allow_vector_stores_for_team_admins", + "disable_key_generate_for_org_admin", + "disable_team_create_for_org_admin", + "disable_model_add_for_org_admin", ] diff --git a/tests/litellm/proxy/common_utils/test_rbac_utils.py b/tests/litellm/proxy/common_utils/test_rbac_utils.py index ebc7c16230..9adc34ed40 100644 --- a/tests/litellm/proxy/common_utils/test_rbac_utils.py +++ b/tests/litellm/proxy/common_utils/test_rbac_utils.py @@ -1,7 +1,8 @@ """ Tests for litellm/proxy/common_utils/rbac_utils.py -Covers check_feature_access_for_user for agents and vector_stores features. +Covers check_feature_access_for_user for agents and vector_stores features, +plus check_org_admin_feature_access for key/team/model creation. """ from unittest.mock import AsyncMock, patch @@ -10,7 +11,10 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user +from litellm.proxy.common_utils.rbac_utils import ( + check_feature_access_for_user, + check_org_admin_feature_access, +) def _make_user(role: str, user_id: str = "user-1") -> UserAPIKeyAuth: @@ -178,3 +182,82 @@ async def test_vector_stores_disabled_non_team_admin_blocked(): with pytest.raises(HTTPException) as exc_info: await check_feature_access_for_user(user, "vector_stores") assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# check_org_admin_feature_access +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "feature_name", + ["key_generate", "team_create", "model_add"], +) +@pytest.mark.asyncio +async def test_org_admin_feature_not_disabled_allows_org_admin(feature_name): + user = _make_user(LitellmUserRoles.ORG_ADMIN.value) + with patch.dict(_GS_PATH, {}, clear=True): + await check_org_admin_feature_access(user, feature_name) + + +@pytest.mark.parametrize( + "feature_name,flag_name", + [ + ("key_generate", "disable_key_generate_for_org_admin"), + ("team_create", "disable_team_create_for_org_admin"), + ("model_add", "disable_model_add_for_org_admin"), + ], +) +@pytest.mark.asyncio +async def test_org_admin_feature_disabled_blocks_org_admin(feature_name, flag_name): + user = _make_user(LitellmUserRoles.ORG_ADMIN.value) + with patch.dict(_GS_PATH, {flag_name: True}, clear=True): + with pytest.raises(HTTPException) as exc_info: + await check_org_admin_feature_access(user, feature_name) + assert exc_info.value.status_code == 403 + + +@pytest.mark.parametrize( + "role", + [ + LitellmUserRoles.PROXY_ADMIN.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + LitellmUserRoles.TEAM.value, + ], +) +@pytest.mark.asyncio +async def test_org_admin_disable_flag_does_not_affect_other_roles(role): + """Only ORG_ADMIN should be gated by these flags — other roles pass through.""" + user = _make_user(role) + with patch.dict( + _GS_PATH, + { + "disable_key_generate_for_org_admin": True, + "disable_team_create_for_org_admin": True, + "disable_model_add_for_org_admin": True, + }, + clear=True, + ): + # No exception expected — these roles are not org admins, so the flag + # should be a no-op. Other auth checks (in the endpoint itself) still + # apply. + await check_org_admin_feature_access(user, "key_generate") + await check_org_admin_feature_access(user, "team_create") + await check_org_admin_feature_access(user, "model_add") + + +@pytest.mark.asyncio +async def test_org_admin_role_enum_and_string_both_blocked(): + """UserAPIKeyAuth.user_role may be either the enum or the string value.""" + with patch.dict(_GS_PATH, {"disable_key_generate_for_org_admin": True}, clear=True): + user_str = _make_user(LitellmUserRoles.ORG_ADMIN.value) + with pytest.raises(HTTPException): + await check_org_admin_feature_access(user_str, "key_generate") + + user_enum = UserAPIKeyAuth( + user_role=LitellmUserRoles.ORG_ADMIN, user_id="user-1" + ) + with pytest.raises(HTTPException): + await check_org_admin_feature_access(user_enum, "key_generate") diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index d7ce66f1d7..1535acaadd 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1098,6 +1098,60 @@ class TestProxySettingEndpoints: assert response.status_code == 200 assert general_settings.get("forward_llm_provider_auth_headers") is True + @pytest.mark.parametrize( + "flag_name", + [ + "disable_key_generate_for_org_admin", + "disable_team_create_for_org_admin", + "disable_model_add_for_org_admin", + ], + ) + def test_update_ui_settings_persists_and_syncs_org_admin_flags( + self, mock_auth, monkeypatch, flag_name + ): + """Org-admin restriction flags must be allowlisted, persisted, and synced to general_settings.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + general_settings: dict = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", general_settings + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + payload = {flag_name: True} + + try: + response = client.patch("/update/ui_settings", json=payload) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + data = response.json() + assert data["settings"][flag_name] is True + + # Persisted in DB + call_args = mock_prisma.db.litellm_uisettings.upsert.call_args + stored_settings = json.loads(call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored_settings[flag_name] is True + + # Synced into general_settings so the enforcement helper sees it + assert general_settings.get(flag_name) is True + def test_get_sso_settings_from_database( self, mock_proxy_config, mock_auth, monkeypatch ):