From 41f95cc5dd5253c102597fffc41828ab0d78eea3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 2 Mar 2026 17:41:00 +0000 Subject: [PATCH] fix(proxy): allow virtual keys with object_permission.vector_stores to access vector store search Fixes #22577 The _check_vector_store_access function only checked team_id membership, causing virtual keys with vector_store permissions via object_permission to be denied access to /v1/vector_stores/{id}/search. Changes: - Added admin bypass: proxy admins always have access - Added object_permission.vector_stores check: if the key's object_permission lists the vector store ID, access is granted even when team_id doesn't match - Added _is_vector_store_id_in_object_permissions helper function - Applied fix to both copies (endpoints.py and management_endpoints.py) - Added comprehensive tests for all new access paths Co-authored-by: Ishaan Jaff --- .../proxy/vector_store_endpoints/endpoints.py | 59 ++++- .../management_endpoints.py | 59 ++++- .../test_vector_store_access_control.py | 211 +++++++++++++++++- .../test_vector_store_endpoints.py | 8 + 4 files changed, 313 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 30cabd3eeff..2dfa48365d8 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -23,31 +23,68 @@ def _check_vector_store_access( user_api_key_dict: UserAPIKeyAuth, ) -> bool: """ - Check if the user has access to the vector store based on team membership. - + Check if the user has access to the vector store. + Args: vector_store: The vector store to check access for user_api_key_dict: User API key authentication info - + Returns: True if user has access, False otherwise - - Access rules: - - If vector store has no team_id, it's accessible to all (legacy behavior) - - If user's team_id matches the vector store's team_id, access is granted - - Otherwise, access is denied + + Access rules (checked in order): + 1. Proxy admins always have access + 2. If vector store has no team_id, it's accessible to all (legacy behavior) + 3. If user's team_id matches the vector store's team_id, access is granted + 4. If the key's object_permission.vector_stores includes the vector store id, access is granted + 5. Otherwise, access is denied """ + from litellm.proxy._types import LitellmUserRoles + + if user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN.value, + ): + return True + + vector_store_id = vector_store.get("vector_store_id") vector_store_team_id = vector_store.get("team_id") - + # If vector store has no team_id, it's accessible to all (legacy behavior) if vector_store_team_id is None: return True - + # Check if user's team matches the vector store's team user_team_id = user_api_key_dict.team_id if user_team_id == vector_store_team_id: return True - + + # Check if the key's object_permission grants access to this vector store + if vector_store_id and _is_vector_store_id_in_object_permissions( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ): + return True + + return False + + +def _is_vector_store_id_in_object_permissions( + vector_store_id: str, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: + """ + Check if vector_store_id is listed in the key's object_permission.vector_stores. + + An empty vector_stores list means no restriction (access to all), so returns True. + """ + object_permission = user_api_key_dict.object_permission + if object_permission is not None: + allowed = object_permission.vector_stores + if allowed is not None: + if len(allowed) == 0: + return True + return vector_store_id in allowed return False diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index cccbb51f47b..891dca25699 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -280,31 +280,68 @@ def _check_vector_store_access( user_api_key_dict: UserAPIKeyAuth, ) -> bool: """ - Check if the user has access to the vector store based on team membership. - + Check if the user has access to the vector store. + Args: vector_store: The vector store to check access for user_api_key_dict: User API key authentication info - + Returns: True if user has access, False otherwise - - Access rules: - - If vector store has no team_id, it's accessible to all (legacy behavior) - - If user's team_id matches the vector store's team_id, access is granted - - Otherwise, access is denied + + Access rules (checked in order): + 1. Proxy admins always have access + 2. If vector store has no team_id, it's accessible to all (legacy behavior) + 3. If user's team_id matches the vector store's team_id, access is granted + 4. If the key's object_permission.vector_stores includes the vector store id, access is granted + 5. Otherwise, access is denied """ + from litellm.proxy._types import LitellmUserRoles + + if user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN.value, + ): + return True + + vector_store_id = vector_store.get("vector_store_id") vector_store_team_id = vector_store.get("team_id") - + # If vector store has no team_id, it's accessible to all (legacy behavior) if vector_store_team_id is None: return True - + # Check if user's team matches the vector store's team user_team_id = user_api_key_dict.team_id if user_team_id == vector_store_team_id: return True - + + # Check if the key's object_permission grants access to this vector store + if vector_store_id and _is_vector_store_id_in_object_permissions( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ): + return True + + return False + + +def _is_vector_store_id_in_object_permissions( + vector_store_id: str, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: + """ + Check if vector_store_id is listed in the key's object_permission.vector_stores. + + An empty vector_stores list means no restriction (access to all), so returns True. + """ + object_permission = user_api_key_dict.object_permission + if object_permission is not None: + allowed = object_permission.vector_stores + if allowed is not None: + if len(allowed) == 0: + return True + return vector_store_id in allowed return False diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py index 74d2a0d66b2..e80430f650d 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py @@ -1,9 +1,13 @@ """ -Test vector store access control based on team membership. +Test vector store access control based on team membership, object permissions, +and admin role. Core tests: 1. Access control logic works correctly for different team scenarios 2. Delete endpoint enforces team access control +3. Admin users always have access +4. Virtual keys with object_permission.vector_stores have access (issue #22577) +5. Both endpoints.py and management_endpoints.py copies behave identically """ from unittest.mock import AsyncMock, MagicMock, patch @@ -11,7 +15,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, + UserAPIKeyAuth, +) from litellm.proxy.vector_store_endpoints.management_endpoints import ( _check_vector_store_access, ) @@ -49,6 +57,130 @@ def test_check_vector_store_access(): assert _check_vector_store_access(vector_store, user) is False +def test_check_vector_store_access_admin_bypass(): + """Test that proxy admin users always have access regardless of team_id. + + Regression test for https://github.com/BerriAI/litellm/issues/22577 + """ + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_team", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + admin_user = UserAPIKeyAuth( + team_id="different_team", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + assert _check_vector_store_access(vector_store, admin_user) is True + + +def test_check_vector_store_access_object_permission_grants_access(): + """Test that a virtual key with object_permission.vector_stores can access + the listed vector store even when the team_id does not match. + + This is the core fix for https://github.com/BerriAI/litellm/issues/22577 + """ + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_123", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + user = UserAPIKeyAuth( + team_id="team_789", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm_1", + vector_stores=["vs_123", "vs_other"], + ), + ) + assert _check_vector_store_access(vector_store, user) is True + + +def test_check_vector_store_access_object_permission_empty_list_grants_all(): + """Test that an empty object_permission.vector_stores list means + unrestricted access (can access all vector stores).""" + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_any", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + user = UserAPIKeyAuth( + team_id="team_789", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm_1", + vector_stores=[], + ), + ) + assert _check_vector_store_access(vector_store, user) is True + + +def test_check_vector_store_access_object_permission_does_not_include_store(): + """Test that access is denied when the key has object_permission but the + specific vector store is not listed.""" + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_restricted", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + user = UserAPIKeyAuth( + team_id="team_789", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm_1", + vector_stores=["vs_other"], + ), + ) + assert _check_vector_store_access(vector_store, user) is False + + +def test_check_vector_store_access_no_object_permission(): + """Test that access is denied when there is no object_permission + and the team_id does not match.""" + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_123", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + user = UserAPIKeyAuth( + team_id="team_789", + object_permission=None, + ) + assert _check_vector_store_access(vector_store, user) is False + + +def test_endpoints_check_vector_store_access_matches_management(): + """Verify that the endpoints.py copy of _check_vector_store_access + behaves identically to the management_endpoints.py copy.""" + from litellm.proxy.vector_store_endpoints.endpoints import ( + _check_vector_store_access as endpoints_check, + ) + + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_123", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + + # Virtual key with object_permission should be granted access + user_with_perm = UserAPIKeyAuth( + team_id="team_789", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm_1", + vector_stores=["vs_123"], + ), + ) + assert endpoints_check(vector_store, user_with_perm) is True + + # Admin should be granted access + admin_user = UserAPIKeyAuth( + team_id="other_team", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + assert endpoints_check(vector_store, admin_user) is True + + # Wrong team, no permission → denied + user_denied = UserAPIKeyAuth(team_id="team_789") + assert endpoints_check(vector_store, user_denied) is False + + @pytest.mark.asyncio async def test_delete_vector_store_checks_access(): """Test that delete endpoint enforces team access control""" @@ -85,3 +217,78 @@ async def test_delete_vector_store_checks_access(): assert exc_info.value.status_code == 403 assert "Access denied" in exc_info.value.detail + + +def test_search_endpoint_check_with_object_permission(): + """Test that _update_request_data_with_litellm_managed_vector_store_registry + allows access when a key has object_permission granting vector store access, + even when team_id doesn't match. + + End-to-end test for the search path from issue #22577. + """ + import litellm + from litellm.proxy.vector_store_endpoints.endpoints import ( + _update_request_data_with_litellm_managed_vector_store_registry, + ) + + mock_vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_pgvector_1", + "custom_llm_provider": "openai", + "team_id": "team_owner", + "litellm_params": {"some_param": "value"}, + } + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( + mock_vector_store + ) + + user_with_perm = UserAPIKeyAuth( + team_id="team_consumer", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm_1", + vector_stores=["vs_pgvector_1"], + ), + ) + + with patch.object(litellm, "vector_store_registry", mock_registry): + result = _update_request_data_with_litellm_managed_vector_store_registry( + data={"query": "test"}, + vector_store_id="vs_pgvector_1", + user_api_key_dict=user_with_perm, + ) + + assert result["custom_llm_provider"] == "openai" + assert result["some_param"] == "value" + + +def test_search_endpoint_denied_without_permission(): + """Test that _update_request_data_with_litellm_managed_vector_store_registry + denies access when the key has no matching object_permission and wrong team.""" + import litellm + from litellm.proxy.vector_store_endpoints.endpoints import ( + _update_request_data_with_litellm_managed_vector_store_registry, + ) + + mock_vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_pgvector_1", + "custom_llm_provider": "openai", + "team_id": "team_owner", + } + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( + mock_vector_store + ) + + user_no_perm = UserAPIKeyAuth(team_id="team_consumer") + + with patch.object(litellm, "vector_store_registry", mock_registry): + with pytest.raises(HTTPException) as exc_info: + _update_request_data_with_litellm_managed_vector_store_registry( + data={"query": "test"}, + vector_store_id="vs_pgvector_1", + user_api_key_dict=user_no_perm, + ) + + assert exc_info.value.status_code == 403 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 b24f0004f22..85428a42526 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 @@ -1668,6 +1668,8 @@ class TestCheckVectorStoreAccess: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.team_id = "team-123" + mock_user_api_key.user_role = None + mock_user_api_key.object_permission = None result = _check_vector_store_access(vector_store, mock_user_api_key) assert result is True @@ -1682,6 +1684,8 @@ class TestCheckVectorStoreAccess: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.team_id = "team-123" + mock_user_api_key.user_role = None + mock_user_api_key.object_permission = None result = _check_vector_store_access(vector_store, mock_user_api_key) assert result is True @@ -1696,6 +1700,8 @@ class TestCheckVectorStoreAccess: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.team_id = "team-456" + mock_user_api_key.user_role = None + mock_user_api_key.object_permission = None result = _check_vector_store_access(vector_store, mock_user_api_key) assert result is False @@ -1710,6 +1716,8 @@ class TestCheckVectorStoreAccess: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.team_id = None + mock_user_api_key.user_role = None + mock_user_api_key.object_permission = None result = _check_vector_store_access(vector_store, mock_user_api_key) assert result is False