Merge pull request #31 from nickroci/feat/daemon-restart-on-update

fix(daemon): restart on version change so updates actually go live

Author: nickroci

daemon/agent_mem_daemon/__main__.py

@@ -118,9 +118,20 @@ def release_pid_file(path: Path) -> None:
 def write_daemon_state(phase: str) -> None:
     """Publish the lifecycle flag other processes read (see
     ``paths.daemon_state_path``). Best-effort — a failed write must never
-    block startup; consumers fall back to pid/socket inference."""
+    block startup; consumers fall back to pid/socket inference.
+
+    Records the running code's ``version`` (read from on-disk metadata at this
+    process's startup, where it equals the code actually loaded). The ``ultan``
+    spawn path compares it against the live installed version to detect a daemon
+    left running OLD code after an ``uv tool install`` update, and restarts it."""
     import json  # noqa: PLC0415 — tiny, cold path
     from datetime import datetime, timezone  # noqa: PLC0415
+    from importlib.metadata import PackageNotFoundError, version  # noqa: PLC0415
+
+    try:
+        running_version: str | None = version("agent-mem-daemon")
+    except PackageNotFoundError:
+        running_version = None
 
     try:
         daemon_state_path().write_text(
@@ -129,6 +140,7 @@ def write_daemon_state(phase: str) -> None:
                     "phase": phase,
                     "pid": os.getpid(),
                     "since": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+                    "version": running_version,
                 }
             ),
             encoding="utf-8",

daemon/tests/test_daemon_state.py

@@ -0,0 +1,23 @@
+"""The daemon stamps the version of the code it is running into daemon.state, so
+the `ultan` spawn path can detect a daemon left on old code after an update and
+restart it. ``agent_mem_home`` isolates AGENT_MEM_HOME (opt-in, per conftest) so
+the write never touches the user's real ~/.agent-mem."""
+
+from __future__ import annotations
+
+import importlib.metadata as md
+import json
+import os
+from pathlib import Path
+
+from agent_mem_daemon import __main__ as dm
+
+
+def test_write_daemon_state_records_running_version(agent_mem_home: Path) -> None:
+    dm.write_daemon_state("ready")
+    state = json.loads((agent_mem_home / "daemon.state").read_text(encoding="utf-8"))
+    assert state["phase"] == "ready"
+    assert state["pid"] == os.getpid()
+    # The version of the code this process is running, read from on-disk metadata
+    # at startup (where it equals the loaded code).
+    assert state["version"] == md.version("agent-mem-daemon")

tests/test_daemon_restart.py

@@ -0,0 +1,169 @@
+"""A daemon left running OLD code after an `ultan` tool update must be detected
+and restarted — a long-running process keeps the code it loaded at spawn, so
+`uv tool install` swapping the venv underneath it doesn't take effect until the
+process is replaced. The version-aware restart lives in the session-start path
+(NOT the per-turn hot path) and the mismatch is surfaced by `ultan doctor`.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from ultan import _daemon, _doctor, _hooks
+
+
+@pytest.fixture()
+def home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
+    h = tmp_path / "agent-mem"
+    h.mkdir()
+    monkeypatch.setenv("AGENT_MEM_HOME", str(h))
+    return h
+
+
+def _write_state(home: Path, *, pid: int, version: Any) -> None:
+    state: dict[str, Any] = {"phase": "ready", "pid": pid, "since": "x"}
+    if version is not None:
+        state["version"] = version
+    (home / "daemon.state").write_text(json.dumps(state), encoding="utf-8")
+
+
+# ── restart_if_stale: only restarts on a genuine version mismatch ──────
+
+
+def test_no_restart_when_no_state(home: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: "0.3.0")
+    assert _daemon.restart_if_stale() is False
+
+
+def test_no_restart_when_pid_dead(home: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    _write_state(home, pid=99999999, version="0.2.0")
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: "0.3.0")
+    assert _daemon.restart_if_stale() is False
+
+
+def test_no_restart_when_versions_match(home: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    _write_state(home, pid=os.getpid(), version="0.3.0")
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: "0.3.0")
+    stopped: list[int] = []
+    monkeypatch.setattr(_daemon, "_stop_daemon", lambda pid, **k: stopped.append(pid))
+    assert _daemon.restart_if_stale() is False
+    assert stopped == []
+
+
+def test_no_restart_when_installed_version_unknown(
+    home: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _write_state(home, pid=os.getpid(), version="0.2.0")
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: None)
+    stopped: list[int] = []
+    monkeypatch.setattr(_daemon, "_stop_daemon", lambda pid, **k: stopped.append(pid))
+    assert _daemon.restart_if_stale() is False
+    assert stopped == []
+
+
+def test_restart_legacy_daemon_with_no_version_stamp(
+    home: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    # A legacy daemon (pre-this-feature) wrote no version. Anchored on a KNOWN
+    # installed version, that's unambiguously old code → restart. This is the
+    # migration case: the daemon a user already has running when they first adopt
+    # the feature has no stamp, and must still be replaced.
+    _write_state(home, pid=os.getpid(), version=None)
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: "0.3.0")
+    stopped: list[int] = []
+    monkeypatch.setattr(_daemon, "_stop_daemon", lambda pid, **k: stopped.append(pid))
+    assert _daemon.restart_if_stale() is True
+    assert stopped == [os.getpid()]
+
+
+def test_restart_when_stale_stops_and_clears_backoff(
+    home: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _write_state(home, pid=os.getpid(), version="0.2.0")
+    (home / ".daemon-spawn-attempt").write_text("x", encoding="utf-8")
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: "0.3.0")
+    stopped: list[int] = []
+    monkeypatch.setattr(_daemon, "_stop_daemon", lambda pid, **k: stopped.append(pid))
+
+    assert _daemon.restart_if_stale() is True
+    assert stopped == [os.getpid()]  # stopped the stale daemon
+    # backoff stamp cleared so ensure_running() respawns immediately, not throttled
+    assert not (home / ".daemon-spawn-attempt").exists()
+
+
+# ── _installed_daemon_version ─────────────────────────────────────────
+
+
+def test_installed_version_matches_metadata_or_none() -> None:
+    import importlib.metadata as md
+
+    # Env-agnostic: in a full `ultan[retrieval]` install the daemon dist resolves;
+    # in the thin root/CI venv (no extras) it doesn't and the helper returns None.
+    # Either way it must agree with importlib.metadata — never invent a version.
+    try:
+        expected: str | None = md.version("agent-mem-daemon")
+    except md.PackageNotFoundError:
+        expected = None
+    assert _daemon._installed_daemon_version() == expected
+
+
+def test_installed_version_none_when_dist_absent(monkeypatch: pytest.MonkeyPatch) -> None:
+    import importlib.metadata as md
+
+    def _raise(_name: str) -> str:
+        raise md.PackageNotFoundError
+
+    monkeypatch.setattr(md, "version", _raise)
+    assert _daemon._installed_daemon_version() is None
+
+
+# ── session-start wires the restart in (before warming) ───────────────
+
+
+def test_session_start_calls_restart_if_stale(monkeypatch: pytest.MonkeyPatch) -> None:
+    order: list[str] = []
+    monkeypatch.setattr(
+        _hooks._daemon, "restart_if_stale", lambda: order.append("restart") or False
+    )
+    monkeypatch.setattr(_hooks._daemon, "ensure_running", lambda: order.append("ensure") or True)
+    monkeypatch.setattr(_hooks._events, "append_event", lambda *a, **k: None)
+
+    rc = _hooks.dispatch("session-start", {"session_id": "s", "source": "startup"})
+
+    assert rc == 0
+    assert order == ["restart", "ensure"]  # stale-check BEFORE warming
+
+
+# ── doctor surfaces the mismatch even before the restart lands ─────────
+
+
+def test_doctor_flags_stale_daemon(
+    home: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
+    monkeypatch.setattr(_doctor, "_report_runtime", lambda home: False)
+    (home / "daemon.pid").write_text(f"{os.getpid()}\n", encoding="utf-8")
+    _write_state(home, pid=os.getpid(), version="0.2.0")
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: "0.3.0")
+
+    _doctor.run()
+    out = capsys.readouterr().out
+
+    assert "daemon code: 0.2.0 — STALE (installed 0.3.0)" in out
+
+
+def test_doctor_no_stale_line_when_versions_match(
+    home: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
+    monkeypatch.setattr(_doctor, "_report_runtime", lambda home: False)
+    (home / "daemon.pid").write_text(f"{os.getpid()}\n", encoding="utf-8")
+    _write_state(home, pid=os.getpid(), version="0.3.0")
+    monkeypatch.setattr(_daemon, "_installed_daemon_version", lambda: "0.3.0")
+
+    _doctor.run()
+    out = capsys.readouterr().out
+
+    assert "daemon code: 0.3.0" in out
+    assert "STALE" not in out

ultan/_daemon.py

@@ -73,6 +73,34 @@ def _pid_alive(pid: object) -> bool:
     return True
 
 
+def _read_daemon_state() -> dict[str, Any] | None:
+    """The daemon's lifecycle flag ``{phase, pid, since, version}`` (written
+    between pid-acquire and exit), or ``None`` when absent/unreadable."""
+    try:
+        raw: Any = json.loads((_home() / "daemon.state").read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    return cast("dict[str, Any]", raw) if isinstance(raw, dict) else None
+
+
+def _installed_daemon_version() -> str | None:
+    """The agent-mem-daemon version installed in THIS venv right now, read live
+    from on-disk metadata — so it reflects an ``uv tool install`` update even
+    while an older daemon process keeps serving the previous code. ``None`` when
+    the daemon dist isn't installed (a thin/uvx env) or metadata is unreadable.
+
+    Cold path only (session-start / doctor): reading dist metadata is stdlib but
+    does I/O, so it must never land on the per-turn hook hot path."""
+    from importlib.metadata import PackageNotFoundError, version  # noqa: PLC0415 — cold path
+
+    try:
+        return version(_DAEMON_ENTRYPOINT)  # dist name == console-script name here
+    except PackageNotFoundError:
+        return None
+    except Exception:  # noqa: BLE001 — a broken install must not crash the caller
+        return None
+
+
 def status() -> str:
     """Coarse daemon state for consumers that must phrase fallbacks honestly:
     ``"ready"`` (socket answering), ``"warming"`` (daemon alive or freshly
@@ -82,12 +110,9 @@ def status() -> str:
         return "ready"
     home = _home()
     # The daemon's own lifecycle flag (written between pid-acquire and exit).
-    try:
-        raw: Any = json.loads((home / "daemon.state").read_text(encoding="utf-8"))
-        if isinstance(raw, dict) and _pid_alive(cast("dict[str, Any]", raw).get("pid")):
-            return "warming"
-    except (OSError, ValueError):
-        pass
+    state = _read_daemon_state()
+    if state is not None and _pid_alive(state.get("pid")):
+        return "warming"
     # A spawn was attempted moments ago (pre-flag window, or older daemon).
     try:
         if (time.time() - (home / ".daemon-spawn-attempt").stat().st_mtime) < 300.0:
@@ -97,6 +122,68 @@ def status() -> str:
     return "down"
 
 
+def _stop_daemon(pid: int, *, timeout_s: float = 3.0) -> None:
+    """Stop a running daemon: SIGTERM, wait briefly for a clean exit, SIGKILL as
+    a last resort. The daemon installs a SIGTERM handler that flips its stop
+    event and shuts down gracefully (releasing the socket), so the common case is
+    a clean stop well under ``timeout_s``."""
+    import signal  # noqa: PLC0415 — cold path (session-start restart only)
+
+    try:
+        os.kill(pid, signal.SIGTERM)
+    except OSError:
+        return  # already gone
+    deadline = time.monotonic() + timeout_s
+    while _pid_alive(pid) and time.monotonic() < deadline:
+        time.sleep(0.1)
+    if _pid_alive(pid):
+        try:
+            os.kill(pid, signal.SIGKILL)
+        except OSError:
+            pass
+
+
+def restart_if_stale() -> bool:
+    """If a daemon is running OLDER code than what's now installed in the venv,
+    stop it so a fresh daemon starts on current code. Returns True when a stale
+    daemon was stopped — the caller's :func:`ensure_running` then spawns the new
+    one (the stopped daemon no longer answers the socket).
+
+    SESSION-START PATH ONLY — never the per-turn hook hot path: it reads package
+    metadata (cold-cheap, hot-needless). The running version is what the daemon
+    stamped into ``daemon.state`` at ITS startup; the installed version is read
+    live, so an ``uv tool install`` update is detected even though the old
+    process keeps serving the previous code.
+
+    Anchored on the INSTALLED version: if we can't read it (a thin/uvx env with
+    no daemon dist, or unreadable metadata) we never restart — there's nothing to
+    compare against. Given a readable installed version, a daemon whose recorded
+    version differs is stale; so is one with NO recorded version, because that is
+    a legacy daemon from before this stamp existed (a current daemon always
+    stamps a version when the metadata it reads is itself readable). Restarting
+    that legacy daemon is the whole point — it's exactly the process left on old
+    code that this feature exists to replace."""
+    state = _read_daemon_state()
+    if not state:
+        return False
+    pid = state.get("pid")
+    if not _pid_alive(pid):
+        return False
+    installed = _installed_daemon_version()
+    if not installed:
+        return False  # can't tell what's installed — never restart blindly
+    if state.get("version") == installed:
+        return False  # up to date
+    _stop_daemon(cast("int", pid))
+    # Clear the spawn backoff so ensure_running() respawns immediately rather
+    # than throttling this intentional restart as if it were a crash loop.
+    try:
+        (_home() / ".daemon-spawn-attempt").unlink()
+    except OSError:
+        pass
+    return True
+
+
 def run_foreground(extra_args: list[str]) -> int:
     """`ultan daemon` — run the installed agent-mem-daemon in the foreground."""
     daemon = _daemon_bin()

ultan/_doctor.py

@@ -90,6 +90,19 @@ def _report_daemon(home: Path) -> tuple[bool, bool]:
     if phase is not None and alive and phase.get("pid") == pid:
         print(f"daemon phase: {phase.get('phase')} (since {phase.get('since')})")
 
+    # Running daemon code vs what's installed now — surfaces a daemon left on old
+    # code after an update (it restarts on the next session; see _daemon).
+    running_ver = phase.get("version") if phase is not None else None
+    installed_ver = _daemon._installed_daemon_version()  # pyright: ignore[reportPrivateUsage]
+    if alive and running_ver:
+        if installed_ver and running_ver != installed_ver:
+            print(
+                f"daemon code: {running_ver} — STALE (installed {installed_ver}); "
+                "restarts on the next session"
+            )
+        else:
+            print(f"daemon code: {running_ver}")
+
     socket_ok = _daemon._socket_answering()  # pyright: ignore[reportPrivateUsage]  # intra-package
     if socket_ok:
         print("priming socket: answering")

ultan/_hooks.py

@@ -87,6 +87,11 @@ def _user_prompt_submit(payload: dict[str, Any]) -> int:
 
 
 def _session_start(payload: dict[str, Any]) -> int:
+    # If a daemon is still running OLDER code than what's now installed (e.g.
+    # right after an `ultan` tool update), stop it so a fresh one starts on the
+    # current code — a long-running process keeps the code it loaded at spawn.
+    # Session-start, not the per-turn hot path, so the version check is fine here.
+    _daemon.restart_if_stale()
     # Warm the daemon at session start so the first prompt is already hot.
     _daemon.ensure_running()
     # Capture the session boundary (legacy parity: src/hooks/session_start.py