[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 <noreply@anthropic.com>
This commit is contained in:
parent
dd183a7fcb
commit
96b75be03d
@ -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 <your_api_key>"
|
||||
```
|
||||
"""
|
||||
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:
|
||||
|
||||
126
litellm/proxy/common_utils/rbac_utils.py
Normal file
126
litellm/proxy/common_utils/rbac_utils.py
Normal file
@ -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
|
||||
@ -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",
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
84
tests/litellm/proxy/agent_endpoints/test_agent_rbac.py
Normal file
84
tests/litellm/proxy/agent_endpoints/test_agent_rbac.py
Normal file
@ -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
|
||||
156
tests/litellm/proxy/common_utils/test_rbac_utils.py
Normal file
156
tests/litellm/proxy/common_utils/test_rbac_utils.py
Normal file
@ -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
|
||||
@ -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"
|
||||
@ -15,6 +15,8 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side
|
||||
const { accessToken } = useAuthorized();
|
||||
const [enabledPagesInternalUsers, setEnabledPagesInternalUsers] = useState<string[] | null>(null);
|
||||
const [enableProjectsUI, setEnableProjectsUI] = useState<boolean>(false);
|
||||
const [disableAgentsForInternalUsers, setDisableAgentsForInternalUsers] = useState<boolean>(false);
|
||||
const [disableVectorStoresForInternalUsers, setDisableVectorStoresForInternalUsers] = useState<boolean>(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}
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
||||
@ -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 (
|
||||
<Card title="UI Settings">
|
||||
{isLoading ? (
|
||||
@ -211,6 +273,80 @@ export default function UISettings() {
|
||||
|
||||
<Divider />
|
||||
|
||||
{/* Agents access control */}
|
||||
<Space align="start" size="middle">
|
||||
<Switch
|
||||
checked={isAgentsDisabled}
|
||||
disabled={isUpdating}
|
||||
loading={isUpdating}
|
||||
onChange={handleToggleDisableAgents}
|
||||
aria-label={disableAgentsProperty?.description ?? "Disable agents for internal users"}
|
||||
/>
|
||||
<Space direction="vertical" size={4}>
|
||||
<Typography.Text strong>Disable agents for internal users</Typography.Text>
|
||||
{disableAgentsProperty?.description && (
|
||||
<Typography.Text type="secondary">{disableAgentsProperty.description}</Typography.Text>
|
||||
)}
|
||||
</Space>
|
||||
</Space>
|
||||
|
||||
<Space align="start" size="middle" style={{ marginLeft: 32 }}>
|
||||
<Switch
|
||||
checked={Boolean(values.allow_agents_for_team_admins)}
|
||||
disabled={isUpdating || !isAgentsDisabled}
|
||||
loading={isUpdating}
|
||||
onChange={handleToggleAllowAgentsTeamAdmins}
|
||||
aria-label={allowAgentsTeamAdminsProperty?.description ?? "Allow agents for team admins"}
|
||||
/>
|
||||
<Space direction="vertical" size={4}>
|
||||
<Typography.Text strong type={!isAgentsDisabled ? "secondary" : undefined}>
|
||||
Allow agents for team admins
|
||||
</Typography.Text>
|
||||
{allowAgentsTeamAdminsProperty?.description && (
|
||||
<Typography.Text type="secondary">{allowAgentsTeamAdminsProperty.description}</Typography.Text>
|
||||
)}
|
||||
</Space>
|
||||
</Space>
|
||||
|
||||
<Divider />
|
||||
|
||||
{/* Vector Stores access control */}
|
||||
<Space align="start" size="middle">
|
||||
<Switch
|
||||
checked={isVectorStoresDisabled}
|
||||
disabled={isUpdating}
|
||||
loading={isUpdating}
|
||||
onChange={handleToggleDisableVectorStores}
|
||||
aria-label={disableVectorStoresProperty?.description ?? "Disable vector stores for internal users"}
|
||||
/>
|
||||
<Space direction="vertical" size={4}>
|
||||
<Typography.Text strong>Disable vector stores for internal users</Typography.Text>
|
||||
{disableVectorStoresProperty?.description && (
|
||||
<Typography.Text type="secondary">{disableVectorStoresProperty.description}</Typography.Text>
|
||||
)}
|
||||
</Space>
|
||||
</Space>
|
||||
|
||||
<Space align="start" size="middle" style={{ marginLeft: 32 }}>
|
||||
<Switch
|
||||
checked={Boolean(values.allow_vector_stores_for_team_admins)}
|
||||
disabled={isUpdating || !isVectorStoresDisabled}
|
||||
loading={isUpdating}
|
||||
onChange={handleToggleAllowVectorStoresTeamAdmins}
|
||||
aria-label={allowVectorStoresTeamAdminsProperty?.description ?? "Allow vector stores for team admins"}
|
||||
/>
|
||||
<Space direction="vertical" size={4}>
|
||||
<Typography.Text strong type={!isVectorStoresDisabled ? "secondary" : undefined}>
|
||||
Allow vector stores for team admins
|
||||
</Typography.Text>
|
||||
{allowVectorStoresTeamAdminsProperty?.description && (
|
||||
<Typography.Text type="secondary">{allowVectorStoresTeamAdminsProperty.description}</Typography.Text>
|
||||
)}
|
||||
</Space>
|
||||
</Space>
|
||||
|
||||
<Divider />
|
||||
|
||||
{/* Page Visibility for Internal Users */}
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={values.enabled_ui_pages_internal_users}
|
||||
|
||||
@ -42,6 +42,8 @@ interface SidebarProps {
|
||||
collapsed?: boolean;
|
||||
enabledPagesInternalUsers?: string[] | null;
|
||||
enableProjectsUI?: boolean;
|
||||
disableAgentsForInternalUsers?: boolean;
|
||||
disableVectorStoresForInternalUsers?: boolean;
|
||||
}
|
||||
|
||||
// Menu item configuration
|
||||
@ -354,7 +356,7 @@ const menuGroups: MenuGroup[] = [
|
||||
},
|
||||
];
|
||||
|
||||
const Sidebar: React.FC<SidebarProps> = ({ setPage, defaultSelectedKey, collapsed = false, enabledPagesInternalUsers, enableProjectsUI }) => {
|
||||
const Sidebar: React.FC<SidebarProps> = ({ 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<SidebarProps> = ({ 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;
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user