diff --git a/src/create-hooks.ts b/src/create-hooks.ts index 67b75fbde..0e40ad480 100644 --- a/src/create-hooks.ts +++ b/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: { diff --git a/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test-support.ts b/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test-support.ts index e394a0040..bb04412bd 100644 --- a/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test-support.ts +++ b/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 + runScheduledTimeout: (index: number) => void restore: () => void getClearTimeoutCalls: () => Array> getScheduledTimeouts: () => Array> @@ -83,6 +84,7 @@ export function setupDelayedTimeoutMocks(): { const originalClearTimeout = globalThis.clearTimeout const clearTimeoutCalls: Array> = [] const scheduledTimeouts: Array> = [] + const scheduledCallbacks: Array<() => void> = [] function createTimeoutHandle(): ReturnType { 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 diff --git a/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test.ts b/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test.ts index 4291bb754..6300ab55e 100644 --- a/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.test.ts +++ b/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) => { + 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() + } + }) + }) diff --git a/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.ts b/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.ts index 2be7569ea..0a80d63bc 100644 --- a/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.ts +++ b/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>() - 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 | 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) + }, } } diff --git a/src/hooks/anthropic-context-window-limit-recovery/session-timeout-map.ts b/src/hooks/anthropic-context-window-limit-recovery/session-timeout-map.ts new file mode 100644 index 000000000..80712d03f --- /dev/null +++ b/src/hooks/anthropic-context-window-limit-recovery/session-timeout-map.ts @@ -0,0 +1,20 @@ +export function clearSessionTimeout( + timeoutBySession: Map>, + sessionID: string, +): void { + const timeoutID = timeoutBySession.get(sessionID) + if (timeoutID !== undefined) { + clearTimeout(timeoutID) + timeoutBySession.delete(sessionID) + } +} + +export function clearAllSessionTimeouts( + timeoutBySession: Map>, +): void { + for (const timeoutID of timeoutBySession.values()) { + clearTimeout(timeoutID) + } + + timeoutBySession.clear() +} diff --git a/src/hooks/anthropic-context-window-limit-recovery/summarize-retry-strategy.test.ts b/src/hooks/anthropic-context-window-limit-recovery/summarize-retry-strategy.test.ts index 2c0137c8b..332aeda20 100644 --- a/src/hooks/anthropic-context-window-limit-recovery/summarize-retry-strategy.test.ts +++ b/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 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 + const handle = timeoutCalls.length + 1 as unknown as ReturnType + 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 () => {