Merge branch 'fix/perf-q05' into fix/perf-omo-in-tree

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<Parameters<typeof clearInterval>[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")
 

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,

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(() => {
+  let pruneIntervalStarted = false
-    const now = Date.now()
+
-    for (const [sessionID, tracked] of sessions.entries()) {
+  function startPruneInterval(): void {
-      if (now - tracked.lastAccessedAt > SESSION_STATE_TTL_MS) {
+    if (pruneIntervalStarted) {
-        cancelCountdown(sessionID)
+      return
-        sessions.delete(sessionID)
+    }
-      }
+
+    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,

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