ai-agent/backend/hooks.py

"""Conversation hooks: follow-up prompts that fire after a conversation finishes.

A hook is the post-task twin of a cron job (``cron.py``): instead of a schedule
it has an **event** (``conversation-done``) and a **delay**, and instead of
spawning a *new* session it **resumes the conversation that just finished** with
the content of its prompt file. The flagship use is the seeded ``post-task``
hook — ten seconds after an agent says ``DONE``, the same session is woken up to
save memories, update the README/CLAUDE.md, copy artefacts into
``.ai/artefacts/``, tear down worktrees and update the changelog.

Storage is a single JSON sidecar (``/data/hooks.json``) holding both the hooks
and the **pending queue**, so a firing survives a backend restart (or the
blue-green cutover where two backends briefly share ``/data``):

    { "hooks": { "<hookId>": {
        "id":            "post-task",
        "name":          "Post-task follow-up",
        "event":         "conversation-done",
        "delaySeconds":  10,
        "promptFile":    ".claude/agents/hooks/post-task.md",
        "enabled":       true,
        "skipCron":      true,      # don't follow up cron-spawned sessions
        "skipSubagents": true,      # …nor subagent (sidechain) conversations
        "maxContextK":   0,         # 0 = no transcript-size cap
        "createdAt":     iso,
        "lastStatus":    "ok" | "error: …" | "skipped: …",
        "history":       { iso: sessionId, … },
      } },
      "pending": [ {"hookId": …, "sessionId": …, "dueAt": iso, "retries": 0}, … ]
    }

**The loop guard is the point.** The follow-up run ends with ``COMPLETED``
(not ``DONE`` — a distinct marker so the viewer can tell "the task finished"
apart from "the task finished AND was tidied up after") and also exits, which
enqueues the hook again. Each firing therefore stamps
``meta.hooks[hookId] = iso`` on the session (``meta.py`` merges that mapping
rather than replacing it), and both the enqueue and the fire-time preconditions
refuse an already-stamped (hook, session) pair — **once per hook per session**,
full stop.

The scheduler is a daemon thread waking every ~10s. A due entry is *claimed*
(removed from the queue, atomically, re-reading the file first) before it fires,
so a parallel backend can't double-resume; a precondition that is merely not-yet
(the run is still live, the transcript hasn't synced) re-enqueues the entry a
minute later, up to ``MAX_RETRIES``.
"""

import json
import os
import pathlib
import threading
import time
from datetime import datetime, timedelta

import cron as cron_mod

# Prompt files live in their own subdirectory of the cron prompt dir, so a hook
# prompt is never mistaken for a schedulable agent and vice versa.
PROMPT_DIR = f"{cron_mod.PROMPT_DIR}/hooks"

# The only event v1 fires on. Kept as a field (and validated) so adding
# "conversation-started" or "commit" later doesn't change the stored shape.
EVENTS = ("conversation-done",)

# A precondition that may still come true (run still live, transcript not
# synced yet) postpones the entry by this many seconds, this many times.
RETRY_DELAY_S = 60
MAX_RETRIES = 10

# Re-exported so main.py has one import for the prompt-file helpers.
strip_frontmatter = cron_mod.strip_frontmatter
run_config = cron_mod.run_config
frontmatter_block = cron_mod.frontmatter_block
slugify = cron_mod.slugify


def _now_iso() -> str:
    return datetime.now().astimezone().strftime("%Y-%m-%dT%H:%M:%S%z")


def _parse_iso(s: str) -> datetime | None:
    try:
        dt = datetime.fromisoformat(s)
    except (TypeError, ValueError):
        return None
    return dt if dt.tzinfo else dt.astimezone()


def valid_prompt_file(rel: str) -> str:
    """Normalize + validate a hook's prompt-file path: a ``.md`` directly under
    ``.claude/agents/hooks/``. Raises ValueError otherwise."""
    rel = (rel or "").strip().lstrip("/")
    p = pathlib.PurePosixPath(rel)
    if ".." in p.parts:
        raise ValueError("promptFile must not contain '..'")
    if p.suffix.lower() != ".md":
        raise ValueError("promptFile must be a .md file")
    if str(p.parent) != PROMPT_DIR:
        raise ValueError(f"promptFile must live in {PROMPT_DIR}/")
    return str(p)


def valid_event(event: str) -> str:
    ev = (event or "").strip() or EVENTS[0]
    if ev not in EVENTS:
        raise ValueError(f"unknown event {ev!r}")
    return ev


def valid_delay(v) -> int:
    try:
        d = int(v)
    except (TypeError, ValueError):
        raise ValueError("delaySeconds must be a whole number of seconds")
    if not 0 <= d <= 24 * 60 * 60:
        raise ValueError("delaySeconds must be between 0 and 86400")
    return d


# ── the seeded hook ──────────────────────────────────────────────────────────
# Enabled on first run: every finished task gets the follow-up pass. Disable it
# in Settings → Hooks if a stretch of work shouldn't pay for one.
DEFAULT_HOOKS = {
    "post-task": {
        "id": "post-task",
        "name": "Post-task follow-up",
        "event": "conversation-done",
        "delaySeconds": 10,
        "promptFile": f"{PROMPT_DIR}/post-task.md",
        "enabled": True,
        "skipCron": True,
        "skipSubagents": True,
        "maxContextK": 0,
        "createdAt": None,   # stamped at seed time
        "lastStatus": None,
        "history": {},
    },
}

# The prompt the seeded hook resumes a finished conversation with. Written to
# the workspace on first run (only if the file doesn't exist), so a fresh
# install has a working post-task pass without any setup.
DEFAULT_POST_TASK_PROMPT = """---
name: post-task
description: Follow-up pass ten seconds after a task finishes — save memories, update docs, collect artefacts, clean up worktrees and processes, update the changelog.
---

Follow-up pass on the task you just finished. Work through this list, do only
what genuinely applies, and keep it short — this is tidying, not new work.

1. **Memory.** Was anything here non-obvious — a gotcha, a wrong turn, a
   constraint that isn't visible in the code? Save it to your memory directory
   (one fact per file, with frontmatter) and add its line to `MEMORY.md`. Skip
   anything the repo already records.
2. **Docs.** If the change altered how something is used, built or deployed,
   update the affected `README.md` / `CLAUDE.md` / `GOAL.md`. No changelog-style
   prose in CLAUDE.md — describe the current state.
3. **Artefacts.** Copy anything worth keeping (screenshots, diagrams, generated
   samples, reports) into the canonical project's `.ai/artefacts/<sessionId>/`
   folder so it renders inline in the conversation viewer. Leave scratch files
   behind.
4. **Cleanup.** Remove worktrees you created once their branch is merged
   (`scripts/new-worktree.sh -r <path>`), and kill background processes, dev
   servers or standby containers you started and no longer need.
5. **Changelog.** If the project keeps one, add the user-facing entry for what
   shipped.

Do **not** re-notify (no `notify-done`) and do not start new feature work. If
nothing on the list applies, say so in one line.

End your reply with `COMPLETED` on its own line (not `DONE` — this pass has its
own marker so the viewer can show that the task was tidied up, not just finished).
"""


class HookStore:
    """The JSON store — hooks + the persistent pending queue.

    Same discipline as ``CronStore``/``MetaStore``: one lock, atomic saves and
    an mtime-based reload so the other container during a deploy cutover (or an
    out-of-band edit) is picked up before any decision is taken."""

    def __init__(self, path: str):
        self.path = pathlib.Path(path)
        self._lock = threading.Lock()
        self._mtime: float | None = None
        self._hooks: dict[str, dict] = {}
        self._pending: list[dict] = []
        self._load()
        if self._mtime is None:  # no file yet → seed the defaults
            with self._lock:
                self._hooks = json.loads(json.dumps(DEFAULT_HOOKS))
                for h in self._hooks.values():
                    h["createdAt"] = _now_iso()
                self._pending = []
                self._save()

    # ── persistence ──────────────────────────────────────────────────────────
    def _load(self) -> None:
        try:
            raw = json.loads(self.path.read_text(encoding="utf-8")) or {}
            self._hooks = raw.get("hooks") or {}
            self._pending = raw.get("pending") or []
            self._mtime = self.path.stat().st_mtime
            self._migrate_delay_unit()
        except (OSError, ValueError):
            self._hooks = {}
            self._pending = []
            self._mtime = None

    def _migrate_delay_unit(self) -> None:
        """Pre-existing stores keep the old ``delayMinutes`` field. Convert it
        in place (minutes → seconds) on load so an already-configured hook's
        delay survives the rename instead of silently resetting to 0."""
        for hook in self._hooks.values():
            if "delaySeconds" not in hook and "delayMinutes" in hook:
                hook["delaySeconds"] = int(hook.pop("delayMinutes") or 0) * 60

    def _save(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".json.tmp")
        tmp.write_text(json.dumps({"hooks": self._hooks,
                                   "pending": self._pending},
                                  indent=2, sort_keys=True), encoding="utf-8")
        os.replace(tmp, self.path)
        try:
            self._mtime = self.path.stat().st_mtime
        except OSError:
            pass

    def _reload_if_changed(self) -> None:
        try:
            m = self.path.stat().st_mtime
        except OSError:
            return
        if m != self._mtime:
            self._load()

    # ── hooks CRUD ───────────────────────────────────────────────────────────
    def list(self) -> "list[dict]":
        with self._lock:
            self._reload_if_changed()
            return json.loads(json.dumps(list(self._hooks.values())))

    def get(self, hook_id: str) -> dict | None:
        with self._lock:
            self._reload_if_changed()
            h = self._hooks.get(hook_id)
            return json.loads(json.dumps(h)) if h else None

    def create(self, name: str, *, event: str = EVENTS[0],
               delay_seconds: int = 10, prompt_file: str,
               enabled: bool = True, skip_cron: bool = True,
               skip_subagents: bool = True, max_context_k: int = 0) -> dict:
        event = valid_event(event)
        delay_seconds = valid_delay(delay_seconds)
        prompt_file = valid_prompt_file(prompt_file)
        with self._lock:
            self._reload_if_changed()
            base = slugify(name)
            hook_id, n = base, 2
            while hook_id in self._hooks:
                hook_id, n = f"{base}-{n}", n + 1
            hook = {"id": hook_id, "name": name or hook_id, "event": event,
                    "delaySeconds": delay_seconds, "promptFile": prompt_file,
                    "enabled": bool(enabled), "skipCron": bool(skip_cron),
                    "skipSubagents": bool(skip_subagents),
                    "maxContextK": max(0, int(max_context_k or 0)),
                    "createdAt": _now_iso(), "lastStatus": None, "history": {}}
            self._hooks[hook_id] = hook
            self._save()
            return json.loads(json.dumps(hook))

    def update(self, hook_id: str, patch: dict) -> dict | None:
        if patch.get("event") is not None:
            patch["event"] = valid_event(patch["event"])
        if patch.get("delaySeconds") is not None:
            patch["delaySeconds"] = valid_delay(patch["delaySeconds"])
        if patch.get("promptFile") is not None:
            patch["promptFile"] = valid_prompt_file(patch["promptFile"])
        if patch.get("maxContextK") is not None:
            patch["maxContextK"] = max(0, int(patch["maxContextK"]))
        with self._lock:
            self._reload_if_changed()
            hook = self._hooks.get(hook_id)
            if not hook:
                return None
            for k in ("name", "event", "delaySeconds", "promptFile", "enabled",
                      "skipCron", "skipSubagents", "maxContextK"):
                if patch.get(k) is not None:
                    hook[k] = patch[k]
            self._save()
            return json.loads(json.dumps(hook))

    def delete(self, hook_id: str) -> bool:
        with self._lock:
            self._reload_if_changed()
            if hook_id not in self._hooks:
                return False
            del self._hooks[hook_id]
            # Drop anything queued for it — a deleted hook must not fire.
            self._pending = [p for p in self._pending
                             if p.get("hookId") != hook_id]
            self._save()
            return True

    def record_run(self, hook_id: str, session_id: str,
                   status: str = "ok") -> None:
        """Append ``{now: sessionId}`` to the hook's history and stamp the
        outcome (what the card's status chip shows)."""
        with self._lock:
            self._reload_if_changed()
            hook = self._hooks.get(hook_id)
            if not hook:
                return
            if session_id:
                hist = dict(hook.get("history") or {})
                hist[_now_iso()] = session_id
                if len(hist) > 200:      # keep the most recent 200 firings
                    for k in sorted(hist)[:len(hist) - 200]:
                        del hist[k]
                hook["history"] = hist
            hook["lastStatus"] = status
            self._save()

    # ── the pending queue ────────────────────────────────────────────────────
    # NB: annotations below are strings — the ``list`` method above shadows the
    # builtin inside the class body, so a bare ``list[dict]`` would evaluate
    # against it and raise at import time.
    def pending(self) -> "list[dict]":
        with self._lock:
            self._reload_if_changed()
            return json.loads(json.dumps(self._pending))

    def is_pending(self, hook_id: str, session_id: str) -> bool:
        with self._lock:
            self._reload_if_changed()
            return any(p.get("hookId") == hook_id
                       and p.get("sessionId") == session_id
                       for p in self._pending)

    def enqueue(self, hook_id: str, session_id: str, delay_seconds: int,
                retries: int = 0) -> dict | None:
        """Queue one firing, due ``delay_seconds`` from now. Returns the entry,
        or None when this (hook, session) pair is already queued — the caller
        checks the ``meta.hooks`` stamp for pairs that already *fired*."""
        due = (datetime.now().astimezone()
               + timedelta(seconds=max(0, delay_seconds)))
        with self._lock:
            self._reload_if_changed()
            if any(p.get("hookId") == hook_id
                   and p.get("sessionId") == session_id
                   for p in self._pending):
                return None
            entry = {"hookId": hook_id, "sessionId": session_id,
                     "dueAt": due.strftime("%Y-%m-%dT%H:%M:%S%z"),
                     "queuedAt": _now_iso(), "retries": int(retries)}
            self._pending.append(entry)
            self._save()
            return json.loads(json.dumps(entry))

    def claim(self, hook_id: str, session_id: str) -> bool:
        """Atomically take this entry off the queue. False = someone else got
        it first (the other backend during a deploy cutover)."""
        with self._lock:
            self._reload_if_changed()
            keep = [p for p in self._pending
                    if not (p.get("hookId") == hook_id
                            and p.get("sessionId") == session_id)]
            if len(keep) == len(self._pending):
                return False
            self._pending = keep
            self._save()
            return True

    def due_entries(self, now: datetime | None = None) -> "list[dict]":
        """Queued firings whose ``dueAt`` has passed (an unparseable one is
        treated as due, so a corrupt stamp can't strand an entry forever)."""
        now = now or datetime.now().astimezone()
        out = []
        for p in self.pending():
            due = _parse_iso(p.get("dueAt") or "")
            if due is None or due <= now:
                out.append(p)
        return out


class HookScheduler(threading.Thread):
    """Fires due queue entries. ``fire(hook, sessionId, entry)`` is injected by
    main.py (it needs the resume path, the summaries and the meta store)."""

    def __init__(self, store: HookStore, fire, interval: float = 10.0):
        super().__init__(daemon=True, name="hook-scheduler")
        self.store = store
        self.fire = fire
        self.interval = interval

    def run(self) -> None:
        while True:
            time.sleep(self.interval)
            try:
                self._tick()
            except Exception:
                pass

    def _tick(self) -> None:
        for entry in self.store.due_entries():
            hook_id = entry.get("hookId") or ""
            sid = entry.get("sessionId") or ""
            hook = self.store.get(hook_id)
            if not hook or not hook.get("enabled"):
                # Deleted or switched off while queued: drop it silently.
                self.store.claim(hook_id, sid)
                continue
            if not self.store.claim(hook_id, sid):
                continue  # a parallel backend took it
            try:
                self.fire(hook, sid, entry)
            except Exception as e:  # never kill the loop on one bad firing
                self.store.record_run(hook_id, "", status=f"error: {e}")