diff --git a/src/features/background-agent/fallback-retry-handler.test.ts b/src/features/background-agent/fallback-retry-handler.test.ts index 825f72a56..7309bb526 100644 --- a/src/features/background-agent/fallback-retry-handler.test.ts +++ b/src/features/background-agent/fallback-retry-handler.test.ts @@ -23,6 +23,21 @@ import { selectFallbackProvider } from "../../shared/model-error-classifier" import { readProviderModelsCache } from "../../shared" import type { BackgroundTask } from "./types" import type { ConcurrencyManager } from "./concurrency" +import type { OpencodeClient, QueueItem } from "./constants" + +function createDeferredPromise(): { + promise: Promise + resolve: () => void +} { + let resolvePromise = () => {} + const promise = new Promise((resolve) => { + resolvePromise = resolve + }) + return { + promise, + resolve: resolvePromise, + } +} function createMockTask(overrides: Partial = {}): BackgroundTask { return { @@ -53,20 +68,27 @@ function createMockConcurrencyManager(): ConcurrencyManager { } as unknown as ConcurrencyManager } -function createMockClient() { +function createMockClient(): { + client: OpencodeClient + abortMock: ReturnType +} { + const abortMock = mock(async () => ({})) return { - session: { - abort: mock(async () => ({})), - }, - } as any + client: { + session: { + abort: abortMock, + }, + } as unknown as OpencodeClient, + abortMock, + } } function createDefaultArgs(taskOverrides: Partial = {}) { const processKeyFn = mock(() => {}) - const queuesByKey = new Map>() + const queuesByKey = new Map() const idleDeferralTimers = new Map>() const concurrencyManager = createMockConcurrencyManager() - const client = createMockClient() + const { client, abortMock } = createMockClient() const task = createMockTask(taskOverrides) return { @@ -75,6 +97,7 @@ function createDefaultArgs(taskOverrides: Partial = {}) { source: "polling", concurrencyManager, client, + abortMock, idleDeferralTimers, queuesByKey, processKey: processKeyFn, @@ -93,97 +116,118 @@ describe("tryFallbackRetry", () => { }) describe("#given retryable error with fallback chain", () => { - test("returns true and enqueues retry", () => { + test("returns true and enqueues retry", async () => { const args = createDefaultArgs() - const result = tryFallbackRetry(args) + const result = await tryFallbackRetry(args) expect(result).toBe(true) }) - test("resets task status to pending", () => { + test("resets task status to pending", async () => { const args = createDefaultArgs() - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.status).toBe("pending") }) - test("increments attemptCount", () => { + test("increments attemptCount", async () => { const args = createDefaultArgs() - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.attemptCount).toBe(1) }) - test("updates task model to fallback", () => { + test("updates task model to fallback", async () => { const args = createDefaultArgs() - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.model?.modelID).toBe("fallback-model-1") expect(args.task.model?.providerID).toBe("provider-a") }) - test("clears sessionID and startedAt", () => { + test("clears sessionID and startedAt", async () => { const args = createDefaultArgs({ sessionID: "old-session", startedAt: new Date(), }) - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.sessionID).toBeUndefined() expect(args.task.startedAt).toBeUndefined() }) - test("clears error field", () => { + test("clears error field", async () => { const args = createDefaultArgs({ error: "previous error" }) - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.error).toBeUndefined() }) - test("sets new queuedAt", () => { + test("sets new queuedAt", async () => { const args = createDefaultArgs() - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.queuedAt).toBeInstanceOf(Date) }) - test("releases concurrency slot", () => { + test("releases concurrency slot", async () => { const args = createDefaultArgs() - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.concurrencyManager.release).toHaveBeenCalledWith("provider-a/original-model") }) - test("clears concurrencyKey after release", () => { + test("clears concurrencyKey after release", async () => { const args = createDefaultArgs() - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.concurrencyKey).toBeUndefined() }) - test("aborts existing session", () => { + test("aborts existing session", async () => { const args = createDefaultArgs({ sessionID: "session-to-abort" }) - tryFallbackRetry(args) + await tryFallbackRetry(args) - expect(args.client.session.abort).toHaveBeenCalledWith({ + expect(args.abortMock).toHaveBeenCalledWith({ path: { id: "session-to-abort" }, }) }) - test("adds retry input to queue and calls processKey", () => { + test("waits for session abort before resolving", async () => { + const args = createDefaultArgs({ sessionID: "session-to-abort" }) + const deferred = createDeferredPromise() + args.abortMock.mockImplementationOnce(() => deferred.promise) + + const retryPromise = tryFallbackRetry(args) + let settled = false + void retryPromise.then(() => { + settled = true + }) + + await Promise.resolve() + + expect(settled).toBe(false) + + deferred.resolve() + await retryPromise + + expect(settled).toBe(true) + }) + + test("adds retry input to queue and calls processKey", async () => { const args = createDefaultArgs() - tryFallbackRetry(args) + await tryFallbackRetry(args) const key = `${args.task.model!.providerID}/${args.task.model!.modelID}` const queue = args.queuesByKey.get(key) @@ -195,81 +239,81 @@ describe("tryFallbackRetry", () => { }) describe("#given non-retryable error", () => { - test("returns false when shouldRetryError returns false", () => { + test("returns false when shouldRetryError returns false", async () => { ;(shouldRetryError as any).mockImplementation(() => false) const args = createDefaultArgs() - const result = tryFallbackRetry(args) + const result = await tryFallbackRetry(args) expect(result).toBe(false) }) }) describe("#given no fallback chain", () => { - test("returns false when fallbackChain is undefined", () => { + test("returns false when fallbackChain is undefined", async () => { const args = createDefaultArgs({ fallbackChain: undefined }) - const result = tryFallbackRetry(args) + const result = await tryFallbackRetry(args) expect(result).toBe(false) }) - test("returns false when fallbackChain is empty", () => { + test("returns false when fallbackChain is empty", async () => { const args = createDefaultArgs({ fallbackChain: [] }) - const result = tryFallbackRetry(args) + const result = await tryFallbackRetry(args) expect(result).toBe(false) }) }) describe("#given exhausted fallbacks", () => { - test("returns false when attemptCount exceeds chain length", () => { + test("returns false when attemptCount exceeds chain length", async () => { const args = createDefaultArgs({ attemptCount: 5 }) - const result = tryFallbackRetry(args) + const result = await tryFallbackRetry(args) expect(result).toBe(false) }) }) describe("#given task without concurrency key", () => { - test("skips concurrency release", () => { + test("skips concurrency release", async () => { const args = createDefaultArgs({ concurrencyKey: undefined }) - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.concurrencyManager.release).not.toHaveBeenCalled() }) }) describe("#given task without session", () => { - test("skips session abort", () => { + test("skips session abort", async () => { const args = createDefaultArgs({ sessionID: undefined }) - tryFallbackRetry(args) + await tryFallbackRetry(args) - expect(args.client.session.abort).not.toHaveBeenCalled() + expect(args.abortMock).not.toHaveBeenCalled() }) }) describe("#given active idle deferral timer", () => { - test("clears the timer and removes from map", () => { + test("clears the timer and removes from map", async () => { const args = createDefaultArgs() const timerId = setTimeout(() => {}, 10000) args.idleDeferralTimers.set("test-task-1", timerId) - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.idleDeferralTimers.has("test-task-1")).toBe(false) }) }) describe("#given second attempt", () => { - test("uses second fallback in chain", () => { + test("uses second fallback in chain", async () => { const args = createDefaultArgs({ attemptCount: 1 }) - tryFallbackRetry(args) + await tryFallbackRetry(args) expect(args.task.model?.modelID).toBe("fallback-model-2") expect(args.task.attemptCount).toBe(2) @@ -277,7 +321,7 @@ describe("tryFallbackRetry", () => { }) describe("#given disconnected fallback providers with connected preferred provider", () => { - test("keeps fallback entry and selects connected preferred provider", () => { + test("keeps fallback entry and selects connected preferred provider", async () => { ;(readProviderModelsCache as any).mockReturnValueOnce({ connected: ["provider-a"] }) ;(selectFallbackProvider as any).mockImplementationOnce( (_providers: string[], preferredProviderID?: string) => preferredProviderID ?? "provider-b", @@ -288,7 +332,7 @@ describe("tryFallbackRetry", () => { model: { providerID: "provider-a", modelID: "original-model" }, }) - const result = tryFallbackRetry(args) + const result = await tryFallbackRetry(args) expect(result).toBe(true) expect(args.task.model?.providerID).toBe("provider-a") diff --git a/src/features/background-agent/fallback-retry-handler.ts b/src/features/background-agent/fallback-retry-handler.ts index 58c828e82..f169fa4eb 100644 --- a/src/features/background-agent/fallback-retry-handler.ts +++ b/src/features/background-agent/fallback-retry-handler.ts @@ -11,7 +11,7 @@ import { } from "../../shared/model-error-classifier" import { transformModelForProvider } from "../../shared/provider-model-id-transform" -export function tryFallbackRetry(args: { +export async function tryFallbackRetry(args: { task: BackgroundTask errorInfo: { name?: string; message?: string } source: string @@ -20,7 +20,7 @@ export function tryFallbackRetry(args: { idleDeferralTimers: Map> queuesByKey: Map processKey: (key: string) => void -}): boolean { +}): Promise { const { task, errorInfo, source, concurrencyManager, client, idleDeferralTimers, queuesByKey, processKey } = args const fallbackChain = task.fallbackChain const canRetry = @@ -84,16 +84,14 @@ export function tryFallbackRetry(args: { task.concurrencyKey = undefined } - if (task.sessionID) { - client.session.abort({ path: { id: task.sessionID } }).catch(() => {}) - } - const idleTimer = idleDeferralTimers.get(task.id) if (idleTimer) { clearTimeout(idleTimer) idleDeferralTimers.delete(task.id) } + const previousSessionID = task.sessionID + task.attemptCount = selectedAttemptCount const transformedModelId = transformModelForProvider(providerID, nextFallback.model) task.model = { @@ -123,6 +121,11 @@ export function tryFallbackRetry(args: { category: task.category, isUnstableAgent: task.isUnstableAgent, } + + if (previousSessionID) { + await client.session.abort({ path: { id: previousSessionID } }).catch(() => {}) + } + queue.push({ task, input: retryInput }) queuesByKey.set(key, queue) processKey(key)