From 96b75be03d1db6e4957183061fb20e97163318ee Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Mar 2026 20:13:14 -0800 Subject: [PATCH] [Feature] RBAC for Vector Stores and Agents Add proxy-admin-configurable toggles to restrict internal users (and optionally team admins) from accessing agent and vector store management features. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/agent_endpoints/endpoints.py | 20 ++- litellm/proxy/common_utils/rbac_utils.py | 126 ++++++++++++++ .../proxy_setting_endpoints.py | 60 +++++-- .../management_endpoints.py | 11 ++ .../proxy/agent_endpoints/test_agent_rbac.py | 84 ++++++++++ .../proxy/common_utils/test_rbac_utils.py | 156 ++++++++++++++++++ .../test_vector_store_rbac.py | 121 ++++++++++++++ .../components/SidebarProvider.tsx | 12 ++ .../AdminSettings/UISettings/UISettings.tsx | 136 +++++++++++++++ .../src/components/leftnav.tsx | 8 +- 10 files changed, 720 insertions(+), 14 deletions(-) create mode 100644 litellm/proxy/common_utils/rbac_utils.py create mode 100644 tests/litellm/proxy/agent_endpoints/test_agent_rbac.py create mode 100644 tests/litellm/proxy/common_utils/test_rbac_utils.py create mode 100644 tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 65674d01be7..80c55f634f7 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -16,6 +16,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.types.agents import ( AgentConfig, @@ -69,6 +70,8 @@ async def get_agents( Returns: List[AgentResponse] """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, @@ -179,6 +182,8 @@ async def create_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -233,7 +238,10 @@ async def create_agent( dependencies=[Depends(user_api_key_auth)], response_model=AgentResponse, ) -async def get_agent_by_id(agent_id: str): +async def get_agent_by_id( + agent_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Get a specific agent by ID @@ -243,6 +251,8 @@ async def get_agent_by_id(agent_id: str): -H "Authorization: Bearer " ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -319,6 +329,8 @@ async def update_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -410,6 +422,8 @@ async def patch_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -484,6 +498,8 @@ async def delete_agent( } ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -763,6 +779,8 @@ async def get_agent_daily_activity( """ Get daily activity for specific agents or all accessible agents. """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: diff --git a/litellm/proxy/common_utils/rbac_utils.py b/litellm/proxy/common_utils/rbac_utils.py new file mode 100644 index 00000000000..2b187d18065 --- /dev/null +++ b/litellm/proxy/common_utils/rbac_utils.py @@ -0,0 +1,126 @@ +""" +RBAC utility helpers for feature-level access control. + +These helpers are used by agent and vector store endpoints to enforce +proxy-admin-configurable toggles that restrict access for internal users. +""" + +from typing import TYPE_CHECKING + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth + +if TYPE_CHECKING: + pass + + +def _is_user_team_admin_for_any_team( + user_api_key_dict: UserAPIKeyAuth, + teams: list, +) -> bool: + """ + Return True if the user is an admin member in at least one of the given teams. + + Args: + user_api_key_dict: The authenticated user. + teams: List of Prisma team records (from litellm_teamtable.find_many). + """ + for team in teams: + team_obj = LiteLLM_TeamTable(**team.model_dump()) + for member in team_obj.members_with_roles: + if ( + member.user_id is not None + and member.user_id == user_api_key_dict.user_id + and member.role == "admin" + ): + return True + return False + + +async def check_feature_access_for_user( + user_api_key_dict: UserAPIKeyAuth, + feature_name: str, +) -> None: + """ + Raise HTTP 403 if the user's role is blocked from accessing the given feature + by the UI settings stored in general_settings. + + Args: + user_api_key_dict: The authenticated user. + feature_name: Either "agents" or "vector_stores". + """ + # Proxy admins (and view-only admins) are never blocked. + if user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.PROXY_ADMIN.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ): + return + + from litellm.proxy.proxy_server import general_settings + + disable_flag = f"disable_{feature_name}_for_internal_users" + allow_team_admins_flag = f"allow_{feature_name}_for_team_admins" + + if not general_settings.get(disable_flag, False): + # Feature is not disabled — allow all authenticated users. + return + + # Feature is disabled. Check if team admins are exempted. + if general_settings.get(allow_team_admins_flag, False): + is_team_admin = await _check_if_team_admin(user_api_key_dict) + if is_team_admin: + return + + raise HTTPException( + status_code=403, + detail={ + "error": f"Access to {feature_name} is disabled for your role. Contact your proxy admin." + }, + ) + + +async def _check_if_team_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Return True if the user is a team admin in any team. + Mirrors the logic in management_endpoints/common_utils._user_has_admin_privileges + but scoped to team-admin check only. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None or user_api_key_dict.user_id is None: + return False + + from litellm.caching import DualCache + from litellm.proxy.auth.auth_checks import get_user_object + + try: + user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + user_id_upsert=False, + proxy_logging_obj=None, + ) + + if user_obj is None: + return False + + if user_obj.teams is None or len(user_obj.teams) == 0: + return False + + teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": user_obj.teams}} + ) + + return _is_user_team_admin_for_any_team(user_api_key_dict, teams) + + except Exception as e: + verbose_proxy_logger.debug( + f"rbac_utils: error checking team admin status for user " + f"{user_api_key_dict.user_id}: {e}" + ) + return False diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index ceda08d520a..8991dc5fd5c 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -104,6 +104,26 @@ class UISettings(BaseModel): description="If enabled, shows the Projects feature in the UI sidebar and the project field in key management.", ) + disable_agents_for_internal_users: bool = Field( + default=False, + description="If true, internal users cannot access agent management endpoints or the Agents page in the UI.", + ) + + allow_agents_for_team_admins: bool = Field( + default=False, + description="If true, team admins are exempt from the agents disable restriction (only takes effect when disable_agents_for_internal_users is true).", + ) + + disable_vector_stores_for_internal_users: bool = Field( + default=False, + description="If true, internal users cannot access vector store management endpoints or the Vector Stores page in the UI.", + ) + + allow_vector_stores_for_team_admins: bool = Field( + default=False, + description="If true, team admins are exempt from the vector stores disable restriction (only takes effect when disable_vector_stores_for_internal_users is true).", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -119,6 +139,10 @@ ALLOWED_UI_SETTINGS_FIELDS = { "require_auth_for_public_ai_hub", "forward_client_headers_to_llm_api", "enable_projects_ui", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", } @@ -976,14 +1000,20 @@ async def get_ui_settings(): k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS } - # Sync forward_client_headers_to_llm_api into general_settings so the proxy - # picks it up at runtime (covers server restart scenarios). - if "forward_client_headers_to_llm_api" in ui_settings: + # Sync runtime flags into general_settings so the proxy picks them up + # at runtime (covers server restart scenarios). + _runtime_flags = [ + "forward_client_headers_to_llm_api", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", + ] + _flags_to_sync = {k: ui_settings[k] for k in _runtime_flags if k in ui_settings} + if _flags_to_sync: from litellm.proxy.proxy_server import general_settings - general_settings["forward_client_headers_to_llm_api"] = ui_settings[ - "forward_client_headers_to_llm_api" - ] + general_settings.update(_flags_to_sync) # Build config-like object for schema helper config: Dict[str, Any] = {"litellm_settings": {"ui_settings": ui_settings}} @@ -1048,14 +1078,20 @@ async def update_ui_settings( }, ) - # Sync forward_client_headers_to_llm_api to general_settings so the proxy - # picks it up at runtime (general_settings is checked in pre-call utils). - if "forward_client_headers_to_llm_api" in ui_settings: + # Sync runtime flags to general_settings so the proxy picks them up + # at runtime (general_settings is checked in pre-call utils). + _runtime_flags = [ + "forward_client_headers_to_llm_api", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", + ] + _flags_to_sync = {k: ui_settings[k] for k in _runtime_flags if k in ui_settings} + if _flags_to_sync: from litellm.proxy.proxy_server import general_settings - general_settings["forward_client_headers_to_llm_api"] = ui_settings[ - "forward_client_headers_to_llm_api" - ] + general_settings.update(_flags_to_sync) return { "message": "UI settings updated successfully", diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index cccbb51f47b..068f4217e0f 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -24,6 +24,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 decrypt_value_helper +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, @@ -439,6 +440,8 @@ async def new_vector_store( - vector_store_description: Optional[str] - Description of the vector store - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client try: @@ -506,6 +509,8 @@ async def list_vector_stores( - page: int - Page number for pagination (default: 1) - page_size: int - Number of items per page (default: 100) """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client vector_store_map: Dict[str, LiteLLM_ManagedVectorStore] = {} @@ -605,6 +610,8 @@ async def delete_vector_store( Parameters: - vector_store_id: str - ID of the vector store to delete """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -687,6 +694,8 @@ async def get_vector_store_info( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return a single vector store's details""" + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -770,6 +779,8 @@ async def update_vector_store( Update vector store details in both database and in-memory registry. The updated data is immediately synchronized to the in-memory registry. """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client from litellm.types.router import GenericLiteLLMParams diff --git a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py new file mode 100644 index 00000000000..a863201ddb5 --- /dev/null +++ b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py @@ -0,0 +1,84 @@ +""" +Tests for RBAC enforcement on agent endpoints. + +Verifies that check_feature_access_for_user is called and that a 403 is +raised when agents are disabled for internal users. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + user_id=user_id, + ) + + +def _make_admin_user(user_id: str = "admin-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_id=user_id, + ) + + +# --------------------------------------------------------------------------- +# get_agents +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_get_agents_blocked_for_internal_user_when_disabled(): + """get_agents should raise 403 when agents are disabled for internal users.""" + from litellm.proxy.agent_endpoints.endpoints import get_agents + + user = _make_internal_user() + gs = {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False} + + request_mock = MagicMock() + with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True): + with pytest.raises(HTTPException) as exc_info: + await get_agents(request=request_mock, user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_agents_allowed_when_not_disabled(): + """get_agents should not raise RBAC 403 when agents are not disabled.""" + from litellm.proxy.agent_endpoints.endpoints import get_agents + + user = _make_internal_user() + request_mock = MagicMock() + + with patch.dict("litellm.proxy.proxy_server.general_settings", {}, clear=True): + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + MagicMock(get_agent_list=MagicMock(return_value=[])), + ): + with patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + new=AsyncMock(return_value=[]), + ): + result = await get_agents(request=request_mock, user_api_key_dict=user) + assert result == [] + + +# --------------------------------------------------------------------------- +# get_agent_daily_activity +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_get_agent_daily_activity_blocked_when_disabled(): + from litellm.proxy.agent_endpoints.endpoints import get_agent_daily_activity + + user = _make_internal_user() + gs = {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False} + + with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True): + with pytest.raises(HTTPException) as exc_info: + await get_agent_daily_activity(user_api_key_dict=user) + assert exc_info.value.status_code == 403 diff --git a/tests/litellm/proxy/common_utils/test_rbac_utils.py b/tests/litellm/proxy/common_utils/test_rbac_utils.py new file mode 100644 index 00000000000..997a2e19b77 --- /dev/null +++ b/tests/litellm/proxy/common_utils/test_rbac_utils.py @@ -0,0 +1,156 @@ +""" +Tests for litellm/proxy/common_utils/rbac_utils.py + +Covers check_feature_access_for_user for agents and vector_stores features. +""" + +from unittest.mock import AsyncMock, patch + +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 + + +def _make_user(role: str, user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=role, user_id=user_id) + + +# general_settings is imported from litellm.proxy.proxy_server inside the +# function, so we patch it via patch.dict on the original dict. +_GS_PATH = "litellm.proxy.proxy_server.general_settings" + + +# --------------------------------------------------------------------------- +# Proxy admin is always allowed +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_proxy_admin_always_allowed(): + user = _make_user(LitellmUserRoles.PROXY_ADMIN.value) + with patch.dict(_GS_PATH, {"disable_agents_for_internal_users": True}): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_proxy_admin_view_only_always_allowed(): + user = _make_user(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) + with patch.dict(_GS_PATH, {"disable_agents_for_internal_users": True}): + await check_feature_access_for_user(user, "agents") + + +# --------------------------------------------------------------------------- +# Feature not disabled — everyone allowed +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_feature_not_disabled_allows_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict(_GS_PATH, {}, clear=True): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_feature_not_disabled_allows_vector_stores(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict(_GS_PATH, {"disable_vector_stores_for_internal_users": False}, clear=True): + await check_feature_access_for_user(user, "vector_stores") + + +# --------------------------------------------------------------------------- +# Feature disabled, team-admin exemption OFF — internal user blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_agents_disabled_blocks_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False}, + clear=True, + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "agents") + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_blocks_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": False}, + clear=True, + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "vector_stores") + assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# Feature disabled, allow_team_admins ON — team admin allowed, non-admin blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_agents_disabled_team_admin_allowed(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="team-admin-user") + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=True), + ): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_agents_disabled_non_team_admin_blocked(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="regular-user") + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=False), + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "agents") + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_team_admin_allowed(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="team-admin-user") + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=True), + ): + await check_feature_access_for_user(user, "vector_stores") + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_non_team_admin_blocked(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="regular-user") + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=False), + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "vector_stores") + assert exc_info.value.status_code == 403 diff --git a/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py b/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py new file mode 100644 index 00000000000..3eb49bdf114 --- /dev/null +++ b/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py @@ -0,0 +1,121 @@ +""" +Tests for RBAC enforcement on vector store management endpoints. + +Verifies that check_feature_access_for_user is called and that a 403 is +raised when vector stores are disabled for internal users. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + user_id=user_id, + ) + + +_DISABLED_GS = { + "disable_vector_stores_for_internal_users": True, + "allow_vector_stores_for_team_admins": False, +} + +_ENABLED_GS: dict = {} + + +# --------------------------------------------------------------------------- +# list_vector_stores +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_list_vector_stores_blocked_when_disabled(): + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + user = _make_internal_user() + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with pytest.raises(HTTPException) as exc_info: + await list_vector_stores(user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_list_vector_stores_allowed_when_not_disabled(): + """list_vector_stores should not raise 403 when vector stores are not disabled.""" + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + import litellm + user = _make_internal_user() + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock(return_value=[]) + + raised_403 = False + with patch.dict("litellm.proxy.proxy_server.general_settings", _ENABLED_GS, clear=True): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch.object(litellm, "vector_store_registry", None): + with patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[]), + ): + try: + await list_vector_stores(user_api_key_dict=user) + except HTTPException as e: + if e.status_code == 403: + raised_403 = True + assert not raised_403, "Should not raise 403 when vector stores are not disabled" + + +# --------------------------------------------------------------------------- +# new_vector_store +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_new_vector_store_blocked_when_disabled(): + from litellm.proxy.vector_store_endpoints.management_endpoints import new_vector_store + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore + + user = _make_internal_user() + vs = LiteLLM_ManagedVectorStore(vector_store_id="vs-1", custom_llm_provider="openai") # type: ignore[call-arg] + + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with pytest.raises(HTTPException) as exc_info: + await new_vector_store(vector_store=vs, user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# Admin user is never blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_list_vector_stores_admin_not_blocked(): + """Proxy admin should never be blocked, even when vector stores are disabled.""" + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + import litellm + admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_id="admin-1", + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock(return_value=[]) + + raised_403 = False + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch.object(litellm, "vector_store_registry", None): + with patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[]), + ): + try: + await list_vector_stores(user_api_key_dict=admin) + except HTTPException as e: + if e.status_code == 403: + raised_403 = True + assert not raised_403, "Admin should not be blocked even when vector stores are disabled" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx index 17f62a20f7d..7dcc3fa8a1a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx @@ -15,6 +15,8 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side const { accessToken } = useAuthorized(); const [enabledPagesInternalUsers, setEnabledPagesInternalUsers] = useState(null); const [enableProjectsUI, setEnableProjectsUI] = useState(false); + const [disableAgentsForInternalUsers, setDisableAgentsForInternalUsers] = useState(false); + const [disableVectorStoresForInternalUsers, setDisableVectorStoresForInternalUsers] = useState(false); useEffect(() => { const fetchUISettings = async () => { @@ -39,6 +41,14 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side if (settings?.values?.enable_projects_ui !== undefined) { setEnableProjectsUI(Boolean(settings.values.enable_projects_ui)); } + + if (settings?.values?.disable_agents_for_internal_users !== undefined) { + setDisableAgentsForInternalUsers(Boolean(settings.values.disable_agents_for_internal_users)); + } + + if (settings?.values?.disable_vector_stores_for_internal_users !== undefined) { + setDisableVectorStoresForInternalUsers(Boolean(settings.values.disable_vector_stores_for_internal_users)); + } } catch (error) { console.error("[SidebarProvider] Failed to fetch UI settings:", error); } @@ -54,6 +64,8 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side collapsed={sidebarCollapsed} enabledPagesInternalUsers={enabledPagesInternalUsers} enableProjectsUI={enableProjectsUI} + disableAgentsForInternalUsers={disableAgentsForInternalUsers} + disableVectorStoresForInternalUsers={disableVectorStoresForInternalUsers} /> ); }; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 5d99dd2969d..dfc66d3484d 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -19,9 +19,15 @@ export default function UISettings() { const forwardClientHeadersProperty = schema?.properties?.forward_client_headers_to_llm_api; const enableProjectsUIProperty = schema?.properties?.enable_projects_ui; const enabledPagesProperty = schema?.properties?.enabled_ui_pages_internal_users; + const disableAgentsProperty = schema?.properties?.disable_agents_for_internal_users; + const allowAgentsTeamAdminsProperty = schema?.properties?.allow_agents_for_team_admins; + const disableVectorStoresProperty = schema?.properties?.disable_vector_stores_for_internal_users; + const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); + const isAgentsDisabled = Boolean(values.disable_agents_for_internal_users); + const isVectorStoresDisabled = Boolean(values.disable_vector_stores_for_internal_users); const handleToggle = (checked: boolean) => { updateSettings( @@ -105,6 +111,62 @@ export default function UISettings() { ); }; + const handleToggleDisableAgents = (checked: boolean) => { + updateSettings( + { disable_agents_for_internal_users: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleAllowAgentsTeamAdmins = (checked: boolean) => { + updateSettings( + { allow_agents_for_team_admins: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleDisableVectorStores = (checked: boolean) => { + updateSettings( + { disable_vector_stores_for_internal_users: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleAllowVectorStoresTeamAdmins = (checked: boolean) => { + updateSettings( + { allow_vector_stores_for_team_admins: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + return ( {isLoading ? ( @@ -211,6 +273,80 @@ export default function UISettings() { + {/* Agents access control */} + + + + Disable agents for internal users + {disableAgentsProperty?.description && ( + {disableAgentsProperty.description} + )} + + + + + + + + Allow agents for team admins + + {allowAgentsTeamAdminsProperty?.description && ( + {allowAgentsTeamAdminsProperty.description} + )} + + + + + + {/* Vector Stores access control */} + + + + Disable vector stores for internal users + {disableVectorStoresProperty?.description && ( + {disableVectorStoresProperty.description} + )} + + + + + + + + Allow vector stores for team admins + + {allowVectorStoresTeamAdminsProperty?.description && ( + {allowVectorStoresTeamAdminsProperty.description} + )} + + + + + {/* Page Visibility for Internal Users */} = ({ setPage, defaultSelectedKey, collapsed = false, enabledPagesInternalUsers, enableProjectsUI }) => { +const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapsed = false, enabledPagesInternalUsers, enableProjectsUI, disableAgentsForInternalUsers, disableVectorStoresForInternalUsers }) => { const { userId, accessToken, userRole } = useAuthorized(); const { data: organizations } = useOrganizations(); @@ -450,6 +452,10 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse // Hide Projects page if enableProjectsUI is not enabled if (item.key === "projects" && !enableProjectsUI) return false; + // Hide agents and vector-stores pages for non-admin users when disabled + if (!isAdmin && item.key === "agents" && disableAgentsForInternalUsers) return false; + if (!isAdmin && item.key === "vector-stores" && disableVectorStoresForInternalUsers) return false; + // Existing role check if (item.roles && !item.roles.includes(userRole)) return false;