fix migration task scheduling and stabilize regression tests

Only create startup migration background tasks when migration is actually needed, and align tests with the new contract and stricter access filtering behavior. This also restores sparse fallback table resolution compatibility and hardens warning-log assertions to avoid flaky failures.

Author: sqhyz55

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

@@ -10,7 +10,6 @@
     SearchWarning,
     SparseSearchResponse,
 )
-from ..LanceDB.model_tag_utils import to_model_tag
 from ..LanceDB.schema_manager import _safe_close_table
 from ..storage.contracts import FilterCondition, FilterExpression, FilterOperator
 from ..storage.factory import (
@@ -223,9 +222,10 @@ def _substring_fallback(
     if filters:
         query_filters.update(filters)
 
+    _table = None
     try:
-        # Avoid opening an extra table handle here; iter_batches will open internally.
-        table_name = f"embeddings_{to_model_tag(model_tag)}"
+        # Resolve embeddings table name with legacy fallback support.
+        _table, table_name = vector_store.open_embeddings_table(model_tag)
 
         for batch in vector_store.iter_batches(
             table_name=table_name,
@@ -289,6 +289,8 @@ def _substring_fallback(
 
     except Exception as exc:
         logger.error("Substring fallback failed: %s", exc)
+    finally:
+        _safe_close_table(_table)
 
     return results
 

src/xagent/web/services/rag_storage_migration_service.py

@@ -16,10 +16,50 @@
 class RAGStorageMigrationService:
     """Run storage compatibility checks and background migrations."""
 
-    async def start_background_migrations(self) -> asyncio.Task[None]:
-        """Create a non-blocking background task for storage migrations."""
+    async def start_background_migrations(self) -> Optional[asyncio.Task[None]]:
+        """Create a non-blocking background task only when migration is needed."""
+        if not self._should_start_background_task():
+            logger.info("Skipping background LanceDB migration task (not needed)")
+            return None
         return asyncio.create_task(self._run_migrations())
 
+    def _should_start_background_task(self) -> bool:
+        """Return whether startup should create a migration background task."""
+        from ...core.tools.core.RAG_tools.LanceDB.schema_manager import (
+            check_table_needs_migration,
+        )
+        from ...core.tools.core.RAG_tools.utils.lancedb_query_utils import (
+            list_embeddings_table_names,
+        )
+        from ...providers.vector_store.lancedb import get_connection_from_env
+
+        auto_migrate = os.getenv("LANCEDB_AUTO_MIGRATE", "true").lower() == "true"
+        if not auto_migrate:
+            return False
+
+        conn = get_connection_from_env()
+        tables_to_check = [
+            "chunks",
+            "documents",
+            "parses",
+            "ingestion_runs",
+            "prompt_templates",
+        ]
+
+        if any(check_table_needs_migration(conn, table) for table in tables_to_check):
+            return True
+
+        try:
+            return any(
+                check_table_needs_migration(conn, table_name)
+                for table_name in list_embeddings_table_names(conn)
+            )
+        except Exception as e:  # noqa: BLE001
+            logger.warning(
+                "Could not inspect embeddings tables for migration need: %s", e
+            )
+            return False
+
     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/generate/test_format_generation_prompt.py

@@ -8,6 +8,7 @@
 )
 
 logger = logging.getLogger(__name__)
+PROMPT_LOGGER = "xagent.core.tools.core.RAG_tools.generate.format_generation_prompt"
 
 
 @pytest.fixture
@@ -78,7 +79,7 @@ def test_format_generation_prompt_empty_contexts_produces_warning_and_formats(
         caplog: pytest.LogCaptureFixture,
     ) -> None:
         """Test that empty formatted contexts produce a warning but still format."""
-        with caplog.at_level(logging.WARNING):
+        with caplog.at_level(logging.WARNING, logger=PROMPT_LOGGER):
             result = format_generation_prompt(
                 prompt_template=sample_prompt_template_plain,
                 formatted_contexts="",

tests/core/tools/core/RAG_tools/pipelines/test_document_search.py

@@ -190,6 +190,8 @@ def test_document_search_end_to_end(
     search_result = document_search.search_documents(
         collection=collection,
         query_text=query_text,
+        user_id=1,
+        is_admin=True,
         config=SearchConfig(
             search_type=SearchType.SPARSE,
             top_k=3,

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

@@ -7,7 +7,9 @@
 
 import pytest
 
-from xagent.core.tools.core.RAG_tools.core.exceptions import DatabaseOperationError, VersionManagementError
+from xagent.core.tools.core.RAG_tools.core.exceptions import (
+    VersionManagementError,
+)
 from xagent.core.tools.core.RAG_tools.core.schemas import StepType
 from xagent.core.tools.core.RAG_tools.version_management.list_candidates import (
     list_candidates,

tests/web/test_rag_storage_migration_service.py

@@ -20,8 +20,10 @@ async def _fake_run() -> None:
         called["run"] += 1
 
     monkeypatch.setattr(service, "_run_migrations", _fake_run)
+    monkeypatch.setattr(service, "_should_start_background_task", lambda: True)
 
     task = await service.start_background_migrations()
+    assert task is not None
     await task
 
     assert isinstance(task, asyncio.Task)