fix(memory): serialize MCP backend initialization (#2309)

## 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 <chopratejas@gmail.com>

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()

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()