From b1dcfa0651a10958289345277c904bc7e66ce972 Mon Sep 17 00:00:00 2001 From: Trevin Chow Date: Thu, 29 Jan 2026 09:17:31 -0800 Subject: [PATCH] feat(hooks): add hook-run-registry core module --- src/gateway/hook-run-registry.test.ts | 161 ++++++++++++++++++++++++++ src/gateway/hook-run-registry.ts | 127 ++++++++++++++++++++ 2 files changed, 288 insertions(+) create mode 100644 src/gateway/hook-run-registry.test.ts create mode 100644 src/gateway/hook-run-registry.ts diff --git a/src/gateway/hook-run-registry.test.ts b/src/gateway/hook-run-registry.test.ts new file mode 100644 index 000000000..94bf79403 --- /dev/null +++ b/src/gateway/hook-run-registry.test.ts @@ -0,0 +1,161 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("./hook-run-registry.store.js", () => ({ + loadHookRunRegistryFromDisk: vi.fn(() => new Map()), + saveHookRunRegistryToDisk: vi.fn(), +})); + +vi.mock("./call.js", () => ({ + callGateway: vi.fn(), +})); + +describe("hook-run-registry", () => { + beforeEach(() => { + vi.resetModules(); + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + vi.resetAllMocks(); + }); + + describe("registerHookRun", () => { + it("registers a new hook run with cleanup=delete", async () => { + const { registerHookRun, getHookRun, clearHookRuns } = await import("./hook-run-registry.js"); + clearHookRuns(); + + registerHookRun({ + runId: "run-1", + sessionKey: "hook:test:1", + jobName: "Test Hook", + cleanup: "delete", + cleanupDelayMinutes: 0, + }); + + const run = getHookRun("run-1"); + expect(run).toBeDefined(); + expect(run?.sessionKey).toBe("hook:test:1"); + expect(run?.cleanup).toBe("delete"); + expect(run?.cleanupHandled).toBe(false); + + clearHookRuns(); + }); + + it("does not register runs with cleanup=keep", async () => { + const { registerHookRun, getHookRun, clearHookRuns } = await import("./hook-run-registry.js"); + clearHookRuns(); + + registerHookRun({ + runId: "run-1", + sessionKey: "hook:test:1", + jobName: "Test Hook", + cleanup: "keep", + cleanupDelayMinutes: 0, + }); + + const run = getHookRun("run-1"); + expect(run).toBeUndefined(); + + clearHookRuns(); + }); + + it("does not register runs with cleanup=undefined", async () => { + const { registerHookRun, getHookRun, clearHookRuns } = await import("./hook-run-registry.js"); + clearHookRuns(); + + registerHookRun({ + runId: "run-1", + sessionKey: "hook:test:1", + jobName: "Test Hook", + cleanup: undefined, + cleanupDelayMinutes: undefined, + }); + + const run = getHookRun("run-1"); + expect(run).toBeUndefined(); + + clearHookRuns(); + }); + }); + + describe("markHookRunComplete", () => { + it("sets endedAt, cleanupAtMs, and cleanupHandled", async () => { + const { registerHookRun, markHookRunComplete, getHookRun, clearHookRuns } = + await import("./hook-run-registry.js"); + clearHookRuns(); + + registerHookRun({ + runId: "run-1", + sessionKey: "hook:test:1", + jobName: "Test Hook", + cleanup: "delete", + cleanupDelayMinutes: 5, + }); + + const now = Date.now(); + vi.setSystemTime(now); + markHookRunComplete("run-1"); + + const run = getHookRun("run-1"); + expect(run?.endedAt).toBe(now); + expect(run?.cleanupAtMs).toBe(now + 5 * 60 * 1000); + expect(run?.cleanupHandled).toBe(true); + + clearHookRuns(); + }); + + it("handles immediate cleanup (cleanupDelayMinutes=0)", async () => { + const { registerHookRun, markHookRunComplete, getHookRun, clearHookRuns } = + await import("./hook-run-registry.js"); + clearHookRuns(); + + registerHookRun({ + runId: "run-1", + sessionKey: "hook:test:1", + jobName: "Test Hook", + cleanup: "delete", + cleanupDelayMinutes: 0, + }); + + const now = Date.now(); + vi.setSystemTime(now); + markHookRunComplete("run-1"); + + const run = getHookRun("run-1"); + expect(run?.cleanupAtMs).toBe(now); + + clearHookRuns(); + }); + }); + + describe("initHookRunRegistry", () => { + it("restores runs from disk on init", async () => { + const { loadHookRunRegistryFromDisk } = await import("./hook-run-registry.store.js"); + const existingRun = { + runId: "restored-run", + sessionKey: "hook:restored:1", + jobName: "Restored", + cleanup: "delete" as const, + cleanupDelayMinutes: 0, + createdAt: Date.now(), + cleanupHandled: true, + cleanupAtMs: Date.now(), + }; + vi.mocked(loadHookRunRegistryFromDisk).mockReturnValue( + new Map([["restored-run", existingRun]]), + ); + + const { initHookRunRegistry, getHookRun, clearHookRuns } = + await import("./hook-run-registry.js"); + clearHookRuns(); + initHookRunRegistry(); + + const run = getHookRun("restored-run"); + expect(run).toBeDefined(); + expect(run?.sessionKey).toBe("hook:restored:1"); + + clearHookRuns(); + }); + }); +}); diff --git a/src/gateway/hook-run-registry.ts b/src/gateway/hook-run-registry.ts new file mode 100644 index 000000000..6264e5ad8 --- /dev/null +++ b/src/gateway/hook-run-registry.ts @@ -0,0 +1,127 @@ +import { callGateway } from "./call.js"; +import { + type HookRunRecord, + loadHookRunRegistryFromDisk, + saveHookRunRegistryToDisk, +} from "./hook-run-registry.store.js"; + +const hookRuns = new Map(); +let sweeper: NodeJS.Timeout | null = null; +let restoreAttempted = false; + +function persistHookRuns() { + try { + saveHookRunRegistryToDisk(hookRuns); + } catch { + // Ignore persistence failures + } +} + +function restoreHookRunsOnce() { + if (restoreAttempted) return; + restoreAttempted = true; + try { + const restored = loadHookRunRegistryFromDisk(); + for (const [runId, entry] of restored.entries()) { + if (!hookRuns.has(runId)) { + hookRuns.set(runId, entry); + } + } + if (hookRuns.size > 0) startSweeper(); + } catch { + // Ignore restore failures + } +} + +function startSweeper() { + if (sweeper) return; + sweeper = setInterval(() => { + void sweepHookRuns(); + }, 60_000); + sweeper.unref?.(); +} + +function stopSweeper() { + if (!sweeper) return; + clearInterval(sweeper); + sweeper = null; +} + +async function sweepHookRuns() { + const now = Date.now(); + let mutated = false; + for (const [runId, entry] of hookRuns.entries()) { + // Skip if not ready for cleanup + if (!entry.cleanupAtMs || entry.cleanupAtMs > now) continue; + // Skip if cleanup not yet marked (endedAt hasn't been set) + if (!entry.cleanupHandled) continue; + + try { + await callGateway({ + method: "sessions.delete", + params: { key: entry.sessionKey, deleteTranscript: true }, + timeoutMs: 10_000, + }); + // Only delete from registry after successful RPC + hookRuns.delete(runId); + mutated = true; + } catch { + // Log and retry on next sweep (entry stays in registry) + } + } + if (mutated) persistHookRuns(); + if (hookRuns.size === 0) stopSweeper(); +} + +export function registerHookRun(params: { + runId: string; + sessionKey: string; + jobName: string; + cleanup: "delete" | "keep" | undefined; + cleanupDelayMinutes: number | undefined; +}) { + restoreHookRunsOnce(); + + // Only track runs that need cleanup + if (params.cleanup !== "delete") return; + + const now = Date.now(); + hookRuns.set(params.runId, { + runId: params.runId, + sessionKey: params.sessionKey, + jobName: params.jobName, + cleanup: "delete", + cleanupDelayMinutes: params.cleanupDelayMinutes ?? 0, + createdAt: now, + cleanupHandled: false, + }); + persistHookRuns(); + startSweeper(); +} + +export function markHookRunComplete(runId: string) { + const entry = hookRuns.get(runId); + if (!entry) return; + + const now = Date.now(); + entry.endedAt = now; + entry.cleanupAtMs = now + entry.cleanupDelayMinutes * 60 * 1000; + entry.cleanupHandled = true; + persistHookRuns(); +} + +export function getHookRun(runId: string): HookRunRecord | undefined { + return hookRuns.get(runId); +} + +/** Initialize registry on gateway startup - restores pending cleanups */ +export function initHookRunRegistry() { + restoreHookRunsOnce(); +} + +/** For testing only */ +export function clearHookRuns() { + hookRuns.clear(); + stopSweeper(); + restoreAttempted = false; +}