diff --git a/src/iai_mcp/daemon/__init__.py b/src/iai_mcp/daemon/__init__.py index c40bc2ac..11c67d31 100644 --- a/src/iai_mcp/daemon/__init__.py +++ b/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 diff --git a/src/iai_mcp/daemon/_boot_warmup.py b/src/iai_mcp/daemon/_boot_warmup.py index 7a2d672d..aeabdaec 100644 --- a/src/iai_mcp/daemon/_boot_warmup.py +++ b/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) diff --git a/tests/test_daemon.py b/tests/test_daemon.py index a611a2a2..ea3195ae 100644 --- a/tests/test_daemon.py +++ b/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']}" )