Merge pull request #25 from nickroci/fix/startup-ux

Startup UX: lifecycle flag + warming-aware skill/doctor

Author: nickroci

README.md

@@ -29,6 +29,7 @@
 from .ingest import DEFAULT_POLL_INTERVAL, JsonlTailer
 from .logging_setup import configure as configure_logging
 from .paths import (
+    daemon_state_path,
     ensure_home,
     events_path,
     knowledge_dir,
@@ -114,6 +115,42 @@ def release_pid_file(path: Path) -> None:
         pass
 
 
+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."""
+    import json  # noqa: PLC0415 — tiny, cold path
+    from datetime import datetime, timezone  # noqa: PLC0415
+
+    try:
+        daemon_state_path().write_text(
+            json.dumps(
+                {
+                    "phase": phase,
+                    "pid": os.getpid(),
+                    "since": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+                }
+            ),
+            encoding="utf-8",
+        )
+    except OSError:
+        pass
+
+
+def clear_daemon_state() -> None:
+    """Best-effort remove of the lifecycle flag. Never raises. Only clears a
+    flag this process wrote — a duplicate daemon exiting must not erase the
+    live daemon's state."""
+    import json  # noqa: PLC0415
+
+    try:
+        raw = json.loads(daemon_state_path().read_text(encoding="utf-8"))
+        if raw.get("pid") == os.getpid():
+            daemon_state_path().unlink()
+    except (OSError, ValueError):
+        pass
+
+
 # ---- args -----------------------------------------------------------
 
 
@@ -341,6 +378,9 @@ def run(args: argparse.Namespace) -> int:
     # files on its way out. acquire_pid_file writes only to raw stderr, which
     # the spawner tees to daemon-spawn.log (see ultan/_daemon.py).
     acquire_pid_file(pidfile)
+    # We own the lifecycle now — flag "warming" so doctor/the search skill
+    # can say "starting up" instead of "broken" during the model load.
+    write_daemon_state("warming")
 
     log = configure_logging(
         logfile,
@@ -405,6 +445,9 @@ def run(args: argparse.Namespace) -> int:
     sched.start()
     tailer_thread.start()
     rpc_thread.start()
+    # Socket is being served — flip the lifecycle flag to ready.
+    write_daemon_state("ready")
+    log.info("daemon ready: priming socket serving")
     try:
         while not stop_event.is_set():
             stop_event.wait(timeout=1.0)
@@ -435,6 +478,7 @@ def run(args: argparse.Namespace) -> int:
             sched.stats.queue_high_water,
             sched.stats.scholar_skipped_backpressure,
         )
+        clear_daemon_state()
         release_pid_file(pidfile)
 
     return 0

daemon/agent_mem_daemon/__main__.py

@@ -38,6 +38,16 @@ def pid_path() -> Path:
     return home() / "daemon.pid"
 
 
+def daemon_state_path() -> Path:
+    """Lifecycle flag the daemon maintains for OTHER processes to read:
+    ``{"phase": "warming"|"ready", "pid": N, "since": <iso>}``. Written right
+    after PID acquisition (warming), updated when the priming socket is
+    serving (ready), removed on shutdown. Consumers — `ultan doctor` and the
+    ultan-search skill — use it to say "starting up, retry shortly" instead
+    of "broken" while the first start loads models for minutes."""
+    return home() / "daemon.state"
+
+
 def offset_state_path() -> Path:
     """Where the JSONL tailer persists its last-read offset across daemon
     restarts. See ``ingest._load_offset_state`` / ``_save_offset_state``.

daemon/agent_mem_daemon/paths.py

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import os
 import signal
@@ -343,3 +344,39 @@ def _killer():
     root = logging.getLogger()
     for h in list(root.handlers):
         root.removeHandler(h)
+
+
+# ── daemon.state lifecycle flag ───────────────────────────────────────
+
+
+def test_write_daemon_state_publishes_phase(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("AGENT_MEM_HOME", str(tmp_path))
+    daemon_main.write_daemon_state("warming")
+    raw = json.loads((tmp_path / "daemon.state").read_text(encoding="utf-8"))
+    assert raw["phase"] == "warming"
+    assert raw["pid"] == os.getpid()
+    assert raw["since"]
+
+
+def test_clear_daemon_state_removes_own_flag(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("AGENT_MEM_HOME", str(tmp_path))
+    daemon_main.write_daemon_state("ready")
+    daemon_main.clear_daemon_state()
+    assert not (tmp_path / "daemon.state").exists()
+
+
+def test_clear_daemon_state_leaves_other_daemons_flag(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A rejected duplicate exiting must not erase the live daemon's flag."""
+    monkeypatch.setenv("AGENT_MEM_HOME", str(tmp_path))
+    (tmp_path / "daemon.state").write_text(
+        json.dumps({"phase": "ready", "pid": os.getpid() + 99999, "since": "x"}),
+        encoding="utf-8",
+    )
+    daemon_main.clear_daemon_state()
+    assert (tmp_path / "daemon.state").exists()

daemon/tests/test_main.py

@@ -13,13 +13,77 @@
 import socket
 import struct
 import sys
+import time
 from pathlib import Path
 
 _HOME = Path(os.environ.get("AGENT_MEM_HOME") or (Path.home() / ".agent-mem"))
 SOCKET_PATH = _HOME / "priming.sock"
+# Default plugin data dir ("ultan" plugin @ "ultan" marketplace) — where the
+# background installer leaves its lock/pid while provisioning.
+_PLUGIN_DATA = Path(
+    os.environ.get("CLAUDE_PLUGIN_DATA")
+    or (Path.home() / ".claude" / "plugins" / "data" / "ultan-ultan")
+)
 PATH_CHARS_RE = re.compile(r"^[a-zA-Z0-9_./-]+$")
 RECV_TIMEOUT_S = 5.0
 
+_MSG_WARMING = (
+    "# Ultan is starting up — not broken\n\n"
+    "The daemon is warming (first start loads the retrieval models; allow 1-3\n"
+    "minutes). Retry this search shortly. Live status: run `ultan doctor`.\n\n"
+    "_Priming still works during warmup via the lexical fallback, so the\n"
+    "session is not memory-blind meanwhile._\n"
+)
+_MSG_INSTALLING = (
+    "# Ultan is still installing — not broken\n\n"
+    "The plugin's background installer is provisioning the runtime (torch +\n"
+    "models; minutes on a cold cache). Retry once it finishes. Progress:\n"
+    "run `ultan doctor`.\n"
+)
+_MSG_DOWN = (
+    "# Ultan daemon not running\n\n"
+    f"No socket at `{SOCKET_PATH}` and no startup in progress. With the\n"
+    "plugin installed the daemon lazy-starts on the next prompt — send any\n"
+    "message and retry. Details: run `ultan doctor`. (From a source\n"
+    "checkout: `uv run agent-mem-daemon -v`.)\n"
+)
+
+
+def _pid_alive(pid: object) -> bool:
+    try:
+        os.kill(int(pid), 0)  # type: ignore[arg-type]
+    except (TypeError, ValueError, ProcessLookupError):
+        return False
+    except PermissionError:
+        return True
+    except OSError:
+        return False
+    return True
+
+
+def _startup_message() -> str | None:
+    """A friendly explanation when the socket isn't answering, or None when
+    the daemon is genuinely down (not installing, not warming). Checks the
+    daemon's lifecycle flag first, then spawn/install breadcrumbs."""
+    # 1. The daemon's own phase flag (written between pid-acquire and exit).
+    try:
+        state = json.loads((_HOME / "daemon.state").read_text(encoding="utf-8"))
+        if _pid_alive(state.get("pid")):
+            return _MSG_WARMING  # alive but socket not serving yet (or restarting)
+    except (OSError, ValueError):
+        pass
+    # 2. A spawn was attempted moments ago (pre-flag window, or older daemon).
+    try:
+        age = time.time() - (_HOME / ".daemon-spawn-attempt").stat().st_mtime
+        if age < 300:
+            return _MSG_WARMING
+    except OSError:
+        pass
+    # 3. The plugin's background install is still running.
+    if (_PLUGIN_DATA / ".install.pid").exists() or (_PLUGIN_DATA / ".install.lock").exists():
+        return _MSG_INSTALLING
+    return None
+
 
 def _looks_like_path(s: str) -> bool:
     s = s.strip()
@@ -119,12 +183,7 @@ def main() -> int:
         mode = "search"
 
     if not SOCKET_PATH.exists():
-        print(
-            "# Ultan daemon not running\n\n"
-            f"Socket missing at `{SOCKET_PATH}`. Start the daemon:\n\n"
-            "```\nuv run agent-mem-daemon -v\n```\n",
-            file=sys.stderr,
-        )
+        print(_startup_message() or _MSG_DOWN, file=sys.stderr)
         return 1
 
     try:
@@ -134,6 +193,11 @@ def main() -> int:
         else:
             resp = _send_request({"op": "bm25_search", "query": arg, "k": 8})
             print(_render_search(resp, arg))
+    except (ConnectionError, OSError):
+        # Socket file present but not answering — stale socket or a daemon
+        # mid-restart. Same friendly triage as the missing-socket path.
+        print(_startup_message() or _MSG_DOWN, file=sys.stderr)
+        return 1
     except Exception as e:
         print(f"# Ultan skill failed\n\nError: {e!r}\n", file=sys.stderr)
         return 1

skills/ultan-search/search.py

@@ -0,0 +1,65 @@
+"""The ultan-search skill must triage a missing socket honestly: "starting
+up" / "still installing" / "down" — never a bare failure with from-source
+advice while the first start is loading models for minutes."""
+
+import importlib.util
+import json
+import os
+from pathlib import Path
+
+import pytest
+
+_SKILL = Path(__file__).resolve().parent.parent / "skills" / "ultan-search" / "search.py"
+
+
+@pytest.fixture()
+def skill(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    spec = importlib.util.spec_from_file_location("ultan_search_skill", _SKILL)
+    assert spec is not None and spec.loader is not None
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    home = tmp_path / "agent-mem"
+    home.mkdir()
+    data = tmp_path / "plugin-data"
+    data.mkdir()
+    monkeypatch.setattr(mod, "_HOME", home)
+    monkeypatch.setattr(mod, "_PLUGIN_DATA", data)
+    return mod
+
+
+def test_warming_when_state_flag_pid_alive(skill) -> None:
+    (skill._HOME / "daemon.state").write_text(
+        json.dumps({"phase": "warming", "pid": os.getpid(), "since": "x"}),
+        encoding="utf-8",
+    )
+    msg = skill._startup_message()
+    assert msg is not None
+    assert "starting up" in msg
+    assert "ultan doctor" in msg
+
+
+def test_warming_when_spawn_stamp_is_fresh(skill) -> None:
+    (skill._HOME / ".daemon-spawn-attempt").write_text("x", encoding="utf-8")
+    msg = skill._startup_message()
+    assert msg is not None
+    assert "starting up" in msg
+
+
+def test_installing_when_install_lock_present(skill) -> None:
+    (skill._PLUGIN_DATA / ".install.lock").write_text("token", encoding="utf-8")
+    msg = skill._startup_message()
+    assert msg is not None
+    assert "still installing" in msg
+
+
+def test_none_when_no_startup_signals(skill) -> None:
+    assert skill._startup_message() is None
+
+
+def test_dead_state_flag_does_not_report_warming(skill) -> None:
+    """A flag left by a crashed daemon (dead pid) must not claim warming."""
+    (skill._HOME / "daemon.state").write_text(
+        json.dumps({"phase": "warming", "pid": 99999999, "since": "x"}),
+        encoding="utf-8",
+    )
+    assert skill._startup_message() is None