From 0924755591d0fb625439facb04905e4fc052b445 Mon Sep 17 00:00:00 2001 From: Gautam Sharma <148205237+GautamSharma99@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:08:53 +0530 Subject: [PATCH] fix(memory): serialize MCP backend initialization (#2309) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The Memory MCP server previously assigned its backend before asynchronous embedder and vector-index warm-up completed. A tool call arriving during the handshake could therefore receive a partially initialized backend. Backend initialization is now atomic and shared between concurrent callers. The backend is published only after warm-up succeeds. Failed candidates are closed and discarded so later calls can retry with a fresh backend. ## Type of Change - [x] Bug fix - [ ] New feature - [ ] Breaking change - [ ] Documentation-only change - [ ] Refactoring ## Changes Made - Keep the initializing backend local until warm-up completes successfully. - Share one initialization task between handshake and concurrent tool calls. - Await the shared task before exposing the backend to tool handlers. - Shield shared initialization from cancellation by an individual tool caller. - Close failed or cancelled backend candidates. - Clear failed initialization state so subsequent calls can retry. - Retrieve and log background initialization failures. - Add regression tests for handshake races, failure recovery, and concurrent initialization. - Add an Unreleased changelog entry. ## Testing - [x] Added regression tests - [x] Focused test suite passes - [x] Ruff checks pass - [x] Mypy passes - [x] Changed files pass formatting checks - [ ] Entire repository test suite passes without baseline failures Commands and results: - `uv run --extra dev --frozen pytest tests/test_memory/test_mcp_server.py -q` - `12 passed` - `uv run --extra dev --frozen ruff check .` - Passed - `uv run --extra dev --frozen ruff format --check headroom/memory/mcp_server.py tests/test_memory/test_mcp_server.py` - Passed - `uv run --extra dev --frozen mypy headroom --ignore-missing-imports` - Success across 504 source files - `uv run --extra dev --frozen pytest -q` - `9363 passed, 565 skipped, 4 failed` - The four failures are existing, unrelated failures outside the changed code: - `test_l2_appends_transform_label` - `test_recovery_records_sockets_and_secures_both_backups` - `test_dashboard_uses_cached_stats_and_lazy_history_feed_polling` - `test_smart_crusher_log_fallback_runs_for_valid_json` Repository-wide `ruff format --check .` also identifies pre-existing formatting drift in the untouched `headroom/proxy/handlers/anthropic.py`. ## Real Behavior Proof The regression tests exercise the affected lifecycle directly: 1. Start backend initialization through the MCP handshake. 2. Suspend warm-up before it completes. 3. Issue a memory tool call and verify its handler is not invoked. 4. Release warm-up and verify the tool receives the initialized backend. 5. Force background initialization to fail and verify the candidate is closed. 6. Issue another tool call and verify initialization retries with a fresh backend. 7. Start two tool calls concurrently and verify only one backend is constructed. Observed behavior: - Tool calls remain pending while handshake warm-up is incomplete. - A partially initialized backend never reaches a tool handler. - Failed candidates are closed and discarded. - A later tool call successfully retries initialization. - Concurrent calls share one initialization task and backend. Environment: macOS arm64, CPython 3.12.13. Not tested: a live stdio MCP client using the real ONNX model and database. The affected initialization lifecycle is covered with deterministic asynchronous regression tests. ## Review Readiness - [x] I have performed a self-review before requesting human review. - [x] This PR is ready for human review. ## Checklist - [x] The implementation follows the repository’s existing style and error-handling conventions. - [x] Tests cover the reported race, concurrent initialization, and failure recovery. - [x] Failed initialization does not leave a partially published backend. - [x] Failed backend candidates are closed before retry. - [x] No unrelated files or formatting changes are included. - [x] No temporary logging, debug code, or commented-out code remains. - [x] Public behavior changes are documented in the changelog. - [x] The branch has been rebased from the intended base and is ready for review. ## Additional Notes The four full-suite failures listed above occur outside the changed Memory MCP code and are unrelated to this PR. All tests covering the modified initialization lifecycle pass. --------- Co-authored-by: Tejas Chopra --- headroom/memory/mcp_server.py | 69 +++++++++--- tests/test_memory/test_mcp_server.py | 157 +++++++++++++++++++++++++++ 2 files changed, 212 insertions(+), 14 deletions(-) diff --git a/headroom/memory/mcp_server.py b/headroom/memory/mcp_server.py index 52830ff06..2a42dc7ab 100644 --- a/headroom/memory/mcp_server.py +++ b/headroom/memory/mcp_server.py @@ -161,34 +161,75 @@ def create_memory_server(db_path: str, user_id: str = "default") -> Server: server = Server("headroom-memory") _backend: LocalBackend | None = None - _init_task: asyncio.Task | None = None + _init_task: asyncio.Task[LocalBackend] | None = None async def _init_backend() -> LocalBackend: """Initialize backend with ONNX embedder (fast, no PyTorch).""" - nonlocal _backend + nonlocal _backend, _init_task config = LocalBackendConfig(db_path=db_path, embedder_backend="onnx") - _backend = LocalBackend(config) - await _warm_up_backend(_backend, user_id) + backend = LocalBackend(config) + init_task = asyncio.current_task() + try: + await _warm_up_backend(backend, user_id) + except (Exception, asyncio.CancelledError): + try: + await backend.close() + except Exception as cleanup_error: + logger.warning("Memory MCP: failed backend cleanup: %s", cleanup_error) + finally: + if _init_task is init_task: + _init_task = None + raise + + _backend = backend + if _init_task is init_task: + _init_task = None logger.info(f"Memory MCP: ready (db={db_path}, user={user_id})") - return _backend + return backend + + def _handle_backend_init_done(init_task: asyncio.Task[LocalBackend]) -> None: + """Clear and log failed background initialization tasks.""" + nonlocal _init_task + if init_task.cancelled(): + if _init_task is init_task: + _init_task = None + return + error = init_task.exception() + if error is not None: + if _init_task is init_task: + _init_task = None + logger.warning("Memory MCP: backend initialization failed: %s", error) + + def _start_backend_init() -> asyncio.Task[LocalBackend]: + """Start backend initialization once for all concurrent callers.""" + nonlocal _init_task + if _init_task is None: + _init_task = asyncio.create_task(_init_backend()) + _init_task.add_done_callback(_handle_backend_init_done) + return _init_task async def _get_backend() -> LocalBackend: nonlocal _backend, _init_task if _backend is not None: return _backend - # Wait for background init if it's running - if _init_task is not None: - await _init_task - return _backend # type: ignore[return-value] - # Fallback: init inline (shouldn't normally happen) - return await _init_backend() + + init_task = _start_backend_init() + try: + return await asyncio.shield(init_task) + except asyncio.CancelledError: + if init_task.done() and _init_task is init_task: + _init_task = None + raise + except Exception: + if _init_task is init_task: + _init_task = None + raise @server.list_tools() async def list_tools() -> list[Tool]: # Kick off background init on first list_tools (called at MCP handshake) - nonlocal _init_task - if _backend is None and _init_task is None: - _init_task = asyncio.create_task(_init_backend()) + if _backend is None: + _start_backend_init() return _TOOLS @server.call_tool() diff --git a/tests/test_memory/test_mcp_server.py b/tests/test_memory/test_mcp_server.py index b86e15179..cd2779627 100644 --- a/tests/test_memory/test_mcp_server.py +++ b/tests/test_memory/test_mcp_server.py @@ -10,6 +10,27 @@ from tests._mcp_stub import import_module_with_mcp_stub mcp_server_mod = import_module_with_mcp_stub("headroom.memory.mcp_server") +class _CapturingServer: + def __init__(self, name: str) -> None: + self.name = name + self.list_tools_handler = None + self.call_tool_handler = None + + def list_tools(self): + def decorator(handler): + self.list_tools_handler = handler + return handler + + return decorator + + def call_tool(self): + def decorator(handler): + self.call_tool_handler = handler + return handler + + return decorator + + def test_warm_up_backend_batches_embedding_and_indexing() -> None: """Warm-up should batch missing embeddings and vector indexing.""" warmup_embedding = np.ones(384, dtype=np.float32) @@ -62,6 +83,142 @@ def test_warm_up_backend_batches_embedding_and_indexing() -> None: assert np.array_equal(memory_without_embedding_b.embedding, batch_embeddings[1]) +def test_tool_call_waits_for_handshake_warm_up(monkeypatch) -> None: + async def scenario() -> None: + warm_up_started = asyncio.Event() + release_warm_up = asyncio.Event() + backend = SimpleNamespace() + + async def warm_up(candidate, user_id: str) -> None: + assert candidate is backend + assert user_id == "alice" + warm_up_started.set() + await release_warm_up.wait() + + handle_search = AsyncMock(return_value=["search result"]) + monkeypatch.setattr(mcp_server_mod, "Server", _CapturingServer) + monkeypatch.setattr(mcp_server_mod, "LocalBackend", lambda config: backend) + monkeypatch.setattr(mcp_server_mod, "_warm_up_backend", warm_up) + monkeypatch.setattr(mcp_server_mod, "_handle_search", handle_search) + + server = mcp_server_mod.create_memory_server("memory.db", user_id="alice") + await server.list_tools_handler() + await warm_up_started.wait() + + tool_call = asyncio.create_task( + server.call_tool_handler("memory_search", {"query": "preferences"}) + ) + await asyncio.sleep(0) + + handle_search.assert_not_awaited() + assert not tool_call.done() + + release_warm_up.set() + assert await tool_call == ["search result"] + handle_search.assert_awaited_once_with( + backend, + {"query": "preferences"}, + "alice", + ) + + asyncio.run(scenario()) + + +def test_failed_handshake_init_is_discarded_and_retried(monkeypatch) -> None: + async def scenario() -> None: + first_warm_up_started = asyncio.Event() + fail_first_warm_up = asyncio.Event() + failed_backend_closed = asyncio.Event() + + async def close_failed_backend() -> None: + failed_backend_closed.set() + + failed_backend = SimpleNamespace(close=AsyncMock(side_effect=close_failed_backend)) + ready_backend = SimpleNamespace(close=AsyncMock()) + backends = iter([failed_backend, ready_backend]) + + async def warm_up(candidate, user_id: str) -> None: + assert user_id == "alice" + if candidate is failed_backend: + first_warm_up_started.set() + await fail_first_warm_up.wait() + raise RuntimeError("warm-up failed") + + handle_search = AsyncMock(return_value=["search result"]) + monkeypatch.setattr(mcp_server_mod, "Server", _CapturingServer) + monkeypatch.setattr(mcp_server_mod, "LocalBackend", lambda config: next(backends)) + monkeypatch.setattr(mcp_server_mod, "_warm_up_backend", warm_up) + monkeypatch.setattr(mcp_server_mod, "_handle_search", handle_search) + + server = mcp_server_mod.create_memory_server("memory.db", user_id="alice") + await server.list_tools_handler() + await first_warm_up_started.wait() + + fail_first_warm_up.set() + await failed_backend_closed.wait() + await asyncio.sleep(0) + + failed_backend.close.assert_awaited_once() + handle_search.assert_not_awaited() + + assert await server.call_tool_handler("memory_search", {"query": "preferences"}) == [ + "search result" + ] + handle_search.assert_awaited_once_with( + ready_backend, + {"query": "preferences"}, + "alice", + ) + ready_backend.close.assert_not_awaited() + + asyncio.run(scenario()) + + +def test_concurrent_tool_calls_share_backend_initialization(monkeypatch) -> None: + async def scenario() -> None: + warm_up_started = asyncio.Event() + release_warm_up = asyncio.Event() + backend = SimpleNamespace() + created_backends = 0 + + def create_backend(config): + nonlocal created_backends + created_backends += 1 + return backend + + async def warm_up(candidate, user_id: str) -> None: + assert candidate is backend + assert user_id == "alice" + warm_up_started.set() + await release_warm_up.wait() + + handle_search = AsyncMock(return_value=["search result"]) + monkeypatch.setattr(mcp_server_mod, "Server", _CapturingServer) + monkeypatch.setattr(mcp_server_mod, "LocalBackend", create_backend) + monkeypatch.setattr(mcp_server_mod, "_warm_up_backend", warm_up) + monkeypatch.setattr(mcp_server_mod, "_handle_search", handle_search) + + server = mcp_server_mod.create_memory_server("memory.db", user_id="alice") + calls = [ + asyncio.create_task( + server.call_tool_handler("memory_search", {"query": f"query-{index}"}) + ) + for index in range(2) + ] + await warm_up_started.wait() + await asyncio.sleep(0) + + assert created_backends == 1 + handle_search.assert_not_awaited() + + release_warm_up.set() + assert await asyncio.gather(*calls) == [["search result"], ["search result"]] + assert handle_search.await_count == 2 + assert all(call.args[0] is backend for call in handle_search.await_args_list) + + asyncio.run(scenario()) + + def test_memory_mcp_startup_context_reports_dynamic_project_db(tmp_path) -> None: project_dir = tmp_path / "project-a" project_dir.mkdir()