fix(background-agent): await retry abort before requeueing
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
@@ -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<void>
|
||||
resolve: () => void
|
||||
} {
|
||||
let resolvePromise = () => {}
|
||||
const promise = new Promise<void>((resolve) => {
|
||||
resolvePromise = resolve
|
||||
})
|
||||
return {
|
||||
promise,
|
||||
resolve: resolvePromise,
|
||||
}
|
||||
}
|
||||
|
||||
function createMockTask(overrides: Partial<BackgroundTask> = {}): BackgroundTask {
|
||||
return {
|
||||
@@ -53,20 +68,27 @@ function createMockConcurrencyManager(): ConcurrencyManager {
|
||||
} as unknown as ConcurrencyManager
|
||||
}
|
||||
|
||||
function createMockClient() {
|
||||
function createMockClient(): {
|
||||
client: OpencodeClient
|
||||
abortMock: ReturnType<typeof mock>
|
||||
} {
|
||||
const abortMock = mock(async () => ({}))
|
||||
return {
|
||||
session: {
|
||||
abort: mock(async () => ({})),
|
||||
},
|
||||
} as any
|
||||
client: {
|
||||
session: {
|
||||
abort: abortMock,
|
||||
},
|
||||
} as unknown as OpencodeClient,
|
||||
abortMock,
|
||||
}
|
||||
}
|
||||
|
||||
function createDefaultArgs(taskOverrides: Partial<BackgroundTask> = {}) {
|
||||
const processKeyFn = mock(() => {})
|
||||
const queuesByKey = new Map<string, Array<{ task: BackgroundTask; input: any }>>()
|
||||
const queuesByKey = new Map<string, QueueItem[]>()
|
||||
const idleDeferralTimers = new Map<string, ReturnType<typeof setTimeout>>()
|
||||
const concurrencyManager = createMockConcurrencyManager()
|
||||
const client = createMockClient()
|
||||
const { client, abortMock } = createMockClient()
|
||||
const task = createMockTask(taskOverrides)
|
||||
|
||||
return {
|
||||
@@ -75,6 +97,7 @@ function createDefaultArgs(taskOverrides: Partial<BackgroundTask> = {}) {
|
||||
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")
|
||||
|
||||
Reference in New Issue
Block a user