diff --git a/src/hooks/todo-continuation-enforcer/dispose.test.ts b/src/hooks/todo-continuation-enforcer/dispose.test.ts index 5423c8068..37971bc6d 100644 --- a/src/hooks/todo-continuation-enforcer/dispose.test.ts +++ b/src/hooks/todo-continuation-enforcer/dispose.test.ts @@ -8,6 +8,7 @@ declare module "bun:test" { import { afterAll, afterEach, describe, expect, it, mock } from "bun:test" +import type { BackgroundManager } from "../../features/background-agent" import * as actualSessionStateModule from "./session-state" import type { SessionStateStore } from "./session-state" @@ -37,6 +38,12 @@ function createMockPluginInput(): PluginInput { } as PluginInput } +function createMockBackgroundManager(): BackgroundManager { + return { + getTasksByParentSession: () => [{ status: "running" }], + } as BackgroundManager +} + function getCreatedSessionStateStore(): SessionStateStore { if (!createdSessionStateStore) { throw new Error("expected session state store to be created") @@ -68,7 +75,7 @@ describe("todo-continuation-enforcer dispose", () => { enforcer.dispose() }) - it("#given enforcer with active session states #when dispose is called #then internal session state store is shut down", () => { + it("#given enforcer with active session states #when dispose is called #then internal session state store is shut down", async () => { // given const originalClearInterval = globalThis.clearInterval const clearIntervalCalls: Array[0]> = [] @@ -78,9 +85,13 @@ describe("todo-continuation-enforcer dispose", () => { }) as typeof clearInterval try { - const enforcer = createTodoContinuationEnforcer(createMockPluginInput()) + const enforcer = createTodoContinuationEnforcer(createMockPluginInput(), { + backgroundManager: createMockBackgroundManager(), + }) const sessionStateStore = getCreatedSessionStateStore() + await enforcer.handler({ event: { type: "session.idle", properties: { sessionID: "session-1" } } }) + enforcer.markRecovering("session-1") enforcer.markRecovering("session-2") diff --git a/src/hooks/todo-continuation-enforcer/handler.ts b/src/hooks/todo-continuation-enforcer/handler.ts index 3347ee666..7136dda44 100644 --- a/src/hooks/todo-continuation-enforcer/handler.ts +++ b/src/hooks/todo-continuation-enforcer/handler.ts @@ -61,6 +61,7 @@ export function createTodoContinuationHandler(args: { const sessionID = props?.sessionID as string | undefined if (!sessionID) return + sessionStateStore.startPruneInterval() await handleSessionIdle({ ctx, sessionID, diff --git a/src/hooks/todo-continuation-enforcer/session-state.ts b/src/hooks/todo-continuation-enforcer/session-state.ts index a87472b7a..dcd88629e 100644 --- a/src/hooks/todo-continuation-enforcer/session-state.ts +++ b/src/hooks/todo-continuation-enforcer/session-state.ts @@ -31,6 +31,7 @@ export interface ContinuationProgressUpdate { export interface SessionStateStore { getState: (sessionID: string) => SessionState getExistingState: (sessionID: string) => SessionState | undefined + startPruneInterval: () => void recordActivity: (sessionID: string) => void trackContinuationProgress: ( sessionID: string, @@ -76,18 +77,26 @@ export function createSessionStateStore(): SessionStateStore { // Periodic pruning of stale session states to prevent unbounded Map growth let pruneInterval: TimerHandle | undefined - pruneInterval = setInterval(() => { - const now = Date.now() - for (const [sessionID, tracked] of sessions.entries()) { - if (now - tracked.lastAccessedAt > SESSION_STATE_TTL_MS) { - cancelCountdown(sessionID) - sessions.delete(sessionID) - } + let pruneIntervalStarted = false + + function startPruneInterval(): void { + if (pruneIntervalStarted) { + return + } + + pruneIntervalStarted = true + pruneInterval = setInterval(() => { + const now = Date.now() + for (const [sessionID, tracked] of sessions.entries()) { + if (now - tracked.lastAccessedAt > SESSION_STATE_TTL_MS) { + cancelCountdown(sessionID) + sessions.delete(sessionID) + } + } + }, SESSION_STATE_PRUNE_INTERVAL_MS) + if (typeof pruneInterval === "object" && typeof pruneInterval.unref === "function") { + pruneInterval.unref() } - }, SESSION_STATE_PRUNE_INTERVAL_MS) - // Allow process to exit naturally even if interval is running - if (typeof pruneInterval === "object" && typeof pruneInterval.unref === "function") { - pruneInterval.unref() } function getTrackedSession(sessionID: string): TrackedSessionState { @@ -272,6 +281,7 @@ export function createSessionStateStore(): SessionStateStore { return { getState, getExistingState, + startPruneInterval, recordActivity, trackContinuationProgress, resetContinuationProgress, diff --git a/src/hooks/todo-continuation-enforcer/todo-continuation-enforcer.test.ts b/src/hooks/todo-continuation-enforcer/todo-continuation-enforcer.test.ts index 9c5a35f5c..5315b0842 100644 --- a/src/hooks/todo-continuation-enforcer/todo-continuation-enforcer.test.ts +++ b/src/hooks/todo-continuation-enforcer/todo-continuation-enforcer.test.ts @@ -249,6 +249,33 @@ describe("todo-continuation-enforcer", () => { _resetForTesting() }) + test("given the first idle event, starts the prune interval lazily", async () => { + // given + const originalSetInterval = globalThis.setInterval + let setIntervalCalls = 0 + globalThis.setInterval = ((callback: TimerCallback, delay?: number, ...args: any[]) => { + setIntervalCalls += 1 + return originalSetInterval(callback, delay, ...args) + }) as typeof setInterval + + try { + const sessionID = "main-lazy-prune" + setMainSession(sessionID) + const hook = createTodoContinuationEnforcer(createMockPluginInput(), { + backgroundManager: createMockBackgroundManager(true), + }) + + // when + await hook.handler({ event: { type: "session.idle", properties: { sessionID } } }) + await hook.handler({ event: { type: "session.idle", properties: { sessionID } } }) + + // then + expect(setIntervalCalls).toBe(1) + } finally { + globalThis.setInterval = originalSetInterval + } + }) + test("should inject continuation when idle with incomplete todos", async () => { fakeTimers.restore() // given - main session with incomplete todos