diff --git a/src/hooks/runtime-fallback/auto-retry.ts b/src/hooks/runtime-fallback/auto-retry.ts index dda3a3b6e..4b966c438 100644 --- a/src/hooks/runtime-fallback/auto-retry.ts +++ b/src/hooks/runtime-fallback/auto-retry.ts @@ -1,4 +1,4 @@ -import type { HookDeps } from "./types" +import type { HookDeps, RuntimeFallbackTimeout } from "./types" import { HOOK_NAME } from "./constants" import { log } from "../../shared/logger" import { normalizeAgentName, resolveAgentForSession } from "./agent-resolver" @@ -9,8 +9,8 @@ import { SessionCategoryRegistry } from "../../shared/session-category-registry" const SESSION_TTL_MS = 30 * 60 * 1000 -declare function setTimeout(callback: () => void | Promise, delay?: number): ReturnType -declare function clearTimeout(timeout: ReturnType): void +declare function setTimeout(callback: () => void | Promise, delay?: number): RuntimeFallbackTimeout +declare function clearTimeout(timeout: RuntimeFallbackTimeout): void export function createAutoRetryHelpers(deps: HookDeps) { const { ctx, config, options, sessionStates, sessionLastAccess, sessionRetryInFlight, sessionAwaitingFallbackResult, sessionFallbackTimeouts, pluginConfig } = deps diff --git a/src/hooks/runtime-fallback/dispose.test.ts b/src/hooks/runtime-fallback/dispose.test.ts new file mode 100644 index 000000000..4810bfb95 --- /dev/null +++ b/src/hooks/runtime-fallback/dispose.test.ts @@ -0,0 +1,160 @@ +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test" +import type { HookDeps, RuntimeFallbackPluginInput } from "./types" + +let capturedDeps: HookDeps | undefined + +const mockCreateAutoRetryHelpers = mock((deps: HookDeps) => { + capturedDeps = deps + + return { + abortSessionRequest: async () => {}, + clearSessionFallbackTimeout: () => {}, + scheduleSessionFallbackTimeout: () => {}, + autoRetryWithFallback: async () => {}, + resolveAgentForSessionFromContext: async () => undefined, + cleanupStaleSessions: () => {}, + } +}) + +const mockCreateEventHandler = mock(() => async () => {}) +const mockCreateMessageUpdateHandler = mock(() => async () => {}) +const mockCreateChatMessageHandler = mock(() => async () => {}) + +mock.module("./auto-retry", () => ({ + createAutoRetryHelpers: mockCreateAutoRetryHelpers, +})) + +mock.module("./event-handler", () => ({ + createEventHandler: mockCreateEventHandler, +})) + +mock.module("./message-update-handler", () => ({ + createMessageUpdateHandler: mockCreateMessageUpdateHandler, +})) + +mock.module("./chat-message-handler", () => ({ + createChatMessageHandler: mockCreateChatMessageHandler, +})) + +const { createRuntimeFallbackHook } = await import("./hook") + +function createMockContext(): RuntimeFallbackPluginInput { + return { + client: { + session: { + abort: async () => ({}), + messages: async () => ({}), + promptAsync: async () => ({}), + }, + tui: { + showToast: async () => ({}), + }, + }, + directory: "/test", + } +} + +describe("createRuntimeFallbackHook dispose", () => { + const originalSetInterval = globalThis.setInterval + const originalClearInterval = globalThis.clearInterval + const originalClearTimeout = globalThis.clearTimeout + const createdIntervals: Array> = [] + const clearedIntervals: Array[0]> = [] + const clearedTimeouts: Array[0]> = [] + const timeoutMapSizesDuringClear: number[] = [] + + beforeEach(() => { + capturedDeps = undefined + createdIntervals.length = 0 + clearedIntervals.length = 0 + clearedTimeouts.length = 0 + timeoutMapSizesDuringClear.length = 0 + + mockCreateAutoRetryHelpers.mockClear() + mockCreateEventHandler.mockClear() + mockCreateMessageUpdateHandler.mockClear() + mockCreateChatMessageHandler.mockClear() + + const wrappedSetInterval = ((handler: () => void, timeout?: number) => { + const interval = originalSetInterval(handler, timeout) + createdIntervals.push(interval) + return interval + }) as typeof globalThis.setInterval + + const wrappedClearInterval = ((interval?: Parameters[0]) => { + clearedIntervals.push(interval) + return originalClearInterval(interval) + }) as typeof globalThis.clearInterval + + const wrappedClearTimeout = ((timeout?: Parameters[0]) => { + timeoutMapSizesDuringClear.push(capturedDeps?.sessionFallbackTimeouts.size ?? -1) + clearedTimeouts.push(timeout) + return originalClearTimeout(timeout) + }) as typeof globalThis.clearTimeout + + globalThis.setInterval = wrappedSetInterval + globalThis.clearInterval = wrappedClearInterval + globalThis.clearTimeout = wrappedClearTimeout + }) + + afterEach(() => { + globalThis.setInterval = originalSetInterval + globalThis.clearInterval = originalClearInterval + globalThis.clearTimeout = originalClearTimeout + }) + + test("#given runtime-fallback hook created #when dispose() is called #then cleanup interval is cleared", () => { + // given + const hook = createRuntimeFallbackHook(createMockContext(), { pluginConfig: {} }) + + // when + hook.dispose?.() + + // then + expect(createdIntervals).toHaveLength(1) + expect(clearedIntervals).toEqual([createdIntervals[0]]) + }) + + test("#given hook with session state data #when dispose() is called #then all Maps and Sets are empty", () => { + // given + const hook = createRuntimeFallbackHook(createMockContext(), { pluginConfig: {} }) + const fallbackTimeout = setTimeout(() => {}, 60_000) + + capturedDeps?.sessionStates.set("session-1", { + originalModel: "anthropic/claude-opus-4-6", + currentModel: "openai/gpt-5.4", + fallbackIndex: 1, + failedModels: new Map([["anthropic/claude-opus-4-6", 1]]), + attemptCount: 1, + }) + capturedDeps?.sessionLastAccess.set("session-1", Date.now()) + capturedDeps?.sessionRetryInFlight.add("session-1") + capturedDeps?.sessionAwaitingFallbackResult.add("session-1") + capturedDeps?.sessionFallbackTimeouts.set("session-1", fallbackTimeout) + + // when + hook.dispose?.() + + // then + expect(capturedDeps?.sessionStates.size).toBe(0) + expect(capturedDeps?.sessionLastAccess.size).toBe(0) + expect(capturedDeps?.sessionRetryInFlight.size).toBe(0) + expect(capturedDeps?.sessionAwaitingFallbackResult.size).toBe(0) + expect(capturedDeps?.sessionFallbackTimeouts.size).toBe(0) + }) + + test("#given hook with pending fallback timeouts #when dispose() is called #then timeouts are cleared before Map is emptied", () => { + // given + const hook = createRuntimeFallbackHook(createMockContext(), { pluginConfig: {} }) + const fallbackTimeout = setTimeout(() => {}, 60_000) + capturedDeps?.sessionFallbackTimeouts.set("session-1", fallbackTimeout) + + // when + hook.dispose?.() + + // then + expect(clearedTimeouts).toEqual([fallbackTimeout]) + expect(timeoutMapSizesDuringClear).toEqual([1]) + expect(capturedDeps?.sessionFallbackTimeouts.size).toBe(0) + }) +}) diff --git a/src/hooks/runtime-fallback/hook.ts b/src/hooks/runtime-fallback/hook.ts index b37887990..a3dfc9334 100644 --- a/src/hooks/runtime-fallback/hook.ts +++ b/src/hooks/runtime-fallback/hook.ts @@ -1,5 +1,4 @@ -import type { PluginInput } from "@opencode-ai/plugin" -import type { HookDeps, RuntimeFallbackHook, RuntimeFallbackOptions } from "./types" +import type { HookDeps, RuntimeFallbackHook, RuntimeFallbackInterval, RuntimeFallbackOptions, RuntimeFallbackPluginInput, RuntimeFallbackTimeout } from "./types" import { DEFAULT_CONFIG, HOOK_NAME } from "./constants" import { log } from "../../shared/logger" import { loadPluginConfig } from "../../plugin-config" @@ -8,8 +7,12 @@ import { createEventHandler } from "./event-handler" import { createMessageUpdateHandler } from "./message-update-handler" import { createChatMessageHandler } from "./chat-message-handler" +declare function setInterval(callback: () => void, delay?: number): RuntimeFallbackInterval +declare function clearInterval(interval: RuntimeFallbackInterval): void +declare function clearTimeout(timeout: RuntimeFallbackTimeout): void + export function createRuntimeFallbackHook( - ctx: PluginInput, + ctx: RuntimeFallbackPluginInput, options?: RuntimeFallbackOptions ): RuntimeFallbackHook { const config = { @@ -60,8 +63,23 @@ export function createRuntimeFallbackHook( await baseEventHandler({ event }) } + const dispose = () => { + clearInterval(cleanupInterval) + + for (const fallbackTimeout of deps.sessionFallbackTimeouts.values()) { + clearTimeout(fallbackTimeout) + } + + deps.sessionStates.clear() + deps.sessionLastAccess.clear() + deps.sessionRetryInFlight.clear() + deps.sessionAwaitingFallbackResult.clear() + deps.sessionFallbackTimeouts.clear() + } + return { event: eventHandler, "chat.message": chatMessageHandler, + dispose, } as RuntimeFallbackHook } diff --git a/src/hooks/runtime-fallback/types.ts b/src/hooks/runtime-fallback/types.ts index 500715b9e..e3ef4a31b 100644 --- a/src/hooks/runtime-fallback/types.ts +++ b/src/hooks/runtime-fallback/types.ts @@ -1,6 +1,40 @@ -import type { PluginInput } from "@opencode-ai/plugin" import type { RuntimeFallbackConfig, OhMyOpenCodeConfig } from "../../config" +export interface RuntimeFallbackInterval { + unref: () => void +} + +export type RuntimeFallbackTimeout = object | number + +export interface RuntimeFallbackPluginInput { + client: { + session: { + abort: (input: { path: { id: string } }) => Promise + messages: (input: { path: { id: string }; query: { directory: string } }) => Promise + promptAsync: (input: { + path: { id: string } + body: { + agent?: string + model: { providerID: string; modelID: string } + parts: Array<{ type: "text"; text: string }> + } + query: { directory: string } + }) => Promise + } + tui: { + showToast: (input: { + body: { + title: string + message: string + variant: "success" | "error" | "info" | "warning" + duration: number + } + }) => Promise + } + } + directory: string +} + export interface FallbackState { originalModel: string currentModel: string @@ -26,10 +60,11 @@ export interface RuntimeFallbackOptions { export interface RuntimeFallbackHook { event: (input: { event: { type: string; properties?: unknown } }) => Promise "chat.message"?: (input: { sessionID: string; agent?: string; model?: { providerID: string; modelID: string } }, output: { message: { model?: { providerID: string; modelID: string } }; parts?: Array<{ type: string; text?: string }> }) => Promise + dispose?: () => void } export interface HookDeps { - ctx: PluginInput + ctx: RuntimeFallbackPluginInput config: Required options: RuntimeFallbackOptions | undefined pluginConfig: OhMyOpenCodeConfig | undefined @@ -37,5 +72,5 @@ export interface HookDeps { sessionLastAccess: Map sessionRetryInFlight: Set sessionAwaitingFallbackResult: Set - sessionFallbackTimeouts: Map> + sessionFallbackTimeouts: Map }