[Feat] Enforce Vector Store Access Controls on LiteLLM Auth (#11281)

* fix LiteLLM_ObjectPermissionTable

* fix include object_permission for list key

* fix key list to inclue obj permissions

* fix object permissions for vector stores on key info

* add key edit view with vector stores

* allow editing vector stores permissions

* fixes obj permissions

* feat: add obj permission on UI

* fix: add object_permission:true

* ui show org vector stores on org info

* fix: show object permissions on /org/info

* feat: allow updating obj permissions for keys

* fixes: key object permissions

* fixes: team object permissions

* fixes: org object permissions

* fix vector store selector for Orgs

* feat: add auth checks for vector store permissions

* feat: working auth checks for vector store permissions

* test vector stores auth checks

* Update litellm/proxy/_types.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* fix: linting

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

litellm/proxy/_types.py

@@ -1254,8 +1254,8 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase):
     """Represents a LiteLLM_ObjectPermissionTable record"""
 
     object_permission_id: str
-    mcp_servers: List[str]
-    vector_stores: List[str]
+    mcp_servers: Optional[List[str]] = []
+    vector_stores: Optional[List[str]] = []
 
 
 class LiteLLM_TeamTable(TeamBase):
@@ -2364,6 +2364,21 @@ class ProxyErrorTypes(str, enum.Enum):
     Team member permission error
     """
 
+    key_vector_store_access_denied = "key_vector_store_access_denied"
+    """
+    Key does not have access to the vector store
+    """
+
+    team_vector_store_access_denied = "team_vector_store_access_denied"
+    """
+    Team does not have access to the vector store
+    """
+
+    org_vector_store_access_denied = "org_vector_store_access_denied"
+    """
+    Organization does not have access to the vector store
+    """
+
     @classmethod
     def get_model_access_error_type_for_object(
         cls, object_type: Literal["key", "user", "team", "org"]
@@ -2380,6 +2395,20 @@ class ProxyErrorTypes(str, enum.Enum):
         elif object_type == "org":
             return cls.org_model_access_denied
 
+    @classmethod
+    def get_vector_store_access_error_type_for_object(
+        cls, object_type: Literal["key", "team", "org"]
+    ) -> "ProxyErrorTypes":
+        """
+        Get the vector store access error type for object_type
+        """
+        if object_type == "key":
+            return cls.key_vector_store_access_denied
+        elif object_type == "team":
+            return cls.team_vector_store_access_denied
+        elif object_type == "org":
+            return cls.org_vector_store_access_denied
+
 
 DB_CONNECTION_ERROR_TYPES = (
     httpx.ConnectError,

litellm/proxy/auth/auth_checks.py

@@ -32,6 +32,7 @@ from litellm.proxy._types import (
     LiteLLM_EndUserTable,
     Litellm_EntityType,
     LiteLLM_JWTAuth,
+    LiteLLM_ObjectPermissionTable,
     LiteLLM_OrganizationMembershipTable,
     LiteLLM_OrganizationTable,
     LiteLLM_TeamTable,
@@ -94,6 +95,7 @@ async def common_checks(
     8. [OPTIONAL] If guardrails modified - is request allowed to change this
     9. Check if request body is safe
     10. [OPTIONAL] Organization checks - is user_object.organization_id is set, run these checks
+    11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
     """
     _model: Optional[Union[str, List[str]]] = get_model_from_request(
         request_body, route
@@ -219,6 +221,13 @@ async def common_checks(
         valid_token=valid_token,
     )
 
+    # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
+    await vector_store_access_check(
+        request_body=request_body,
+        team_object=team_object,
+        valid_token=valid_token,
+    )
+
     return True
 
 
@@ -1576,3 +1585,104 @@ def _is_wildcard_pattern(allowed_model_pattern: str) -> bool:
     Checks if `*` is in the pattern.
     """
     return "*" in allowed_model_pattern
+
+
+async def vector_store_access_check(
+    request_body: dict,
+    team_object: Optional[LiteLLM_TeamTable],
+    valid_token: Optional[UserAPIKeyAuth],
+):
+    """
+    Checks if the object (key, team, org) has access to the vector store.
+
+    Raises ProxyException if the object (key, team, org) cannot access the specific vector store.
+    """
+    from litellm.proxy.proxy_server import prisma_client
+
+    #########################################################
+    # Get the vector store the user is trying to access
+    #########################################################
+    if prisma_client is None:
+        verbose_proxy_logger.debug(
+            "Prisma client not found, skipping vector store access check"
+        )
+        return True
+
+    if litellm.vector_store_registry is None:
+        verbose_proxy_logger.debug(
+            "Vector store registry not found, skipping vector store access check"
+        )
+        return True
+
+    vector_store_ids_to_run = litellm.vector_store_registry.get_vector_store_ids_to_run(
+        non_default_params=request_body, tools=request_body.get("tools", None)
+    )
+    if vector_store_ids_to_run is None:
+        verbose_proxy_logger.debug(
+            "Vector store to run not found, skipping vector store access check"
+        )
+        return True
+
+    #########################################################
+    # Check if the object (key, team, org) has access to the vector store
+    #########################################################
+    # Check if the key can access the vector store
+    if valid_token is not None and valid_token.object_permission_id is not None:
+        key_object_permission = (
+            await prisma_client.db.litellm_objectpermissiontable.find_unique(
+                where={"object_permission_id": valid_token.object_permission_id},
+            )
+        )
+        if key_object_permission is not None:
+            _can_object_call_vector_stores(
+                object_type="key",
+                vector_store_ids_to_run=vector_store_ids_to_run,
+                object_permissions=key_object_permission,
+            )
+
+    # Check if the team can access the vector store
+    if team_object is not None and team_object.object_permission_id is not None:
+        team_object_permission = (
+            await prisma_client.db.litellm_objectpermissiontable.find_unique(
+                where={"object_permission_id": team_object.object_permission_id},
+            )
+        )
+        if team_object_permission is not None:
+            _can_object_call_vector_stores(
+                object_type="team",
+                vector_store_ids_to_run=vector_store_ids_to_run,
+                object_permissions=team_object_permission,
+            )
+    return True
+
+
+def _can_object_call_vector_stores(
+    object_type: Literal["key", "team", "org"],
+    vector_store_ids_to_run: List[str],
+    object_permissions: Optional[LiteLLM_ObjectPermissionTable],
+):
+    """
+    Raises ProxyException if the object (key, team, org) cannot access the specific vector store.
+    """
+    if object_permissions is None:
+        return True
+
+    if object_permissions.vector_stores is None:
+        return True
+
+    # If length is 0, then the object has access to all vector stores.
+    if len(object_permissions.vector_stores) == 0:
+        return True
+
+    for vector_store_id in vector_store_ids_to_run:
+        if vector_store_id not in object_permissions.vector_stores:
+            raise ProxyException(
+                message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}",
+                type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(
+                    object_type
+                ),
+                param="vector_store",
+                code=status.HTTP_401_UNAUTHORIZED,
+            )
+
+    return True

litellm/proxy/proxy_config.yaml

@@ -3,16 +3,5 @@ model_list:
     litellm_params:
       model: anthropic/*
   
-
 general_settings:
   store_prompts_in_spend_logs: true
-
-guardrails:
-  - guardrail_name: "bedrock-pre-guard"
-    litellm_params:
-      guardrail: bedrock  # supported values: "aporia", "bedrock", "lakera"
-      mode: "post_call"
-      guardrailIdentifier: wf0hkdb5x07f # your guardrail ID on bedrock
-      guardrailVersion: "DRAFT"         # your guardrail version on bedrock
-      default_on: true
-  

tests/test_litellm/proxy/auth/test_auth_checks.py

@@ -14,11 +14,21 @@ import pytest
 
 import litellm
 from litellm.proxy._types import (
+    LiteLLM_ObjectPermissionTable,
+    LiteLLM_TeamTable,
     LiteLLM_UserTable,
     LitellmUserRoles,
+    ProxyErrorTypes,
+    ProxyException,
     SSOUserDefinedValues,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.auth.auth_checks import (
+    ExperimentalUIJWTToken,
+    _can_object_call_vector_stores,
+    get_user_object,
+    vector_store_access_check,
 )
-from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object
 from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
 from litellm.utils import get_utc_datetime
 
@@ -178,3 +188,149 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch):
     assert creation_args["models"] == ["gpt-4", "claude-3-opus"]
     assert creation_args["max_budget"] == 200.0
     assert creation_args["user_role"] == "internal_user"
+
+
+# Vector Store Auth Check Tests
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "prisma_client,vector_store_registry,expected_result",
+    [
+        (None, MagicMock(), True),  # No prisma client
+        (MagicMock(), None, True),  # No vector store registry
+        (MagicMock(), MagicMock(), True),  # No vector stores to run
+    ],
+)
+async def test_vector_store_access_check_early_returns(
+    prisma_client, vector_store_registry, expected_result
+):
+    """Test vector_store_access_check returns True for early exit conditions"""
+    request_body = {"messages": [{"role": "user", "content": "test"}]}
+
+    if vector_store_registry:
+        vector_store_registry.get_vector_store_ids_to_run.return_value = None
+
+    with patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch(
+        "litellm.vector_store_registry", vector_store_registry
+    ):
+        result = await vector_store_access_check(
+            request_body=request_body,
+            team_object=None,
+            valid_token=None,
+        )
+
+    assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    "object_permissions,vector_store_ids,should_raise,error_type",
+    [
+        (None, ["store-1"], False, None),  # None permissions - should pass
+        (
+            {"vector_stores": []},
+            ["store-1"],
+            False,
+            None,
+        ),  # Empty vector_stores - should pass (access to all)
+        (
+            {"vector_stores": ["store-1", "store-2"]},
+            ["store-1"],
+            False,
+            None,
+        ),  # Has access
+        (
+            {"vector_stores": ["store-1", "store-2"]},
+            ["store-3"],
+            True,
+            ProxyErrorTypes.key_vector_store_access_denied,
+        ),  # No access
+        (
+            {"vector_stores": ["store-1"]},
+            ["store-1", "store-3"],
+            True,
+            ProxyErrorTypes.team_vector_store_access_denied,
+        ),  # Partial access
+    ],
+)
+def test_can_object_call_vector_stores_scenarios(
+    object_permissions, vector_store_ids, should_raise, error_type
+):
+    """Test _can_object_call_vector_stores with various permission scenarios"""
+    # Convert dict to object if not None
+    if object_permissions is not None:
+        mock_permissions = MagicMock()
+        mock_permissions.vector_stores = object_permissions["vector_stores"]
+        object_permissions = mock_permissions
+
+    object_type = (
+        "key"
+        if error_type == ProxyErrorTypes.key_vector_store_access_denied
+        else "team"
+    )
+
+    if should_raise:
+        with pytest.raises(ProxyException) as exc_info:
+            _can_object_call_vector_stores(
+                object_type=object_type,
+                vector_store_ids_to_run=vector_store_ids,
+                object_permissions=object_permissions,
+            )
+        assert exc_info.value.type == error_type
+    else:
+        result = _can_object_call_vector_stores(
+            object_type=object_type,
+            vector_store_ids_to_run=vector_store_ids,
+            object_permissions=object_permissions,
+        )
+        assert result is True
+
+
+@pytest.mark.asyncio
+async def test_vector_store_access_check_with_permissions():
+    """Test vector_store_access_check with actual permission checking"""
+    request_body = {"tools": [{"type": "function", "function": {"name": "test"}}]}
+
+    # Test with valid token that has access
+    valid_token = UserAPIKeyAuth(
+        token="test-token",
+        object_permission_id="perm-123",
+        models=["gpt-4"],
+        max_budget=100.0,
+    )
+
+    mock_prisma_client = MagicMock()
+    mock_permissions = MagicMock()
+    mock_permissions.vector_stores = ["store-1", "store-2"]
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
+        return_value=mock_permissions
+    )
+
+    mock_vector_store_registry = MagicMock()
+    mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-1"]
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
+        "litellm.vector_store_registry", mock_vector_store_registry
+    ):
+        result = await vector_store_access_check(
+            request_body=request_body,
+            team_object=None,
+            valid_token=valid_token,
+        )
+
+    assert result is True
+
+    # Test with denied access
+    mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-3"]
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
+        "litellm.vector_store_registry", mock_vector_store_registry
+    ):
+        with pytest.raises(ProxyException) as exc_info:
+            await vector_store_access_check(
+                request_body=request_body,
+                team_object=None,
+                valid_token=valid_token,
+            )
+
+        assert exc_info.value.type == ProxyErrorTypes.key_vector_store_access_denied