From 672f4c66374f10714e62252e1c207b427a3e5fb3 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Fri, 24 Jul 2026 11:41:34 -0400 Subject: [PATCH] fix: apply session scoping to all working-representation query paths (#881) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: apply session scoping to all working-representation query paths session_name was only applied to the recent-documents query in RepresentationManager; the semantic and most-derived paths ignored it, so limit_to_session leaked cross-session conclusions into perspectives. - Thread a session allowlist (session_names) uniformly through all three query paths; pushed down to pgvector and external vector stores - Accept a list so the upcoming session-allowlist API reuses this path - Fail closed on an empty allowlist (downstream stores drop empty IN clauses, which would silently widen scope) Fixes DEV-1994 Co-Authored-By: Claude Fable 5 * feat: bare-list membership sugar in the filter DSL {"session_id": ["s1", "s2"]} is now shorthand for {"session_id": {"in": [...]}} on regular columns, generically (peer_id, etc.). JSONB metadata columns are excluded — a bare list there keeps JSONB containment semantics, unchanged. Previously a bare list on a regular column compiled to a type-mismatched equality that matched nothing, so this is strictly additive. Also translates the same shape in the turbopuffer/lancedb filter builders, and fixes lancedb dropping empty IN clauses (fail-open) — an empty membership list now emits an always-false condition. Groundwork for DEV-1995 (session allowlist via the existing filters DSL, no new API params) Co-Authored-By: Claude Fable 5 * fix: fail closed on empty session allowlist across all filter builders Empty session allowlists relied solely on the early-return guard in _get_working_representation_internal. The layers below it were inconsistent, so a future direct caller (the DEV-1995 allowlist API) would silently widen scope instead of failing closed: - _build_filter_conditions used a truthiness check; an empty list was treated like None and dropped the filter. Now uses `is not None`, matching the recent/most-derived SQL paths. - turbopuffer emitted a bare `In []` with undocumented (possibly fail-open) semantics. Now emits an explicit always-false predicate, mirroring lancedb's `1 = 0`. Also extract the duplicated JSONB column tuple in filter.py to a JSONB_COLUMNS constant. Tests exercise each fail-closed guarantee at the layer it lives, rather than masking it behind the early-return guard. --------- Co-authored-by: Claude Fable 5 --- src/crud/representation.py | 59 ++++-- src/routers/peers.py | 6 +- src/routers/sessions.py | 8 +- src/utils/filter.py | 17 +- src/vector_store/lancedb.py | 18 +- src/vector_store/turbopuffer.py | 28 ++- tests/crud/test_representation_manager.py | 218 ++++++++++++++++++++++ tests/test_advanced_filters.py | 80 ++++++++ tests/vector_store/test_lancedb.py | 21 +++ tests/vector_store/test_turbopuffer.py | 24 +++ 10 files changed, 446 insertions(+), 33 deletions(-) diff --git a/src/crud/representation.py b/src/crud/representation.py index 3061c2d9..69c9edc9 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -213,7 +213,7 @@ class RepresentationManager: self, *, db: AsyncSession | None = None, - session_name: str | None = None, + session_names: list[str] | None = None, include_semantic_query: str | None = None, embedding: list[float] | None = None, semantic_search_top_k: int | None = None, @@ -229,7 +229,10 @@ class RepresentationManager: Args: db: Optional database session. If provided, uses it directly; otherwise creates a new session via tracked_db. - session_name: Optional session to filter by + session_names: Optional session allowlist to filter by. Applied + uniformly to every query path (semantic, most-derived, and + recent). None means no session restriction; an empty list + fail-closes to an empty representation. include_semantic_query: Query for semantic search embedding: Pre-computed embedding for the semantic query. semantic_search_top_k: Number of semantic results @@ -267,7 +270,7 @@ class RepresentationManager: if db is not None: return await self._get_working_representation_internal( db, - session_name=session_name, + session_names=session_names, include_semantic_query=include_semantic_query, embedding=embedding, semantic_search_top_k=semantic_search_top_k, @@ -281,7 +284,7 @@ class RepresentationManager: ) as new_db: return await self._get_working_representation_internal( new_db, - session_name=session_name, + session_names=session_names, include_semantic_query=include_semantic_query, embedding=embedding, semantic_search_top_k=semantic_search_top_k, @@ -296,7 +299,7 @@ class RepresentationManager: self, db: AsyncSession, *, - session_name: str | None = None, + session_names: list[str] | None = None, include_semantic_query: str | None = None, embedding: list[float] | None = None, semantic_search_top_k: int | None = None, @@ -305,6 +308,12 @@ class RepresentationManager: max_observations: int = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS, ) -> Representation: """Internal implementation of get_working_representation.""" + # Fail closed on an empty allowlist. This must short-circuit before + # any query: downstream stores drop an `IN ()` clause with an empty + # list (lancedb), which would silently widen the scope instead. + if session_names is not None and not session_names: + return Representation() + total = max_observations # Calculate how many observations to get from each source @@ -345,6 +354,7 @@ class RepresentationManager: top_k=semantic_observations, max_distance=semantic_search_max_distance, embedding=embedding, + session_names=session_names, ) representation.merge_representation( Representation.from_documents(semantic_docs) @@ -353,7 +363,7 @@ class RepresentationManager: # Get most derived observations if requested if include_most_derived: derived_docs = await self._query_documents_most_derived( - db, top_k=top_observations + db, top_k=top_observations, session_names=session_names ) representation.merge_representation( Representation.from_documents(derived_docs) @@ -361,7 +371,7 @@ class RepresentationManager: # Get recent observations recent_docs = await self._query_documents_recent( - db, top_k=recent_observations, session_name=session_name + db, top_k=recent_observations, session_names=session_names ) representation.merge_representation(Representation.from_documents(recent_docs)) @@ -376,6 +386,7 @@ class RepresentationManager: max_distance: float | None = None, level: str | None = None, embedding: list[float] | None = None, + session_names: list[str] | None = None, ) -> list[models.Document]: """Query documents by semantic similarity.""" try: @@ -387,6 +398,7 @@ class RepresentationManager: top_k, max_distance, embedding=embedding, + session_names=session_names, ) else: documents = await crud.query_documents( @@ -398,6 +410,8 @@ class RepresentationManager: max_distance=max_distance, top_k=top_k, embedding=embedding, + filters=self._build_filter_conditions(session_names=session_names) + or None, ) db.expunge_all() return list(documents) @@ -407,7 +421,7 @@ class RepresentationManager: return [] async def _query_documents_recent( - self, db: AsyncSession, top_k: int, session_name: str | None = None + self, db: AsyncSession, top_k: int, session_names: list[str] | None = None ) -> list[models.Document]: """Query most recent documents.""" stmt = ( @@ -419,8 +433,8 @@ class RepresentationManager: models.Document.observed == self.observed, models.Document.deleted_at.is_(None), *( - [models.Document.session_name == session_name] - if session_name is not None + [models.Document.session_name.in_(session_names)] + if session_names is not None else [] ), ) @@ -433,7 +447,7 @@ class RepresentationManager: return list(documents) async def _query_documents_most_derived( - self, db: AsyncSession, top_k: int + self, db: AsyncSession, top_k: int, session_names: list[str] | None = None ) -> list[models.Document]: """Query most derived documents.""" stmt = ( @@ -444,6 +458,11 @@ class RepresentationManager: models.Document.observer == self.observer, models.Document.observed == self.observed, models.Document.deleted_at.is_(None), + *( + [models.Document.session_name.in_(session_names)] + if session_names is not None + else [] + ), ) .order_by( models.Document.times_derived.desc(), @@ -480,6 +499,7 @@ class RepresentationManager: count: int, max_distance: float | None = None, embedding: list[float] | None = None, + session_names: list[str] | None = None, ) -> list[models.Document]: """Query documents for a specific level.""" documents = await crud.query_documents( @@ -490,7 +510,7 @@ class RepresentationManager: query=query, max_distance=max_distance, top_k=count, - filters=self._build_filter_conditions(level), + filters=self._build_filter_conditions(level, session_names=session_names), embedding=embedding, ) @@ -503,17 +523,28 @@ class RepresentationManager: def _build_filter_conditions( self, level: str | None = None, + session_names: list[str] | None = None, ) -> dict[str, Any]: """ Build filter conditions for document queries. Returns a flat dict of key-value pairs for vector store filtering. + Callers must not pass an empty session_names list — empty allowlists + fail closed before any query is issued (see + _get_working_representation_internal). """ filters: dict[str, Any] = {} if level: filters["level"] = level + # `is not None` (not truthiness): an explicit empty allowlist must emit + # an empty `in` so downstream stores fail closed, matching + # _query_documents_recent / _query_documents_most_derived. Truthiness + # here would silently drop the filter and widen scope. + if session_names is not None: + filters["session_name"] = {"in": session_names} + return filters @@ -526,7 +557,7 @@ async def get_working_representation( db: AsyncSession | None = None, observer: str, observed: str, - session_name: str | None = None, + session_names: list[str] | None = None, include_semantic_query: str | None = None, embedding: list[float] | None = None, semantic_search_top_k: int | None = None, @@ -559,7 +590,7 @@ async def get_working_representation( ) return await manager.get_working_representation( db=db, - session_name=session_name, + session_names=session_names, include_semantic_query=include_semantic_query, embedding=embedding, semantic_search_top_k=semantic_search_top_k, diff --git a/src/routers/peers.py b/src/routers/peers.py index b61bf442..d6e35269 100644 --- a/src/routers/peers.py +++ b/src/routers/peers.py @@ -312,7 +312,9 @@ async def get_representation( workspace_id, observer=peer_id, observed=options.target if options.target is not None else peer_id, - session_name=options.session_id, + session_names=[options.session_id] + if options.session_id is not None + else None, include_semantic_query=options.search_query, embedding=embedding, semantic_search_top_k=options.search_top_k, @@ -475,7 +477,7 @@ async def get_peer_context( workspace_id, observer=peer_id, observed=observed, - session_name=None, # Peer context is global, not session-scoped + session_names=None, # Peer context is global, not session-scoped include_semantic_query=search_query, embedding=embedding, semantic_search_top_k=search_top_k, diff --git a/src/routers/sessions.py b/src/routers/sessions.py index 89ce665c..3cc34043 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -43,7 +43,7 @@ async def _get_working_representation_task( *, observer: str, observed: str, - session_name: str | None, + session_names: list[str] | None, search_top_k: int | None, search_max_distance: float | None, include_most_derived: bool, @@ -59,7 +59,7 @@ async def _get_working_representation_task( last_message: Optional last message for semantic query observer: Name of the observer peer observed: Name of the observed peer - session_name: Optional session to filter by + session_names: Optional session allowlist to filter by search_top_k: Number of semantic-search-retrieved observations to include in the representation search_max_distance: Maximum distance to search for semantically relevant observations include_most_derived: Whether to include the most derived observations in the representation @@ -74,7 +74,7 @@ async def _get_working_representation_task( db=db, observer=observer, observed=observed, - session_name=session_name, + session_names=session_names, include_semantic_query=last_message, semantic_search_top_k=search_top_k, semantic_search_max_distance=search_max_distance, @@ -765,7 +765,7 @@ async def get_session_context( search_query, observer=observer, observed=observed, - session_name=session_id if limit_to_session else None, + session_names=[session_id] if limit_to_session else None, search_top_k=search_top_k, search_max_distance=search_max_distance, include_most_derived=include_most_frequent, diff --git a/src/utils/filter.py b/src/utils/filter.py index 1394a8c6..36e34f4d 100644 --- a/src/utils/filter.py +++ b/src/utils/filter.py @@ -1,7 +1,8 @@ import datetime -from collections.abc import Callable +from collections.abc import Callable, Sequence from logging import getLogger from typing import Any, TypeVar +from typing import cast as typing_cast from sqlalchemy import ColumnElement, Select, and_, case, cast, literal, not_, or_ from sqlalchemy.types import Numeric @@ -28,6 +29,10 @@ COMPARISON_OPERATORS = { NUMERIC_OPERATORS = {"gte", "lte", "gt", "lt", "ne"} +# JSONB columns keep containment semantics: bare lists are not membership +# sugar, and dict values map to nested-metadata conditions rather than IN/Eq. +JSONB_COLUMNS = ("h_metadata", "configuration", "internal_metadata") + ALLOWED_EXTERNAL_TO_INTERNAL_COLUMN_MAPPING = { "id": "name", "created_at": "created_at", @@ -238,6 +243,12 @@ def _build_field_condition( if value == "*": return None + # Bare-list sugar on regular columns: {"session_id": ["a", "b"]} is + # shorthand for {"session_id": {"in": ["a", "b"]}}. JSONB columns are + # excluded — a bare list there keeps JSONB containment semantics. + if isinstance(value, list | tuple | set) and column_name not in JSONB_COLUMNS: + value = {"in": list(typing_cast(Sequence[Any], value))} + # Handle comparison operators vs regular values if isinstance(value, dict): # Check if this is a comparison operators dict by looking for known operators @@ -248,12 +259,12 @@ def _build_field_condition( else: # This is a regular value that happens to be a dict # For JSONB fields (metadata, configuration), check if it contains nested comparison operators - if column_name in ("h_metadata", "configuration", "internal_metadata"): + if column_name in JSONB_COLUMNS: return _build_nested_metadata_conditions(column, value) # pyright: ignore else: return column == value else: - if column_name in ("h_metadata", "configuration", "internal_metadata"): + if column_name in JSONB_COLUMNS: return column.contains(value) else: return column == value diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index 0b621971..1b4c4880 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -298,10 +298,15 @@ class LanceDBVectorStore(VectorStore): if not _VALID_IDENTIFIER_PATTERN.match(key): raise ValueError(f"Invalid filter key: {key!r}") - # Check if value is a dict with "in" operator - if isinstance(value, dict) and "in" in value: - # IN clause for list membership - in_values = cast(Sequence[Any], value["in"]) + # Membership: dict form {"in": [...]} or bare-list sugar + if (isinstance(value, dict) and "in" in value) or isinstance( + value, list | tuple | set + ): + in_values = ( + cast(Sequence[Any], value["in"]) + if isinstance(value, dict) + else list(cast(Sequence[Any], value)) + ) if in_values: escaped_values = [ f"'{str(v).replace(chr(39), chr(39) + chr(39))}'" @@ -310,6 +315,11 @@ class LanceDBVectorStore(VectorStore): for v in in_values ] conditions.append(f"{key} IN ({', '.join(escaped_values)})") + else: + # An empty membership list matches nothing. Emitting no + # condition would silently widen the result set + # (fail-open); force an always-false condition instead. + conditions.append("1 = 0") # Handle string values with proper quoting elif isinstance(value, str): # Escape single quotes in the value diff --git a/src/vector_store/turbopuffer.py b/src/vector_store/turbopuffer.py index ded0c691..e7825310 100644 --- a/src/vector_store/turbopuffer.py +++ b/src/vector_store/turbopuffer.py @@ -20,9 +20,7 @@ from . import VectorQueryResult, VectorRecord, VectorStore logger = logging.getLogger(__name__) -# Type aliases for Turbopuffer's filter formats -EqFilter = tuple[str, Literal["Eq"], Any] -InFilter = tuple[str, Literal["In"], Sequence[Any]] +# Type alias for Turbopuffer's AND filter format AndFilter = tuple[Literal["And"], Sequence[Filter]] DISTANCE_METRIC = "cosine_distance" @@ -245,13 +243,17 @@ class TurbopufferVectorStore(VectorStore): if not filters: return None - filter_list: list[EqFilter | InFilter] = [] + filter_list: list[Filter] = [] for key, value in filters.items(): # Check if value is a dict with "in" operator if isinstance(value, dict) and "in" in value: # Membership filter using "In" operator - in_values = cast(Sequence[Any], value["in"]) - filter_list.append((key, "In", in_values)) + in_values = list(cast(Sequence[Any], value["in"])) + filter_list.append(self._membership_filter(key, in_values)) + elif isinstance(value, list | tuple | set): + # Bare-list sugar: same membership semantics as {"in": [...]} + in_values = list(cast(Sequence[Any], value)) + filter_list.append(self._membership_filter(key, in_values)) else: # Simple equality filter using "Eq" operator filter_list.append((key, "Eq", cast(Any, value))) @@ -266,6 +268,20 @@ class TurbopufferVectorStore(VectorStore): and_filter: AndFilter = ("And", filter_list) return and_filter + @staticmethod + def _membership_filter(key: str, values: list[Any]) -> Filter: + """Build an "In" membership filter, failing closed on an empty list. + + Turbopuffer's empty-"In" semantics are undocumented, so an empty + allowlist emits an explicit contradiction (`Eq(x) AND NotEq(x)` is + false for every document) rather than risk a fail-open widening. + Mirrors lancedb's `1 = 0` guard. + """ + if not values: + never: AndFilter = ("And", [(key, "Eq", ""), (key, "NotEq", "")]) + return never + return (key, "In", values) + async def delete_many(self, namespace: str, ids: list[str]) -> None: """ Delete multiple vectors from Turbopuffer. diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index f4be22dd..b1c24c01 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -233,6 +233,224 @@ class TestRepresentationManagerSoftDelete: assert contents[1:] == ["tie 2", "tie 1", "tie 0"] +class TestRepresentationManagerSessionScoping: + """Tests that the session allowlist is applied uniformly to every query path. + + Regression for DEV-1994: session_name used to be applied only to the + recent-documents query; the semantic and most-derived paths ignored it, + so limit_to_session leaked cross-session conclusions. + """ + + async def _setup( + self, + db_session: AsyncSession, + test_workspace: models.Workspace, + test_peer: models.Peer, + ) -> tuple[models.Session, models.Session, RepresentationManager]: + """Create two sessions and documents in each, plus a session-less doc.""" + test_peer2 = models.Peer( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add(test_peer2) + await db_session.flush() + + session_a = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + session_b = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add_all([session_a, session_b]) + await db_session.flush() + + collection = models.Collection( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + db_session.add(collection) + await db_session.flush() + + db_session.add_all( + [ + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="in-scope observation", + session_name=session_a.name, + times_derived=1, + ), + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="out-of-scope observation", + session_name=session_b.name, + times_derived=100, + ), + # Dream-produced documents have no session_name; a session + # allowlist must exclude them (fail-closed). + models.Document( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + content="sessionless dream observation", + session_name=None, + times_derived=50, + ), + ] + ) + await db_session.flush() + + manager = RepresentationManager( + test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + return session_a, session_b, manager + + @pytest.mark.asyncio + async def test_recent_respects_session_allowlist( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + results = await manager._query_documents_recent( # pyright: ignore[reportPrivateUsage] + db_session, top_k=10, session_names=[session_a.name] + ) + + contents = [doc.content for doc in results] + assert contents == ["in-scope observation"] + + @pytest.mark.asyncio + async def test_most_derived_respects_session_allowlist( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """The out-of-scope doc has far higher times_derived; it must still be excluded.""" + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + results = await manager._query_documents_most_derived( # pyright: ignore[reportPrivateUsage] + db_session, top_k=10, session_names=[session_a.name] + ) + + contents = [doc.content for doc in results] + assert contents == ["in-scope observation"] + + @pytest.mark.asyncio + async def test_semantic_passes_session_allowlist_as_filters( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """The semantic path must push the allowlist down to query_documents.""" + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await manager._query_documents_semantic( # pyright: ignore[reportPrivateUsage] + db_session, + query="anything", + top_k=5, + embedding=[0.1], + session_names=[session_a.name], + ) + + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] == { + "session_name": {"in": [session_a.name]} + } + + @pytest.mark.asyncio + async def test_semantic_passes_no_filters_when_unscoped( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + test_workspace, test_peer = sample_data + _, _, manager = await self._setup(db_session, test_workspace, test_peer) + + with patch( + "src.crud.query_documents", new=AsyncMock(return_value=[]) + ) as mock_query: + await manager._query_documents_semantic( # pyright: ignore[reportPrivateUsage] + db_session, + query="anything", + top_k=5, + embedding=[0.1], + ) + + assert mock_query.await_args is not None + assert mock_query.await_args.kwargs["filters"] is None + + @pytest.mark.asyncio + async def test_working_representation_scoped_end_to_end( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """All blended paths active: only in-scope content may appear.""" + test_workspace, test_peer = sample_data + session_a, _, manager = await self._setup(db_session, test_workspace, test_peer) + + representation = await manager.get_working_representation( + db=db_session, + session_names=[session_a.name], + include_most_derived=True, + ) + + contents = [obs.content for obs in representation.explicit] + assert "in-scope observation" in contents + assert "out-of-scope observation" not in contents + assert "sessionless dream observation" not in contents + + @pytest.mark.asyncio + async def test_empty_allowlist_fails_closed( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """An empty allowlist must return an empty representation, not fall + back to unscoped behavior (downstream stores drop empty IN clauses).""" + test_workspace, test_peer = sample_data + _, _, manager = await self._setup(db_session, test_workspace, test_peer) + + representation = await manager.get_working_representation( + db=db_session, + session_names=[], + include_most_derived=True, + ) + + assert representation.explicit == [] + assert representation.deductive == [] + + def test_build_filter_conditions_empty_allowlist_fails_closed(self): + """The filter-builder layer itself must fail closed, independent of the + early-return guard in _get_working_representation_internal. An empty + allowlist emits an empty `in` (renders as always-false downstream), not + an omitted filter.""" + manager = RepresentationManager( + "workspace", observer="observer", observed="observed" + ) + + assert manager._build_filter_conditions(session_names=[]) == { # pyright: ignore[reportPrivateUsage] + "session_name": {"in": []} + } + # None means unscoped — no session filter emitted. + assert manager._build_filter_conditions(session_names=None) == {} # pyright: ignore[reportPrivateUsage] + assert manager._build_filter_conditions(session_names=["s1"]) == { # pyright: ignore[reportPrivateUsage] + "session_name": {"in": ["s1"]} + } + + class TestRepresentationManagerSave: @pytest.mark.asyncio async def test_save_representation_filters_blank_observations_before_embedding( diff --git a/tests/test_advanced_filters.py b/tests/test_advanced_filters.py index 768f29bc..3b187470 100644 --- a/tests/test_advanced_filters.py +++ b/tests/test_advanced_filters.py @@ -235,6 +235,86 @@ async def test_comparison_operators_filters( ), f"Unexpected message '{message_config['content']}' found in results for {description}" +@pytest.mark.asyncio +async def test_bare_list_membership_sugar( + client: TestClient, + sample_data: tuple[Workspace, Peer], +): + """A bare list on a regular column is shorthand for {"in": [...]}. + + JSONB metadata columns are excluded from the sugar: a bare list there + keeps JSONB containment semantics. + """ + test_workspace, test_peer = sample_data + + # Second peer so peer_id membership has something to exclude + peer2_name = str(generate_nanoid()) + client.post( + f"/v3/workspaces/{test_workspace.name}/peers", + json={"id": peer2_name}, + ) + + session_id = str(generate_nanoid()) + session_response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions", + json={ + "id": session_id, + "peer_names": {test_peer.name: {}, peer2_name: {}}, + }, + ) + assert session_response.status_code == 201 + + message_configs = [ + { + "content": "From peer one", + "peer_id": test_peer.name, + "metadata": {"tags": ["important", "urgent"]}, + }, + { + "content": "From peer two", + "peer_id": peer2_name, + "metadata": {"tags": ["normal"]}, + }, + ] + messages_response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{session_id}/messages", + json={"messages": message_configs}, + ) + assert messages_response.status_code == 201 + + def list_contents(filter_config: dict[str, Any]) -> list[str]: + response = client.post( + f"/v3/workspaces/{test_workspace.name}/sessions/{session_id}/messages/list", + json={"filters": filter_config}, + ) + assert response.status_code == 200 + return [item["content"] for item in response.json()["items"]] + + # Bare list == membership on a regular column + assert list_contents({"peer_id": [test_peer.name]}) == ["From peer one"] + + # Multiple values + assert sorted(list_contents({"peer_id": [test_peer.name, peer2_name]})) == [ + "From peer one", + "From peer two", + ] + + # Equivalent to the explicit {"in": [...]} form + assert list_contents({"peer_id": [test_peer.name]}) == list_contents( + {"peer_id": {"in": [test_peer.name]}} + ) + + # Empty list matches nothing (fail-closed), never everything + assert list_contents({"peer_id": []}) == [] + + # JSONB metadata keeps containment semantics for bare lists: + # matches arrays containing ALL listed elements, not membership. + assert list_contents({"metadata": {"tags": ["important", "urgent"]}}) == [ + "From peer one" + ] + assert list_contents({"metadata": {"tags": ["important", "missing"]}}) == [] + + @pytest.mark.asyncio async def test_wildcard_filters( client: TestClient, sample_data: tuple[Workspace, Peer] diff --git a/tests/vector_store/test_lancedb.py b/tests/vector_store/test_lancedb.py index 0c502e4c..9b52708a 100644 --- a/tests/vector_store/test_lancedb.py +++ b/tests/vector_store/test_lancedb.py @@ -37,6 +37,27 @@ def store() -> LanceDBVectorStore: return LanceDBVectorStore() +def test_build_where_clause_membership(store: LanceDBVectorStore) -> None: + """Both the dict `in` form and the bare-list sugar produce an IN clause.""" + assert ( + store._build_where_clause({"session_name": {"in": ["s1", "s2"]}}) # pyright: ignore[reportPrivateUsage] + == "session_name IN ('s1', 's2')" + ) + assert ( + store._build_where_clause({"session_name": ["s1", "s2"]}) # pyright: ignore[reportPrivateUsage] + == "session_name IN ('s1', 's2')" + ) + + +def test_build_where_clause_empty_membership_fails_closed( + store: LanceDBVectorStore, +) -> None: + """An empty membership list must emit an always-false predicate, never an + omitted condition that would widen scope (fail-open).""" + assert store._build_where_clause({"session_name": {"in": []}}) == "1 = 0" # pyright: ignore[reportPrivateUsage] + assert store._build_where_clause({"session_name": []}) == "1 = 0" # pyright: ignore[reportPrivateUsage] + + @pytest.mark.asyncio async def test_query_returns_empty_when_table_missing( store: LanceDBVectorStore, diff --git a/tests/vector_store/test_turbopuffer.py b/tests/vector_store/test_turbopuffer.py index ca8c4ff3..cdae44d5 100644 --- a/tests/vector_store/test_turbopuffer.py +++ b/tests/vector_store/test_turbopuffer.py @@ -53,6 +53,30 @@ async def test_upsert_many_raises_vector_store_error_on_5xx( namespace_mock.write.assert_awaited_once() +def test_build_filters_membership(store: TurbopufferVectorStore) -> None: + """Both the dict `in` form and the bare-list sugar produce an In filter.""" + assert store._build_filters({"session_name": {"in": ["s1", "s2"]}}) == ( # pyright: ignore[reportPrivateUsage] + "session_name", + "In", + ["s1", "s2"], + ) + assert store._build_filters({"session_name": ["s1", "s2"]}) == ( # pyright: ignore[reportPrivateUsage] + "session_name", + "In", + ["s1", "s2"], + ) + + +def test_build_filters_empty_membership_fails_closed( + store: TurbopufferVectorStore, +) -> None: + """An empty membership list must produce an always-false filter, never an + omitted/empty In that could widen scope (fail-open).""" + never = ("And", [("session_name", "Eq", ""), ("session_name", "NotEq", "")]) + assert store._build_filters({"session_name": {"in": []}}) == never # pyright: ignore[reportPrivateUsage] + assert store._build_filters({"session_name": []}) == never # pyright: ignore[reportPrivateUsage] + + @pytest.mark.asyncio async def test_upsert_many_short_circuits_on_empty( store: TurbopufferVectorStore,