fix(daemon): avoid duplicate startup work and clean cancellation

src/iai_mcp/daemon/__init__.py

@@ -761,7 +761,6 @@ def _install_warm_embedder_override(store) -> tuple[object, bool]:
     orig_efs = _embed_mod.embedder_for_store
     try:
         warm = orig_efs(store)
-        warm.embed("warmup")
 
         def _held_embedder_for_store(_store):
             return warm
@@ -1094,8 +1093,14 @@ def _capture_handler(record: dict) -> None:
         try:
             from iai_mcp.daemon._boot_warmup import warm_dispatch_surface
 
+            # Shield keeps the worker alive after the bounded pre-bind wait:
+            # asyncio cannot cancel a running thread, and retaining the task
+            # makes that continuation explicit and observable until shutdown.
+            _wds_task = asyncio.create_task(
+                asyncio.to_thread(warm_dispatch_surface, store)
+            )
             _wds_summary = await asyncio.wait_for(
-                asyncio.to_thread(warm_dispatch_surface, store), timeout=20.0,
+                asyncio.shield(_wds_task), timeout=20.0,
             )
             log.info(
                 "dispatch surface warmed pre-bind in %.0fms",
@@ -1113,7 +1118,7 @@ def _capture_handler(record: dict) -> None:
 
             async def _boot_warmup_task() -> None:
                 try:
-                    await asyncio.to_thread(run_boot_warmup, store)
+                    await asyncio.to_thread(run_boot_warmup, store, warm_dispatch=False)
                 except Exception as _exc:  # noqa: BLE001 -- warm-up must never crash the daemon
                     log.debug("boot_warmup failed: %s", _exc, exc_info=True)
 
@@ -1869,6 +1874,16 @@ def _interrupt_check() -> bool:
             except (OSError, RuntimeError) as exc:
                 log.debug("lifecycle_lock release failed: %s", exc)
     finally:
+        try:
+            if getattr(store, "_write_queue", None) is not None:
+                await store.disable_async_writes()
+        except Exception as exc:  # noqa: BLE001 -- cancellation cleanup must finish
+            log.debug("outer async-write cleanup failed: %s", exc, exc_info=True)
+        try:
+            if lifecycle_lock.is_held_by_self():
+                lifecycle_lock.release()
+        except (OSError, RuntimeError) as exc:
+            log.debug("outer lifecycle-lock release failed: %s", exc)
         _restore_embedder_funnel(_orig_efs, _override_installed)
     return 0
 

src/iai_mcp/daemon/_boot_warmup.py

@@ -127,10 +127,10 @@ def warm_dispatch_surface(store: Any) -> dict:
     These are one-time per-process costs that otherwise land on the FIRST
     real recall. The daemon runs this BEFORE binding its socket — a memory
     that answers the socket before its recall surface is resident is not
-    awake yet — and the full warm-up re-invokes it as its first shape (an
-    idempotent no-op then). All read-only: an embed produces a vector and
-    discards it, and load_recall_structural only reads/decodes the on-disk
-    cache. Returns a summary dict; never raises.
+    awake yet — then starts the remaining warm-up shapes without repeating
+    this encode. All read-only: an embed produces a vector and discards it,
+    and load_recall_structural only reads/decodes the on-disk cache. Returns
+    a summary dict; never raises.
     """
     _t0 = time.perf_counter()
     try:
@@ -156,6 +156,7 @@ def run_boot_warmup(
     *,
     probe_count: int = _DEFAULT_PROBE_COUNT,
     k: int = _DEFAULT_K,
+    warm_dispatch: bool = True,
 ) -> dict:
     """Fire the real first-recall read shapes once, write-free.
 
@@ -170,7 +171,8 @@ def run_boot_warmup(
     shapes: dict[str, dict] = {}
     probe_hit_ids: list[str] = []
 
-    shapes["dispatch_surface"] = warm_dispatch_surface(store)
+    if warm_dispatch:
+        shapes["dispatch_surface"] = warm_dispatch_surface(store)
 
     try:
         probes = _sample_probe_embeddings(store, probe_count)

tests/test_daemon.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import plistlib
+import threading
 from pathlib import Path
 
 import pytest
@@ -27,6 +28,7 @@ def _short_socket_paths(tmp_path, monkeypatch):
     sock_dir.mkdir(parents=True, exist_ok=True)
     sock_path = sock_dir / "d.sock"
     monkeypatch.setattr(concurrency, "SOCKET_PATH", sock_path)
+    monkeypatch.setenv("IAI_DAEMON_SOCKET_PATH", str(sock_path))
     return lock_path, sock_path, sock_dir
 
 
@@ -134,6 +136,7 @@ async def runner():
 def test_prewarm_called_once_at_boot(tmp_path, monkeypatch):
     from iai_mcp import daemon as daemon_mod
     from iai_mcp import daemon_state as ds_mod
+    from iai_mcp.daemon import _boot_warmup
 
     monkeypatch.setenv("IAI_MCP_STORE", str(tmp_path / "iai"))
     monkeypatch.setenv("IAI_MCP_EMBED_DIM", "384")
@@ -152,16 +155,31 @@ def _fake_embedder(store):
 
     monkeypatch.setattr("iai_mcp.embed.embedder_for_store", _fake_embedder)
 
+    warmup_finished = threading.Event()
+    real_warm_dispatch_surface = _boot_warmup.warm_dispatch_surface
+
+    def _observed_warm_dispatch_surface(store):
+        try:
+            return real_warm_dispatch_surface(store)
+        finally:
+            warmup_finished.set()
+
+    monkeypatch.setattr(
+        _boot_warmup, "warm_dispatch_surface", _observed_warm_dispatch_surface
+    )
+
     async def runner():
         task = asyncio.create_task(daemon_mod.main())
-        await asyncio.sleep(0.15)
-        task.cancel()
         try:
-            await task
-        except asyncio.CancelledError:
-            pass
+            return await asyncio.to_thread(warmup_finished.wait, 5)
+        finally:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
 
-    asyncio.run(runner())
+    assert asyncio.run(runner()), "daemon did not complete its pre-bind dispatch warmup"
     assert prewarm_calls["n"] == 1, (
         f"prewarm expected once, got {prewarm_calls['n']}"
     )