Merge pull request #3116 from code-yeongyu/fix/prepublish-hook-recovery

fix(anthropic-recovery): add dispose path and fix test isolation

src/create-hooks.ts

@@ -19,14 +19,16 @@ export type DisposableCreatedHooks = {
   runtimeFallback?: DisposableHook
   todoContinuationEnforcer?: DisposableHook
   autoSlashCommand?: DisposableHook
+  anthropicContextWindowLimitRecovery?: DisposableHook
 }
 
 export function disposeCreatedHooks(hooks: DisposableCreatedHooks): void {
-	 hooks.claudeCodeHooks?.dispose?.()
-	 hooks.commentChecker?.dispose?.()
+  hooks.claudeCodeHooks?.dispose?.()
+  hooks.commentChecker?.dispose?.()
   hooks.runtimeFallback?.dispose?.()
   hooks.todoContinuationEnforcer?.dispose?.()
   hooks.autoSlashCommand?.dispose?.()
+  hooks.anthropicContextWindowLimitRecovery?.dispose?.()
 }
 
 export function createHooks(args: {

src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test-support.ts

@@ -75,6 +75,7 @@ export function createMockContext(): PluginInput {
 
 export function setupDelayedTimeoutMocks(): {
   createUntrackedTimeout: () => ReturnType<typeof setTimeout>
+  runScheduledTimeout: (index: number) => void
   restore: () => void
   getClearTimeoutCalls: () => Array<ReturnType<typeof setTimeout>>
   getScheduledTimeouts: () => Array<ReturnType<typeof setTimeout>>
@@ -83,6 +84,7 @@ export function setupDelayedTimeoutMocks(): {
   const originalClearTimeout = globalThis.clearTimeout
   const clearTimeoutCalls: Array<ReturnType<typeof setTimeout>> = []
   const scheduledTimeouts: Array<ReturnType<typeof setTimeout>> = []
+  const scheduledCallbacks: Array<() => void> = []
 
   function createTimeoutHandle(): ReturnType<typeof setTimeout> {
     const timeoutID = originalSetTimeout(() => {}, 60_000)
@@ -90,9 +92,10 @@ export function setupDelayedTimeoutMocks(): {
     return timeoutID
   }
 
-  globalThis.setTimeout = ((_: () => void, _delay?: number) => {
+  globalThis.setTimeout = ((callback: () => void, _delay?: number) => {
     const timeoutID = createTimeoutHandle()
     scheduledTimeouts.push(timeoutID)
+    scheduledCallbacks.push(callback)
     return timeoutID
   }) as typeof setTimeout
 
@@ -103,6 +106,9 @@ export function setupDelayedTimeoutMocks(): {
 
   return {
     createUntrackedTimeout: createTimeoutHandle,
+    runScheduledTimeout: (index: number) => {
+      scheduledCallbacks[index]?.()
+    },
     restore: () => {
       globalThis.setTimeout = originalSetTimeout
       globalThis.clearTimeout = originalClearTimeout

src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test.ts

@@ -94,4 +94,45 @@ describe("createAnthropicContextWindowLimitRecoveryHook", () => {
     }
   })
 
+  test("#given active pending and retry timers #when dispose is called #then it clears both timer maps", async () => {
+    //#given
+    const { createUntrackedTimeout, getClearTimeoutCalls, getScheduledTimeouts, restore, runScheduledTimeout } =
+      setupDelayedTimeoutMocks()
+    executeCompactMock.mockImplementationOnce(async (...args: Parameters<typeof executeCompactMock>) => {
+      const sessionID = args[0]
+      const autoCompactState = args[2]
+
+      autoCompactState.retryTimerBySession.set(sessionID, createUntrackedTimeout())
+    })
+    const hook = createRecoveryHook()
+
+    try {
+      await hook.event({
+        event: {
+          type: "session.error",
+          properties: { sessionID: "session-retry", error: "prompt is too long" },
+        },
+      })
+
+      await hook.event({
+        event: {
+          type: "session.error",
+          properties: { sessionID: "session-pending", error: "prompt is too long" },
+        },
+      })
+
+      runScheduledTimeout(0)
+
+      const [retryTimer, pendingTimer] = getScheduledTimeouts()
+
+      //#when
+      hook.dispose()
+
+      //#then
+      expect(getClearTimeoutCalls()).toEqual(expect.arrayContaining([retryTimer, pendingTimer]))
+    } finally {
+      restore()
+    }
+  })
+
 })

src/hooks/anthropic-context-window-limit-recovery/recovery-hook.ts

@@ -6,6 +6,7 @@ import { parseAnthropicTokenLimitError } from "./parser"
 import { executeCompact, getLastAssistant } from "./executor"
 import { attemptDeduplicationRecovery } from "./deduplication-recovery"
 import { clearSessionState } from "./state"
+import { clearAllSessionTimeouts, clearSessionTimeout } from "./session-timeout-map"
 import { log } from "../../shared/logger"
 
 export interface AnthropicContextWindowLimitRecoveryOptions {
@@ -48,21 +49,13 @@ export function createAnthropicContextWindowLimitRecoveryHook(
   }
   const pendingCompactionTimeoutBySession = new Map<string, ReturnType<typeof setTimeout>>()
 
-  function clearPendingCompactionTimeout(sessionID: string): void {
-    const timeoutID = pendingCompactionTimeoutBySession.get(sessionID)
-    if (timeoutID !== undefined) {
-      clearTimeout(timeoutID)
-      pendingCompactionTimeoutBySession.delete(sessionID)
-    }
-  }
-
   const eventHandler = async ({ event }: { event: { type: string; properties?: unknown } }) => {
     const props = event.properties as Record<string, unknown> | undefined
 
     if (event.type === "session.deleted") {
       const sessionInfo = props?.info as { id?: string } | undefined
       if (sessionInfo?.id) {
-        clearPendingCompactionTimeout(sessionInfo.id)
+        clearSessionTimeout(pendingCompactionTimeoutBySession, sessionInfo.id)
 
         clearSessionState(autoCompactState, sessionInfo.id)
       }
@@ -105,7 +98,7 @@ export function createAnthropicContextWindowLimitRecoveryHook(
           })
           .catch(() => {})
 
-        clearPendingCompactionTimeout(sessionID)
+        clearSessionTimeout(pendingCompactionTimeoutBySession, sessionID)
 
         const timeoutID = setTimeout(() => {
           pendingCompactionTimeoutBySession.delete(sessionID)
@@ -149,7 +142,7 @@ export function createAnthropicContextWindowLimitRecoveryHook(
 
       if (!autoCompactState.pendingCompact.has(sessionID)) return
 
-      clearPendingCompactionTimeout(sessionID)
+      clearSessionTimeout(pendingCompactionTimeoutBySession, sessionID)
 
       const errorData = autoCompactState.errorDataBySession.get(sessionID)
       const lastAssistant = await dependencies.getLastAssistant(
@@ -192,5 +185,9 @@ export function createAnthropicContextWindowLimitRecoveryHook(
 
   return {
     event: eventHandler,
+    dispose: (): void => {
+      clearAllSessionTimeouts(pendingCompactionTimeoutBySession)
+      clearAllSessionTimeouts(autoCompactState.retryTimerBySession)
+    },
   }
 }

src/hooks/anthropic-context-window-limit-recovery/session-timeout-map.ts

@@ -0,0 +1,20 @@
+export function clearSessionTimeout(
+  timeoutBySession: Map<string, ReturnType<typeof setTimeout>>,
+  sessionID: string,
+): void {
+  const timeoutID = timeoutBySession.get(sessionID)
+  if (timeoutID !== undefined) {
+    clearTimeout(timeoutID)
+    timeoutBySession.delete(sessionID)
+  }
+}
+
+export function clearAllSessionTimeouts(
+  timeoutBySession: Map<string, ReturnType<typeof setTimeout>>,
+): void {
+  for (const timeoutID of timeoutBySession.values()) {
+    clearTimeout(timeoutID)
+  }
+
+  timeoutBySession.clear()
+}

src/hooks/anthropic-context-window-limit-recovery/summarize-retry-strategy.test.ts

@@ -4,6 +4,7 @@ import type { AutoCompactState, ParsedTokenLimitError, RetryState } from "./type
 import type { OhMyOpenCodeConfig } from "../../config"
 
 type TimeoutCall = {
+  handle: ReturnType<typeof setTimeout>
   delay: number
 }
 
@@ -94,8 +95,9 @@ describe("runSummarizeRetryStrategy", () => {
     //#given
     const timeoutCalls: TimeoutCall[] = []
     globalThis.setTimeout = ((_: (...args: unknown[]) => void, delay?: number) => {
-      timeoutCalls.push({ delay: delay ?? 0 })
-      return 1 as unknown as ReturnType<typeof setTimeout>
+      const handle = timeoutCalls.length + 1 as unknown as ReturnType<typeof setTimeout>
+      timeoutCalls.push({ handle, delay: delay ?? 0 })
+      return handle
     }) as typeof setTimeout
 
     autoCompactState.pendingCompact.add(sessionID)
@@ -117,9 +119,12 @@ describe("runSummarizeRetryStrategy", () => {
     })
 
     //#then
-    expect(timeoutCalls.length).toBe(1)
-    expect(timeoutCalls[0]!.delay).toBeGreaterThan(0)
-    expect(timeoutCalls[0]!.delay).toBeLessThanOrEqual(2000)
+    const retryTimer = autoCompactState.retryTimerBySession.get(sessionID)
+    const retryTimeoutCall = timeoutCalls.find(({ handle }) => handle === retryTimer)
+
+    expect(retryTimeoutCall).toBeDefined()
+    expect(retryTimeoutCall?.delay).toBeGreaterThan(0)
+    expect(retryTimeoutCall?.delay).toBeLessThanOrEqual(2000)
   })
 
   test("#given pending retry timer after session cleanup #when scheduled callback fires #then it does not recreate retry state", async () => {