scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(proxy): respect object-level permissions for managed vector store endpoints (#26351)

Browse source 
* fix(proxy): honor object_permission for managed vector store access

* perf(proxy): preload team object_permission on UserAPIKeyAuth

Populate team_object_permission during virtual-key and JWT auth when the
team is loaded, so can_user_access_vector_store uses it in memory first
and only falls back to get_object_permission by id when missing.

Made-with: Cursor
This commit is contained in:
Shivam Rawat 2026-04-24 09:21:13 -07:00 • committed by GitHub
parent 863f922be8
commit 9dcb2bd528
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 260 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/_types.py
View file

@ -2569,6 +2569,9 @@ class UserAPIKeyAuth(
None # Expanded created_by user when expand=user is used
)
end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
# Team object_permission preloaded in auth (e.g. get_team_object) to avoid
# per-request object_permission fetches in downstream checks (vector stores, etc.)
team_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
# Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery
# and forwarded into outbound tokens by guardrails such as MCPJWTSigner.
jwt_claims: Optional[Dict] = None

10
litellm/proxy/auth/user_api_key_auth.py
View file

@ -867,6 +867,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
),
jwt_claims=jwt_claims,
)
valid_token.team_object_permission = (
team_object.object_permission
if team_object is not None
else None
)
# Check if model has zero cost - if so, skip all budget checks
model = get_model_from_request(request_data, route)
@ -1452,6 +1457,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
else:
_team_obj = None
if _team_obj is not None:
valid_token.team_object_permission = _team_obj.object_permission
else:
valid_token.team_object_permission = None
await user_api_key_cache.async_set_cache(
key=valid_token.team_id, value=_team_obj
) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py

49
litellm/proxy/vector_store_endpoints/endpoints.py
View file

@ -10,6 +10,7 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.utils import jsonify_object
from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store
from litellm.types.vector_stores import IndexCreateRequest
router = APIRouter()
@ -18,40 +19,25 @@ router = APIRouter()
########################################################
def _check_vector_store_access(
async def _check_vector_store_access(
vector_store: LiteLLM_ManagedVectorStore,
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
Delegates to :func:`can_user_access_vector_store`, which honors:
- PROXY_ADMIN bypass
- legacy vector stores with no team_id
- key-level and team-level ``object_permission.vector_stores`` allowlists
- team_id match between key and store
"""
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
return False
return await can_user_access_vector_store(
vector_store=vector_store, user_api_key_dict=user_api_key_dict
)
def _update_request_data_with_litellm_managed_vector_store_registry(
async def _update_request_data_with_litellm_managed_vector_store_registry(
data: Dict,
vector_store_id: str,
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
@ -74,9 +60,8 @@ def _update_request_data_with_litellm_managed_vector_store_registry(
)
)
if vector_store_to_run is not None:
# Check access control if user_api_key_dict is provided
if user_api_key_dict is not None:
if not _check_vector_store_access(
if not await _check_vector_store_access(
vector_store_to_run, user_api_key_dict
):
raise HTTPException(
@ -140,7 +125,7 @@ async def vector_store_search(
data["vector_store_id"] = vector_store_id
# Check for legacy vector store registry (non-managed vector stores)
data = _update_request_data_with_litellm_managed_vector_store_registry(
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
)
@ -322,7 +307,7 @@ async def vector_store_retrieve(
data = {"vector_store_id": vector_store_id}
data = _update_request_data_with_litellm_managed_vector_store_registry(
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
)
@ -462,7 +447,7 @@ async def vector_store_update(
if "vector_store_id" not in data:
data["vector_store_id"] = vector_store_id
data = _update_request_data_with_litellm_managed_vector_store_registry(
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
)
@ -529,7 +514,7 @@ async def vector_store_delete(
data = {"vector_store_id": vector_store_id}
data = _update_request_data_with_litellm_managed_vector_store_registry(
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
)

55
litellm/proxy/vector_store_endpoints/management_endpoints.py
View file

@ -25,6 +25,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.proxy.vector_store_endpoints.utils import can_user_access_vector_store
from litellm.secret_managers.main import get_secret
from litellm.types.vector_stores import (
LiteLLM_ManagedVectorStore,
@ -274,37 +275,22 @@ async def _resolve_embedding_config(
########################################################
# Helper Functions
########################################################
def _check_vector_store_access(
async def _check_vector_store_access(
vector_store: LiteLLM_ManagedVectorStore,
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
Delegates to :func:`can_user_access_vector_store`, which honors:
- PROXY_ADMIN bypass
- legacy vector stores with no team_id
- key-level and team-level ``object_permission.vector_stores`` allowlists
- team_id match between key and store
"""
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
return False
return await can_user_access_vector_store(
vector_store=vector_store, user_api_key_dict=user_api_key_dict
)
async def create_vector_store_in_db(
@ -565,12 +551,11 @@ async def list_vector_stores(
vector_store_id=vector_store_id, updated_data=vector_store
)
# Filter vector stores based on team access
accessible_vector_stores = [
vs
for vs in vector_store_map.values()
if _check_vector_store_access(vs, user_api_key_dict)
]
# Filter vector stores based on access control
accessible_vector_stores = []
for vs in vector_store_map.values():
if await _check_vector_store_access(vs, user_api_key_dict):
accessible_vector_stores.append(vs)
total_count = len(accessible_vector_stores)
total_pages = (total_count + page_size - 1) // page_size
@ -647,7 +632,7 @@ async def delete_vector_store(
)
# Check access control
if vector_store_to_check and not _check_vector_store_access(
if vector_store_to_check and not await _check_vector_store_access(
vector_store_to_check, user_api_key_dict
):
raise HTTPException(
@ -703,7 +688,9 @@ async def get_vector_store_info(
)
if vector_store is not None:
# Check access control
if not _check_vector_store_access(vector_store, user_api_key_dict):
if not await _check_vector_store_access(
vector_store, user_api_key_dict
):
raise HTTPException(
status_code=403,
detail="Access denied: You do not have permission to access this vector store",
@ -749,7 +736,7 @@ async def get_vector_store_info(
# Check access control for DB vector store
vector_store_dict = vector_store.model_dump() # type: ignore[attr-defined]
vector_store_typed = LiteLLM_ManagedVectorStore(**vector_store_dict)
if not _check_vector_store_access(vector_store_typed, user_api_key_dict):
if not await _check_vector_store_access(vector_store_typed, user_api_key_dict):
raise HTTPException(
status_code=403,
detail="Access denied: You do not have permission to access this vector store",

115
litellm/proxy/vector_store_endpoints/utils.py
View file

@ -2,11 +2,124 @@ from typing import Any, Dict, Literal, Optional
from fastapi import HTTPException, Request
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.types.utils import LlmProviders
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
from litellm.utils import ProviderConfigManager
def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
return (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
def _object_permission_allows_vector_store(
object_permission: Optional[LiteLLM_ObjectPermissionTable],
vector_store_id: str,
) -> bool:
"""Returns True if an object permission explicitly allowlists the vector store."""
if object_permission is None:
return False
allowed = object_permission.vector_stores
if not allowed:
return False
return vector_store_id in allowed
async def _get_object_permission_for_id(
object_permission_id: Optional[str],
) -> Optional[LiteLLM_ObjectPermissionTable]:
"""Load an object permission record by id, using the shared cache/DB helper."""
if not object_permission_id:
return None
from litellm.proxy.auth.auth_checks import get_object_permission
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if prisma_client is None:
return None
try:
return await get_object_permission(
object_permission_id=object_permission_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
except Exception as e:
verbose_proxy_logger.debug(
"Failed to load object_permission id=%s: %s",
object_permission_id,
e,
)
return None
async def can_user_access_vector_store(
vector_store: LiteLLM_ManagedVectorStore,
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
"""
Returns True if the caller is allowed to access this managed vector store.
Access is granted (first match wins) when any of the following is true:
1. The caller's role is PROXY_ADMIN.
2. The vector store has no team_id (legacy behavior - accessible to all).
3. The caller's key-level object_permission.vector_stores explicitly lists
this vector store id.
4. The caller's team-level object_permission.vector_stores explicitly lists
this vector store id.
5. The caller's team_id matches the vector store's team_id.
Otherwise access is denied.
"""
if _is_proxy_admin(user_api_key_dict):
return True
vector_store_team_id = vector_store.get("team_id")
if vector_store_team_id is None:
return True
vector_store_id = vector_store.get("vector_store_id") or ""
key_object_permission = user_api_key_dict.object_permission
if key_object_permission is None:
key_object_permission = await _get_object_permission_for_id(
user_api_key_dict.object_permission_id
)
if _object_permission_allows_vector_store(key_object_permission, vector_store_id):
return True
team_object_permission: Optional[LiteLLM_ObjectPermissionTable] = (
user_api_key_dict.team_object_permission
)
if team_object_permission is None:
team_object_permission = await _get_object_permission_for_id(
user_api_key_dict.team_object_permission_id
)
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 False
def _does_endpoint_match(endpoint_path: str, request_path: str) -> bool:
if endpoint_path in request_path:
return True

65
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py
View file

@ -11,14 +11,19 @@ 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,
)
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
def test_check_vector_store_access():
@pytest.mark.asyncio
async def test_check_vector_store_access():
"""Test core access control logic for team-based vector store access"""
# Test 1: Legacy vector stores (no team_id) are accessible to all
@ -28,7 +33,7 @@ def test_check_vector_store_access():
"team_id": None,
}
user = UserAPIKeyAuth(team_id="team_456")
assert _check_vector_store_access(vector_store, user) is True
assert await _check_vector_store_access(vector_store, user) is True
# Test 2: User can access their team's vector stores
vector_store = {
@ -37,7 +42,7 @@ def test_check_vector_store_access():
"team_id": "team_456",
}
user = UserAPIKeyAuth(team_id="team_456")
assert _check_vector_store_access(vector_store, user) is True
assert await _check_vector_store_access(vector_store, user) is True
# Test 3: User cannot access other teams' vector stores
vector_store = {
@ -46,7 +51,57 @@ def test_check_vector_store_access():
"team_id": "team_456",
}
user = UserAPIKeyAuth(team_id="team_789")
assert _check_vector_store_access(vector_store, user) is False
assert await _check_vector_store_access(vector_store, user) is False
@pytest.mark.asyncio
async def test_check_vector_store_access_proxy_admin_bypass():
"""PROXY_ADMIN can access a vector store even if teams don't match."""
vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "vs_team",
"custom_llm_provider": "openai",
"team_id": "team_456",
}
admin = UserAPIKeyAuth(team_id="team_999", user_role=LitellmUserRoles.PROXY_ADMIN)
assert await _check_vector_store_access(vector_store, admin) is True
@pytest.mark.asyncio
async def test_check_vector_store_access_key_object_permission_grants_access():
"""A key whose object_permission.vector_stores allowlists the store can access it
even if its team_id does not match the store's team_id."""
vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "vs_explicit",
"custom_llm_provider": "openai",
"team_id": "team_456",
}
user = UserAPIKeyAuth(
team_id="team_789",
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="op-1",
vector_stores=["vs_explicit"],
),
)
assert await _check_vector_store_access(vector_store, user) is True
@pytest.mark.asyncio
async def test_check_vector_store_access_key_object_permission_wrong_store_denied():
"""A key whose object_permission.vector_stores lists *other* stores is still denied
when the key has no other reason to access this store."""
vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "vs_target",
"custom_llm_provider": "openai",
"team_id": "team_456",
}
user = UserAPIKeyAuth(
team_id="team_789",
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="op-1",
vector_stores=["vs_other"],
),
)
assert await _check_vector_store_access(vector_store, user) is False
@pytest.mark.asyncio

58
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

@ -115,7 +115,8 @@ def test_router_vector_store_file_delete_passes_correct_args():
assert call_kwargs["custom_llm_provider"] == "openai"
def test_update_request_data_with_litellm_managed_vector_store_registry():
@pytest.mark.asyncio
async def test_update_request_data_with_litellm_managed_vector_store_registry():
"""
Test that _update_request_data_with_litellm_managed_vector_store_registry
correctly updates request data with vector store registry information.
@ -139,7 +140,7 @@ def test_update_request_data_with_litellm_managed_vector_store_registry():
# Test with vector store registry
with patch.object(litellm, "vector_store_registry", mock_registry):
result = _update_request_data_with_litellm_managed_vector_store_registry(
result = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
)
@ -158,7 +159,7 @@ def test_update_request_data_with_litellm_managed_vector_store_registry():
# Test with no vector store registry
with patch.object(litellm, "vector_store_registry", None):
original_data = {"existing_key": "existing_value"}
result = _update_request_data_with_litellm_managed_vector_store_registry(
result = await _update_request_data_with_litellm_managed_vector_store_registry(
data=original_data, vector_store_id=vector_store_id
)
@ -1686,10 +1687,26 @@ async def test_new_vector_store_auto_resolves_from_router():
)
def _stub_user_api_key(
*,
team_id=None,
user_role=None,
object_permission=None,
object_permission_id=None,
team_object_permission_id=None,
):
user = UserAPIKeyAuth(team_id=team_id, user_role=user_role)
user.object_permission = object_permission
user.object_permission_id = object_permission_id
user.team_object_permission_id = team_object_permission_id
return user
class TestCheckVectorStoreAccess:
"""Test suite for _check_vector_store_access function."""
def test_access_granted_when_no_team_id(self):
@pytest.mark.asyncio
async def test_access_granted_when_no_team_id(self):
"""Test that access is granted when vector store has no team_id (legacy behavior)."""
vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "test-store",
@ -1697,13 +1714,12 @@ class TestCheckVectorStoreAccess:
# No team_id field
}
mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.team_id = "team-123"
result = _check_vector_store_access(vector_store, mock_user_api_key)
user = _stub_user_api_key(team_id="team-123")
result = await _check_vector_store_access(vector_store, user)
assert result is True
def test_access_granted_when_team_ids_match(self):
@pytest.mark.asyncio
async def test_access_granted_when_team_ids_match(self):
"""Test that access is granted when user's team_id matches vector store's team_id."""
vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "test-store",
@ -1711,13 +1727,12 @@ class TestCheckVectorStoreAccess:
"team_id": "team-123",
}
mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.team_id = "team-123"
result = _check_vector_store_access(vector_store, mock_user_api_key)
user = _stub_user_api_key(team_id="team-123")
result = await _check_vector_store_access(vector_store, user)
assert result is True
def test_access_denied_when_team_ids_dont_match(self):
@pytest.mark.asyncio
async def test_access_denied_when_team_ids_dont_match(self):
"""Test that access is denied when user's team_id doesn't match vector store's team_id."""
vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "test-store",
@ -1725,13 +1740,12 @@ class TestCheckVectorStoreAccess:
"team_id": "team-123",
}
mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.team_id = "team-456"
result = _check_vector_store_access(vector_store, mock_user_api_key)
user = _stub_user_api_key(team_id="team-456")
result = await _check_vector_store_access(vector_store, user)
assert result is False
def test_access_denied_when_vector_store_has_team_id_but_user_doesnt(self):
@pytest.mark.asyncio
async def test_access_denied_when_vector_store_has_team_id_but_user_doesnt(self):
"""Test that access is denied when vector store has team_id but user doesn't."""
vector_store: LiteLLM_ManagedVectorStore = {
"vector_store_id": "test-store",
@ -1739,10 +1753,8 @@ class TestCheckVectorStoreAccess:
"team_id": "team-123",
}
mock_user_api_key = MagicMock(spec=UserAPIKeyAuth)
mock_user_api_key.team_id = None
result = _check_vector_store_access(vector_store, mock_user_api_key)
user = _stub_user_api_key(team_id=None)
result = await _check_vector_store_access(vector_store, user)
assert result is False