diff --git a/litellm/models/access_group.py b/litellm/models/access_group.py index fb0ab211bbc..cec0c1ef27e 100644 --- a/litellm/models/access_group.py +++ b/litellm/models/access_group.py @@ -20,6 +20,7 @@ class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase): mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] = [] access_skill_names: tuple[str, ...] = () + access_vector_store_ids: tuple[str, ...] = () assigned_team_ids: list[str] = [] assigned_key_ids: list[str] = [] created_at: datetime | None = None diff --git a/litellm/models/team.py b/litellm/models/team.py index 8edf10703b1..6d526528589 100644 --- a/litellm/models/team.py +++ b/litellm/models/team.py @@ -137,6 +137,10 @@ class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): last_refreshed_at: float | None = None +class LiteLLM_TeamRebuiltFromToken(LiteLLM_TeamTableCachedObj): # frozen-ok: same mutability as its cached base + """Auth's stand-in, built from the key's fields when the team row was unreadable. It has no access groups.""" + + class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): """Audit record for deleted teams; mirrors the team plus deletion metadata.""" diff --git a/litellm/proxy-extras/litellm_proxy_extras/migrations/20261111120000_add_access_group_vector_store_ids/migration.sql b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261111120000_add_access_group_vector_store_ids/migration.sql new file mode 100644 index 00000000000..fe06185e1f0 --- /dev/null +++ b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261111120000_add_access_group_vector_store_ids/migration.sql @@ -0,0 +1,4 @@ +SET LOCAL lock_timeout = '10s'; + +-- AlterTable +ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_vector_store_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma index 498f86cc40f..df8134f2b0e 100644 --- a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma @@ -1630,6 +1630,7 @@ model LiteLLM_AccessGroupTable { mcp_tool_permissions Json? access_agent_ids String[] @default([]) access_skill_names String[] @default([]) + access_vector_store_ids String[] @default([]) assigned_team_ids String[] @default([]) assigned_key_ids String[] @default([]) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fa0c423fc25..d89e4a9d9d4 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -939,6 +939,20 @@ ], "title": "Access Skill Names" }, + "access_vector_store_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Vector Store Ids" + }, "assigned_key_ids": { "anyOf": [ { @@ -1128,6 +1142,14 @@ "title": "Access Skill Names", "type": "array" }, + "access_vector_store_ids": { + "default": [], + "items": { + "type": "string" + }, + "title": "Access Vector Store Ids", + "type": "array" + }, "assigned_key_ids": { "items": { "type": "string" @@ -1304,6 +1326,20 @@ ], "title": "Access Skill Names" }, + "access_vector_store_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Vector Store Ids" + }, "assigned_key_ids": { "anyOf": [ { diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c7b698b7f18..8ecfce31076 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -54,6 +54,7 @@ from litellm.llms.base_llm.managed_resources.utils import ( is_base64_encoded_unified_id, ) from litellm.models.project import LiteLLM_ProjectTable +from litellm.models.team import LiteLLM_TeamRebuiltFromToken from litellm.proxy._types import ( MODEL_LIST_ROUTES, RBAC_ROLES, @@ -5807,6 +5808,17 @@ async def can_team_access_model( raise +def _group_grants_key( + group: LiteLLM_AccessGroupTable, valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None +) -> bool: + key_team_id: Final = valid_token.team_id or (team_object.team_id if team_object is not None else None) + return ( + bool(key_team_id and key_team_id in (group.assigned_team_ids or [])) + or bool(valid_token.token and valid_token.token in (group.assigned_key_ids or [])) + or group.access_group_id in valid_token.access_group_rule_grants + ) + + async def get_authorized_resources_from_key_access_groups( valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, @@ -5835,9 +5847,6 @@ async def get_authorized_resources_from_key_access_groups( if _prisma_client is None or _user_api_key_cache is None: return [] - key_team_id: Final = valid_token.team_id or (team_object.team_id if team_object is not None else None) - key_token: Final = valid_token.token - authorized_resources: Final[list[str]] = [] for ag_id in key_access_group_ids: try: @@ -5856,9 +5865,7 @@ async def get_authorized_resources_from_key_access_groups( if check_db_only or strict_errors: raise continue - team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or [])) - key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) - if team_authorized or key_authorized or ag_id in valid_token.access_group_rule_grants: + if _group_grants_key(ag, valid_token, team_object): authorized_resources.extend(getattr(ag, resource_field, []) or []) return list(set(authorized_resources)) @@ -8095,10 +8102,104 @@ def _legacy_requested_vector_store_ids( ) -def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, LiteLLM_ObjectPermissionTable | None]]) -> bool: +def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, _VectorStorePermissionsRow | None]]) -> bool: return any(grant is not None and bool(grant.vector_stores) for _, grant in grants) +async def access_group_vector_stores( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + unreadable: Literal["skip", "raise"] = "skip", +) -> Mapping[GrantLayer, tuple[str, ...]]: + """The stores the key's and its team's access groups list, per layer. + A deleted group lists nothing. So does a group that cannot be read, unless `unreadable` is "raise".""" + key_groups: Final = await get_access_group_tables( + (valid_token.access_group_ids or ()) if valid_token is not None else (), unreadable=unreadable + ) + team_groups: Final = await get_access_group_tables( + (team_object.access_group_ids or ()) if team_object is not None else (), unreadable=unreadable + ) + listed_key_groups: Final = ( + tuple(group for group in key_groups if _group_grants_key(group, valid_token, team_object)) + if valid_token is not None + else () + ) + return MappingProxyType( + { + "key": tuple(chain.from_iterable(group.access_vector_store_ids for group in listed_key_groups)), + "team": tuple(chain.from_iterable(group.access_vector_store_ids for group in team_groups)), + } + ) + + +async def _team_read_again( + valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None +) -> LiteLLM_TeamTable | None: + """Auth's stand-in for a team it could not read has none of the team's groups, so the team is read again.""" + team_id: Final = valid_token.team_id if valid_token is not None else None + if ( + not isinstance(team_object, LiteLLM_TeamRebuiltFromToken) + or valid_token is None + or team_id is None + or team_id == UI_TEAM_ID + ): + return team_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + +@dataclass(frozen=True, slots=True) +class _OwnAndGroupVectorStores: + vector_stores: Sequence[str] + + +def _with_group_stores( + grant: LiteLLM_ObjectPermissionTable | None, group_stores: Sequence[str] +) -> _VectorStorePermissionsRow | None: + if not group_stores: + return grant + own_stores: Final = (grant.vector_stores if grant is not None else None) or () + return _OwnAndGroupVectorStores([*dict.fromkeys((*own_stores, *group_stores))]) + + +async def _vector_store_grants( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + strict_layers: frozenset[GrantLayer] = frozenset(), +) -> tuple[tuple[GrantLayer, _VectorStorePermissionsRow | None], ...]: + """Each identity's own grant. The stores its access groups list count as if they were on its own list. + A strict team layer can only lose grants without its groups, so it keeps auth's stand-in team.""" + try: + team: Final = team_object if "team" in strict_layers else await _team_read_again(valid_token, team_object) + group_stores: Final = await access_group_vector_stores(valid_token, team, unreadable="raise") + except Exception as e: + # Skipping an unreadable team or group could empty a list and lift the restriction. + raise ProxyException( + message="Could not read the team or its access groups to check vector store access", + type=ProxyErrorTypes.internal_server_error, + param="vector_store", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) from e + return tuple( + [ + (layer, _with_group_stores(grant, group_stores.get(layer, ()))) + for layer, grant in await _identity_grants( + valid_token, team_object, user_object, prisma_client, user_api_key_cache + ) + ] + ) + + def _request_drop_policies( request_body: Mapping[str, object], llm_router: Router | None, valid_token: UserAPIKeyAuth | None ) -> tuple[DropPolicy, ...]: @@ -8161,7 +8262,7 @@ async def vector_store_access_check( Requested stores come from `vector_store_ids`, `tools[].vector_store_ids`, the tools of `additional_tools` input items and the RAG `retrieval_config.vector_store_id`. Grants come from each identity's - `object_permission.vector_stores`. + `object_permission.vector_stores`, plus the stores its access groups list, which count as if they were on it. Both modes read the same request fields, with or without a vector store registry, and answer 400 when the top-level `vector_store_ids` is not a list of strings, which legacy mode does only when the @@ -8209,7 +8310,7 @@ async def vector_store_access_check( ) if isinstance(requested_ids, _UnreadableVectorStoreIds): if _restricts_vector_stores( - await _identity_grants(valid_token, team_object, None, prisma_client, user_api_key_cache) + await _vector_store_grants(valid_token, team_object, None, prisma_client, user_api_key_cache) ): raise _malformed_vector_store_ids() return True @@ -8223,12 +8324,13 @@ async def vector_store_access_check( # Check if the object (key, team, org) has access to the vector store ######################################################### strict_layers: Final = _strict_grant_layers(deny_by_default, valid_token, team_object) - grants: Final = await _identity_grants( + grants: Final = await _vector_store_grants( valid_token, team_object, user_object if "user" in strict_layers else None, prisma_client, user_api_key_cache, + strict_layers, ) for layer, grant in grants: if layer in strict_layers: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index fd361e4a539..ada6d90c128 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -43,6 +43,7 @@ from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.models.team import LiteLLM_TeamRebuiltFromToken from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers from litellm.proxy.auth.access_group_rules import access_group_ids_for_identity @@ -3614,13 +3615,13 @@ async def _safe_fetch(label: str, awaitable): return None -def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCachedObj: +def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamRebuiltFromToken: """Reconstruct a cached team object from the fields already on the UserAPIKeyAuth. Only called when valid_token.team_id is known to be non-None (the caller gates on it).""" assert valid_token.team_id is not None token_team_models: Final = _token_team_models(valid_token) - return LiteLLM_TeamTableCachedObj( + return LiteLLM_TeamRebuiltFromToken( team_id=valid_token.team_id, max_budget=valid_token.team_max_budget, soft_budget=valid_token.team_soft_budget, diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 10bdd1d2353..dd412cfac25 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -578,6 +578,7 @@ async def create_access_group( "mcp_tool_permissions": _tool_permissions_json(data.mcp_tool_permissions), "access_agent_ids": data.access_agent_ids or [], "access_skill_names": list(data.access_skill_names or ()), + "access_vector_store_ids": list(data.access_vector_store_ids or ()), "assigned_team_ids": data.assigned_team_ids or [], "assigned_key_ids": data.assigned_key_ids or [], "created_by": user_api_key_dict.user_id, @@ -676,6 +677,7 @@ async def update_access_group( "access_mcp_server_ids", "access_agent_ids", "access_skill_names", + "access_vector_store_ids", ): value = list(iter_items(value or ())) update_data[field] = ( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 498f86cc40f..df8134f2b0e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1630,6 +1630,7 @@ model LiteLLM_AccessGroupTable { mcp_tool_permissions Json? access_agent_ids String[] @default([]) access_skill_names String[] @default([]) + access_vector_store_ids String[] @default([]) assigned_team_ids String[] @default([]) assigned_key_ids String[] @default([]) diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 4044e50f8c1..57dbca25f40 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,6 +1,7 @@ import json import re from collections.abc import Iterable, Mapping +from itertools import chain from types import MappingProxyType from typing import Final, Literal, Protocol @@ -13,6 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth, ) @@ -181,6 +183,8 @@ async def can_user_access_vector_store( 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. + 6. An access group of the caller's key or team lists this vector store id, + which counts exactly like rules 3 and 4. A dashboard session credential is evaluated against the same effective contexts as listing (its own grants plus each real team of the user). @@ -193,7 +197,9 @@ async def can_user_access_vector_store( return True auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict) - return await _is_vector_store_granted_to_any(vector_store, auth_contexts) + return await _is_vector_store_granted_to_any(vector_store, auth_contexts) or ( + vector_store.get("vector_store_id") in await _access_group_vector_store_ids(auth_contexts) + ) async def _is_vector_store_granted( @@ -266,16 +272,54 @@ async def _is_vector_store_granted_to_any( return False +async def _team_or_none(auth_context: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + """A team that cannot be read has no access groups here, as it has no own list.""" + 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 + + if auth_context.team_id is None: + return None + try: + return await get_team_object( + team_id=auth_context.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=auth_context.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # no groups, per the docstring + verbose_proxy_logger.debug("Could not read team id=%s for its access groups: %s", auth_context.team_id, e) + return None + + +async def _access_group_vector_store_ids(auth_contexts: tuple[UserAPIKeyAuth, ...]) -> frozenset[str]: + """Groups only add access here, so a group that cannot be read lists nothing.""" + from litellm.proxy.auth.auth_checks import access_group_vector_stores + + listed: Final = [ + await access_group_vector_stores(auth_context, await _team_or_none(auth_context)) + for auth_context in auth_contexts + ] + return frozenset(chain.from_iterable(chain.from_iterable(stores.values()) for stores in listed)) + + 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.""" + """Non-admins only see stores their key, a team's object_permission, team ownership or access groups grant.""" if _is_proxy_admin(user_api_key_dict): return tuple(vector_stores) auth_contexts: Final = await _vector_store_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)]) + group_store_ids: Final = await _access_group_vector_store_ids(auth_contexts) + return tuple( + [ + vs + for vs in vector_stores + if vs.get("vector_store_id") in group_store_ids or await _is_vector_store_granted_to_any(vs, auth_contexts) + ] + ) async def get_litellm_managed_vector_store( diff --git a/litellm/schema.prisma b/litellm/schema.prisma index 498f86cc40f..df8134f2b0e 100644 --- a/litellm/schema.prisma +++ b/litellm/schema.prisma @@ -1630,6 +1630,7 @@ model LiteLLM_AccessGroupTable { mcp_tool_permissions Json? access_agent_ids String[] @default([]) access_skill_names String[] @default([]) + access_vector_store_ids String[] @default([]) assigned_team_ids String[] @default([]) assigned_key_ids String[] @default([]) diff --git a/litellm/types/access_group.py b/litellm/types/access_group.py index 2fa8cb8e592..d4ad7afc54a 100644 --- a/litellm/types/access_group.py +++ b/litellm/types/access_group.py @@ -12,6 +12,7 @@ class AccessGroupCreateRequest(LiteLLMBaseModel): mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] | None = None access_skill_names: tuple[str, ...] | None = None + access_vector_store_ids: tuple[str, ...] | None = None assigned_team_ids: list[str] | None = None assigned_key_ids: list[str] | None = None @@ -24,6 +25,7 @@ class AccessGroupUpdateRequest(LiteLLMBaseModel): mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] | None = None access_skill_names: tuple[str, ...] | None = None + access_vector_store_ids: tuple[str, ...] | None = None assigned_team_ids: list[str] | None = None assigned_key_ids: list[str] | None = None @@ -44,6 +46,7 @@ class AccessGroupResponse(LiteLLMBaseModel): mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None access_agent_ids: list[str] access_skill_names: tuple[str, ...] = () + access_vector_store_ids: tuple[str, ...] = () assigned_team_ids: list[str] assigned_key_ids: list[str] access_mcp_servers: tuple[AccessGroupResource, ...] diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index f2048ca3ba9..3076fe8a4e2 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26299,6 +26299,8 @@ export interface components { access_model_names?: string[] | null; /** Access Skill Names */ access_skill_names?: string[] | null; + /** Access Vector Store Ids */ + access_vector_store_ids?: string[] | null; /** Assigned Key Ids */ assigned_key_ids?: string[] | null; /** Assigned Team Ids */ @@ -26353,6 +26355,11 @@ export interface components { * @default [] */ access_skill_names: string[]; + /** + * Access Vector Store Ids + * @default [] + */ + access_vector_store_ids: string[]; /** Assigned Key Ids */ assigned_key_ids: string[]; /** Assigned Keys */ @@ -26405,6 +26412,8 @@ export interface components { access_model_names?: string[] | null; /** Access Skill Names */ access_skill_names?: string[] | null; + /** Access Vector Store Ids */ + access_vector_store_ids?: string[] | null; /** Assigned Key Ids */ assigned_key_ids?: string[] | null; /** Assigned Team Ids */ diff --git a/tests/unit/proxy/auth/test_vector_store_access_groups.py b/tests/unit/proxy/auth/test_vector_store_access_groups.py new file mode 100644 index 00000000000..375658a5bee --- /dev/null +++ b/tests/unit/proxy/auth/test_vector_store_access_groups.py @@ -0,0 +1,359 @@ +"""A team or key in an access group that grants stores can use those stores and no others.""" + +import base64 +from collections.abc import Sequence +from typing import Final, Literal +from unittest.mock import AsyncMock + +import pytest + +import litellm +from litellm import Router +from litellm.proxy._types import ( + LiteLLM_TeamTableCachedObj, + ProxyException, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import vector_store_access_check +from litellm.proxy.auth.user_api_key_auth import _team_obj_from_token +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore +from litellm.vector_stores.vector_store_registry import VectorStoreRegistry +from tests.unit.proxy.vector_store_access_group_state import ( + TEAM_ID, + TOKEN, + ProxyState, + make_key, + make_team, + own_permission, +) + + +@pytest.fixture +def proxy(monkeypatch: pytest.MonkeyPatch) -> ProxyState: + return ProxyState(monkeypatch) + + +def _virtual_key(groups: Sequence[str]) -> UserAPIKeyAuth: + key: Final = UserAPIKeyAuth(token=TOKEN, api_key=TOKEN, team_id=TEAM_ID, access_group_ids=list(groups)) + key.via_virtual_key = True + return key + + +async def _request_is_allowed( + key: UserAPIKeyAuth, team: LiteLLM_TeamTableCachedObj | None, store: str, *, deny_by_default: bool = False +) -> bool: + try: + await vector_store_access_check( + request_body={"vector_store_ids": [store]}, + team_object=team, + valid_token=key, + deny_by_default=deny_by_default, + ) + except ProxyException as refused: + if refused.code != "401": + raise + return False + return True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key_stores", "key_groups", "team_stores", "team_groups", "store", "allowed"), + [ + pytest.param(None, ("g1",), None, (), "store-a", True, id="key-group-grants"), + pytest.param(None, ("g1",), None, (), "store-b", False, id="key-group-restricts"), + pytest.param((), ("g1",), None, (), "store-b", False, id="key-group-restricts-with-empty-own-list"), + pytest.param(None, (), None, ("g1",), "store-a", True, id="team-group-grants"), + pytest.param(None, (), None, ("g1",), "store-b", False, id="team-group-restricts"), + pytest.param(None, (), (), ("g1",), "store-b", False, id="team-group-restricts-with-empty-own-list"), + pytest.param(("store-x",), ("g1",), None, (), "store-a", True, id="key-group-widens-own-list"), + pytest.param(("store-x",), ("g1",), None, (), "store-x", True, id="key-own-list-still-counts"), + pytest.param(("store-x",), ("g1",), None, (), "store-b", False, id="key-neither-list-denies"), + pytest.param(None, (), ("store-x",), ("g1",), "store-a", True, id="team-group-widens-own-list"), + pytest.param(None, (), ("store-x",), ("g1",), "store-x", True, id="team-own-list-still-counts"), + pytest.param(("store-x",), (), None, ("g1",), "store-a", False, id="key-own-list-narrows-team-group"), + pytest.param(None, ("g1",), ("store-x",), (), "store-a", False, id="team-own-list-narrows-key-group"), + pytest.param(None, ("g1", "g2"), None, (), "store-c", True, id="stores-from-a-second-group"), + pytest.param(None, ("models-only",), None, (), "store-b", True, id="models-only-group-changes-nothing"), + pytest.param(None, (), None, ("models-only",), "store-b", True, id="team-models-only-group-changes-nothing"), + pytest.param(("store-x",), ("models-only",), None, (), "store-b", False, id="models-only-keeps-own-list"), + pytest.param(None, (), None, (), "store-b", True, id="no-grants-anywhere-changes-nothing"), + ], +) +async def test_request_check_combines_own_list_and_group_stores( + proxy: ProxyState, + key_stores: Sequence[str] | None, + key_groups: Sequence[str], + team_stores: Sequence[str] | None, + team_groups: Sequence[str], + store: str, + allowed: bool, +) -> None: + await proxy.add_group("g1", ("store-a",)) + await proxy.add_group("g2", ("store-c",)) + await proxy.add_group("models-only", models=("gpt-4",)) + + assert ( + await _request_is_allowed(make_key(key_stores, key_groups, TEAM_ID), make_team(team_stores, team_groups), store) + is allowed + ) + + +@pytest.mark.asyncio +async def test_key_group_counts_when_it_lists_only_the_keys_team(proxy: ProxyState) -> None: + await proxy.add_group("g1", ("store-a",), key_ids=()) + + assert await _request_is_allowed(make_key(groups=("g1",), team_id=TEAM_ID), None, "store-a") is True + assert await _request_is_allowed(make_key(groups=("g1",), team_id=TEAM_ID), None, "store-b") is False + + +@pytest.mark.asyncio +async def test_group_granted_by_an_identity_rule_counts_for_the_key(proxy: ProxyState) -> None: + await proxy.add_group("g1", ("store-a",), team_ids=(), key_ids=()) + key: Final = make_key(("store-x",)) + key.apply_access_group_rule_grants(("g1",)) + + assert await _request_is_allowed(key, None, "store-a") is True + assert await _request_is_allowed(key, None, "store-b") is False + + +@pytest.mark.asyncio +async def test_key_cannot_claim_a_group_that_does_not_list_it(proxy: ProxyState) -> None: + await proxy.add_group("g1", ("store-a",), team_ids=(), key_ids=()) + + assert await _request_is_allowed(make_key(("store-x",), ("g1",)), None, "store-a") is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("store", "allowed"), + [pytest.param("store-a", True, id="granted"), pytest.param("store-b", False, id="not-granted")], +) +async def test_deny_by_default_counts_group_grants(proxy: ProxyState, store: str, allowed: bool) -> None: + await proxy.add_group("g1", ("store-a",)) + key: Final = _virtual_key(("g1",)) + + assert await _request_is_allowed(key, make_team(groups=("g1",)), store, deny_by_default=True) is allowed + + +@pytest.mark.asyncio +async def test_deny_by_default_requires_the_key_too_not_only_its_team(proxy: ProxyState) -> None: + await proxy.add_group("g1", ("store-a",)) + key: Final = _virtual_key(()) + + assert await _request_is_allowed(key, make_team(groups=("g1",)), "store-a", deny_by_default=True) is False + + +_MANAGED_A: Final = base64.urlsafe_b64encode( + b"litellm_proxy:vector_store;unified_id,abc;target_model_names,gpt-4;resource_id,kb-shared" +).decode() +_PLAIN_B: Final = LiteLLM_ManagedVectorStore(vector_store_id="kb-shared", custom_llm_provider="bedrock") + + +async def _all_allowed(key: UserAPIKeyAuth, stores: Sequence[str]) -> bool: + try: + await vector_store_access_check( + request_body={"vector_store_ids": list(stores)}, team_object=None, valid_token=key + ) + except ProxyException as refused: + if refused.code != "401": + raise + return False + return True + + +@pytest.mark.asyncio +async def test_an_own_list_still_matches_by_provider_id(proxy: ProxyState, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry(vector_stores=[_PLAIN_B])) + + assert await _all_allowed(make_key((_MANAGED_A,)), ("kb-shared",)) is True + + +async def _answer( + layer: Literal["key", "team"], + own: Sequence[str] | None, + groups: Sequence[str], + requested: Sequence[str], + deny_by_default: bool, +) -> tuple[object, ...] | None: + """The refusal a virtual key gets when `layer` holds the `own` list and the `groups`, or None when allowed. + Under deny-by-default the other layer lists every requested store, so `layer` is the one that decides.""" + other: Final = tuple(requested) if deny_by_default else None + key_stores: Final = own if layer == "key" else other + key: Final = _virtual_key(groups if layer == "key" else ()).model_copy( + update={"object_permission_id": "own", "object_permission": own_permission(key_stores)} + if key_stores is not None + else {} + ) + team: Final = make_team(own if layer == "team" else other, groups if layer == "team" else ()) + try: + await vector_store_access_check( + request_body={"vector_store_ids": list(requested)}, + team_object=team, + valid_token=key, + deny_by_default=deny_by_default, + ) + except ProxyException as refused: + return (refused.code, refused.type, refused.message) + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deny_by_default", [False, True], ids=["legacy", "deny-by-default"]) +@pytest.mark.parametrize("layer", ["key", "team"]) +@pytest.mark.parametrize("own", [None, ("store-x",)], ids=["no-own-list", "own-list"]) +@pytest.mark.parametrize( + "group", [("store-a",), ("kb-shared",), (_MANAGED_A,)], ids=["plain", "provider-id", "managed"] +) +@pytest.mark.parametrize( + "requested", + [("store-a",), ("store-x",), ("store-b",), ("kb-shared",), (_MANAGED_A,), ("store-a", "store-b")], + ids=["store-a", "store-x", "store-b", "provider-id", "managed", "two-stores"], +) +async def test_a_group_store_is_checked_exactly_like_the_same_store_on_the_own_list( + proxy: ProxyState, + monkeypatch: pytest.MonkeyPatch, + deny_by_default: bool, + layer: Literal["key", "team"], + own: Sequence[str] | None, + group: Sequence[str], + requested: Sequence[str], +) -> None: + monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry(vector_stores=[_PLAIN_B])) + await proxy.add_group("g1", group) + + through_the_group: Final = await _answer(layer, own, ("g1",), requested, deny_by_default) + on_the_own_list: Final = await _answer(layer, (*(own or ()), *group), (), requested, deny_by_default) + + assert through_the_group == on_the_own_list + + +@pytest.mark.asyncio +async def test_a_refusal_lists_the_allowed_stores_as_main_does(proxy: ProxyState) -> None: + with pytest.raises(ProxyException) as denied: + await vector_store_access_check( + request_body={"vector_store_ids": ["store-b"]}, team_object=None, valid_token=make_key(("store-x",)) + ) + + assert "Only allowed to access ['store-x']" in denied.value.message + + +@pytest.mark.asyncio +async def test_keyless_team_member_gets_the_teams_group_stores(proxy: ProxyState) -> None: + await proxy.add_group("g1", ("store-a",)) + member: Final = UserAPIKeyAuth(team_id=TEAM_ID) + + assert await _request_is_allowed(member, make_team(groups=("g1",)), "store-a") is True + assert await _request_is_allowed(member, make_team(groups=("g1",)), "store-b") is False + + +@pytest.mark.asyncio +async def test_deny_by_default_still_denies_a_models_only_group(proxy: ProxyState) -> None: + await proxy.add_group("models-only", models=("gpt-4",)) + key: Final = _virtual_key(("models-only",)) + + assert await _request_is_allowed(key, make_team(groups=("models-only",)), "store-a", deny_by_default=True) is False + + +_UNREADABLE: Final = "Could not read the team or its access groups to check vector store access" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own", [None, ("store-x",)], ids=["no-own-list", "own-list"]) +@pytest.mark.parametrize("layer", ["key", "team"]) +async def test_a_group_that_cannot_be_read_blocks_the_request_instead_of_lifting_the_restriction( + monkeypatch: pytest.MonkeyPatch, layer: Literal["key", "team"], own: Sequence[str] | None +) -> None: + ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable"))) + + refusal: Final = await _answer(layer, own, ("g1",), ("store-b",), deny_by_default=False) + + assert refusal is not None + assert refusal[0] == "503" + assert refusal[2] == _UNREADABLE + + +@pytest.mark.asyncio +async def test_auths_stand_in_for_a_team_it_could_not_read_is_read_again_so_the_teams_groups_still_apply( + proxy: ProxyState, +) -> None: + await proxy.add_group("g1", ("store-a",)) + await proxy.add_team(make_team(groups=("g1",))) + key: Final = make_key(team_id=TEAM_ID) + + assert await _request_is_allowed(key, _team_obj_from_token(key), "store-a") is True + assert await _request_is_allowed(key, _team_obj_from_token(key), "store-b") is False + + +@pytest.mark.asyncio +async def test_a_stand_in_team_that_still_cannot_be_read_blocks_the_request(monkeypatch: pytest.MonkeyPatch) -> None: + ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable"))) + key: Final = make_key(team_id=TEAM_ID) + + with pytest.raises(ProxyException) as blocked: + await vector_store_access_check( + request_body={"vector_store_ids": ["store-b"]}, team_object=_team_obj_from_token(key), valid_token=key + ) + + assert (blocked.value.code, blocked.value.message) == ("503", _UNREADABLE) + + +@pytest.mark.asyncio +async def test_a_strict_team_layer_keeps_the_stand_in_team_and_is_refused_as_on_main( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable"))) + key: Final = _virtual_key(()) + + assert await _request_is_allowed(key, _team_obj_from_token(key), "store-b", deny_by_default=True) is False + + +@pytest.mark.asyncio +async def test_deleted_group_grants_nothing_and_does_not_fail_the_request(proxy: ProxyState) -> None: + assert await _request_is_allowed(make_key(team_id=TEAM_ID), make_team(groups=("deleted",)), "store-b") is True + assert ( + await _request_is_allowed(make_key(("store-x",), team_id=TEAM_ID), make_team(groups=("deleted",)), "store-b") + is False + ) + + +@pytest.mark.asyncio +async def test_group_cached_before_the_upgrade_keeps_the_callers_own_list(proxy: ProxyState) -> None: + await proxy.cache.async_set_cache( + key="access_group_id:old", + value={"access_group_id": "old", "access_group_name": "old", "assigned_team_ids": [TEAM_ID]}, + ) + + assert ( + await _request_is_allowed(make_key(("store-x",), team_id=TEAM_ID), make_team(groups=("old",)), "store-b") + is False + ) + assert ( + await _request_is_allowed(make_key(("store-x",), team_id=TEAM_ID), make_team(groups=("old",)), "store-x") + is True + ) + + +@pytest.mark.asyncio +async def test_malformed_store_ids_are_rejected_when_only_a_group_restricts(proxy: ProxyState) -> None: + await proxy.add_group("g1", ("store-a",)) + + with pytest.raises(ProxyException, match="vector_store_ids must be a list of strings"): + await vector_store_access_check( + request_body={"model": "gpt-4o-mini", "vector_store_ids": 123}, + team_object=make_team(groups=("g1",)), + valid_token=make_key(team_id=TEAM_ID), + llm_router=Router(model_list=[{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}]), + ) + + +@pytest.mark.asyncio +async def test_request_naming_no_store_never_reads_groups(monkeypatch: pytest.MonkeyPatch) -> None: + proxy: Final = ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable"))) + + result: Final = await vector_store_access_check( + request_body={"model": "gpt-4"}, team_object=make_team(groups=("g1",)), valid_token=make_key(team_id=TEAM_ID) + ) + + assert result is True + assert proxy.group_reads.await_count == 0 diff --git a/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py b/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py index f06008bb0ef..6908e61c96d 100644 --- a/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py @@ -3,8 +3,10 @@ Tests for access group management endpoints. """ import types +from collections.abc import Iterator, Mapping, Sequence from contextlib import asynccontextmanager from datetime import datetime +from typing import Final, TypeAlias from unittest.mock import AsyncMock, MagicMock import pytest @@ -22,6 +24,9 @@ from litellm.proxy.proxy_server import app from litellm.types.agents import AgentResponse +_ClientAndMocks: TypeAlias = tuple[TestClient, MagicMock, MagicMock, MagicMock, MagicMock] + + def _make_access_group_record( access_group_id: str = "ag-123", access_group_name: str = "test-group", @@ -30,12 +35,13 @@ def _make_access_group_record( access_mcp_server_ids: list | None = None, access_agent_ids: list | None = None, access_skill_names: list | None = None, + access_vector_store_ids: Sequence[str] = (), assigned_team_ids: list | None = None, assigned_key_ids: list | None = None, created_by: str | None = "admin-user", updated_by: str | None = "admin-user", created_at: datetime | None = None, -): +) -> MagicMock: created_at_val = created_at or datetime.now() updated_at_val = datetime.now() data = { @@ -46,6 +52,7 @@ def _make_access_group_record( "access_mcp_server_ids": access_mcp_server_ids or [], "access_agent_ids": access_agent_ids or [], "access_skill_names": access_skill_names or [], + "access_vector_store_ids": list(access_vector_store_ids), "assigned_team_ids": assigned_team_ids or [], "assigned_key_ids": assigned_key_ids or [], "created_at": created_at_val, @@ -78,12 +85,12 @@ def _make_key_record(token: str, key_alias: str | None = None): @pytest.fixture -def client_and_mocks(monkeypatch): +def client_and_mocks(monkeypatch: pytest.MonkeyPatch) -> Iterator[_ClientAndMocks]: """Setup mock prisma and admin auth for access group endpoints.""" mock_access_group_table = MagicMock() mock_prisma = MagicMock() - def _create_side_effect(*, data): + def _create_side_effect(*, data: Mapping[str, object]) -> MagicMock: return _make_access_group_record( access_group_id="ag-new", access_group_name=data.get("access_group_name", "new"), @@ -92,6 +99,7 @@ def client_and_mocks(monkeypatch): access_mcp_server_ids=data.get("access_mcp_server_ids", []), access_agent_ids=data.get("access_agent_ids", []), access_skill_names=data.get("access_skill_names", []), + access_vector_store_ids=data.get("access_vector_store_ids", []), assigned_team_ids=data.get("assigned_team_ids", []), assigned_key_ids=data.get("assigned_key_ids", []), created_by=data.get("created_by"), @@ -110,6 +118,7 @@ def client_and_mocks(monkeypatch): access_mcp_server_ids=data.get("access_mcp_server_ids", []), access_agent_ids=data.get("access_agent_ids", []), access_skill_names=data.get("access_skill_names", []), + access_vector_store_ids=data.get("access_vector_store_ids", []), assigned_team_ids=data.get("assigned_team_ids", []), assigned_key_ids=data.get("assigned_key_ids", []), updated_by=data.get("updated_by"), @@ -1714,6 +1723,41 @@ def test_update_access_group_writes_skill_names(client_and_mocks, update_payload assert mock_table.update.await_args.kwargs["data"]["access_skill_names"] == expected_skills +def test_create_access_group_persists_and_returns_vector_store_ids(client_and_mocks: _ClientAndMocks) -> None: + client, _, mock_table, *_ = client_and_mocks + + resp: Final = client.post( + "/v1/access_group", json={"access_group_name": "stores", "access_vector_store_ids": ["store-a"]} + ) + + assert resp.status_code == 201 + assert resp.json()["access_vector_store_ids"] == ["store-a"] + assert mock_table.create.await_args.kwargs["data"]["access_vector_store_ids"] == ["store-a"] + + +@pytest.mark.parametrize( + ("update_payload", "expected_stores"), + [ + ({"access_vector_store_ids": ["store-a"]}, ["store-a"]), + ({"access_vector_store_ids": []}, []), + ({"access_vector_store_ids": None}, []), + ], +) +def test_update_access_group_writes_vector_store_ids( + client_and_mocks: _ClientAndMocks, update_payload: Mapping[str, object], expected_stores: Sequence[str] +) -> None: + client, _, mock_table, *_ = client_and_mocks + mock_table.find_unique = AsyncMock( + return_value=_make_access_group_record(access_group_id="ag-stores", access_vector_store_ids=["old-store"]) + ) + + resp: Final = client.put("/v1/access_group/ag-stores", json=update_payload) + + assert resp.status_code == 200 + assert resp.json()["access_vector_store_ids"] == expected_stores + assert mock_table.update.await_args.kwargs["data"]["access_vector_store_ids"] == expected_stores + + def test_update_access_group_without_skill_field_leaves_skills_untouched(client_and_mocks): client, _, mock_table, *_ = client_and_mocks mock_table.find_unique = AsyncMock( diff --git a/tests/unit/proxy/vector_store_access_group_state.py b/tests/unit/proxy/vector_store_access_group_state.py new file mode 100644 index 00000000000..c9d462a9ad4 --- /dev/null +++ b/tests/unit/proxy/vector_store_access_group_state.py @@ -0,0 +1,104 @@ +"""Proxy state and identities shared by the vector store access group tests.""" + +from collections.abc import Sequence +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTableCachedObj, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import cache_access_object +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore + +TOKEN: Final = "hashed-key" +TEAM_ID: Final = "team-t" +OTHER_TEAM_ID: Final = "owner-team" + + +class ProxyState: + """The proxy globals the checks read: a real cache in front of a database that holds no groups.""" + + def __init__( + self, + monkeypatch: pytest.MonkeyPatch, + *, + group_reads: AsyncMock | None = None, + team_reads: AsyncMock | None = None, + store_reads: AsyncMock | None = None, + cache: UserApiKeyCache | MagicMock | None = None, + ) -> None: + self.cache: Final = cache if cache is not None else UserApiKeyCache() + self.group_reads: Final = group_reads if group_reads is not None else AsyncMock(return_value=None) + self.prisma: Final = MagicMock( + db=MagicMock( + litellm_accessgrouptable=MagicMock(find_unique=self.group_reads), + litellm_teamtable=MagicMock(find_unique=team_reads or AsyncMock(return_value=None)), + litellm_managedvectorstorestable=MagicMock(find_unique=store_reads or AsyncMock(return_value=None)), + ) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", self.prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", self.cache) + + async def add_group( + self, + group_id: str, + stores: Sequence[str] = (), + *, + models: Sequence[str] = (), + team_ids: Sequence[str] = (TEAM_ID,), + key_ids: Sequence[str] = (TOKEN,), + ) -> None: + await cache_access_object( + access_group_id=group_id, + access_group_table=LiteLLM_AccessGroupTable( + access_group_id=group_id, + access_group_name=group_id, + access_model_names=list(models), + access_vector_store_ids=tuple(stores), + assigned_team_ids=list(team_ids), + assigned_key_ids=list(key_ids), + ), + user_api_key_cache=self.cache, + ) + + async def add_team(self, team: LiteLLM_TeamTableCachedObj) -> None: + await self.cache.async_set_cache(key=f"team_id:{team.team_id}", value=team) + + +def own_permission(stores: Sequence[str] | None) -> LiteLLM_ObjectPermissionTable | None: + if stores is None: + return None + return LiteLLM_ObjectPermissionTable(object_permission_id="own", vector_stores=list(stores)) + + +def make_key( + own_stores: Sequence[str] | None = None, groups: Sequence[str] = (), team_id: str | None = None +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token=TOKEN, + team_id=team_id, + access_group_ids=list(groups), + object_permission_id="own" if own_stores is not None else None, + object_permission=own_permission(own_stores), + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def make_team(own_stores: Sequence[str] | None = None, groups: Sequence[str] = ()) -> LiteLLM_TeamTableCachedObj: + return LiteLLM_TeamTableCachedObj( + team_id=TEAM_ID, + access_group_ids=list(groups), + object_permission_id="own" if own_stores is not None else None, + object_permission=own_permission(own_stores), + ) + + +def team_owned_store(store_id: str) -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore(vector_store_id=store_id, custom_llm_provider="openai", team_id=OTHER_TEAM_ID) diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_access_groups.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_access_groups.py new file mode 100644 index 00000000000..e7df0c5a6d9 --- /dev/null +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_access_groups.py @@ -0,0 +1,310 @@ +"""A store an access group lists is answered exactly like the same store on the key's or team's own list.""" + +import json +from collections.abc import Mapping, Sequence +from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +import respx +from fastapi.testclient import TestClient + +import litellm +from litellm import Router +from litellm.proxy import proxy_server +from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.proxy_server import app +from litellm.proxy.utils import hash_token +from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store, filter_listable_vector_stores +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore +from litellm.vector_stores.vector_store_registry import VectorStoreRegistry +from tests._master_key import MASTER_KEY +from tests.unit.proxy.vector_store_access_group_state import ( + TEAM_ID, + TOKEN, + ProxyState, + make_key, + make_team, + own_permission, + team_owned_store, +) + +Layer = Literal["key", "team"] +_KEY: Final = "sk-group-store" +_STORES: Final = ( + team_owned_store("store-a"), + team_owned_store("store-b"), + LiteLLM_ManagedVectorStore(vector_store_id="store-free", custom_llm_provider="openai"), +) + + +async def _caller(proxy: ProxyState, layer: Layer, how: Literal["own-list", "group"]) -> UserAPIKeyAuth: + """A team key granted store-a on `layer`, either on that layer's own list or through its access group.""" + if how == "own-list": + await proxy.add_team(make_team(("store-a",) if layer == "team" else None)) + key: Final = make_key(("store-a",) if layer == "key" else None, team_id=TEAM_ID) + return ( + key.model_copy(update={"team_object_permission": own_permission(("store-a",))}) if layer == "team" else key + ) + await proxy.add_group("g1", ("store-a",)) + await proxy.add_team(make_team(groups=("g1",) if layer == "team" else ())) + return make_key(groups=("g1",) if layer == "key" else (), team_id=TEAM_ID) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("layer", ["key", "team"]) +@pytest.mark.parametrize("how", ["own-list", "group"]) +async def test_the_store_gate_and_the_list_answer_a_group_store_like_an_own_list_store( + monkeypatch: pytest.MonkeyPatch, layer: Layer, how: Literal["own-list", "group"] +) -> None: + caller: Final = await _caller(ProxyState(monkeypatch), layer, how) + + assert [await can_user_access_vector_store(store, caller) for store in _STORES] == [True, False, True] + assert await filter_listable_vector_stores(_STORES, caller) == (_STORES[0],) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_is_readable", "allowed"), + [pytest.param(True, [True, True, True], id="readable"), pytest.param(False, [True, False, True], id="unreadable")], +) +async def test_a_team_that_cannot_be_read_loses_its_group_stores_and_the_keys_still_count( + monkeypatch: pytest.MonkeyPatch, team_is_readable: bool, allowed: Sequence[bool] +) -> None: + """The key's group lists store-a and names only the key's team. The team's group lists store-b.""" + proxy: Final = ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable"))) + await proxy.add_group("key-group", ("store-a",), key_ids=()) + await proxy.add_group("team-group", ("store-b",)) + if team_is_readable: + await proxy.add_team(make_team(groups=("team-group",))) + caller: Final = make_key(groups=("key-group",), team_id=TEAM_ID) + + assert [await can_user_access_vector_store(store, caller) for store in _STORES] == allowed + + +@pytest.mark.asyncio +async def test_a_caller_without_groups_keeps_its_answers_when_its_team_cannot_be_read( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable"))) + caller: Final = make_key(("store-a",), team_id=TEAM_ID) + + assert [await can_user_access_vector_store(store, caller) for store in _STORES] == [True, False, True] + assert await filter_listable_vector_stores(_STORES, caller) == (_STORES[0],) + + +@pytest.mark.asyncio +async def test_a_group_that_cannot_be_read_grants_nothing_and_own_list_stores_still_pass( + monkeypatch: pytest.MonkeyPatch, +) -> None: + proxy: Final = ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable"))) + await proxy.add_team(make_team(groups=("g1",))) + caller: Final = make_key(("store-b",), ("g1",), TEAM_ID) + + assert [await can_user_access_vector_store(store, caller) for store in _STORES] == [False, True, True] + + +@pytest.mark.asyncio +async def test_a_dashboard_session_gets_the_group_stores_of_each_of_the_users_teams( + monkeypatch: pytest.MonkeyPatch, +) -> None: + proxy: Final = ProxyState(monkeypatch) + await proxy.add_group("g1", ("store-a",), team_ids=("team-1",)) + await proxy.add_group("g2", ("store-c",), team_ids=("team-2",)) + for team_id, group in (("team-1", "g1"), ("team-2", "g2")): + await proxy.cache.async_set_cache( + key=f"team_id:{team_id}", value=LiteLLM_TeamTableCachedObj(team_id=team_id, access_group_ids=[group]) + ) + await proxy.cache.async_set_cache( + key="session-user", value=LiteLLM_UserTable(user_id="session-user", teams=["team-1", "team-2"]) + ) + session: Final = UserAPIKeyAuth(token=TOKEN, user_id="session-user", team_id=UI_TEAM_ID) + stores: Final = tuple(team_owned_store(store_id) for store_id in ("store-a", "store-b", "store-c")) + + assert await filter_listable_vector_stores(stores, session) == (stores[0], stores[2]) + assert [await can_user_access_vector_store(store, session) for store in stores] == [True, False, True] + + +READ_ROUTES: Final = ( + pytest.param("POST", "/v1/vector_stores/{}/search", {"query": "hi"}, id="search"), + pytest.param("GET", "/v1/vector_stores/{}", None, id="retrieve"), + pytest.param("GET", "/v1/vector_stores/{}/files", None, id="file-list"), + pytest.param("GET", "/v1/vector_stores/{}/files/file-1", None, id="file-retrieve"), + pytest.param("GET", "/v1/vector_stores/{}/files/file-1/content", None, id="file-content"), + pytest.param("POST", "/vector_store/info", {"vector_store_id": "{}"}, id="info"), + pytest.param( + "POST", + "/v1/rag/query", + {"model": "m", "messages": [{"role": "user", "content": "hi"}], "retrieval_config": {"vector_store_id": "{}"}}, + id="rag-query", + ), + pytest.param( + "POST", + "/v1/responses", + {"model": "m", "input": "hi", "tools": [{"type": "file_search", "vector_store_ids": ["{}"]}]}, + id="responses", + ), +) + +MANAGE_ROUTES: Final = ( + pytest.param( + "POST", "/vector_store/update", {"vector_store_id": "{}", "vector_store_name": "renamed"}, id="update" + ), + pytest.param("POST", "/vector_store/delete", {"vector_store_id": "{}"}, id="delete"), + pytest.param("POST", "/v1/vector_stores/{}", {"name": "renamed"}, id="provider-update"), + pytest.param("DELETE", "/v1/vector_stores/{}", None, id="provider-delete"), + pytest.param("POST", "/v1/vector_stores/{}/files", {"file_id": "file-1"}, id="file-create"), + pytest.param("POST", "/v1/vector_stores/{}/files/file-1", {"attributes": {"topic": "parks"}}, id="file-update"), + pytest.param("DELETE", "/v1/vector_stores/{}/files/file-1", None, id="file-delete"), +) + + +def _body_for(body: Mapping[str, object] | None, store_id: str) -> object: + return None if body is None else json.loads(json.dumps(body).replace("{}", store_id)) + + +def _openai_answer(request: httpx.Request) -> httpx.Response: + path: Final = request.url.path + if path.endswith("/responses"): + return httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + if path.endswith("/chat/completions"): + return httpx.Response( + 200, + json={ + "id": "chat_1", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + if path.endswith("/search"): + return httpx.Response( + 200, + json={ + "object": "vector_store.search_results.page", + "search_query": "hi", + "data": [], + "has_more": False, + "next_page": None, + }, + ) + if request.method == "DELETE": + return httpx.Response( + 200, json={"id": path.rsplit("/", 1)[-1], "object": "vector_store.deleted", "deleted": True} + ) + if path.endswith("/files"): + return httpx.Response( + 200, json={"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False} + ) + if path.endswith("/content"): + return httpx.Response(200, json={"object": "vector_store.file_content.page", "data": [], "has_more": False}) + if "/files/" in path: + return httpx.Response( + 200, + json={ + "id": "file-1", + "object": "vector_store.file", + "created_at": 1, + "usage_bytes": 0, + "vector_store_id": "store-a", + "status": "completed", + }, + ) + return httpx.Response( + 200, + json={ + "id": "store-a", + "object": "vector_store", + "created_at": 1, + "name": "store-a", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": {}, + }, + ) + + +@pytest.fixture +def openai_stub(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, httpx_transport: None) -> None: + respx_mock.route(host="api.openai.com").mock(side_effect=_openai_answer) + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[{"model_name": "m", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}}] + ), + ) + monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry(vector_stores=list(_STORES))) + + +async def _statuses( + monkeypatch: pytest.MonkeyPatch, + layer: Layer, + how: Literal["own-list", "group"], + method: str, + path: str, + body: Mapping[str, object] | None, +) -> Sequence[int]: + stores: Final = {store["vector_store_id"]: store for store in _STORES} + + def row(where: Mapping[str, str], data: object = None) -> MagicMock: + return MagicMock(model_dump=MagicMock(return_value=stores[where["vector_store_id"]])) + + proxy: Final = ProxyState(monkeypatch, store_reads=AsyncMock(side_effect=row)) + monkeypatch.setattr(proxy.prisma.db.litellm_managedvectorstorestable, "update", AsyncMock(side_effect=row)) + monkeypatch.setattr(proxy.prisma.db.litellm_managedvectorstorestable, "delete", AsyncMock()) + token: Final = hash_token(_KEY) + caller: Final = (await _caller(proxy, layer, how)).model_copy( + update={"token": token, "allowed_routes": [path] if path.startswith("/vector_store/") else []} + ) + await proxy.cache.async_set_cache(key=token, value=caller) + monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + client: Final = TestClient(app, raise_server_exceptions=False) + return [ + client.request( + method, + path.format(store_id), + json=_body_for(body, store_id), + headers={"Authorization": f"Bearer {_KEY}"}, + ).status_code + for store_id in stores + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("layer", ["key", "team"]) +@pytest.mark.parametrize(("method", "path", "body"), [*READ_ROUTES, *MANAGE_ROUTES]) +async def test_every_store_route_answers_a_group_store_like_an_own_list_store( + monkeypatch: pytest.MonkeyPatch, + openai_stub: None, + layer: Layer, + method: str, + path: str, + body: Mapping[str, object] | None, +) -> None: + """Real auth runs the request check and the route's store gate. Statuses are for store-a (granted), + store-b (another team's) and store-free (no team), in that order.""" + listed: Final = await _statuses(monkeypatch, layer, "own-list", method, path, body) + grouped: Final = await _statuses(monkeypatch, layer, "group", method, path, body) + + assert grouped == listed + assert grouped[0] == 200 + assert grouped[1] in (401, 403)