fix(storage): address PR #359 review on tenancy, scans, and migrations

Restore per-embeddings-table cascade predicates, pin route collection in
sparse fallback scans, forward user scope through async FTS, batch unbounded
file-status document lookups, and always schedule migration compatibility
checks while keeping backfill gated by LANCEDB_AUTO_MIGRATE.

Author: sqhyz55

src/xagent/core/tools/core/RAG_tools/retrieval/search_sparse.py

@@ -201,6 +201,18 @@ def search_sparse(
         )
 
 
+def _build_substring_scan_filters(
+    *,
+    collection: str,
+    filters: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+    """Build iter_batches filters with a non-overridable route collection."""
+    scan_filters: Dict[str, Any] = dict(filters) if filters else {}
+    scan_filters.pop("collection", None)
+    scan_filters["collection"] = collection
+    return scan_filters
+
+
 def _substring_fallback(
     *,
     model_tag: str,
@@ -218,9 +230,9 @@ def _substring_fallback(
     vector_store = get_vector_index_store()
     results: List[SearchResult] = []
 
-    query_filters: Dict[str, Any] = {"collection": collection}
-    if filters:
-        query_filters.update(filters)
+    query_filters = _build_substring_scan_filters(
+        collection=collection, filters=filters
+    )
 
     _table = None
     try:
@@ -413,6 +425,8 @@ async def search_sparse_async(
             top_k=top_k,
             filters=filter_expr,
             text_column_name="text",
+            user_id=user_id,
+            is_admin=is_admin,
         )
 
         if not raw_results:
@@ -512,10 +526,9 @@ async def _substring_fallback_async(
     vector_store = get_vector_index_store()
     results: List[SearchResult] = []
 
-    # Build query filters
-    query_filters: Dict[str, Any] = {"collection": collection}
-    if filters:
-        query_filters.update(filters)
+    query_filters = _build_substring_scan_filters(
+        collection=collection, filters=filters
+    )
 
     _table = None
     try:

src/xagent/core/tools/core/RAG_tools/storage/contracts.py

@@ -1043,6 +1043,8 @@ async def search_fts_async(
         top_k: int,
         filters: Optional[FilterExpression] = None,
         text_column_name: str = "text",
+        user_id: Optional[int] = None,
+        is_admin: bool = False,
     ) -> List[Dict[str, Any]]:
         """Execute full-text search (async).
 
@@ -1052,6 +1054,8 @@ async def search_fts_async(
             top_k: Number of top results to return.
             filters: Optional abstract filter expression.
             text_column_name: Name of text column with FTS index (default "text").
+            user_id: Optional user ID for multi-tenancy filtering.
+            is_admin: Whether the user has admin privileges.
 
         Returns:
             List of search result dictionaries with keys:
@@ -1073,6 +1077,8 @@ async def search_fts_by_model_async(
         top_k: int,
         filters: Optional[FilterExpression] = None,
         text_column_name: str = "text",
+        user_id: Optional[int] = None,
+        is_admin: bool = False,
     ) -> List[Dict[str, Any]]:
         """Convenience method: search FTS by model_tag with automatic table resolution.
 
@@ -1085,6 +1091,8 @@ async def search_fts_by_model_async(
             top_k: Number of top results to return.
             filters: Optional abstract filter expression.
             text_column_name: Name of text column with FTS index (default "text").
+            user_id: Optional user ID for multi-tenancy filtering.
+            is_admin: Whether the user has admin privileges.
 
         Returns:
             List of search result dictionaries with keys:
@@ -1105,6 +1113,8 @@ async def search_fts_by_model_async(
                 top_k=top_k,
                 filters=filters,
                 text_column_name=text_column_name,
+                user_id=user_id,
+                is_admin=is_admin,
             )
         finally:
             _release_embeddings_table_probe(_table)

src/xagent/core/tools/core/RAG_tools/storage/lancedb_stores.py

@@ -1867,6 +1867,8 @@ async def search_fts_async(
         top_k: int,
         filters: Optional[FilterExpression] = None,
         text_column_name: str = "text",
+        user_id: Optional[int] = None,
+        is_admin: bool = False,
     ) -> List[Dict[str, Any]]:
         """Execute full-text search using async LanceDB FTS API.
 
@@ -1881,7 +1883,7 @@ async def search_fts_async(
 
             # Build filter expression
             backend_filter = self.build_filter_expression(
-                filters, user_id=None, is_admin=False
+                filters, user_id=user_id, is_admin=is_admin
             )
 
             # Build FTS search query

src/xagent/core/tools/core/RAG_tools/version_management/cascade_cleaner.py

@@ -269,30 +269,25 @@ def cascade_delete(
         if t.startswith("embeddings_")
         and (model_tag is None or t == f"embeddings_{model_tag}")
     ]
-    if target_embeddings_tables:
-        sample_embeddings_table = target_embeddings_tables[0]
+    for t in target_embeddings_tables:
         if target == "collection":
-            embed_filter = _build_collection_filter(
+            predicates[t] = _build_collection_filter(
                 conn=conn,
-                table_name=sample_embeddings_table,
+                table_name=t,
                 collection=collection,
                 user_id=user_id,
                 is_admin=is_admin,
             )
         else:
-            embed_filter = _build_document_filter(
+            predicates[t] = _build_document_filter(
                 conn=conn,
-                table_name=sample_embeddings_table,
+                table_name=t,
                 collection=collection,
                 doc_id=str(doc_id),
                 user_id=user_id,
                 is_admin=is_admin,
             )
 
-        # Reuse one computed embeddings filter across all target embeddings_* tables.
-        for t in target_embeddings_tables:
-            predicates[t] = embed_filter
-
     if preview_only and not confirm:
         return _plan_by_predicates(conn, predicates, model_tag=model_tag)
 

src/xagent/web/services/kb_file_service.py

@@ -8,7 +8,7 @@
 from collections import OrderedDict
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 from sqlalchemy.orm import Session
 
@@ -81,6 +81,33 @@ def clear(self) -> None:
 _file_status_cache = _FileStatusCache(ttl_seconds=5)
 
 
+def _list_document_records_for_file_ids(
+    store: VectorIndexStore,
+    *,
+    file_ids: Sequence[str],
+    user_id: int,
+    is_admin: bool,
+) -> List[DocumentRecord]:
+    """Load document rows for many file IDs without scan-limit truncation."""
+    unique_ids = sorted({file_id for file_id in file_ids if file_id})
+    if not unique_ids:
+        return []
+
+    records: List[DocumentRecord] = []
+    for offset in range(0, len(unique_ids), _FILE_STATUS_BATCH_SIZE):
+        batch = unique_ids[offset : offset + _FILE_STATUS_BATCH_SIZE]
+        records.extend(
+            store.list_document_records(
+                collection_name=None,
+                user_id=user_id,
+                is_admin=is_admin,
+                file_ids=batch,
+                max_results=-1,
+            )
+        )
+    return records
+
+
 def upsert_uploaded_file_record(
     db: Session,
     *,
@@ -376,11 +403,11 @@ def aggregate_uploaded_file_statuses(
 
     # Cache miss - compute from database via abstraction layer
     store = get_vector_index_store()
-    records = store.list_document_records(
-        collection_name=None,
+    records = _list_document_records_for_file_ids(
+        store,
+        file_ids=normalized_file_ids,
         user_id=user_id,
         is_admin=is_admin,
-        file_ids=normalized_file_ids,
     )
 
     doc_refs_by_file_id: Dict[str, List[tuple[str, str]]] = {

src/xagent/web/services/rag_storage_migration_service.py

@@ -16,17 +16,10 @@
 class RAGStorageMigrationService:
     """Run storage compatibility checks and background migrations."""
 
-    async def start_background_migrations(self) -> Optional[asyncio.Task[None]]:
-        """Create a non-blocking background task when auto-migrate is enabled."""
-        if not self._should_start_background_task():
-            logger.info("Skipping background LanceDB migration task (not needed)")
-            return None
+    async def start_background_migrations(self) -> asyncio.Task[None]:
+        """Schedule compatibility checks; backfill execution is env-gated."""
         return asyncio.create_task(self._run_migrations())
 
-    def _should_start_background_task(self) -> bool:
-        """Return whether startup should create a migration background task."""
-        return os.getenv("LANCEDB_AUTO_MIGRATE", "true").lower() == "true"
-
     async def _run_migrations(self) -> None:
         """Run migration checks and backfills in background."""
         auto_migrate = os.getenv("LANCEDB_AUTO_MIGRATE", "true").lower() == "true"

tests/core/tools/core/RAG_tools/retrieval/test_search_sparse.py

@@ -401,6 +401,47 @@ def test_search_sparse_fts_fallback_warning_content(self) -> None:
         assert "Check FTS tokenizer configuration" in warning.message
         assert "update LanceDB to ensure proper tokenisation" in warning.message
 
+    def test_substring_fallback_pins_route_collection_over_caller_override(
+        self,
+    ) -> None:
+        """Route collection must not be replaced by caller filters with the same key."""
+        from xagent.core.tools.core.RAG_tools.retrieval.search_sparse import (
+            _substring_fallback,
+        )
+
+        mock_batch = Mock()
+        mock_batch.to_pandas.return_value = pd.DataFrame(
+            {
+                "doc_id": ["doc1"],
+                "chunk_id": ["chunk1"],
+                "text": ["needle"],
+                "parse_hash": ["hash1"],
+                "created_at": [pd.Timestamp.now()],
+                "metadata": [None],
+            }
+        )
+
+        mock_vector_store = Mock()
+        mock_vector_store.open_embeddings_table.return_value = (Mock(), "embeddings_x")
+        mock_vector_store.iter_batches.return_value = [mock_batch]
+
+        with patch(
+            "xagent.core.tools.core.RAG_tools.retrieval.search_sparse.get_vector_index_store",
+            return_value=mock_vector_store,
+        ):
+            _substring_fallback(
+                model_tag="test_model",
+                collection="route_collection",
+                query_text="needle",
+                top_k=5,
+                filters={"collection": "evil_override", "doc_id": "d1"},
+                current_warnings=[],
+            )
+
+        iter_kwargs = mock_vector_store.iter_batches.call_args.kwargs
+        assert iter_kwargs["filters"]["collection"] == "route_collection"
+        assert iter_kwargs["filters"]["doc_id"] == "d1"
+
 
 @pytest.mark.asyncio
 class TestSearchSparseAsync:
@@ -450,6 +491,8 @@ async def test_search_sparse_async_success_forwards_user_scope(self) -> None:
                 "filters"
             ],
             text_column_name="text",
+            user_id=9,
+            is_admin=False,
         )
 
     async def test_search_sparse_async_triggers_fallback(self) -> None:

tests/core/tools/core/RAG_tools/version_management/test_cascade_cleaner.py

@@ -1177,3 +1177,49 @@ def _delete(_: str) -> None:
         "ingestion_runs",
         "documents",
     ]
+
+
+@patch(
+    "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner._plan_by_predicates"
+)
+@patch(
+    "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner._build_collection_filter"
+)
+@patch(
+    "xagent.core.tools.core.RAG_tools.version_management.cascade_cleaner.get_vector_store_raw_connection"
+)
+def test_cascade_delete_builds_distinct_embeddings_predicates_per_table(
+    mock_get_conn: MagicMock,
+    mock_build_collection_filter: MagicMock,
+    mock_plan: MagicMock,
+) -> None:
+    """Each embeddings_* table must get its own predicate (schema may differ during migration)."""
+    mock_build_collection_filter.side_effect = (
+        lambda *, table_name, **_: f"expr:{table_name}"
+    )
+    mock_plan.return_value = {}
+
+    conn = MagicMock(spec=["table_names"])
+    conn.table_names.return_value = ["embeddings_m1", "embeddings_m2"]
+    mock_get_conn.return_value = conn
+
+    cascade_delete(
+        target="collection",
+        collection="kb",
+        user_id=1,
+        is_admin=False,
+        preview_only=True,
+        confirm=False,
+    )
+
+    predicates = mock_plan.call_args[0][1]
+    assert predicates["embeddings_m1"] == "expr:embeddings_m1"
+    assert predicates["embeddings_m2"] == "expr:embeddings_m2"
+    assert predicates["embeddings_m1"] != predicates["embeddings_m2"]
+
+    embeddings_calls = [
+        call.kwargs["table_name"]
+        for call in mock_build_collection_filter.call_args_list
+        if str(call.kwargs["table_name"]).startswith("embeddings_")
+    ]
+    assert sorted(embeddings_calls) == ["embeddings_m1", "embeddings_m2"]

tests/web/test_kb_file_service.py

@@ -42,3 +42,38 @@ def test_list_documents_for_user_delegates_to_vector_store_without_raw_lancedb()
             "source_path": "/tmp/demo.pdf",
         }
     ]
+
+
+def test_aggregate_uploaded_file_statuses_chunks_file_ids_with_unbounded_limit() -> (
+    None
+):
+    """Large file-id lists must be queried in batches with max_results=-1."""
+    file_ids = [f"file-{index}" for index in range(250)]
+    list_calls: list[dict[str, object]] = []
+
+    def _fake_list_document_records(**kwargs: object) -> list[object]:
+        list_calls.append(dict(kwargs))
+        return []
+
+    fake_store = SimpleNamespace(list_document_records=_fake_list_document_records)
+
+    with (
+        patch.object(
+            kb_file_service, "get_vector_index_store", return_value=fake_store
+        ),
+        patch.object(kb_file_service, "load_ingestion_status", return_value=[]),
+        patch.object(kb_file_service, "_load_indexed_doc_refs", return_value=set()),
+    ):
+        status_map = kb_file_service.aggregate_uploaded_file_statuses(
+            file_ids=file_ids,
+            user_id=7,
+            is_admin=False,
+            use_cache=False,
+        )
+
+    assert len(list_calls) == 2
+    assert len(list_calls[0]["file_ids"]) == 200
+    assert len(list_calls[1]["file_ids"]) == 50
+    assert all(call["max_results"] == -1 for call in list_calls)
+    assert all(call["user_id"] == 7 for call in list_calls)
+    assert status_map == {file_id: "UNKNOWN" for file_id in file_ids}