diff --git a/src/hooks/runtime-fallback/session-status-handler.test.ts b/src/hooks/runtime-fallback/session-status-handler.test.ts new file mode 100644 index 000000000..42316a6ec --- /dev/null +++ b/src/hooks/runtime-fallback/session-status-handler.test.ts @@ -0,0 +1,111 @@ +import { describe, expect, it } from "bun:test" +import type { HookDeps, RuntimeFallbackPluginInput } from "./types" +import type { AutoRetryHelpers } from "./auto-retry" +import { createFallbackState } from "./fallback-state" +import { createSessionStatusHandler } from "./session-status-handler" +import { SessionCategoryRegistry } from "../../shared/session-category-registry" + +function createContext(): RuntimeFallbackPluginInput { + return { + client: { + session: { + abort: async () => ({}), + messages: async () => ({ data: [] }), + promptAsync: async () => ({}), + }, + tui: { + showToast: async () => ({}), + }, + }, + directory: "/test/dir", + } +} + +function createDeps(): HookDeps { + return { + ctx: createContext(), + config: { + enabled: true, + retry_on_errors: [429, 503, 529], + max_fallback_attempts: 4, + cooldown_seconds: 60, + timeout_seconds: 30, + notify_on_fallback: false, + }, + options: undefined, + pluginConfig: { + categories: { + test: { + fallback_models: ["openai/gpt-5.4", "google/gemini-2.5-pro"], + }, + }, + }, + sessionStates: new Map(), + sessionLastAccess: new Map(), + sessionRetryInFlight: new Set(), + sessionAwaitingFallbackResult: new Set(), + sessionFallbackTimeouts: new Map(), + sessionStatusRetryKeys: new Map(), + } +} + +function createHelpers(abortCalls: string[], retryCalls: Array<{ sessionID: string; model: string; source: string }>): AutoRetryHelpers { + return { + abortSessionRequest: async (sessionID: string) => { + abortCalls.push(sessionID) + }, + clearSessionFallbackTimeout: () => {}, + scheduleSessionFallbackTimeout: () => {}, + autoRetryWithFallback: async (sessionID: string, model: string, _resolvedAgent: string | undefined, source: string) => { + retryCalls.push({ sessionID, model, source }) + }, + resolveAgentForSessionFromContext: async () => undefined, + cleanupStaleSessions: () => {}, + } +} + +describe("createSessionStatusHandler", () => { + it("#given a pending fallback model #when a new provider cooldown retry arrives #then the handler overrides the pending fallback and advances the chain", async () => { + // given + SessionCategoryRegistry.clear() + const sessionID = "session-status-pending-fallback" + SessionCategoryRegistry.register(sessionID, "test") + + const deps = createDeps() + const abortCalls: string[] = [] + const retryCalls: Array<{ sessionID: string; model: string; source: string }> = [] + const state = createFallbackState("anthropic/claude-opus-4-6") + state.currentModel = "openai/gpt-5.4" + state.fallbackIndex = 0 + state.attemptCount = 1 + state.pendingFallbackModel = "openai/gpt-5.4" + state.failedModels.set("anthropic/claude-opus-4-6", Date.now()) + deps.sessionStates.set(sessionID, state) + + const handler = createSessionStatusHandler(deps, createHelpers(abortCalls, retryCalls), deps.sessionStatusRetryKeys) + + // when + await handler({ + sessionID, + model: "openai/gpt-5.4", + status: { + type: "retry", + attempt: 2, + message: "All credentials for model gpt-5.4 are cooling down [retrying in 7m 56s attempt #2]", + }, + }) + + // then + expect(abortCalls).toEqual([sessionID]) + expect(retryCalls).toEqual([ + { + sessionID, + model: "google/gemini-2.5-pro", + source: "session.status", + }, + ]) + expect(state.currentModel).toBe("google/gemini-2.5-pro") + expect(state.pendingFallbackModel).toBe("google/gemini-2.5-pro") + SessionCategoryRegistry.clear() + }) +}) diff --git a/src/hooks/runtime-fallback/session-status-handler.ts b/src/hooks/runtime-fallback/session-status-handler.ts index 72aa76d66..2e15d41d5 100644 --- a/src/hooks/runtime-fallback/session-status-handler.ts +++ b/src/hooks/runtime-fallback/session-status-handler.ts @@ -26,6 +26,7 @@ export function createSessionStatusHandler( const status = props?.status as { type?: string; message?: string; attempt?: number } | undefined const agent = props?.agent as string | undefined const model = props?.model as string | undefined + const timeoutEnabled = deps.config.timeout_seconds > 0 if (!sessionID || status?.type !== "retry") return @@ -40,8 +41,17 @@ export function createSessionStatusHandler( sessionStatusRetryKeys.set(sessionID, retryKey) if (sessionRetryInFlight.has(sessionID)) { - log(`[${HOOK_NAME}] session.status retry skipped — retry already in flight`, { sessionID }) - return + if (timeoutEnabled) { + log(`[${HOOK_NAME}] Overriding in-flight retry due to provider auto-retry signal`, { + sessionID, + model, + }) + await helpers.abortSessionRequest(sessionID, "session.status.retry-signal") + sessionRetryInFlight.delete(sessionID) + } else { + log(`[${HOOK_NAME}] session.status retry skipped — retry already in flight`, { sessionID }) + return + } } const resolvedAgent = await helpers.resolveAgentForSessionFromContext(sessionID, agent) @@ -75,11 +85,19 @@ export function createSessionStatusHandler( sessionLastAccess.set(sessionID, Date.now()) if (state.pendingFallbackModel) { - log(`[${HOOK_NAME}] session.status retry skipped (pending fallback in progress)`, { - sessionID, - pendingFallbackModel: state.pendingFallbackModel, - }) - return + if (timeoutEnabled) { + log(`[${HOOK_NAME}] Clearing pending fallback due to provider auto-retry signal`, { + sessionID, + pendingFallbackModel: state.pendingFallbackModel, + }) + state.pendingFallbackModel = undefined + } else { + log(`[${HOOK_NAME}] session.status retry skipped (pending fallback in progress)`, { + sessionID, + pendingFallbackModel: state.pendingFallbackModel, + }) + return + } } log(`[${HOOK_NAME}] Detected provider auto-retry signal in session.status`, {