@andares/pi-toolkits 0.11.2 → 0.12.0

package/src/features/quit-guard/guard.ts

@@ -0,0 +1,214 @@
+/**
+ * quit-guard feature — watchdog state machine.
+ *
+ * Deliberately dependency-free and side-effect free: it is fed stdout chunks
+ * and clock/scheduler callbacks, so every branch is unit-testable without a
+ * real TUI. The wiring to `process.stdout`, `SIGCONT`, the settings store and
+ * the actual `process.exit()` lives in index.ts.
+ *
+ * State machine:
+ *
+ *   idle ──teardown marker──▶ armed ──timer──▶ fired (onFire, then exit)
+ *     ▲                         │
+ *     └──restart marker / resume┘
+ *
+ * Invariants:
+ *  - arming is idempotent while armed (a teardown writes two markers);
+ *  - a TUI restart (mode switch, Ctrl+Z resume) always wins over the timer:
+ *    the fire callback is deferred one turn so a SIGCONT handler can disarm
+ *    first, and the token check drops a fire whose arm was superseded;
+ *  - disabled means inert: it never arms and disarms immediately.
+ */
+import { RESTART_MARKERS, TEARDOWN_MARKERS } from "./constants.js";
+
+/** A chunk as written to `process.stdout` (string, or Buffer/typed array). */
+export type GuardChunk = string | Uint8Array;
+
+export interface GuardFireInfo {
+	/** Clock value captured when the guard armed. */
+	startedAt: number;
+	/** Milliseconds spent armed (i.e. inside the shutdown phase). */
+	elapsedMs: number;
+	/** The teardown sequence that armed the guard. */
+	marker: string;
+}
+
+export interface QuitGuardDeps {
+	now?: () => number;
+	/** Start a timer; returns a handle handed back to `cancel`. */
+	schedule?: (run: () => void, ms: number) => unknown;
+	cancel?: (handle: unknown) => void;
+	/**
+	 * Defer one turn (setImmediate in production). Used so an expired timer
+	 * cannot beat a SIGCONT-driven disarm on the resume path.
+	 */
+	defer?: (run: () => void) => void;
+	/** Called exactly once per arm that reaches the deadline. */
+	onFire: (info: GuardFireInfo) => void;
+	/** Optional debug tap (PI_QUIT_GUARD_DEBUG); never affects behavior. */
+	onEvent?: (event: "arm" | "disarm" | "fire", detail: string) => void;
+	/**
+	 * Second, independent "is a TUI still running" signal, checked right before
+	 * firing. The teardown marker can also arrive inside rendered content (a
+	 * file or tool output that happens to contain the escape sequence); raw
+	 * mode cannot: the TUI leaves raw mode while tearing down. An active
+	 * session must never be killed by a stray marker.
+	 */
+	isTuiActive?: () => boolean;
+}
+
+/** First marker in `markers` that appears in the chunk, if any. */
+export function findMarker(
+	chunk: GuardChunk,
+	markers: readonly string[],
+): string | undefined {
+	if (typeof chunk === "string") {
+		if (chunk.length === 0) return undefined;
+		for (const marker of markers) {
+			if (chunk.includes(marker)) return marker;
+		}
+		return undefined;
+	}
+	if (!(chunk instanceof Uint8Array) || chunk.byteLength === 0) {
+		return undefined;
+	}
+	// Buffers are the normal case; other typed arrays are rare enough that a
+	// one-off copy is cheaper than a hand-rolled byte scan.
+	const buffer: Buffer = Buffer.isBuffer(chunk)
+		? chunk
+		: Buffer.from(chunk);
+	for (const marker of markers) {
+		if (buffer.includes(marker)) return marker;
+	}
+	return undefined;
+}
+
+export class QuitGuard {
+	private readonly now: () => number;
+	private readonly schedule: (run: () => void, ms: number) => unknown;
+	private readonly cancel: (handle: unknown) => void;
+	private readonly defer: (run: () => void) => void;
+	private readonly onFire: (info: GuardFireInfo) => void;
+	private readonly onEvent?: (event: "arm" | "disarm" | "fire", detail: string) => void;
+	private readonly isTuiActive?: () => boolean;
+
+	private enabled = true;
+	private timeoutMs: number;
+	private armed = false;
+	private handle: unknown;
+	private startedAt = 0;
+	private armedBy = "";
+	/** Incremented on every arm; a stale fire callback compares against it. */
+	private armToken = 0;
+
+	constructor(timeoutMs: number, deps: QuitGuardDeps) {
+		this.timeoutMs = timeoutMs;
+		// performance.now() is monotonic: a wall-clock jump (NTP, WSL VM
+		// resume, manual `date`) must not make a 10 s deadline look like 7 s,
+		// and `elapsedMs` in the hang report must measure real time spent.
+		this.now = deps.now ?? (() => performance.now());
+		this.schedule =
+			deps.schedule ??
+			((run, ms) => setTimeout(run, ms));
+		this.cancel =
+			deps.cancel ??
+			((handle) => clearTimeout(handle as ReturnType<typeof setTimeout>));
+		this.defer = deps.defer ?? ((run) => setImmediate(run));
+		this.onFire = deps.onFire;
+		this.onEvent = deps.onEvent;
+		this.isTuiActive = deps.isTuiActive;
+	}
+
+	/** Feed one stdout chunk. Converts markers into arm/disarm transitions. */
+	observeOutput(chunk: GuardChunk): void {
+		if (!this.enabled) return;
+		if (this.armed) {
+			// While armed the only meaningful signal is "the TUI came back".
+			if (findMarker(chunk, RESTART_MARKERS)) {
+				this.disarm();
+			}
+			return;
+		}
+		const marker = findMarker(chunk, TEARDOWN_MARKERS);
+		if (marker) this.arm(marker);
+	}
+
+	/** Process resumed from Ctrl+Z: the TUI is about to re-enter, never kill it. */
+	resume(): void {
+		this.disarm();
+	}
+
+	setEnabled(enabled: boolean): void {
+		this.enabled = enabled;
+		if (!enabled) this.disarm();
+	}
+
+	isEnabled(): boolean {
+		return this.enabled;
+	}
+
+	setTimeoutMs(timeoutMs: number): void {
+		this.timeoutMs = timeoutMs;
+	}
+
+	getTimeoutMs(): number {
+		return this.timeoutMs;
+	}
+
+	isArmed(): boolean {
+		return this.armed;
+	}
+
+	private arm(marker: string): void {
+		const token = (this.armToken += 1);
+		const startedAt = this.now();
+		this.armed = true;
+		this.startedAt = startedAt;
+		this.armedBy = marker;
+		this.onEvent?.("arm", `${marker} timeout=${this.timeoutMs}ms`);
+		this.handle = this.schedule(
+			() => this.onTimeout(token, startedAt),
+			this.timeoutMs,
+		);
+	}
+
+	private disarm(): void {
+		if (!this.armed) return;
+		this.armed = false;
+		this.onEvent?.("disarm", this.armedBy);
+		if (this.handle !== undefined) {
+			this.cancel(this.handle);
+			this.handle = undefined;
+		}
+	}
+
+	private onTimeout(token: number, startedAt: number): void {
+		this.handle = undefined;
+		// One turn later: a pending SIGCONT handler runs before this and can
+		// disarm, so resume always wins over an expired timer.
+		this.defer(() => {
+			if (!this.armed || this.armToken !== token || this.startedAt !== startedAt) {
+				return;
+			}
+			// The liveness probe is best-effort: if it throws (stdin already
+			// destroyed/replaced during teardown), treat the TUI as gone and
+			// still fire. A probe that can throw must not disable the watchdog.
+			let tuiActive = false;
+			try {
+				tuiActive = this.isTuiActive?.() === true;
+			} catch {
+				tuiActive = false;
+			}
+			if (tuiActive) {
+				// A live TUI means the marker was content, not a teardown.
+				this.armed = false;
+				this.onEvent?.("disarm", "tui-still-active");
+				return;
+			}
+			this.armed = false;
+			const elapsedMs = this.now() - startedAt;
+			this.onEvent?.("fire", `elapsed=${Math.round(elapsedMs)}ms`);
+			this.onFire({ startedAt, elapsedMs, marker: this.armedBy });
+		});
+	}
+}

package/src/features/quit-guard/index.ts

@@ -0,0 +1,292 @@
+/**
+ * quit-guard feature — quit watchdog (`/quit-guard`).
+ *
+ * Problem this solves: pi's interactive `/quit` stops the TUI, then awaits
+ * `runtimeHost.dispose()` — every extension's `session_shutdown` handler,
+ * one after another, with no timeout (upstream issue #9997; still unbounded
+ * in 1.5.0 and on main). One handler that never settles leaves the process
+ * alive at 0% CPU, the terminal already restored and the final frame
+ * reprinted: the screen looks like a frozen "exit screen" and the shell
+ * prompt never comes back. The `isShuttingDown` latch in the host then
+ * swallows Ctrl+C / SIGTERM / SIGHUP, so only `kill -9` helps.
+ *
+ * The guard watches the terminal instead of the event stream, because the
+ * event stream is exactly what can hang (a handler loaded before ours never
+ * returns, so our own `session_shutdown` never runs). It arms when the TUI
+ * writes its teardown sequence, and disarms if the TUI comes back (mode
+ * switch, Ctrl+Z suspend/resume). If the deadline passes while pi is still
+ * alive, it writes a JSON hang report and forces `process.exit(0)`.
+ *
+ * Wiring:
+ *  - one global `process.stdout.write` wrapper per process (survives
+ *    `/reload`; the shared state object is the source of truth, so a reload
+ *    only refreshes config/settings, never double-wraps);
+ *  - one `SIGCONT` listener: a resume must never be mistaken for a quit;
+ *  - `/quit-guard` — status / `on` / `off` / `<timeoutMs>`, persisted to
+ *    `<agent-dir>/pi-toolkits-quit-guard.json`.
+ *
+ * Settings file (all fields optional, defaults shown):
+ *   { "enabled": true, "timeoutMs": 10000, "report": true }
+ */
+import { appendFileSync, existsSync, readdirSync, writeSync } from "node:fs";
+import { join } from "node:path";
+import { getAgentDir } from "@earendil-works/pi-coding-agent";
+import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
+import {
+	MAX_TIMEOUT_MS,
+	MIN_TIMEOUT_MS,
+	QUIT_GUARD_FILE,
+	QUIT_GUARD_REPORT_DIR,
+} from "./constants.js";
+import { QuitGuard, type GuardFireInfo } from "./guard.js";
+import {
+	collectHangReport,
+	formatHangReportSummary,
+	writeHangReport,
+} from "./report.js";
+import {
+	QuitGuardStore,
+	normalizeTimeoutMs,
+	type QuitGuardConfig,
+} from "./store.js";
+
+/** Process-wide key; `/reload` re-imports the module but keeps this state. */
+const STATE_KEY = Symbol.for("@andares/pi-toolkits/quit-guard/state");
+
+interface QuitGuardState {
+	store: QuitGuardStore;
+	guard: QuitGuard;
+	config: QuitGuardConfig;
+	reportDir: string;
+	installed: boolean;
+}
+
+function sharedState(): QuitGuardState | undefined {
+	// SAFETY: `globalThis` is an untyped property bag; the symbol key is the
+	// private contract and this module is the only writer.
+	return (
+		globalThis as unknown as Record<symbol, QuitGuardState | undefined>
+	)[STATE_KEY];
+}
+
+function storeSharedState(state: QuitGuardState): void {
+	// SAFETY: same contract as `sharedState`; the value shape is checked by this
+	// module, not by the property bag.
+	(
+		globalThis as unknown as Record<symbol, QuitGuardState | undefined>
+	)[STATE_KEY] = state;
+}
+
+function fire(state: QuitGuardState, info: GuardFireInfo): void {
+	const timeoutMs = state.guard.getTimeoutMs();
+	const elapsedMs = Math.round(info.elapsedMs);
+	let reportPath: string | undefined;
+	let summary: string | undefined;
+	if (state.config.report) {
+		// Diagnostics only: a throwing report (deleted cwd, unreadable home, ...)
+		// must never keep the user from getting the prompt back.
+		try {
+			const report = collectHangReport({
+				timeoutMs,
+				elapsedMs,
+				marker: info.marker,
+			});
+			summary = formatHangReportSummary(report);
+			reportPath = writeHangReport(state.reportDir, report);
+		} catch {
+			summary = undefined;
+			reportPath = undefined;
+		}
+	}
+	const line = summary
+		? `[pi-toolkits/quit-guard] ${summary}${reportPath ? ` — report: ${reportPath}` : ""}; forcing exit.`
+		: `[pi-toolkits/quit-guard] no exit ${elapsedMs}ms after the TUI teardown started (limit ${timeoutMs}ms); forcing exit.`;
+	// stderr write is synchronous on POSIX; on failure the exit still happens.
+	try {
+		writeSync(2, `\n${line}\n`);
+	} catch {
+		// The terminal is gone or blocked; the report (if any) still exists.
+	}
+	// `process.exit(0)` is normally immediate, but every step of it is host
+	// code: an `exit` listener (or a patched `reallyExit`) can turn it into a
+	// no-op or throw — one such case was observed while testing this feature.
+	// The SIGKILL fallback guarantees the prompt comes back; the report is
+	// already on disk and session entries are appended as they happen, so
+	// nothing recoverable is lost.
+	const hardKill = setTimeout(() => process.kill(process.pid, "SIGKILL"), 500);
+	try {
+		process.exit(0);
+	} catch {
+		// Fall through to the hard kill.
+	}
+	clearTimeout(hardKill);
+	process.kill(process.pid, "SIGKILL");
+}
+
+function installProcessHooks(state: QuitGuardState): void {
+	if (state.installed) return;
+	state.installed = true;
+
+	const originalWrite = process.stdout.write;
+	// Capture the overload parameters once, then forward the exact same tuple:
+	// no casts, and every `write(chunk)` / `write(chunk, cb)` / `write(chunk,
+	// encoding, cb)` call site keeps working.
+	const wrapped = (...args: Parameters<typeof process.stdout.write>): boolean => {
+		try {
+			state.guard.observeOutput(args[0]);
+		} catch {
+			// Never let observation break terminal output.
+		}
+		return originalWrite.apply(process.stdout, args);
+	};
+	// SAFETY: the forwarded tuple is exactly what the original function receives
+	// (Node normalizes the (chunk, cb) form itself); the cast only restores
+	// Writable.write's two overloads, which `Parameters<...>` collapses into one.
+	process.stdout.write = wrapped as unknown as typeof process.stdout.write;
+
+	// Ctrl+Z resume: `handleCtrlZ()` calls ui.stop() (teardown marker), then
+	// SIGCONT calls ui.start() (restart marker). Both cancel; this listener
+	// additionally disarms synchronously so an already-expired timer cannot
+	// win the race against the resume. Windows cannot deliver SIGCONT at all.
+	if (process.platform !== "win32") {
+		process.on("SIGCONT", () => state.guard.resume());
+	}
+}
+
+function latestReportName(dir: string): string | undefined {
+	try {
+		if (!existsSync(dir)) return undefined;
+		const names = readdirSync(dir)
+			.filter(
+				(name) => name.startsWith("quit-hang-") && name.endsWith(".json"),
+			)
+			.sort();
+		return names.length > 0 ? names[names.length - 1] : undefined;
+	} catch {
+		return undefined;
+	}
+}
+
+function statusText(state: QuitGuardState): string {
+	const config = state.config;
+	const last = latestReportName(state.reportDir);
+	const parts = [
+		`quit guard: ${config.enabled ? "on" : "off"}`,
+		`timeout ${config.timeoutMs}ms`,
+		`report ${config.report ? "on" : "off"}`,
+		`last hang: ${last ?? "none"}`,
+	];
+	return parts.join(" · ");
+}
+
+function buildState(agentDir: string): QuitGuardState {
+	const store = new QuitGuardStore(join(agentDir, QUIT_GUARD_FILE));
+	const config = store.get();
+	let state: QuitGuardState | undefined;
+	// PI_QUIT_GUARD_DEBUG points at a file that receives one line per
+	// arm/disarm/fire. It exists because the failure being debugged happens
+	// after the TUI is gone: there is no other channel left to observe.
+	const debugPath = process.env.PI_QUIT_GUARD_DEBUG;
+	const guard = new QuitGuard(config.timeoutMs, {
+		onFire: (info) => {
+			if (state) fire(state, info);
+		},
+		// Raw mode is the independent "TUI still running" signal; see guard.ts.
+		isTuiActive: () => process.stdin.isRaw === true,
+		onEvent: debugPath
+			? (event, detail) => {
+					try {
+						appendFileSync(
+							debugPath,
+							`${new Date().toISOString()} ${event} ${detail}\n`,
+						);
+					} catch {
+						// Debug output must never affect the guard.
+					}
+				}
+			: undefined,
+	});
+	state = {
+		store,
+		guard,
+		config,
+		reportDir: join(agentDir, QUIT_GUARD_REPORT_DIR),
+		installed: false,
+	};
+	return state;
+}
+
+export interface QuitGuardOptions {
+	/** pi config directory; defaults to `getAgentDir()` (tests inject a temp dir). */
+	readonly agentDir?: string;
+}
+
+export function registerQuitGuard(
+	pi: ExtensionAPI,
+	options: QuitGuardOptions = {},
+): void {
+	const agentDir = options.agentDir ?? getAgentDir();
+	const state = sharedState() ?? buildState(agentDir);
+	storeSharedState(state);
+
+	// A reload rebuilds the extension instance: refresh settings from disk and
+	// start unarmed. Any watchdog armed by the previous instance belonged to a
+	// teardown that has already been replaced.
+	state.store = new QuitGuardStore(join(agentDir, QUIT_GUARD_FILE));
+	state.config = state.store.get();
+	state.reportDir = join(agentDir, QUIT_GUARD_REPORT_DIR);
+	state.guard.resume();
+	state.guard.setEnabled(state.config.enabled);
+	state.guard.setTimeoutMs(state.config.timeoutMs);
+
+	installProcessHooks(state);
+	if (process.env.PI_QUIT_GUARD_DEBUG) {
+		try {
+			appendFileSync(
+				process.env.PI_QUIT_GUARD_DEBUG,
+				`${new Date().toISOString()} register agentDir=${agentDir} enabled=${state.config.enabled} timeout=${state.config.timeoutMs}\n`,
+			);
+		} catch {
+			// ignore
+		}
+	}
+
+	pi.registerCommand("quit-guard", {
+		description:
+			"Quit watchdog: force-exit when pi's shutdown hangs (/quit-guard [on|off|<timeoutMs>])",
+		handler: async (args: string, ctx: ExtensionCommandContext) => {
+			const raw = args.trim();
+			if (raw.length === 0) {
+				ctx.ui.notify(statusText(state), "info");
+				return;
+			}
+			const token = raw.toLowerCase();
+			if (token === "on" || token === "off") {
+				const enabled = token === "on";
+				state.config = state.store.update({ enabled });
+				state.guard.setEnabled(enabled);
+				ctx.ui.notify(
+					enabled
+						? `quit guard ON — force exits ${state.config.timeoutMs}ms after a stuck shutdown.`
+						: "quit guard OFF — a stuck shutdown will keep hanging until kill -9.",
+					"info",
+				);
+				return;
+			}
+			if (/^[0-9]+$/.test(token)) {
+				const timeoutMs = normalizeTimeoutMs(Number(token));
+				state.config = state.store.update({ timeoutMs });
+				state.guard.setTimeoutMs(timeoutMs);
+				ctx.ui.notify(
+					`quit guard timeout ${timeoutMs}ms (allowed range ${MIN_TIMEOUT_MS}–${MAX_TIMEOUT_MS}ms).`,
+					"info",
+				);
+				return;
+			}
+			ctx.ui.notify(
+				"Usage: /quit-guard (status) | /quit-guard on|off | /quit-guard <timeoutMs>",
+				"error",
+			);
+		},
+	});
+}