From 7ca796beb10fd1f84d673f2a84c41ebeffece06c Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Sat, 30 May 2026 15:10:21 -0700 Subject: [PATCH] fix(proxy): restrict vector store index create/delete to proxy admins (#29202) * fix(proxy): restrict vector store index create/delete to proxy admins Prevent non-admin API keys from registering indexes via POST /v1/indexes or deleting Azure AI Search indexes through managed pass-through routes. Co-authored-by: Cursor * fix(proxy): tighten vector store index lifecycle checks Co-authored-by: Cursor --------- Co-authored-by: Cursor --- .../proxy/vector_store_endpoints/endpoints.py | 6 + litellm/proxy/vector_store_endpoints/utils.py | 84 ++++++- .../test_vector_store_endpoints.py | 216 +++++++++++++++++- 3 files changed, 302 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index ccf15c206b..b3bdfecbe5 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -12,6 +12,7 @@ from litellm.proxy.vector_store_endpoints.management_endpoints import ( _resolve_embedding_config, ) from litellm.proxy.vector_store_endpoints.utils import ( + assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, get_litellm_managed_vector_store, ) @@ -575,6 +576,11 @@ async def index_create( """ from litellm.proxy.proxy_server import prisma_client + assert_proxy_admin_for_vector_store_index_management( + user_api_key_dict, + operation="create", + ) + if prisma_client is None: raise HTTPException( status_code=500, diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 1221ccf119..d4afc54703 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,4 +1,5 @@ import json +import re from typing import Any, Dict, Literal, Optional from fastapi import HTTPException, Request @@ -37,6 +38,64 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: ) +def assert_proxy_admin_for_vector_store_index_management( + user_api_key_dict: UserAPIKeyAuth, + *, + operation: Literal["create", "delete", "update"] = "create", +) -> None: + """Raise 403 unless the caller is a proxy admin.""" + if _is_proxy_admin(user_api_key_dict): + return + raise HTTPException( + status_code=403, + detail=( + f"Only proxy admins can {operation} vector store indexes. " + "Contact your LiteLLM administrator." + ), + ) + + +def _suffix_after_index_name(request_path: str, index_name: str) -> Optional[str]: + """Return the path suffix after ``/indexes/{index_name}``, or None if absent.""" + match = re.search(rf"/indexes/{re.escape(index_name)}(?=$|[/?])", request_path) + if match is None: + return None + return request_path[match.end() :] + + +def _is_vector_store_index_lifecycle_request( + request_method: str, + request_path: str, + index_name: str, +) -> bool: + """ + True when the request creates or deletes a search index itself (not documents). + + Examples (admin-only): + - DELETE /azure_ai/indexes/my-index + - PUT /azure_ai/indexes/my-index + - POST /azure_ai/indexes + """ + if request_method not in ("POST", "PUT", "DELETE", "PATCH"): + return False + + suffix = _suffix_after_index_name(request_path, index_name) + if suffix is not None: + # Document operations live under /indexes/{name}/docs/... + if suffix.startswith("/docs"): + return False + # DELETE/PUT/PATCH on /indexes/{name} itself is index lifecycle. + if suffix == "" or suffix.startswith("?"): + return True + + # POST /indexes (create index at service level; no index name in path). + normalized = request_path.rstrip("/") + if request_method == "POST" and normalized.endswith("/indexes"): + return True + + return False + + def _object_permission_allows_vector_store( object_permission: Optional[LiteLLM_ObjectPermissionTable], vector_store_id: str, @@ -335,6 +394,22 @@ def is_allowed_to_call_vector_store_endpoint( request_route = get_request_route(request) + if _is_vector_store_index_lifecycle_request( + request_method=request.method, + request_path=request_route, + index_name=index_name, + ): + operation_label: Literal["create", "delete", "update"] = "create" + if request.method == "DELETE": + operation_label = "delete" + elif request.method in ("PUT", "PATCH"): + operation_label = "update" + assert_proxy_admin_for_vector_store_index_management( + user_api_key_dict, + operation=operation_label, + ) + return True + # Determine the permission type based on the request permission_type = None for endpoint in provider_vector_store_endpoints["read"]: @@ -353,7 +428,14 @@ def is_allowed_to_call_vector_store_endpoint( break if permission_type is None: - return None + raise HTTPException( + status_code=403, + detail=( + f"User does not have permission to call vector store endpoint " + f"{index_name}. Ask your administrator to add the necessary " + "permissions to your API key/Team." + ), + ) # Check if key has specific permission for allowed_vector_store_indexes has_permission = check_vector_store_permission( diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 81b67e8bc5..882081b33c 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -16,9 +16,10 @@ import litellm from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.vector_store_endpoints.endpoints import ( _update_request_data_with_litellm_managed_vector_store_registry, + index_create, ) from litellm.proxy.vector_store_endpoints.management_endpoints import ( _check_vector_store_access, @@ -33,6 +34,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( is_allowed_to_call_vector_store_endpoint, is_allowed_to_call_vector_store_files_endpoint, ) +from litellm.types.vector_stores import IndexCreateRequest from litellm.types.utils import LlmProviders @@ -674,18 +676,151 @@ class TestIsAllowedToCallVectorStoreEndpoint: "write": [("POST", "/create")], } + with patch( + "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config", + return_value=mock_provider_config, + ): + with pytest.raises(HTTPException) as exc_info: + is_allowed_to_call_vector_store_endpoint( + provider=LlmProviders.OPENAI, + index_name="my-index", + request=mock_request, + user_api_key_dict=mock_user_api_key, + ) + + assert exc_info.value.status_code == 403 + + def test_delete_index_requires_admin(self): + """Non-admin users must not delete managed search indexes via pass-through.""" + mock_request = MagicMock(spec=Request) + mock_request.method = "DELETE" + mock_request.url.path = "/azure_ai/indexes/my-index" + + mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key.user_role = None + mock_user_api_key.metadata = { + "allowed_vector_store_indexes": [ + {"index_name": "my-index", "index_permissions": ["read", "write"]} + ] + } + mock_user_api_key.team_metadata = None + + mock_provider_config = MagicMock() + mock_provider_config.get_vector_store_endpoints_by_type.return_value = { + "read": [("GET", "/docs/search"), ("POST", "/docs/search")], + "write": [("PUT", "/docs")], + } + + with patch( + "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config", + return_value=mock_provider_config, + ): + with pytest.raises(HTTPException) as exc_info: + is_allowed_to_call_vector_store_endpoint( + provider=LlmProviders.AZURE_AI, + index_name="my-index", + request=mock_request, + user_api_key_dict=mock_user_api_key, + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins can delete" in exc_info.value.detail + + def test_delete_index_allowed_for_admin(self): + """Proxy admins can delete managed search indexes via pass-through.""" + mock_request = MagicMock(spec=Request) + mock_request.method = "DELETE" + mock_request.url.path = "/azure_ai/indexes/my-index" + + mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key.user_role = LitellmUserRoles.PROXY_ADMIN + + mock_provider_config = MagicMock() + mock_provider_config.get_vector_store_endpoints_by_type.return_value = { + "read": [("GET", "/docs/search"), ("POST", "/docs/search")], + "write": [("PUT", "/docs")], + } + with patch( "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config", return_value=mock_provider_config, ): result = is_allowed_to_call_vector_store_endpoint( - provider=LlmProviders.OPENAI, + provider=LlmProviders.AZURE_AI, index_name="my-index", request=mock_request, user_api_key_dict=mock_user_api_key, ) - assert result is None + assert result is True + + def test_update_index_requires_admin_with_update_message(self): + """Non-admin users get an update-specific message for index replacement.""" + mock_request = MagicMock(spec=Request) + mock_request.method = "PUT" + mock_request.url.path = "/azure_ai/indexes/my-index" + + mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key.user_role = None + mock_user_api_key.metadata = { + "allowed_vector_store_indexes": [ + {"index_name": "my-index", "index_permissions": ["read", "write"]} + ] + } + mock_user_api_key.team_metadata = None + + mock_provider_config = MagicMock() + mock_provider_config.get_vector_store_endpoints_by_type.return_value = { + "read": [("GET", "/docs/search"), ("POST", "/docs/search")], + "write": [("PUT", "/docs")], + } + + with patch( + "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config", + return_value=mock_provider_config, + ): + with pytest.raises(HTTPException) as exc_info: + is_allowed_to_call_vector_store_endpoint( + provider=LlmProviders.AZURE_AI, + index_name="my-index", + request=mock_request, + user_api_key_dict=mock_user_api_key, + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins can update" in exc_info.value.detail + + def test_index_name_prefix_does_not_match_lifecycle_request(self): + """An index name that is only a path prefix must not trigger lifecycle checks.""" + mock_request = MagicMock(spec=Request) + mock_request.method = "DELETE" + mock_request.url.path = "/azure_ai/indexes/my-index-archive" + + mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key.user_role = None + mock_user_api_key.metadata = None + mock_user_api_key.team_metadata = None + + mock_provider_config = MagicMock() + mock_provider_config.get_vector_store_endpoints_by_type.return_value = { + "read": [], + "write": [], + } + + with patch( + "litellm.proxy.vector_store_endpoints.utils.ProviderConfigManager.get_provider_vector_stores_config", + return_value=mock_provider_config, + ): + with pytest.raises(HTTPException) as exc_info: + is_allowed_to_call_vector_store_endpoint( + provider=LlmProviders.AZURE_AI, + index_name="my-index", + request=mock_request, + user_api_key_dict=mock_user_api_key, + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins" not in exc_info.value.detail def test_team_metadata_permissions(self): """Test that team metadata permissions work.""" @@ -800,6 +935,81 @@ class TestIsAllowedToCallVectorStoreEndpoint: assert exc_info.value.status_code == 403 +class TestIndexCreate: + @pytest.mark.asyncio + async def test_index_create_requires_admin(self): + """Non-admin users must not register managed vector store indexes.""" + request = IndexCreateRequest( + index_name="test-index", + litellm_params={ + "vector_store_index": "real-index", + "vector_store_name": "azure-ai-search", + }, + ) + mock_request = MagicMock(spec=Request) + mock_response = MagicMock() + + with pytest.raises(HTTPException) as exc_info: + await index_create( + request=mock_request, + index_create_request=request, + fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth( + token="sk-test", + key_name="sk-...test", + user_role=LitellmUserRoles.INTERNAL_USER, + ), + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins can create" in exc_info.value.detail + + @pytest.mark.asyncio + async def test_index_create_allowed_for_admin(self): + """Proxy admins can register managed vector store indexes.""" + create_request = IndexCreateRequest( + index_name="test-index", + litellm_params={ + "vector_store_index": "real-index", + "vector_store_name": "azure-ai-search", + }, + ) + mock_request = MagicMock(spec=Request) + mock_response = MagicMock() + mock_row = MagicMock() + mock_row.model_dump.return_value = { + "index_name": "test-index", + "litellm_params": create_request.litellm_params.model_dump(), + } + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstoreindextable.find_unique = AsyncMock( + return_value=None + ) + mock_prisma.db.litellm_managedvectorstoreindextable.create = AsyncMock( + return_value=mock_row + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ): + result = await index_create( + request=mock_request, + index_create_request=create_request, + fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth( + token="sk-test", + key_name="sk-...test", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + ), + ) + + assert result["index_name"] == "test-index" + mock_prisma.db.litellm_managedvectorstoreindextable.create.assert_awaited_once() + + class TestIsAllowedToCallVectorStoreFilesEndpoint: def _mock_provider_config(self): provider_config = MagicMock()