fix(core): own startup index tasks and remove flaky test causes

src/basic_memory/mcp/container.py

@@ -10,6 +10,8 @@
 - File indexing decisions are centralized here
 """
 
+from collections.abc import Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -89,22 +91,48 @@ def create_watch_coordinator(self) -> "WatchCoordinator":
         )
 
 
-# Module-level container instance (set by lifespan)
+# Container set for the whole process by an embedding host (set_container).
 _container: McpContainer | None = None
+# Containers of the lifespans running now, oldest first.
+_lifespan_containers: list[McpContainer] = []
 
 
 def get_container() -> McpContainer:
-    """Get the current MCP container.
+    """Get the current MCP container: the newest running lifespan's, else the host's.
 
     Raises:
         RuntimeError: If container hasn't been initialized
     """
+    if _lifespan_containers:
+        return _lifespan_containers[-1]
     if _container is None:
         raise RuntimeError("MCP container not initialized. Call set_container() first.")
     return _container
 
 
 def set_container(container: McpContainer) -> None:
-    """Set the MCP container (called by lifespan)."""
+    """Set the process-wide MCP container for a host that runs no lifespan."""
     global _container
     _container = container
+
+
+@contextmanager
+def installed_container(container: McpContainer) -> Iterator[None]:
+    """Expose a container for one server lifespan and withdraw it when that lifespan ends.
+
+    Trigger: an in-process MCP session (``Client(mcp)``) enters and exits the lifespan.
+    Why: a lifespan that set a global and never reset it left its config installed
+      after the session ended, and later code in the same process, such as the
+      default search mode, read a config that no longer applied (#1955). Saving and
+      restoring a single slot is not enough either: two lifespans that overlap and
+      exit out of order would restore an exited container.
+    Outcome: each lifespan adds its container on entry and removes that same
+      container, by identity, on exit; the newest still-running one is current.
+    """
+    _lifespan_containers.append(container)
+    try:
+        yield
+    finally:
+        newest_first = reversed(range(len(_lifespan_containers)))
+        own_entry = next(i for i in newest_first if _lifespan_containers[i] is container)
+        del _lifespan_containers[own_entry]

src/basic_memory/mcp/server.py

@@ -18,7 +18,7 @@
 from basic_memory.db import scoped_session
 from basic_memory.index.local_schedulers import drain_background_tasks
 from basic_memory.mcp.client_info import MCPClientInfoMiddleware
-from basic_memory.mcp.container import McpContainer, set_container
+from basic_memory.mcp.container import McpContainer, installed_container as installed_mcp_container
 from basic_memory.mcp.empty_results import EmptyListResultMiddleware
 from basic_memory.read_cache import ReadCacheUnavailable
 from basic_memory.read_cache.redis import RedisReadCache
@@ -116,12 +116,11 @@ async def lifespan(app: FastMCP):
         max_connections=config.redis_max_connections,
     ) as read_cache:
         container.read_cache = read_cache
-        set_container(container)
         api_container = ApiContainer(config=config, mode=container.mode, read_cache=read_cache)
 
         # Local MCP requests use FastAPI through ASGITransport without entering its lifespan.
         # Installing the same container keeps request reads and watcher invalidation on one cache.
-        with installed_container(api_container):
+        with installed_mcp_container(container), installed_container(api_container):
             with logfire.span(
                 "mcp.lifecycle.startup",
                 entrypoint="mcp",

src/basic_memory/services/initialization.py

@@ -212,11 +212,6 @@ async def reconcile_projects_with_config_once(app_config: BasicMemoryConfig) ->
         _reconciled_database_paths.add(database_path)
 
 
-# Strong references for fire-and-forget startup index tasks; the event loop
-# alone would hold only weak references (asyncio.create_task docs).
-_initial_index_tasks: set[asyncio.Task[None]] = set()
-
-
 async def initialize_file_indexing(
     app_config: BasicMemoryConfig,
     quiet: bool = True,
@@ -316,7 +311,6 @@ async def initialize_file_indexing(
     if recovery_complete is not None:
         recovery_complete.set()
 
-    # Start indexing for all projects as background tasks (non-blocking)
     async def index_project_background(project: Project):
         """Index a single project in the background."""
         logger.info(f"Starting background project index for project: {project.name}")
@@ -326,29 +320,32 @@ async def index_project_background(project: Project):
                 runtime_factory=project_index_runtime_factory,
             )
             logger.info(f"Background project index completed for project: {project.name}")
-        except Exception as e:  # pragma: no cover
+        except Exception as e:
+            # Trigger: one project's startup index failed (unreadable root, bad file).
+            # Why: an exception escaping this task would make the task group below
+            #   cancel the watcher, stopping file watching for every project.
+            # Outcome: log it and keep watching; the next startup or reindex retries it.
+            #   CancelledError is not an Exception, so shutdown still propagates.
             logger.error(f"Error in background project index for project {project.name}: {e}")
 
-    # Create background tasks for all project indexes (non-blocking). The event
-    # loop keeps only weak task references, so hold them in the module-level set
-    # to keep GC from cancelling an index mid-flight.
-    for project in active_projects:
-        index_task = asyncio.create_task(index_project_background(project))
-        _initial_index_tasks.add(index_task)
-        index_task.add_done_callback(_initial_index_tasks.discard)
-    logger.info(f"Created {len(active_projects)} background indexing tasks")
-
-    # Don't await the tasks - let them run in background while we continue
-
-    # Then start the watch service in the background
-    if constrained_project:
-        logger.info(f"Starting watch service constrained to project: {constrained_project}")
-    else:
-        logger.info("Starting watch service for all projects")
-
-    # run the watch service
-    await watch_service.run()
-    logger.info("Watch service started")
+    # Trigger: startup indexes every active project while the watcher starts.
+    # Why: the indexes used to be fire-and-forget tasks in a module-level set, so
+    #   cancelling the watcher at shutdown left them running against a database
+    #   the lifespan was about to close, and the set leaked between event loops.
+    # Outcome: this function owns them. The watcher and the indexes run in one
+    #   task group, so cancelling the watcher cancels and awaits the indexes too.
+    async with asyncio.TaskGroup() as startup_tasks:
+        for project in active_projects:
+            startup_tasks.create_task(index_project_background(project))
+        logger.info(f"Created {len(active_projects)} background indexing tasks")
+
+        if constrained_project:
+            logger.info(f"Starting watch service constrained to project: {constrained_project}")
+        else:
+            logger.info("Starting watch service for all projects")
+
+        await watch_service.run()
+        logger.info("Watch service started")
 
     return None
 

test-int/conftest.py

@@ -86,6 +86,7 @@ async def test_my_mcp_tool(mcp_server, app):
 from basic_memory.models.base import Base
 from basic_memory.repository.postgres_search_repository import PostgresSearchRepository
 from basic_memory.repository.project_repository import ProjectRepository
+from basic_memory.repository.search_repository import create_search_repository
 from fastapi import FastAPI
 
 from basic_memory.deps import get_engine_factory, get_app_config
@@ -318,8 +319,6 @@ async def engine_factory(
     None,
 ]:
     """Create engine and session factory for the configured database backend."""
-    from basic_memory.models.search import CREATE_SEARCH_INDEX
-    from basic_memory import db
 
     if db_backend == "postgres":
         assert postgres_engine is not None
@@ -371,12 +370,23 @@ async def engine_factory(
             async with engine.begin() as conn:
                 await conn.run_sync(Base.metadata.create_all)
 
-            # Drop any SearchIndex ORM table, then create FTS5 virtual table
+            # create_all made search_index a plain table from the ORM model; drop it so
+            # the search repository creates the FTS5 virtual table instead.
             async with db.scoped_session(session_maker) as session:
                 await session.execute(text("DROP TABLE IF EXISTS search_index"))
-                await session.execute(CREATE_SEARCH_INDEX)
                 await session.commit()
 
+            # Trigger: production initializes search storage (the FTS5 table and the
+            # vector storage) once, when the database is opened, before any request.
+            # Why: left to the first embedding, the vector-table DDL runs while other
+            #   connections search; each schema change makes those connections reload
+            #   the schema and reconnect search_index, and a reconnect that overlaps
+            #   the DDL fails with "vtable constructor failed: search_index" (#1955).
+            # Outcome: tests open the database the way production does.
+            await create_search_repository(
+                session_maker, project_id=1, app_config=app_config
+            ).init_search_index()
+
             yield engine, session_maker
 
 

test-int/mcp/test_new_project_watch_lifecycle.py

@@ -11,7 +11,6 @@
 from basic_memory.config import BasicMemoryConfig, ConfigManager
 from basic_memory.index.local_schedulers import drain_background_tasks
 from basic_memory.mcp.server import mcp
-from basic_memory.services.initialization import _initial_index_tasks
 
 
 pytestmark = [pytest.mark.asyncio, pytest.mark.real_embedder]
@@ -103,7 +102,3 @@ async def test_new_project_copied_files_and_live_changes_are_searchable(
                 noise.cancel()
                 with suppress(asyncio.CancelledError):
                     await noise
-                # Startup scans have explicit task ownership outside the watch task.
-                # Let them settle before fixture teardown closes the database.
-                await asyncio.gather(*_initial_index_tasks)
-                await drain_background_tasks()

test-int/test_sqlite_lifecycle_integration.py

@@ -48,27 +48,39 @@ def count_connection(connection: object, record: object) -> None:
             await connection.execute(text("CREATE TABLE notes (id INTEGER PRIMARY KEY, body TEXT)"))
             await connection.execute(text("CREATE VIRTUAL TABLE notes_fts USING fts5(body)"))
 
-        async def write_notes(writer: int) -> None:
-            for index in range(40):
-                note_id = writer * 40 + index
-                async with engine.begin() as connection:
-                    await connection.execute(
-                        text("INSERT INTO notes VALUES (:id, :body)"),
-                        {"id": note_id, "body": "searchable note"},
-                    )
+        async def write_note(note_id: int) -> None:
+            async with engine.begin() as connection:
+                await connection.execute(
+                    text("INSERT INTO notes VALUES (:id, :body)"),
+                    {"id": note_id, "body": "searchable note"},
+                )
+                await connection.execute(
+                    text("INSERT INTO notes_fts(rowid, body) VALUES (:id, :body)"),
+                    {"id": note_id, "body": "searchable note"},
+                )
+                assert (
                     await connection.execute(
-                        text("INSERT INTO notes_fts(rowid, body) VALUES (:id, :body)"),
-                        {"id": note_id, "body": "searchable note"},
+                        text("SELECT body FROM notes WHERE id = :id"), {"id": note_id}
                     )
-                    assert (
-                        await connection.execute(
-                            text("SELECT body FROM notes WHERE id = :id"), {"id": note_id}
-                        )
-                    ).scalar_one() == "searchable note"
+                ).scalar_one() == "searchable note"
 
-        async with asyncio.timeout(90), asyncio.TaskGroup() as writers:
-            for writer in range(64):
-                writers.create_task(write_notes(writer))
+        # Trigger: 64 writers share the pool's connections, and every round they all
+        #   want SQLite's single write lock at once.
+        # Why: SQLite's busy handler polls with growing sleeps (up to 100 ms) instead of
+        #   queueing waiters. When writers ran 40 transactions back to back, the other
+        #   pooled connections could hand the lock between themselves for the whole run
+        #   while one waiter kept sleeping through every free moment; on a slow Windows
+        #   runner that waiter passed busy_timeout (10 s) and raised "database is
+        #   locked" (#1955). Nothing in production holds the write lock saturated for
+        #   20 s, and the busy handler never promised fairness.
+        # Outcome: each round ends when its 64 writes commit, so a waiting connection
+        #   competes only with that round's work (well under a second), never with the
+        #   rest of the run. Contention per round is unchanged.
+        async with asyncio.timeout(90):
+            for index in range(40):
+                async with asyncio.TaskGroup() as writers:
+                    for writer in range(64):
+                        writers.create_task(write_note(writer * 40 + index))
         async with engine.connect() as connection:
             assert (await connection.execute(text("SELECT count(*) FROM notes"))).scalar() == 2560
             assert (

tests/mcp/test_mcp_container.py

@@ -5,6 +5,7 @@
 from basic_memory.mcp.container import (
     McpContainer,
     get_container,
+    installed_container,
     set_container,
 )
 from basic_memory.runtime.mode import RuntimeMode
@@ -92,3 +93,37 @@ def test_set_and_get_container(self, app_config, monkeypatch):
 
         set_container(container)
         assert get_container() is container
+
+    def test_overlapping_lifespans_exit_out_of_order(self, app_config, monkeypatch):
+        """Each lifespan withdraws only its own container, whatever the exit order."""
+        import basic_memory.mcp.container as container_module
+
+        monkeypatch.setattr(container_module, "_container", None)
+        first = McpContainer(config=app_config, mode=RuntimeMode.LOCAL)
+        second = McpContainer(config=app_config, mode=RuntimeMode.LOCAL)
+
+        first_lifespan = installed_container(first)
+        second_lifespan = installed_container(second)
+        first_lifespan.__enter__()
+        second_lifespan.__enter__()
+        assert get_container() is second
+
+        first_lifespan.__exit__(None, None, None)
+        assert get_container() is second
+
+        second_lifespan.__exit__(None, None, None)
+        with pytest.raises(RuntimeError, match="MCP container not initialized"):
+            get_container()
+
+    def test_lifespan_container_wins_over_the_host_container_while_it_runs(
+        self, app_config, monkeypatch
+    ):
+        import basic_memory.mcp.container as container_module
+
+        host = McpContainer(config=app_config, mode=RuntimeMode.LOCAL)
+        lifespan = McpContainer(config=app_config, mode=RuntimeMode.LOCAL)
+        monkeypatch.setattr(container_module, "_container", host)
+
+        with installed_container(lifespan):
+            assert get_container() is lifespan
+        assert get_container() is host

tests/mcp/test_server_lifespan_branches.py

@@ -42,6 +42,20 @@ async def test_mcp_lifespan_sync_enabled_branch(config_manager):
         pass
 
 
+@pytest.mark.asyncio
+async def test_mcp_lifespan_uninstalls_its_container_on_exit(config_manager):
+    """A session's container, and the config it read, end with the session (#1955).
+
+    Later work in the same process, such as picking the default search mode, must
+    not read a config that a finished in-process session captured.
+    """
+    async with lifespan(mcp):
+        assert get_container().config.default_project == config_manager.config.default_project
+
+    with pytest.raises(RuntimeError, match="MCP container not initialized"):
+        get_container()
+
+
 @pytest.mark.asyncio
 async def test_mcp_lifespan_shuts_down_db_when_engine_was_none(config_manager):
     db._engine = None