pi-wakeup 0.2.0 → 0.2.2

package/README.md

@@ -27,7 +27,7 @@ Every tmux / UI call is best-effort: a failure in a timer callback is caught and
 
 ## Design
 
-- **Session-local.** A wakeup is an in-process `setTimeout`, not a cron job. It lives and dies with the session. This is deliberate: subagents and pacing in pi are in-process (the muster-teammate shape is shelved), and cross-session scheduling is a different product (`cron` + `pi -p`), not this tool.
+- **Session-local.** A wakeup is an in-process `setTimeout`, not a cron job. It lives and dies with the session, and survives `/reload`: the pending wakeups are handed to the reloaded extension, which re-arms each one for the time it had left. This is deliberate: subagents and pacing in pi are in-process (the muster-teammate shape is shelved), and cross-session scheduling is a different product (`cron` + `pi -p`), not this tool.
 - **Busy-gated delivery.** `agent_start` / `agent_settled` track whether a turn is running; delivery picks `steer` (idle) or `followUp` (busy) accordingly, so a wakeup never interrupts a turn in flight.
 - **Best-effort.** Every handler and the timer callback swallow-and-log on error — a missed wakeup must never take the session process down. The pure decisions (clamp, delivery mode, message shape) live in `src/schedule.ts` and are unit-tested with zero mocking; `src/index.ts` is the thin pi wiring, in the same split the guardrails package uses.
 - **Persistence.** Each schedule / fire / cancel is written with `pi.appendEntry` so the session transcript records the wakeup lifecycle.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-wakeup",
-  "version": "0.2.0",
+  "version": "0.2.2",
   "description": "schedule_wakeup / cancel_wakeup for the pi coding agent: session-local self-pacing, the pi analog of Claude Code's ScheduleWakeup dynamic mode.",
   "type": "module",
   "license": "MIT",

package/src/index.test.ts

@@ -21,6 +21,9 @@ function harness(
     // Make a tmux call or setWidget throw, to prove containment.
     tmuxThrows?: boolean;
     widgetThrows?: boolean;
+    // The session_start event (a reload re-arms handed-over wakeups) and clock.
+    startEvent?: unknown;
+    nowMs?: number;
   } = {},
 ) {
   const idle = opts.idle ?? true;
@@ -40,6 +43,7 @@ function harness(
       return idle;
     },
     ui: {
+      theme: { fg: (color: string, text: string) => `<${color}>${text}` },
       notify: (text: string, level: string) => notifies.push({ text, level }),
       setWidget: (key: string, lines: string[] | undefined) => {
         if (opts.widgetThrows) throw new Error("stale widget ctx");
@@ -73,7 +77,7 @@ function harness(
     clearTimer: (handle: unknown) => {
       (handle as Timer).cancelled = true;
     },
-    now: () => 1000,
+    now: () => opts.nowMs ?? 1000,
     setInterval: (fn: () => void, ms: number) => {
       const i: Interval = { fn, ms, cancelled: false };
       intervals.push(i);
@@ -97,7 +101,7 @@ function harness(
 
   createWakeup(pi, deps);
   // Prime the captured ctx as a real session would via session_start.
-  handlers.get("session_start")?.(undefined, ctx);
+  handlers.get("session_start")?.(opts.startEvent, ctx);
 
   const schedule = (params: any) =>
     tools.get("schedule_wakeup").execute("call", params, undefined, undefined, ctx);
@@ -242,6 +246,38 @@ test("session_shutdown clears every timer and no wakeup can still deliver", asyn
   assert.equal(h.sent.length, 0);
 });
 
+test("a /reload hands pending wakeups to the new copy, which re-arms them for the time left", async () => {
+  const before = harness();
+  await before.schedule({ delay_seconds: 120, prompt: "a", reason: "ci" });
+  await before.schedule({ delay_seconds: 200, prompt: "b" });
+  await before.cancel({ id: "wake-2" });
+  before.handlers.get("session_shutdown")!({ reason: "reload" });
+  assert.ok(before.timers.every((t) => t.cancelled));
+
+  const after = harness({ startEvent: { reason: "reload" }, nowMs: 31_000 });
+  assert.equal(after.timers.length, 1);
+  assert.equal(after.timers[0].ms, 90_000, "120s scheduled at t=1s, reloaded at t=31s");
+  after.fire(after.timers[0]);
+  assert.equal(after.sent.length, 1);
+  assert.match(after.sent[0].msg.content, /a/);
+  await after.schedule({ delay_seconds: 60, prompt: "c" });
+  assert.match(after.entries.at(-1)!.data.id, /wake-3/, "ids continue, not reused");
+});
+
+test("a quit hands nothing over, and a fresh start re-arms nothing", async () => {
+  const before = harness();
+  await before.schedule({ delay_seconds: 120, prompt: "a" });
+  before.handlers.get("session_shutdown")!({ reason: "quit" });
+  const after = harness({ startEvent: { reason: "reload" } });
+  assert.equal(after.timers.length, 0);
+
+  const reloaded = harness();
+  await reloaded.schedule({ delay_seconds: 120, prompt: "a" });
+  reloaded.handlers.get("session_shutdown")!({ reason: "reload" });
+  const fresh = harness({ startEvent: { reason: "startup" } });
+  assert.equal(fresh.timers.length, 0);
+});
+
 // --- Countdown display -----------------------------------------------------
 
 test("widget mode: scheduling shows a line above the editor and arms a refresh", async () => {
@@ -252,6 +288,7 @@ test("widget mode: scheduling shows a line above the editor and arms a refresh",
   assert.equal(last.lines!.length, 1);
   assert.match(last.lines![0], /2m/); // 120s from now()=1000 -> 2m remaining
   assert.match(last.lines![0], /CI/);
+  assert.match(last.lines![0], /^<warning>⏰ /); // styled like the tmux bar
   assert.equal(h.intervals.length, 1);
   assert.equal(h.intervals[0].cancelled, false);
   // The tmux runner is never touched in widget mode.

package/src/index.ts

@@ -18,7 +18,9 @@
 // pi.appendEntry so the session transcript records the wakeup lifecycle.
 // Re-arming a pending wakeup across a session RESUME is a documented follow-up
 // (the extension API exposes no entry read-back). Within a live session,
-// wakeups fire.
+// wakeups fire, including across /reload: a reload shuts this extension down
+// and loads a fresh copy in the same process, so shutdown hands the pending
+// wakeups to the new copy through a global, which re-arms them on start.
 
 import { execFileSync } from "node:child_process";
 import {
@@ -33,6 +35,11 @@ import {
 // how long is left; it's the same value persisted in the wakeup_scheduled entry.
 type Pending = { prompt: string; reason: string; fireAtMs: number };
 
+// Where a shutdown for /reload leaves the pending wakeups (and the id counter,
+// so new ids don't collide) for the freshly loaded copy to pick up.
+const HANDOFF_KEY = Symbol.for("pi-wakeup:reload-handoff");
+type Handoff = { counter: number; pending: Array<[string, Pending]> };
+
 // How often the countdown re-renders while anything is pending. The interval is
 // unref'd so it never holds the process open on its own.
 const REFRESH_MS = 30_000;
@@ -194,6 +201,19 @@ export function createWakeup(pi: any, deps: Deps = {}): void {
     }
   }
 
+  // The widget line matches the tmux bar: a clock and the theme's warning
+  // (yellow) colour, so it stands out above the editor. The theme read is
+  // guarded like every other ctx.ui access; without one the line is plain.
+  function widgetLine(text: string): string {
+    const line = `⏰ ${text}`;
+    try {
+      const theme = ctx?.ui?.theme;
+      return theme && typeof theme.fg === "function" ? theme.fg("warning", line) : line;
+    } catch {
+      return line;
+    }
+  }
+
   // Render the current pending list to the chosen target, or clear it when
   // nothing is pending. Runs inside timer/interval callbacks, so every tmux and
   // ctx.ui call is guarded — a throw here must never escape (best-effort).
@@ -214,7 +234,7 @@ export function createWakeup(pi: any, deps: Deps = {}): void {
       // Widget mode. Guard ctx.ui: a ctx captured at session_start goes stale
       // after /reload or a session switch, and setWidget would then throw.
       if (ctx?.ui?.setWidget) {
-        ctx.ui.setWidget("wakeup", text === undefined ? undefined : [text]);
+        ctx.ui.setWidget("wakeup", text === undefined ? undefined : [widgetLine(text)]);
       }
     } catch (error) {
       log(`display refresh failed: ${String(error)}`);
@@ -303,14 +323,27 @@ export function createWakeup(pi: any, deps: Deps = {}): void {
     }
   }
 
-  pi.on("session_start", (_event: unknown, sessionCtx: any) => {
+  pi.on("session_start", (event: any, sessionCtx: any) => {
     ctx = sessionCtx;
     // Decide the display target once, now, from the env + tmux status strings.
     decideDisplay();
-    // Nothing is pending at a fresh start, so this just clears any stale display.
-    refresh();
+    const store = globalThis as Record<symbol, Handoff | undefined>;
+    const handoff = event?.reason === "reload" ? store[HANDOFF_KEY] : undefined;
+    delete store[HANDOFF_KEY];
+    if (handoff) {
+      counter = handoff.counter;
+      for (const [id, w] of handoff.pending) {
+        pending.set(id, w);
+        timers.set(id, setTimer(() => deliver(id), Math.max(0, w.fireAtMs - now())));
+      }
+    }
+    // Shows the re-armed wakeups, or clears any stale display.
+    syncDisplay();
   });
-  pi.on("session_shutdown", () => {
+  pi.on("session_shutdown", (event: any) => {
+    if (event?.reason === "reload") {
+      (globalThis as Record<symbol, Handoff | undefined>)[HANDOFF_KEY] = { counter, pending: [...pending] };
+    }
     for (const handle of timers.values()) clearTimer(handle);
     timers.clear();
     pending.clear();