diff --git a/litellm/litellm_core_utils/error_normalization.py b/litellm/litellm_core_utils/error_normalization.py index be0098ec34b..afb7a0e9b65 100644 --- a/litellm/litellm_core_utils/error_normalization.py +++ b/litellm/litellm_core_utils/error_normalization.py @@ -106,6 +106,7 @@ _PROXY_ERROR_TYPE_MAP: Final[Mapping[str, str]] = MappingProxyType( "key_vector_store_access_denied": PERMISSION_DENIED, "team_vector_store_access_denied": PERMISSION_DENIED, "org_vector_store_access_denied": PERMISSION_DENIED, + "user_vector_store_access_denied": PERMISSION_DENIED, "tool_access_denied": PERMISSION_DENIED, "team_member_permission_error": PERMISSION_DENIED, "not_found_error": RESOURCE_NOT_FOUND, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4526765edcf..fe505ad6b06 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2968,7 +2968,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): ) vector_store_deny_by_default: bool = Field( default=False, - description="When True, a virtual key may only use vector stores explicitly listed in its object_permission.vector_stores, and a key on a team also needs the team to list them. A missing permission record, an empty list, or an unresolved team grants nothing. Dashboard session keys and non-key callers are not yet covered", + description="When True, a vector store must be explicitly listed in object_permission.vector_stores: a virtual key needs its own grant plus its team's, a keyless team member needs the team's, and a user with neither needs their own. A missing permission record, an empty list, or an unresolved team grants nothing. Dashboard session keys are not yet covered", ) missing_session_id: Literal["generate", "reject", "omit"] | None = Field( None, @@ -4514,6 +4514,11 @@ class ProxyErrorTypes(str, enum.Enum): Organization does not have access to the vector store """ + user_vector_store_access_denied = "user_vector_store_access_denied" + """ + User does not have access to the vector store + """ + team_member_already_in_team = "team_member_already_in_team" """ Team member is already in team @@ -4546,7 +4551,7 @@ class ProxyErrorTypes(str, enum.Enum): @classmethod def get_vector_store_access_error_type_for_object( - cls, object_type: Literal["key", "team", "org"] + cls, object_type: Literal["key", "team", "org", "user"] ) -> "ProxyErrorTypes": """ Get the vector store access error type for object_type @@ -4557,6 +4562,8 @@ class ProxyErrorTypes(str, enum.Enum): return cls.team_vector_store_access_denied elif object_type == "org": return cls.org_vector_store_access_denied + elif object_type == "user": + return cls.user_vector_store_access_denied DB_CONNECTION_ERROR_TYPES: Final = ( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 04003dc8187..9a4f72441cb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1284,6 +1284,7 @@ async def common_checks( request_body=request_body, team_object=team_object, valid_token=valid_token, + user_object=user_object, deny_by_default=ConfigGeneralSettings.model_validate( MappingProxyType( { @@ -6634,17 +6635,16 @@ def _get_rag_query_vector_store_id(request_body: Mapping[str, object]) -> str | return vector_store_id if isinstance(vector_store_id, str) and vector_store_id else None -def _is_strict_vector_store_virtual_key(valid_token: UserAPIKeyAuth | None) -> bool: +def _is_strict_vector_store_identity(valid_token: UserAPIKeyAuth | None) -> bool: return ( valid_token is not None - and valid_token.via_virtual_key and valid_token.api_key != LITELLM_PROXY_MASTER_KEY_ALIAS and valid_token.team_id != UI_TEAM_ID ) def _require_vector_store_grant( - object_type: Literal["key", "team"], + object_type: Literal["key", "team", "user"], vector_store_ids_to_run: Sequence[str], object_permission: _VectorStorePermissionsRow | None, ) -> None: @@ -6667,6 +6667,7 @@ async def vector_store_access_check( team_object: LiteLLM_TeamTable | None, valid_token: UserAPIKeyAuth | None, *, + user_object: LiteLLM_UserTable | None = None, deny_by_default: bool = False, ): """ @@ -6708,7 +6709,11 @@ async def vector_store_access_check( if valid_token is not None and valid_token.object_permission_id is not None else None ) - strict_key: Final = deny_by_default and _is_strict_vector_store_virtual_key(valid_token) + strict_identity: Final = deny_by_default and _is_strict_vector_store_identity(valid_token) + strict_key: Final = ( + strict_identity and valid_token is not None and valid_token.via_virtual_key and not valid_token.is_session_token + ) + has_team: Final = team_object is not None or (valid_token is not None and valid_token.team_id is not None) if strict_key: _require_vector_store_grant("key", vector_store_ids_to_run, key_object_permission) elif key_object_permission is not None: @@ -6726,7 +6731,7 @@ async def vector_store_access_check( if team_object is not None and team_object.object_permission_id is not None else None ) - if strict_key and (team_object is not None or (valid_token is not None and valid_token.team_id is not None)): + if strict_identity and has_team: _require_vector_store_grant("team", vector_store_ids_to_run, team_object_permission) elif team_object_permission is not None: _can_object_call_vector_stores( @@ -6734,11 +6739,21 @@ async def vector_store_access_check( vector_store_ids_to_run=vector_store_ids_to_run, object_permissions=team_object_permission, ) + + if strict_identity and not strict_key and not has_team: + user_object_permission: Final = ( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique( + where={"object_permission_id": user_object.object_permission_id}, + ) + if user_object is not None and user_object.object_permission_id is not None + else None + ) + _require_vector_store_grant("user", vector_store_ids_to_run, user_object_permission) return True def _can_object_call_vector_stores( - object_type: Literal["key", "team", "org"], + object_type: Literal["key", "team", "org", "user"], vector_store_ids_to_run: Sequence[str], object_permissions: _VectorStorePermissionsRow | None, ): diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index f7e58612b84..436a773b85f 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -1818,6 +1818,7 @@ async def test_vector_store_access_check_enforces_team_allowlist_for_rag_query( _KEY_DENIED: Final = ProxyErrorTypes.key_vector_store_access_denied _TEAM_DENIED: Final = ProxyErrorTypes.team_vector_store_access_denied +_USER_DENIED: Final = ProxyErrorTypes.user_vector_store_access_denied def _virtual_key( @@ -1844,8 +1845,13 @@ async def _common_checks_for_rag_query( team_object: LiteLLM_TeamTable | None = None, team_permission: SimpleNamespace | None = None, user_object: LiteLLM_UserTable | None = None, + user_permission: SimpleNamespace | LiteLLM_ObjectPermissionTable | None = None, ) -> bool: - permissions: Final = {"key-permission": key_permission, "team-permission": team_permission} + permissions: Final = { + "key-permission": key_permission, + "team-permission": team_permission, + "user-permission": user_permission, + } mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( side_effect=lambda where: permissions.get(where["object_permission_id"]) @@ -2028,6 +2034,143 @@ async def test_user_owned_team_key_without_key_grant_is_denied_despite_team_memb ) +def _user_row(user_vector_stores: list[str] | None, teams: tuple[str, ...] = ()) -> LiteLLM_UserTable: + return LiteLLM_UserTable( + user_id="user-1", + teams=list(teams), + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission_id=None if user_vector_stores is None else "user-permission", + ) + + +async def _keyless_rag_query( + general_settings: Mapping[str, object], + user_vector_stores: list[str] | None, + team_vector_stores: list[str] | None = None, + team_id: str | None = None, + user_permission: LiteLLM_ObjectPermissionTable | None = None, +) -> bool: + return await _common_checks_for_rag_query( + general_settings, + UserAPIKeyAuth(user_id="user-1", team_id=team_id, user_role=LitellmUserRoles.INTERNAL_USER), + None, + None + if team_id is None + else LiteLLM_TeamTable( + team_id=team_id, object_permission_id=None if team_vector_stores is None else "team-permission" + ), + None if team_vector_stores is None else SimpleNamespace(vector_stores=team_vector_stores), + _user_row(user_vector_stores, teams=() if team_id is None else (team_id,)), + user_permission + or (None if user_vector_stores is None else SimpleNamespace(vector_stores=user_vector_stores)), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_vector_stores", "denied_by"), + [(["KBSTOREA"], None), (["KBSTOREB"], _USER_DENIED), ([], _USER_DENIED), (None, _USER_DENIED)], + ids=["user-grants", "user-excludes", "user-empty", "user-no-record"], +) +async def test_keyless_user_without_team_needs_user_grant_under_deny_by_default( + user_vector_stores: list[str] | None, denied_by: ProxyErrorTypes | None +): + await _assert_rag_query_outcome( + _keyless_rag_query({"vector_store_deny_by_default": True}, user_vector_stores), denied_by + ) + + +@pytest.mark.asyncio +async def test_keyless_user_null_vector_store_list_grants_nothing_under_deny_by_default(): + null_list: Final = LiteLLM_ObjectPermissionTable(object_permission_id="user-permission", vector_stores=None) + + await _assert_rag_query_outcome( + _keyless_rag_query({"vector_store_deny_by_default": True}, [], user_permission=null_list), _USER_DENIED + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("general_settings", "user_vector_stores"), + [({}, ["KBSTOREB"]), ({"vector_store_deny_by_default": False}, ["KBSTOREB"]), ({}, None)], + ids=["omitted-user-excludes", "false-user-excludes", "omitted-user-no-record"], +) +async def test_keyless_user_keeps_legacy_vector_store_behavior_when_flag_off( + general_settings: Mapping[str, object], user_vector_stores: list[str] | None +): + await _assert_rag_query_outcome(_keyless_rag_query(general_settings, user_vector_stores), None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_vector_stores", "user_vector_stores", "denied_by"), + [ + (["KBSTOREA"], [], None), + (["KBSTOREB"], ["KBSTOREA"], _TEAM_DENIED), + ([], ["KBSTOREA"], _TEAM_DENIED), + (None, ["KBSTOREA"], _TEAM_DENIED), + ], + ids=["team-grants-user-empty", "team-excludes-user-grants", "team-empty-user-grants", "team-no-record-user-grants"], +) +async def test_keyless_team_member_needs_only_team_grant_under_deny_by_default( + team_vector_stores: list[str] | None, user_vector_stores: list[str] | None, denied_by: ProxyErrorTypes | None +): + await _assert_rag_query_outcome( + _keyless_rag_query({"vector_store_deny_by_default": True}, user_vector_stores, team_vector_stores, "team-1"), + denied_by, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "denied_by"), + [("team-1", None), (None, None)], + ids=["session-on-team-uses-team-grant", "session-without-team-uses-user-grant"], +) +async def test_session_token_is_not_a_virtual_key_under_deny_by_default( + team_id: str | None, denied_by: ProxyErrorTypes | None +): + session = UserAPIKeyAuth( + api_key="hashed-session-token", + user_id="user-1", + team_id=team_id, + user_role=LitellmUserRoles.INTERNAL_USER, + is_session_token=True, + ) + session.via_virtual_key = True + + await _assert_rag_query_outcome( + _common_checks_for_rag_query( + {"vector_store_deny_by_default": True}, + session, + None, + None if team_id is None else LiteLLM_TeamTable(team_id=team_id, object_permission_id="team-permission"), + SimpleNamespace(vector_stores=["KBSTOREA"]), + _user_row(["KBSTOREA"] if team_id is None else []), + SimpleNamespace(vector_stores=["KBSTOREA"] if team_id is None else []), + ), + denied_by, + ) + + +@pytest.mark.asyncio +async def test_user_owned_standalone_key_cannot_use_owner_grants_under_deny_by_default(): + owner: Final = LiteLLM_UserTable( + user_id="key-owner", user_role=LitellmUserRoles.INTERNAL_USER, object_permission_id="user-permission" + ) + + await _assert_rag_query_outcome( + _common_checks_for_rag_query( + {"vector_store_deny_by_default": True}, + _virtual_key(), + None, + user_object=owner, + user_permission=SimpleNamespace(vector_stores=["KBSTOREA"]), + ), + _KEY_DENIED, + ) + + def test_can_object_call_model_with_alias(): """Test that can_object_call_model works with model aliases""" from litellm import Router diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 2fe435e733e..90078fa836f 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -4916,6 +4916,87 @@ async def test_team_key_vector_store_access_when_team_cannot_be_resolved( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_lookup", "denied"), + [ + ({"team-1": "team-1-grants-a"}, False), + (RuntimeError("team cache unavailable"), True), + ({"team-1": "team-1-grants-none", "team-2": "team-2-grants-a"}, True), + ], + ids=["resolved-team-grants", "team-lookup-swallowed-no-personal-fallback", "other-member-team-grants-ignored"], +) +async def test_keyless_team_member_vector_store_access_uses_only_the_resolved_team( + monkeypatch: pytest.MonkeyPatch, team_lookup: dict[str, str] | Exception, denied: bool +): + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_models=["gpt-4o-mini"], + user_role=LitellmUserRoles.INTERNAL_USER, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/rag/query") + grants = { + "team-1-grants-a": ["KBSTOREA"], + "team-1-grants-none": [], + "team-2-grants-a": ["KBSTOREA"], + "user-grants-a": ["KBSTOREA"], + } + database = MagicMock() + database.db.litellm_objectpermissiontable.find_unique = AsyncMock( + side_effect=lambda where: SimpleNamespace(vector_stores=grants[where["object_permission_id"]]) + ) + attrs = { + **_proxy_attrs_for_centralized_checks(), + "prisma_client": database, + "general_settings": {"vector_store_deny_by_default": True}, + } + for name, value in attrs.items(): + monkeypatch.setattr(litellm.proxy.proxy_server, name, value) + monkeypatch.setattr(litellm, "vector_store_registry", None) + + async def get_team(team_id: str, **_: object) -> LiteLLM_TeamTableCachedObj: + if isinstance(team_lookup, Exception): + raise team_lookup + return LiteLLM_TeamTableCachedObj(team_id=team_id, models=["gpt-4o-mini"], object_permission_id=team_lookup[team_id]) + + monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.get_team_object", get_team) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_membership", AsyncMock(return_value=None)) + monkeypatch.setattr( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + AsyncMock( + return_value=LiteLLM_UserTable( + user_id="user-1", teams=["team-1", "team-2"], object_permission_id="user-grants-a" + ) + ), + ) + + checks = _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "what is in this KB?"}], + "retrieval_config": {"vector_store_id": "KBSTOREA", "custom_llm_provider": "bedrock"}, + }, + route="/v1/rag/query", + ) + if not denied: + await checks + return + with pytest.raises(ProxyException) as exc_info: + await checks + assert (exc_info.value.type, exc_info.value.param, exc_info.value.code) == ( + ProxyErrorTypes.team_vector_store_access_denied, + "vector_store", + "401", + ) + + @pytest.mark.asyncio async def test_centralized_common_checks_propagates_end_user_budget_error(): """Regression: ``get_end_user_object`` raises ``litellm.BudgetExceededError``