From 3e475aa8673a0c502b8a1c231e33152bda8c1eb4 Mon Sep 17 00:00:00 2001 From: phernandez Date: Thu, 29 Oct 2026 14:20:06 -0500 Subject: [PATCH 1/2] fix(core): own startup index tasks and remove flaky test causes - SQLite integration fixture now initializes search storage (FTS5 and vector tables) at database open, as production does. Lazy vector-table DDL during the first embedding raced concurrent searches and failed them with "vtable constructor failed: search_index". - initialize_file_indexing runs startup indexes and the watcher in one TaskGroup, replacing the module-level _initial_index_tasks set. Shutdown cancels and awaits them; a failed index is logged and the watcher keeps running. - MCP lifespan installs its container for the lifespan only and restores the previous one on exit, so a finished in-process session no longer leaves a stale config behind. - SQLite contention test writes in rounds so a busy-handler waiter cannot be starved for the whole run. - Drop the Postgres refresh negative controls: the pkey collision needs two upserts to pass the arbiter check in the same instant, which no barrier can force; the tolerance is covered deterministically elsewhere. - Linear-scaling tokenizer tests measure thread CPU time with gc paused. Closes #1955 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CMYVL1x91B8kJxy7nn2GU8 Signed-off-by: phernandez --- src/basic_memory/mcp/container.py | 21 +++++ src/basic_memory/mcp/server.py | 5 +- src/basic_memory/services/initialization.py | 51 +++++----- test-int/conftest.py | 18 +++- .../mcp/test_new_project_watch_lifecycle.py | 5 - test-int/test_sqlite_lifecycle_integration.py | 48 ++++++---- tests/mcp/test_server_lifespan_branches.py | 14 +++ .../test_search_relaxation_unicode_classes.py | 55 +++++++---- tests/services/test_initialization.py | 94 +++++++++++++------ .../test_search_refresh_concurrency.py | 64 ++----------- 10 files changed, 216 insertions(+), 159 deletions(-) diff --git a/src/basic_memory/mcp/container.py b/src/basic_memory/mcp/container.py index 4fbb7591e..e0e65c523 100644 --- a/src/basic_memory/mcp/container.py +++ b/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 @@ -108,3 +110,22 @@ def set_container(container: McpContainer) -> None: """Set the MCP container (called by lifespan).""" global _container _container = container + + +@contextmanager +def installed_container(container: McpContainer) -> Iterator[None]: + """Expose a container for one server lifespan, then restore the previous one. + + Trigger: an in-process MCP session (``Client(mcp)``) enters and exits the lifespan. + Why: set_container() alone left the last lifespan's container, and its config, + installed after the session ended. Later code in the same process, such as + the default search mode, then read a config that no longer applied. + Outcome: the container lives exactly as long as the lifespan that created it. + """ + global _container + previous = _container + _container = container + try: + yield + finally: + _container = previous diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index a2535f25c..9e61b030d 100644 --- a/src/basic_memory/mcp/server.py +++ b/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", diff --git a/src/basic_memory/services/initialization.py b/src/basic_memory/services/initialization.py index 195686a2e..80de64082 100644 --- a/src/basic_memory/services/initialization.py +++ b/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 diff --git a/test-int/conftest.py b/test-int/conftest.py index 5afdbbe0b..c5b013090 100644 --- a/test-int/conftest.py +++ b/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 diff --git a/test-int/mcp/test_new_project_watch_lifecycle.py b/test-int/mcp/test_new_project_watch_lifecycle.py index a5ece2fef..9c8c5cad7 100644 --- a/test-int/mcp/test_new_project_watch_lifecycle.py +++ b/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() diff --git a/test-int/test_sqlite_lifecycle_integration.py b/test-int/test_sqlite_lifecycle_integration.py index b02337542..41b0ece6e 100644 --- a/test-int/test_sqlite_lifecycle_integration.py +++ b/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 ( diff --git a/tests/mcp/test_server_lifespan_branches.py b/tests/mcp/test_server_lifespan_branches.py index 6175eac4e..e3f594c9b 100644 --- a/tests/mcp/test_server_lifespan_branches.py +++ b/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 diff --git a/tests/repository/test_search_relaxation_unicode_classes.py b/tests/repository/test_search_relaxation_unicode_classes.py index 5fb725b2a..e11d53c9c 100644 --- a/tests/repository/test_search_relaxation_unicode_classes.py +++ b/tests/repository/test_search_relaxation_unicode_classes.py @@ -13,6 +13,7 @@ rather than identifiers. """ +import gc import time import unicodedata @@ -244,6 +245,30 @@ def test_structural_punctuation_keeps_splitting(structural: str) -> None: assert len(relaxation_word_tokens(f"сло{structural}во доступа")) == 3 +def _tokenizing_cpu_seconds(text: str) -> float: + """CPU time this thread spends tokenizing ``text``, best of five runs. + + The ratios below measure algorithmic growth, so the clock must not include time + the test spent descheduled. Wall-clock time (perf_counter) does: on a loaded CI + runner another process can take the CPU in the middle of one measurement and + turn a linear ×2.0 into a "quadratic" ×2.5 (#1955). thread_time() counts only + this thread's CPU work. The collector is paused because a collection that lands + in one measurement and not the other is the same kind of outside noise. + """ + gc_was_enabled = gc.isenabled() + gc.disable() + try: + samples = [] + for _ in range(5): + start = time.thread_time() + relaxation_word_tokens(text) + samples.append(time.thread_time() - start) + finally: + if gc_was_enabled: + gc.enable() + return min(samples) + + def test_tokenizing_scales_linearly_with_query_length() -> None: """Doubling the query must roughly double the work, not more. @@ -252,19 +277,16 @@ def test_tokenizing_scales_linearly_with_query_length() -> None: quadratic: an unbounded query full of apostrophes then ties up the worker that tokenizes it. The ratio is asserted rather than a duration, so the test does not depend on how fast the machine is. + + There is no operation count to assert instead: the slicing version makes the + same calls and inspects the same characters, and only the bytes each slice + copies grow, inside C. Measured here, the indexed reader scales at ×2.0 and the + sliced one at ×3.2. """ query = "a'b " * (128 * 1024 // 4) - def elapsed(text: str) -> float: - start = time.perf_counter() - relaxation_word_tokens(text) - return time.perf_counter() - start - - # Fastest of three runs each: the ratio, not the duration, is what is being - # asserted, so a slow or busy machine does not turn this red. Measured here, - # the indexed reader scales at ×2.0 and the sliced one at ×3.2. - single = min(elapsed(query) for _ in range(3)) - double = min(elapsed(query * 2) for _ in range(3)) + single = _tokenizing_cpu_seconds(query) + double = _tokenizing_cpu_seconds(query * 2) assert double < single * 2.5, ( f"tokenizing scaled worse than linearly: {single * 1000:.1f} ms then {double * 1000:.1f} ms" @@ -278,20 +300,13 @@ def test_trailing_format_trim_scales_linearly() -> None: trailing trim runs. Trimming one character at a time copies the shrinking token at every step, which is quadratic: an unbounded query ending in enough soft hyphens then ties up the worker that tokenizes it. As above, the ratio - is asserted rather than a duration. + is asserted rather than a duration, in CPU time. """ - - def elapsed(count: int) -> float: - text = "a" + "\u00ad" * count - start = time.perf_counter() - relaxation_word_tokens(text) - return time.perf_counter() - start - # Sized so the trim dominates the measurement: at shorter lengths the # tokenizer's linear per-character work dilutes the quadratic term below # the ratio threshold and a regression would pass unnoticed. - single = min(elapsed(256 * 1024) for _ in range(3)) - double = min(elapsed(512 * 1024) for _ in range(3)) + single = _tokenizing_cpu_seconds("a" + "\u00ad" * (256 * 1024)) + double = _tokenizing_cpu_seconds("a" + "\u00ad" * (512 * 1024)) assert double < single * 2.5, ( f"trailing trim scaled worse than linearly: {single * 1000:.1f} ms then {double * 1000:.1f} ms" diff --git a/tests/services/test_initialization.py b/tests/services/test_initialization.py index c4d0532cc..3dacde0f5 100644 --- a/tests/services/test_initialization.py +++ b/tests/services/test_initialization.py @@ -9,6 +9,7 @@ import asyncio import pytest +from loguru import logger from basic_memory import db from basic_memory.config import BasicMemoryConfig, DatabaseBackend @@ -17,7 +18,6 @@ from basic_memory.models import Project from basic_memory.repository.project_repository import ProjectRepository from basic_memory.services.initialization import ( - _initial_index_tasks, ensure_initialization, initialize_database, initialize_file_indexing, @@ -265,23 +265,6 @@ async def test_initialize_file_indexing_uses_project_index_runtime_for_initial_s monkeypatch.setattr("basic_memory.index.watch_service.WatchService", _FakeWatchService) monkeypatch.setenv("BASIC_MEMORY_MCP_PROJECT", "Event Startup") - original_create_task = asyncio.create_task - created_coroutines = [] - - class _CapturedTask: - """Stand-in task: the scheduler holds a strong ref and registers a - done callback, so the fake must accept both.""" - - def add_done_callback(self, callback) -> None: # noqa: ANN001 - del callback - - def capture_task(coro): - coroutine_name = getattr(getattr(coro, "cr_code", None), "co_name", "") - if coroutine_name == "index_project_background": - created_coroutines.append(coro) - return _CapturedTask() - return original_create_task(coro) - class RecordingProjectIndexRuntimeFactory: def __init__(self, *, read_cache: object) -> None: self.read_cache = read_cache @@ -303,9 +286,6 @@ async def run_project_index_for_project( # noqa: ANN001 deleted_files=0, ) - monkeypatch.setattr( - "basic_memory.services.initialization.asyncio.create_task", capture_task - ) monkeypatch.setattr( "basic_memory.index.local_project.LocalProjectIndexRuntimeFactory", RecordingProjectIndexRuntimeFactory, @@ -315,6 +295,8 @@ async def run_project_index_for_project( # noqa: ANN001 run_project_index_for_project, ) + # The fake watcher returns at once; initialize_file_indexing still returns only + # after the startup index it owns has finished. await initialize_file_indexing(updated, quiet=True) watch_kwargs = _FakeWatchService.last_kwargs @@ -326,9 +308,6 @@ async def run_project_index_for_project( # noqa: ANN001 assert len(initial_projects) == 1 assert isinstance(initial_projects[0], Project) assert initial_projects[0].name == "event-startup" - assert len(created_coroutines) == 1 - await created_coroutines[0] - assert project_index_calls == [("event-startup", "runtime:event-startup", False)] finally: await db.shutdown_db() @@ -344,7 +323,6 @@ async def test_initialize_file_indexing_skips_project_with_non_absolute_path( recognized as cloud, yet Path("") resolves to the process cwd. Indexing it would inject frontmatter into unrelated files, so it must be skipped. """ - initial_index_tasks = _initial_index_tasks.copy() await db.shutdown_db() try: from basic_memory.config import ProjectEntry @@ -383,9 +361,69 @@ async def test_initialize_file_indexing_skips_project_with_non_absolute_path( assert "empty-path" in skip_logs[0] assert "good" not in skip_logs[0] finally: - # The fake watcher returns immediately; let this test's startup index - # finish before closing its database and event loop. - await asyncio.gather(*(_initial_index_tasks - initial_index_tasks)) + await db.shutdown_db() + + +@pytest.mark.asyncio +async def test_a_failed_startup_index_leaves_the_watcher_running( + app_config: BasicMemoryConfig, config_manager, config_home, monkeypatch +): + """One project's failed startup index must not stop file watching for every project. + + The startup indexes and the watcher share one task group, so an index failure that + escaped its task would cancel the watcher. The index of a project whose root is a + regular file fails for real; the watcher keeps running until it finishes on its own. + """ + await db.shutdown_db() + try: + from basic_memory.config import ProjectEntry + + watched = config_home / "watched" + watched.mkdir(parents=True, exist_ok=True) + not_a_directory = config_home / "not-a-directory" + not_a_directory.write_text("a file where a project root should be\n") + updated = app_config.model_copy( + update={ + "projects": { + "watched": ProjectEntry(path=str(watched)), + "broken": ProjectEntry(path=str(not_a_directory)), + }, + "default_project": "watched", + } + ) + config_manager.save_config(updated) + await initialize_database(updated) + await reconcile_projects_with_config(updated) + + index_failed = asyncio.Event() + + def record_index_failure(message) -> None: # noqa: ANN001 + if "Error in background project index for project broken" in str(message): + index_failed.set() + + sink_id = logger.add(record_index_failure, level="ERROR") + + class WatchUntilTheIndexFails(_FakeWatchService): + finished = False + + @override + async def run(self) -> None: + await index_failed.wait() + # Still running after the failure: the task group did not cancel it. + WatchUntilTheIndexFails.finished = True + + _disable_test_env_short_circuit(monkeypatch) + monkeypatch.setattr( + "basic_memory.index.watch_service.WatchService", WatchUntilTheIndexFails + ) + try: + async with asyncio.timeout(60): + await initialize_file_indexing(updated, quiet=True) + finally: + logger.remove(sink_id) + + assert WatchUntilTheIndexFails.finished + finally: await db.shutdown_db() diff --git a/tests/services/test_search_refresh_concurrency.py b/tests/services/test_search_refresh_concurrency.py index fa5e52005..5c076daa9 100644 --- a/tests/services/test_search_refresh_concurrency.py +++ b/tests/services/test_search_refresh_concurrency.py @@ -9,9 +9,16 @@ A real 80 KB note opens the race window by itself: indexing it takes longer than the gap between saves. Here every refresh also pauses between its delete and its insert, and the collision rounds start from an entity whose rows another writer has just deleted, so the -deletes have nothing to wait on and the inserts arrive together. The negative controls -remove the tolerance and show the same refreshes still collide underneath, so the -positive tests are not passing by luck. +deletes have nothing to wait on and the inserts arrive together. + +Whether two such inserts actually collide is Postgres's call, not this file's: the upsert +absorbs a conflict it sees in its arbiter index, and ``search_index_pkey`` fires only when +two inserts pass that check in the same instant. No barrier outside the INSERT can force +that instant, so these tests assert only that overlapping refreshes never fail. The lost- +race branch itself is pinned deterministically by +``test_search_refresh_atomicity.py`` (an injected ``deadlock detected`` is logged and the +previous projection kept; any other database error still raises) and by +``tests/repository/test_search_refresh_race.py`` (which errors count as a lost race). """ from __future__ import annotations @@ -31,7 +38,6 @@ ) from basic_memory.models import Entity from basic_memory.repository.accepted_note_repositories import AcceptedNoteRepositories -from basic_memory.services import search_service as search_service_module pytestmark = [pytest.mark.asyncio, pytest.mark.postgres] @@ -67,10 +73,6 @@ async def pause_then_insert(search_index_rows, session=None): monkeypatch.setattr(repository, "bulk_index_items", pause_then_insert) -def _remove_the_race_tolerance(monkeypatch) -> None: - monkeypatch.setattr(search_service_module, "lost_search_refresh_race", lambda error: None) - - async def _projection(search_service, entity: Entity) -> list[tuple[str, int]]: rows = await search_service.repository.get_entity_search_rows(entity.id) return sorted((row.type, row.id) for row in rows if row.entity_id == entity.id) @@ -159,34 +161,6 @@ async def test_concurrent_saves_of_a_long_note_all_succeed_and_converge( ) == await file_service.read_entity_content(entity) -async def test_negative_control_without_the_tolerance_overlapping_saves_fail( - monkeypatch, - db_backend, - search_service, - file_service, - full_entity, - entity_repository, - session_maker, -): - _skip_unless_postgres(db_backend) - entity = await _reload(session_maker, entity_repository, full_entity) - await search_service.index_entity_data(entity, content="before") - _widen_the_window(monkeypatch, search_service.repository) - _remove_the_race_tolerance(monkeypatch) - - failures: list[BaseException] = [] - # The collision is a race, so a round can miss; ten rounds have not all missed. - for round_start in range(0, 1000, 100): - results = await _save_round(search_service, file_service, entity, round_start) - failures += [r for r in results if isinstance(r, BaseException)] - if failures: - break - - messages = " ".join(str(failure) for failure in failures) - assert failures, "without the tolerance, overlapping saves were expected to collide" - assert "search_index_pkey" in messages or "deadlock detected" in messages - - async def _refresh_relations_concurrently(search_service, entity: Entity) -> list[BaseException]: """Refresh with the content in hand, as the batch indexer does; relation rows ride along. @@ -220,24 +194,6 @@ async def test_concurrent_refreshes_keep_relation_rows_single( assert await _projection(search_service, entity) == before -async def test_negative_control_relation_rows_collide_without_the_tolerance( - monkeypatch, db_backend, search_service, full_entity, entity_repository, session_maker -): - _skip_unless_postgres(db_backend) - entity = await _reload(session_maker, entity_repository, full_entity) - await search_service.index_entity_data(entity, content="before") - _widen_the_window(monkeypatch, search_service.repository) - _remove_the_race_tolerance(monkeypatch) - - failures: list[BaseException] = [] - for _ in range(10): - failures += await _refresh_relations_concurrently(search_service, entity) - if failures: - break - assert failures, "without the tolerance, overlapping refreshes were expected to collide" - assert "search_index_pkey" in " ".join(str(failure) for failure in failures) - - async def test_an_accepted_write_and_an_index_refresh_overlap_without_failing( monkeypatch, db_backend, search_service, full_entity, entity_repository, session_maker ): From c4af332b7fbd4bb3107d162b659d29b752a1ddcd Mon Sep 17 00:00:00 2001 From: phernandez Date: Thu, 29 Oct 2026 14:40:25 -0500 Subject: [PATCH 2/2] fix(mcp): withdraw only a lifespan's own container on exit Overlapping lifespans that exit out of order restored an exited container with the save/restore slot. Each lifespan now adds its container to a list and removes that same entry on exit; the newest running one is current. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CMYVL1x91B8kJxy7nn2GU8 Signed-off-by: phernandez --- src/basic_memory/mcp/container.py | 31 ++++++++++++++++----------- tests/mcp/test_mcp_container.py | 35 +++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 12 deletions(-) diff --git a/src/basic_memory/mcp/container.py b/src/basic_memory/mcp/container.py index e0e65c523..6359210bf 100644 --- a/src/basic_memory/mcp/container.py +++ b/src/basic_memory/mcp/container.py @@ -91,41 +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, then restore the previous one. + """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: set_container() alone left the last lifespan's container, and its config, - installed after the session ended. Later code in the same process, such as - the default search mode, then read a config that no longer applied. - Outcome: the container lives exactly as long as the lifespan that created it. + 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. """ - global _container - previous = _container - _container = container + _lifespan_containers.append(container) try: yield finally: - _container = previous + 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] diff --git a/tests/mcp/test_mcp_container.py b/tests/mcp/test_mcp_container.py index f6e3f9a66..1b852a97f 100644 --- a/tests/mcp/test_mcp_container.py +++ b/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