diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index c928398a87f..fe4732c6492 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -30,7 +30,10 @@ from litellm.proxy._types import ( ) 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.vector_store_endpoints.utils import can_user_access_vector_store +from litellm.proxy.vector_store_endpoints.utils import ( + can_user_access_vector_store, + filter_listable_vector_stores, +) from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.types.vector_stores import ( @@ -390,11 +393,10 @@ async def list_vector_stores( # Filter vector stores based on access control accessible_vector_stores: Final = [] - for vs in vector_store_map.values(): - if await _check_vector_store_access(vs, user_api_key_dict): - redacted = LiteLLM_ManagedVectorStore(**vs) - redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) - accessible_vector_stores.append(redacted) + for vs in await filter_listable_vector_stores(vector_store_map.values(), user_api_key_dict): + redacted = LiteLLM_ManagedVectorStore(**vs) + redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) + accessible_vector_stores.append(redacted) total_count: Final = len(accessible_vector_stores) total_pages: Final = (total_count + page_size - 1) // page_size diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 93f1510bf22..6e94a5a88ac 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,11 +1,17 @@ import json import re +from collections.abc import Iterable +from types import MappingProxyType from typing import Any, Final, Literal from fastapi import HTTPException, Request import litellm from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + is_ui_session_credential, + resolve_ui_session_team_ids, +) from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LitellmUserRoles, @@ -160,10 +166,16 @@ async def can_user_access_vector_store( if _is_proxy_admin(user_api_key_dict): return True - vector_store_team_id: Final = vector_store.get("team_id") - if vector_store_team_id is None: + if vector_store.get("team_id") is None: return True + return await _is_vector_store_granted(vector_store, user_api_key_dict) + + +async def _is_vector_store_granted( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: vector_store_id: Final = vector_store.get("vector_store_id") or "" key_object_permission = user_api_key_dict.object_permission @@ -178,12 +190,70 @@ async def can_user_access_vector_store( if _object_permission_allows_vector_store(team_object_permission, vector_store_id): return True - if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store_team_id: - return True + return user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store.get("team_id") + +async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth: + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + team: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return user_api_key_dict.model_copy( + update=MappingProxyType( + { + "team_id": team_id, + "team_object_permission": team.object_permission, + "team_object_permission_id": team.object_permission_id, + } + ) + ) + + +async def _vector_store_listing_auth_contexts( + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[UserAPIKeyAuth, ...]: + if not is_ui_session_credential(user_api_key_dict): + return (user_api_key_dict,) + session_key_context: Final = user_api_key_dict.model_copy( + update=MappingProxyType({"team_id": None, "team_object_permission": None, "team_object_permission_id": None}) + ) + team_ids: Final = await resolve_ui_session_team_ids(user_api_key_dict) + team_contexts: Final = tuple([await _team_auth_context(team_id, user_api_key_dict) for team_id in team_ids]) + return (session_key_context, *team_contexts) + + +async def _is_vector_store_granted_to_any( + vector_store: LiteLLM_ManagedVectorStore, + auth_contexts: tuple[UserAPIKeyAuth, ...], +) -> bool: + for auth_context in auth_contexts: + if await _is_vector_store_granted(vector_store, auth_context): + return True return False +async def filter_listable_vector_stores( + vector_stores: Iterable[LiteLLM_ManagedVectorStore], + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[LiteLLM_ManagedVectorStore, ...]: + """Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants.""" + if _is_proxy_admin(user_api_key_dict): + return tuple(vector_stores) + + auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict) + return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)]) + + async def get_litellm_managed_vector_store( vector_store_id: str, ) -> LiteLLM_ManagedVectorStore | None: 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 7d72121456a..93049b21460 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 @@ -120,9 +120,7 @@ async def test_delete_vector_store_checks_access(): "team_id": "team_456", } ) - mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=mock_vector_store - ) + mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=mock_vector_store) # User from different team should get 403 user_api_key_dict = UserAPIKeyAuth(team_id="team_789") @@ -134,9 +132,115 @@ async def test_delete_vector_store_checks_access(): ): with patch("litellm.vector_store_registry", None): with pytest.raises(HTTPException) as exc_info: - await delete_vector_store( - data=request, user_api_key_dict=user_api_key_dict - ) + await delete_vector_store(data=request, user_api_key_dict=user_api_key_dict) assert exc_info.value.status_code == 403 assert "Access denied" in exc_info.value.detail + + +_UNSCOPED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_unscoped", + "custom_llm_provider": "openai", + "team_id": None, +} +_TEAM_A_OWNED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_team_a", + "custom_llm_provider": "openai", + "team_id": "team_a", +} +_UI_CREATED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_ui_created", + "custom_llm_provider": "openai", + "team_id": "litellm-dashboard", +} + + +async def _listed_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + list_vector_stores, + ) + + with patch( # test-quality-ok: the list route reads rows through this module-level DB helper, no injection seam + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[_UNSCOPED, _TEAM_A_OWNED, _UI_CREATED]), + ): + response = await list_vector_stores(user_api_key_dict=user_api_key_dict) + return sorted(vs["vector_store_id"] for vs in response["data"]) + + +@pytest.mark.asyncio +async def test_list_vector_stores_hides_ungranted_stores_from_non_admin_keys(): + """A store with no team_id and no allowlist entry is not listed for a key it was never granted to; + only team ownership or an explicit object_permission grant makes a store visible.""" + assert await _listed_ids(UserAPIKeyAuth()) == [] + assert await _listed_ids(UserAPIKeyAuth(team_id="team_a")) == ["vs_team_a"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", vector_stores=["vs_unscoped"]), + ) + ) == ["vs_unscoped"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + team_object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-2", vector_stores=["vs_unscoped"] + ), + ) + ) == ["vs_unscoped"] + assert await _listed_ids(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) == [ + "vs_team_a", + "vs_ui_created", + "vs_unscoped", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_team_ids", "session_key_grants", "expected"), + [ + ([], None, []), + ([], ["vs_unscoped"], ["vs_unscoped"]), + (["team_a"], None, ["vs_team_a"]), + (["team_a", "team_granted"], None, ["vs_team_a", "vs_unscoped"]), + ], +) +async def test_list_vector_stores_dashboard_session_resolves_real_teams( + user_team_ids: list[str], session_key_grants: list[str] | None, expected: list[str] +): + """A dashboard session lists through the user's real teams plus the session key's own grants: stores created + from the dashboard (team_id litellm-dashboard) are not visible just because every session shares that team id, + while stores owned by or granted to one of the user's teams, or granted to the session key itself, are.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + + alice = UserAPIKeyAuth( + team_id="litellm-dashboard", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission=( + LiteLLM_ObjectPermissionTable(object_permission_id="op-4", vector_stores=session_key_grants) + if session_key_grants is not None + else None + ), + ) + teams = { + "team_a": LiteLLM_TeamTableCachedObj(team_id="team_a"), + "team_granted": LiteLLM_TeamTableCachedObj( + team_id="team_granted", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-3", vector_stores=["vs_unscoped"]), + ), + } + + async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj: + return teams[team_id] + + with ( + patch( # test-quality-ok: team rows come from the module-level prisma client, no injection seam + "litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object + ), + patch( # test-quality-ok: the user row comes from the module-level prisma client, no injection seam + "litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids", + new=AsyncMock(return_value=user_team_ids), + ), + ): + assert await _listed_ids(alice) == expected