From 806b7cccb7212e7c8281d1dc31679c7973ab98b1 Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Thu, 9 Jul 2026 20:01:46 +0000 Subject: [PATCH 1/3] Retry transient provider errors --- packages/agent/src/agent.ts | 7 +- packages/agent/src/provider-retry.ts | 159 +++++++++++++ packages/agent/test/agent.test.ts | 38 ++++ packages/agent/test/provider-retry.test.ts | 252 +++++++++++++++++++++ packages/cli/test/action-runner.test.ts | 40 +++- 5 files changed, 493 insertions(+), 3 deletions(-) create mode 100644 packages/agent/src/provider-retry.ts create mode 100644 packages/agent/test/provider-retry.test.ts diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 38ca37b..e0a7183 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -26,6 +26,7 @@ import { type SimpleStreamOptions, } from "@onkernel/cua-ai"; import type Kernel from "@onkernel/sdk"; +import { withProviderRetry, withProviderRetryModels } from "./provider-retry"; import { buildCuaComputerTools } from "./tools"; import { InternalComputerTranslator, type KernelBrowser } from "./translator/translator"; @@ -308,13 +309,14 @@ export class CuaAgent extends Agent { playwright, onPayload, }); + const retryingStream = withProviderRetry(streamFn ?? defaultCuaStream); const wrappedStreamFn: StreamFn = (model, context, streamOptions) => { const optionsWithCuaRuntime: CuaSimpleStreamOptions = { ...streamOptions, onPayload: runtime.onPayload(), keepToolNames: runtime.keepToolNames(), }; - return (streamFn ?? defaultCuaStream)(model, context, optionsWithCuaRuntime); + return retryingStream(model, context, optionsWithCuaRuntime); }; super({ @@ -459,11 +461,12 @@ export class CuaAgentHarness< onPayload, }); const resolvedTools = runtime.tools(); + const retryingModels = withProviderRetryModels(models ?? cuaModels()); super({ ...harnessOptions, model: runtime.model, - models: models ?? cuaModels(), + models: retryingModels, tools: resolvedTools, systemPrompt: systemPrompt ?? (() => runtime.systemPrompt), activeToolNames: activeToolNames ?? resolvedTools.map((tool) => tool.name), diff --git a/packages/agent/src/provider-retry.ts b/packages/agent/src/provider-retry.ts new file mode 100644 index 0000000..9f85259 --- /dev/null +++ b/packages/agent/src/provider-retry.ts @@ -0,0 +1,159 @@ +import type { StreamFn } from "@earendil-works/pi-agent-core"; +import { + createAssistantMessageEventStream, + isRetryableAssistantError, + type Api, + type AssistantMessage, + type AssistantMessageEvent, + type AssistantMessageEventStream, + type Context, + type Model, + type Models, + type SimpleStreamOptions, +} from "@earendil-works/pi-ai"; + +const MAX_ATTEMPTS = 4; +const DEFAULT_MAX_RETRY_DELAY_MS = 60_000; +const MAX_TIMER_DELAY_MS = 2_147_483_647; + +type RetryStreamFn = ( + model: Model, + context: Context, + options?: SimpleStreamOptions, +) => AssistantMessageEventStream; + +export function withProviderRetry(streamFn: StreamFn): RetryStreamFn { + return (model, context, options) => { + const output = createAssistantMessageEventStream(); + void (async () => { + for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt += 1) { + if (options?.signal?.aborted) { + replayAborted(output, model); + return; + } + + const input = await streamFn(model, context, options); + const events: AssistantMessageEvent[] = []; + for await (const event of input) events.push(structuredClone(event)); + const terminal = events.at(-1); + const message = terminalMessage(terminal); + if ( + !message || + message.stopReason !== "error" || + !isRetryableAssistantError(message) || + attempt === MAX_ATTEMPTS - 1 + ) { + replay(output, events, message); + return; + } + + const providerDelay = parseProviderDelay(message.errorMessage); + const cap = options?.maxRetryDelayMs ?? DEFAULT_MAX_RETRY_DELAY_MS; + if (providerDelay !== undefined && cap > 0 && providerDelay > cap) { + replay(output, events, message); + return; + } + const delay = Math.max(2_000 * 2 ** attempt, providerDelay ?? 0); + if (!(await wait(delay, options?.signal))) { + replayAborted(output, model); + return; + } + } + })(); + return output; + }; +} + +export function withProviderRetryModels(models: Models): Models { + const streamSimple = withProviderRetry((model, context, options) => models.streamSimple(model, context, options)); + return { + getProviders: () => models.getProviders(), + getProvider: (id) => models.getProvider(id), + getModels: (provider) => models.getModels(provider), + getModel: (provider, id) => models.getModel(provider, id), + refresh: (provider) => models.refresh(provider), + getAuth: (model) => models.getAuth(model), + // AgentHarness uses the simple paths; generic API dispatch remains untouched. + stream: (model, context, options) => models.stream(model, context, options), + complete: (model, context, options) => models.complete(model, context, options), + streamSimple, + completeSimple: (model, context, options) => streamSimple(model, context, options).result(), + }; +} + +function terminalMessage(event: AssistantMessageEvent | undefined): AssistantMessage | undefined { + if (event?.type === "done") return event.message; + if (event?.type === "error") return event.error; + return undefined; +} + +function replay( + output: ReturnType, + events: AssistantMessageEvent[], + message: AssistantMessage | undefined, +): void { + for (const event of events) output.push(event); + if (message) output.end(message); + else output.end(); +} + +function replayAborted(output: ReturnType, model: Model): void { + const message: AssistantMessage = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "aborted", + errorMessage: "Request aborted", + timestamp: Date.now(), + }; + output.push({ type: "start", partial: structuredClone(message) }); + output.push({ type: "error", reason: "aborted", error: message }); + output.end(message); +} + +function parseProviderDelay(message: string | undefined): number | undefined { + if (!message) return undefined; + const match = /(?:retry\s+in|retry\s+delay(?:\s+of|\s*:)?)[\s:]*(-?(?:\d+(?:\.\d*)?|\.\d+|infinity))\s*(ms|s)\b/i.exec( + message, + ); + if (!match) return undefined; + const value = Number(match[1]); + if (!Number.isFinite(value) || value < 0) return undefined; + const delay = value * (match[2].toLowerCase() === "s" ? 1_000 : 1); + return Number.isFinite(delay) ? delay : undefined; +} + +async function wait(delayMs: number, signal: AbortSignal | undefined): Promise { + let remaining = delayMs; + while (remaining > 0) { + const chunk = Math.min(remaining, MAX_TIMER_DELAY_MS); + if (!(await waitChunk(chunk, signal))) return false; + remaining -= chunk; + } + return !signal?.aborted; +} + +function waitChunk(delayMs: number, signal: AbortSignal | undefined): Promise { + if (signal?.aborted) return Promise.resolve(false); + return new Promise((resolve) => { + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(true); + }, delayMs); + const onAbort = () => { + clearTimeout(timer); + resolve(false); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + }); +} diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index d7d2273..27720aa 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -231,6 +231,44 @@ describe("CuaAgent", () => { expect(agent.state.systemPrompt).toBe("custom prompt"); }); + it("retries transient errors from a custom stream function", async () => { + vi.useFakeTimers(); + try { + let calls = 0; + const streamFn: StreamFn = (model) => { + calls += 1; + const stream = createAssistantMessageEventStream(); + const message = createAssistantMessage(model); + stream.push({ type: "start", partial: message }); + if (calls === 1) { + message.stopReason = "error"; + message.errorMessage = "HTTP 429: rate limited"; + message.content.push({ type: "text", text: "discarded" }); + stream.push({ type: "error", reason: "error", error: message }); + } else { + message.content.push({ type: "text", text: "done" }); + stream.push({ type: "done", reason: "stop", message }); + } + return stream; + }; + const agent = new CuaAgent({ + browser, + client, + streamFn, + initialState: { model: "openai:gpt-5.5" }, + }); + const prompt = agent.prompt("hello"); + await vi.advanceTimersByTimeAsync(2_000); + await prompt; + + expect(calls).toBe(2); + expect(JSON.stringify(agent.state.messages)).toContain("done"); + expect(JSON.stringify(agent.state.messages)).not.toContain("discarded"); + } finally { + vi.useRealTimers(); + } + }); + it("composes payload hooks for custom stream functions", async () => { const payloads: unknown[] = []; const streamFn: StreamFn = (model, _context, options) => { diff --git a/packages/agent/test/provider-retry.test.ts b/packages/agent/test/provider-retry.test.ts new file mode 100644 index 0000000..143dc39 --- /dev/null +++ b/packages/agent/test/provider-retry.test.ts @@ -0,0 +1,252 @@ +import { + createAssistantMessageEventStream, + type Api, + type AssistantMessage, + type AssistantMessageEvent, + type Model, + type Models, +} from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { withProviderRetry, withProviderRetryModels } from "../src/provider-retry"; + +const model = { + id: "test", + provider: "google", + api: "google-generative-ai", +} as Model; +const context = { messages: [] }; + +afterEach(() => vi.useRealTimers()); + +describe("provider retry", () => { + it("discards a retryable attempt and returns the successful retry", async () => { + vi.useFakeTimers(); + let calls = 0; + const retrying = withProviderRetry(() => { + calls += 1; + return calls === 1 ? errorStream("HTTP 429: rate limited") : textStream("done"); + }); + + const eventsPromise = collect(retrying(model, context)); + await vi.advanceTimersByTimeAsync(1_999); + expect(calls).toBe(1); + await vi.advanceTimersByTimeAsync(1); + const events = await eventsPromise; + + expect(calls).toBe(2); + expect(events.some((event) => event.type === "error")).toBe(false); + expect(events.filter((event) => event.type === "text_delta").map((event) => event.delta)).toEqual(["done"]); + }); + + it("uses exponential backoff and exposes the fourth error", async () => { + vi.useFakeTimers(); + let calls = 0; + const retrying = withProviderRetry(() => { + calls += 1; + return errorStream("HTTP 429: rate limited"); + }); + const result = retrying(model, context).result(); + + await vi.advanceTimersByTimeAsync(2_000); + expect(calls).toBe(2); + await vi.advanceTimersByTimeAsync(4_000); + expect(calls).toBe(3); + await vi.advanceTimersByTimeAsync(8_000); + expect(calls).toBe(4); + expect((await result).stopReason).toBe("error"); + }); + + it.each([ + ["Please retry in 10.5s", 10_500], + ["retry delay of 2500 ms", 2_500], + ["Retry Delay: 3 s", 3_000], + ])("honors provider delay %s", async (message, delay) => { + vi.useFakeTimers(); + let calls = 0; + const retrying = withProviderRetry(() => (++calls === 1 ? errorStream(`HTTP 429: ${message}`) : textStream("ok"))); + const result = retrying(model, context).result(); + await vi.advanceTimersByTimeAsync(delay - 1); + expect(calls).toBe(1); + await vi.advanceTimersByTimeAsync(1); + expect((await result).stopReason).toBe("stop"); + expect(calls).toBe(2); + }); + + it.each([ + ["HTTP 429: retry in 61s", undefined], + ["HTTP 429: retry in 10s", 5_000], + ])("does not retry a provider delay above the cap", async (message, maxRetryDelayMs) => { + let calls = 0; + const retrying = withProviderRetry(() => { + calls += 1; + return errorStream(message); + }); + expect((await retrying(model, context, { maxRetryDelayMs }).result()).stopReason).toBe("error"); + expect(calls).toBe(1); + }); + + it.each(["retry in -3s", "retry delay: Infinity ms", "retry soon"])( + "falls back to exponential delay for malformed hint %s", + async (hint) => { + vi.useFakeTimers(); + let calls = 0; + const retrying = withProviderRetry(() => (++calls === 1 ? errorStream(`HTTP 429: ${hint}`) : textStream("ok"))); + const result = retrying(model, context).result(); + await vi.advanceTimersByTimeAsync(1_999); + expect(calls).toBe(1); + await vi.advanceTimersByTimeAsync(1); + expect((await result).stopReason).toBe("stop"); + }, + ); + + it("allows an uncapped provider delay when maxRetryDelayMs is zero", async () => { + vi.useFakeTimers(); + let calls = 0; + const retrying = withProviderRetry(() => (++calls === 1 ? errorStream("HTTP 429: retry in 61s") : textStream("ok"))); + const result = retrying(model, context, { maxRetryDelayMs: 0 }).result(); + await vi.advanceTimersByTimeAsync(61_000); + expect((await result).stopReason).toBe("stop"); + expect(calls).toBe(2); + }); + + it("waits the full uncapped delay when it exceeds the runtime timer limit", async () => { + vi.useFakeTimers(); + let calls = 0; + const delay = 3_000_000_000; + const retrying = withProviderRetry(() => + ++calls === 1 ? errorStream(`HTTP 429: retry in ${delay}ms`) : textStream("ok"), + ); + const result = retrying(model, context, { maxRetryDelayMs: 0 }).result(); + + await vi.advanceTimersByTimeAsync(delay - 1); + expect(calls).toBe(1); + await vi.advanceTimersByTimeAsync(1); + expect((await result).stopReason).toBe("stop"); + expect(calls).toBe(2); + }); + + it("does not retry non-retryable errors or aborted responses", async () => { + for (const message of [errorMessage("invalid API key"), errorMessage("aborted", "aborted")]) { + let calls = 0; + const retrying = withProviderRetry(() => { + calls += 1; + return message; + }); + await retrying(model, context).result(); + expect(calls).toBe(1); + } + }); + + it("emits only an aborted logical attempt when aborted during backoff", async () => { + vi.useFakeTimers(); + const controller = new AbortController(); + let calls = 0; + const retrying = withProviderRetry(() => { + calls += 1; + return errorStream("HTTP 429: retry in 10s", "discarded"); + }); + const eventsPromise = collect(retrying(model, context, { signal: controller.signal })); + await vi.advanceTimersByTimeAsync(1); + controller.abort(); + const events = await eventsPromise; + + expect(calls).toBe(1); + expect(events.map((event) => event.type)).toEqual(["start", "error"]); + expect(events.at(-1)?.type === "error" && events.at(-1).error.stopReason).toBe("aborted"); + expect(JSON.stringify(events)).not.toContain("discarded"); + }); + + it("supports promise-returning streams and snapshots mutable partials", async () => { + const retrying = withProviderRetry(async () => mutableTextStream()); + const events = await collect(retrying(model, context)); + const deltas = events.filter((event) => event.type === "text_delta"); + expect(deltas[0].partial.content).toEqual([{ type: "text", text: "a" }]); + expect(deltas[1].partial.content).toEqual([{ type: "text", text: "ab" }]); + }); + + it("applies retries to Models.completeSimple", async () => { + vi.useFakeTimers(); + let calls = 0; + const models = { + streamSimple: () => (++calls === 1 ? errorStream("HTTP 429: rate limited") : textStream("done")), + } as Models; + const retryingModels = withProviderRetryModels(models); + const completion = retryingModels.completeSimple(model, context); + await vi.advanceTimersByTimeAsync(2_000); + + expect((await completion).content).toEqual([{ type: "text", text: "done" }]); + expect(calls).toBe(2); + }); +}); + +function assistant(stopReason: AssistantMessage["stopReason"], errorMessage?: string): AssistantMessage { + return { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason, + errorMessage, + timestamp: 0, + }; +} + +function errorStream(message: string, text?: string) { + const stream = createAssistantMessageEventStream(); + const value = assistant("error", message); + stream.push({ type: "start", partial: value }); + if (text) { + value.content.push({ type: "text", text }); + stream.push({ type: "text_delta", contentIndex: 0, delta: text, partial: value }); + } + stream.push({ type: "error", reason: "error", error: value }); + return stream; +} + +function errorMessage(message: string, reason: "error" | "aborted" = "error") { + const stream = createAssistantMessageEventStream(); + const value = assistant(reason, message); + stream.push({ type: "start", partial: value }); + stream.push({ type: "error", reason, error: value }); + return stream; +} + +function textStream(text: string) { + const stream = createAssistantMessageEventStream(); + const value = assistant("stop"); + stream.push({ type: "start", partial: value }); + value.content.push({ type: "text", text }); + stream.push({ type: "text_delta", contentIndex: 0, delta: text, partial: value }); + stream.push({ type: "done", reason: "stop", message: value }); + return stream; +} + +function mutableTextStream() { + const stream = createAssistantMessageEventStream(); + const value = assistant("stop"); + void (async () => { + stream.push({ type: "start", partial: value }); + value.content.push({ type: "text", text: "a" }); + stream.push({ type: "text_delta", contentIndex: 0, delta: "a", partial: value }); + await new Promise((resolve) => setTimeout(resolve, 0)); + value.content[0] = { type: "text", text: "ab" }; + stream.push({ type: "text_delta", contentIndex: 0, delta: "b", partial: value }); + stream.push({ type: "done", reason: "stop", message: value }); + })(); + return stream; +} + +async function collect(stream: AsyncIterable): Promise { + const events: AssistantMessageEvent[] = []; + for await (const event of stream) events.push(event); + return events; +} diff --git a/packages/cli/test/action-runner.test.ts b/packages/cli/test/action-runner.test.ts index 6833e79..ce5d3ba 100644 --- a/packages/cli/test/action-runner.test.ts +++ b/packages/cli/test/action-runner.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { runAction } from "../src/action/harness-runner"; import { buildTestHarness, type TestHarnessFixture } from "./fixtures/harness"; @@ -61,6 +61,44 @@ describe("action harness-runner", () => { expect(res.result.text).toContain("boom"); }); + it("retries a transient provider error before completing the action", async () => { + vi.useFakeTimers(); + try { + fixture = await buildTestHarness({ + turns: [ + { + steps: [ + { type: "text", text: "discarded" }, + { type: "tool_call", toolName: "click", args: { x: 9, y: 9 } }, + { type: "error", message: "HTTP 429: Please retry in 10.367614288s" }, + ], + }, + { steps: [{ type: "text", text: "done" }] }, + ], + }); + const resultPromise = runAction( + { action: "do", text: "recover" }, + { harness: fixture.harness, browserHandle: handleFor(fixture), session: fixture.session, maxTurns: 3 }, + ); + + await vi.advanceTimersByTimeAsync(10_366); + expect(fixture.provider.callCount()).toBe(1); + await vi.advanceTimersByTimeAsync(1); + const res = await resultPromise; + + expect(fixture.provider.callCount()).toBe(2); + expect(res.exitCode).toBe(0); + expect(res.result.status).toBe("ok"); + expect(fixture.kernel.batchCalls).toHaveLength(0); + const messages = (await fixture.session.getBranch()).filter((entry) => entry.type === "message"); + expect(messages).toHaveLength(2); + expect(JSON.stringify(messages)).toContain("done"); + expect(JSON.stringify(messages)).not.toContain("discarded"); + } finally { + vi.useRealTimers(); + } + }); + it("invokes harness.abort once the turn cap is reached", async () => { const toolCall = { steps: [ From c4be964d194bd4243b2d6d9514304f7afe86b6a5 Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Thu, 9 Jul 2026 21:21:25 +0000 Subject: [PATCH 2/3] Make provider retries configurable --- packages/agent/src/agent.ts | 23 +- packages/agent/src/index.ts | 1 + packages/agent/src/provider-retry.ts | 273 +++++++++++----- packages/agent/test/agent.test.ts | 22 ++ packages/agent/test/provider-retry.test.ts | 347 ++++++++++++++------- packages/cli/src/harness.ts | 3 + packages/cli/test/action-runner.test.ts | 14 +- packages/cli/test/fixtures/harness.ts | 2 + 8 files changed, 475 insertions(+), 210 deletions(-) diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index e0a7183..cde7de4 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -26,7 +26,12 @@ import { type SimpleStreamOptions, } from "@onkernel/cua-ai"; import type Kernel from "@onkernel/sdk"; -import { withProviderRetry, withProviderRetryModels } from "./provider-retry"; +import { + type CuaRetryOptions, + resolveProviderRetryPolicy, + withProviderRetry, + withProviderRetryModels, +} from "./provider-retry"; import { buildCuaComputerTools } from "./tools"; import { InternalComputerTranslator, type KernelBrowser } from "./translator/translator"; @@ -79,6 +84,8 @@ export type CuaAgentOptions = Omit & { nativeTool?: CuaNativeToolSpec; /** Expose a tool that runs Playwright code against the browser session. */ playwright?: boolean; + /** Optional CUA-level retries around each provider request. Disabled by default. */ + retry?: CuaRetryOptions; }; /** @@ -117,6 +124,8 @@ export type CuaAgentHarnessOptions< playwright?: boolean; /** Optional payload hook composed after the provider-specific CUA payload hook. */ onPayload?: SimpleStreamOptions["onPayload"]; + /** Optional CUA-level retries around each provider request. Disabled by default. */ + retry?: CuaRetryOptions; }; /** @@ -297,6 +306,7 @@ export class CuaAgent extends Agent { mode, nativeTool, playwright, + retry, ...agentOptions } = options; const runtime = new CuaRuntimeController({ @@ -309,7 +319,10 @@ export class CuaAgent extends Agent { playwright, onPayload, }); - const retryingStream = withProviderRetry(streamFn ?? defaultCuaStream); + const retryingStream = withProviderRetry( + streamFn ?? defaultCuaStream, + resolveProviderRetryPolicy(retry), + ); const wrappedStreamFn: StreamFn = (model, context, streamOptions) => { const optionsWithCuaRuntime: CuaSimpleStreamOptions = { ...streamOptions, @@ -448,6 +461,7 @@ export class CuaAgentHarness< systemPrompt, onPayload, activeToolNames, + retry, ...harnessOptions } = options; const runtime = new CuaRuntimeController({ @@ -461,7 +475,10 @@ export class CuaAgentHarness< onPayload, }); const resolvedTools = runtime.tools(); - const retryingModels = withProviderRetryModels(models ?? cuaModels()); + const retryingModels = withProviderRetryModels( + models ?? cuaModels(), + resolveProviderRetryPolicy(retry), + ); super({ ...harnessOptions, diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 2f79b50..a27747f 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -17,3 +17,4 @@ export type { } from "./tools"; export { CuaAgent, CuaAgentHarness } from "./agent"; export type { CuaAgentHarnessOptions, CuaAgentOptions, CuaAgentState } from "./agent"; +export type { CuaRetryOptions } from "./provider-retry"; diff --git a/packages/agent/src/provider-retry.ts b/packages/agent/src/provider-retry.ts index 9f85259..bdc4880 100644 --- a/packages/agent/src/provider-retry.ts +++ b/packages/agent/src/provider-retry.ts @@ -1,71 +1,87 @@ import type { StreamFn } from "@earendil-works/pi-agent-core"; import { createAssistantMessageEventStream, + isContextOverflow, isRetryableAssistantError, type Api, type AssistantMessage, type AssistantMessageEvent, - type AssistantMessageEventStream, type Context, type Model, type Models, type SimpleStreamOptions, } from "@earendil-works/pi-ai"; -const MAX_ATTEMPTS = 4; -const DEFAULT_MAX_RETRY_DELAY_MS = 60_000; +const DEFAULT_MAX_RETRIES = 3; +const DEFAULT_BASE_DELAY_MS = 2_000; const MAX_TIMER_DELAY_MS = 2_147_483_647; +/** + * Controls CUA retries around a single provider request. + * Enabling provider request retries as well can multiply network attempts. + */ +export interface CuaRetryOptions { + /** Enable CUA-level retries. Disabled by default. */ + enabled?: boolean; + /** Number of additional provider requests. Defaults to 3. */ + maxRetries?: number; + /** Initial delay in milliseconds; subsequent delays double. Defaults to 2000. */ + baseDelayMs?: number; +} + +interface ResolvedRetryPolicy { + enabled: boolean; + maxRetries: number; + baseDelayMs: number; +} + type RetryStreamFn = ( model: Model, context: Context, options?: SimpleStreamOptions, -) => AssistantMessageEventStream; +) => ReturnType; -export function withProviderRetry(streamFn: StreamFn): RetryStreamFn { - return (model, context, options) => { - const output = createAssistantMessageEventStream(); - void (async () => { - for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt += 1) { - if (options?.signal?.aborted) { - replayAborted(output, model); - return; - } +export function resolveProviderRetryPolicy(options?: CuaRetryOptions): ResolvedRetryPolicy { + const enabled = options?.enabled ?? false; + const maxRetries = options?.maxRetries ?? DEFAULT_MAX_RETRIES; + const baseDelayMs = options?.baseDelayMs ?? DEFAULT_BASE_DELAY_MS; - const input = await streamFn(model, context, options); - const events: AssistantMessageEvent[] = []; - for await (const event of input) events.push(structuredClone(event)); - const terminal = events.at(-1); - const message = terminalMessage(terminal); - if ( - !message || - message.stopReason !== "error" || - !isRetryableAssistantError(message) || - attempt === MAX_ATTEMPTS - 1 - ) { - replay(output, events, message); - return; - } + if (typeof enabled !== "boolean") throw new TypeError("retry.enabled must be a boolean"); + if (!Number.isSafeInteger(maxRetries) || maxRetries < 0) { + throw new RangeError("retry.maxRetries must be a non-negative safe integer"); + } + if (!Number.isFinite(baseDelayMs) || baseDelayMs < 0) { + throw new RangeError("retry.baseDelayMs must be a finite non-negative number"); + } + if (enabled && maxRetries > 0 && baseDelayMs > 0) { + const largestDelay = baseDelayMs * 2 ** (maxRetries - 1); + if (!Number.isFinite(largestDelay) || largestDelay > MAX_TIMER_DELAY_MS) { + throw new RangeError(`retry delay must not exceed ${MAX_TIMER_DELAY_MS}ms`); + } + } - const providerDelay = parseProviderDelay(message.errorMessage); - const cap = options?.maxRetryDelayMs ?? DEFAULT_MAX_RETRY_DELAY_MS; - if (providerDelay !== undefined && cap > 0 && providerDelay > cap) { - replay(output, events, message); - return; - } - const delay = Math.max(2_000 * 2 ** attempt, providerDelay ?? 0); - if (!(await wait(delay, options?.signal))) { - replayAborted(output, model); - return; - } - } - })(); + return { enabled, maxRetries, baseDelayMs }; +} + +export function withProviderRetry(streamFn: StreamFn, policy: ResolvedRetryPolicy): StreamFn { + if (!policy.enabled || policy.maxRetries === 0) return streamFn; + + return (model, context, options) => { + const output = createAssistantMessageEventStream(); + void runWithRetry(output, streamFn, model, context, options, policy); return output; }; } -export function withProviderRetryModels(models: Models): Models { - const streamSimple = withProviderRetry((model, context, options) => models.streamSimple(model, context, options)); +export function withProviderRetryModels(models: Models, policy: ResolvedRetryPolicy): Models { + if (!policy.enabled || policy.maxRetries === 0) return models; + + const retryingStream = withProviderRetry( + (model, context, options) => models.streamSimple(model, context, options), + policy, + ); + const streamSimple: RetryStreamFn = (model, context, options) => + retryingStream(model, context, options) as ReturnType; return { getProviders: () => models.getProviders(), getProvider: (id) => models.getProvider(id), @@ -73,7 +89,6 @@ export function withProviderRetryModels(models: Models): Models { getModel: (provider, id) => models.getModel(provider, id), refresh: (provider) => models.refresh(provider), getAuth: (model) => models.getAuth(model), - // AgentHarness uses the simple paths; generic API dispatch remains untouched. stream: (model, context, options) => models.stream(model, context, options), complete: (model, context, options) => models.complete(model, context, options), streamSimple, @@ -81,24 +96,129 @@ export function withProviderRetryModels(models: Models): Models { }; } -function terminalMessage(event: AssistantMessageEvent | undefined): AssistantMessage | undefined { - if (event?.type === "done") return event.message; - if (event?.type === "error") return event.error; - return undefined; +async function runWithRetry( + output: ReturnType, + streamFn: StreamFn, + model: Model, + context: Context, + options: SimpleStreamOptions | undefined, + policy: ResolvedRetryPolicy, +): Promise { + let settled = false; + const settle = (message: AssistantMessage, events?: AssistantMessageEvent[]) => { + if (settled) return; + settled = true; + let terminalEmitted = false; + try { + if (events) { + for (const event of events) { + output.push(event); + if (terminalMessage(event)) terminalEmitted = true; + } + } + if (!terminalEmitted) { + output.push({ type: "start", partial: structuredClone(message) }); + output.push({ type: "error", reason: message.stopReason === "aborted" ? "aborted" : "error", error: message }); + } + } catch (error) { + if (!terminalEmitted) { + const failure = createFailureMessage(model, `Failed to replay provider response: ${errorMessage(error)}`); + try { + output.push({ type: "error", reason: "error", error: failure }); + } catch { + // end(result) still settles result() and any waiting iterator. + } + output.end(failure); + return; + } + } finally { + output.end(message); + } + }; + + for (let attempt = 0; attempt <= policy.maxRetries; attempt += 1) { + if (options?.signal?.aborted) { + settle(createFailureMessage(model, "Request aborted", "aborted")); + return; + } + + // Buffering is opt-in because pi has no way to retract deltas from a failed attempt. + let events: AssistantMessageEvent[] = []; + let message: AssistantMessage | undefined; + try { + const input = await streamFn(model, context, options); + try { + for await (const event of input) { + try { + events.push(structuredClone(event)); + } catch (error) { + throw new SnapshotError(errorMessage(error)); + } + message = terminalMessage(event); + if (message) break; + } + } catch (error) { + if (error instanceof SnapshotError) { + settle(createFailureMessage(model, error.message)); + return; + } + throw error; + } + if (!message) { + settle(createFailureMessage(model, "Provider stream ended without a terminal event")); + return; + } + } catch (error) { + message = createFailureMessage(model, errorMessage(error)); + events = []; + } + + if (!shouldRetry(message, model, attempt, policy, options?.signal)) { + settle(message, events); + return; + } + + events = []; + try { + if (!(await wait(policy.baseDelayMs * 2 ** attempt, options?.signal))) { + settle(createFailureMessage(model, "Request aborted", "aborted")); + return; + } + } catch (error) { + settle(createFailureMessage(model, errorMessage(error))); + return; + } + } } -function replay( - output: ReturnType, - events: AssistantMessageEvent[], - message: AssistantMessage | undefined, -): void { - for (const event of events) output.push(event); - if (message) output.end(message); - else output.end(); +function shouldRetry( + message: AssistantMessage, + model: Model, + attempt: number, + policy: ResolvedRetryPolicy, + signal: AbortSignal | undefined, +): boolean { + return ( + message.stopReason === "error" && + !isContextOverflow(message, model.contextWindow) && + isRetryableAssistantError(message) && + attempt < policy.maxRetries && + !signal?.aborted + ); } -function replayAborted(output: ReturnType, model: Model): void { - const message: AssistantMessage = { +function terminalMessage(event: AssistantMessageEvent): AssistantMessage | undefined { + if (event.type === "done") return event.message; + if (event.type === "error") return event.error; + return undefined; +} + +function createFailureMessage( + model: Model, + message: string, + stopReason: "error" | "aborted" = "error", +): AssistantMessage { + return { role: "assistant", content: [], api: model.api, @@ -112,48 +232,29 @@ function replayAborted(output: ReturnType { - let remaining = delayMs; - while (remaining > 0) { - const chunk = Math.min(remaining, MAX_TIMER_DELAY_MS); - if (!(await waitChunk(chunk, signal))) return false; - remaining -= chunk; - } - return !signal?.aborted; -} +class SnapshotError extends Error {} -function waitChunk(delayMs: number, signal: AbortSignal | undefined): Promise { +function wait(delayMs: number, signal: AbortSignal | undefined): Promise { if (signal?.aborted) return Promise.resolve(false); return new Promise((resolve) => { - const timer = setTimeout(() => { - signal?.removeEventListener("abort", onAbort); - resolve(true); - }, delayMs); const onAbort = () => { clearTimeout(timer); resolve(false); }; + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(true); + }, delayMs); signal?.addEventListener("abort", onAbort, { once: true }); }); } diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 27720aa..b699f7c 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -231,6 +231,27 @@ describe("CuaAgent", () => { expect(agent.state.systemPrompt).toBe("custom prompt"); }); + it("does not retry transient errors by default", async () => { + let calls = 0; + const agent = new CuaAgent({ + browser, + client, + streamFn: (model) => { + calls += 1; + const stream = createAssistantMessageEventStream(); + const message = createAssistantMessage(model); + message.stopReason = "error"; + message.errorMessage = "HTTP 429: rate limited"; + stream.push({ type: "error", reason: "error", error: message }); + return stream; + }, + initialState: { model: "openai:gpt-5.5" }, + }); + + await agent.prompt("hello"); + expect(calls).toBe(1); + }); + it("retries transient errors from a custom stream function", async () => { vi.useFakeTimers(); try { @@ -255,6 +276,7 @@ describe("CuaAgent", () => { browser, client, streamFn, + retry: { enabled: true }, initialState: { model: "openai:gpt-5.5" }, }); const prompt = agent.prompt("hello"); diff --git a/packages/agent/test/provider-retry.test.ts b/packages/agent/test/provider-retry.test.ts index 143dc39..ed18180 100644 --- a/packages/agent/test/provider-retry.test.ts +++ b/packages/agent/test/provider-retry.test.ts @@ -5,177 +5,264 @@ import { type AssistantMessageEvent, type Model, type Models, + type SimpleStreamOptions, } from "@earendil-works/pi-ai"; +import type { StreamFn } from "@earendil-works/pi-agent-core"; import { afterEach, describe, expect, it, vi } from "vitest"; -import { withProviderRetry, withProviderRetryModels } from "../src/provider-retry"; +import { + resolveProviderRetryPolicy, + withProviderRetry, + withProviderRetryModels, +} from "../src/provider-retry"; const model = { id: "test", provider: "google", api: "google-generative-ai", + contextWindow: 128_000, } as Model; const context = { messages: [] }; -afterEach(() => vi.useRealTimers()); +const enabled = resolveProviderRetryPolicy({ enabled: true }); -describe("provider retry", () => { - it("discards a retryable attempt and returns the successful retry", async () => { - vi.useFakeTimers(); - let calls = 0; - const retrying = withProviderRetry(() => { - calls += 1; - return calls === 1 ? errorStream("HTTP 429: rate limited") : textStream("done"); - }); +afterEach(() => vi.useRealTimers()); - const eventsPromise = collect(retrying(model, context)); - await vi.advanceTimersByTimeAsync(1_999); - expect(calls).toBe(1); - await vi.advanceTimersByTimeAsync(1); - const events = await eventsPromise; +describe("provider retry policy", () => { + it("is disabled by default and preserves stream and Models identity", () => { + const streamFn = (() => textStream("ok")) as StreamFn; + const models = {} as Models; + expect(withProviderRetry(streamFn, resolveProviderRetryPolicy())).toBe(streamFn); + expect(withProviderRetry(streamFn, resolveProviderRetryPolicy({ enabled: false }))).toBe(streamFn); + expect(withProviderRetry(streamFn, resolveProviderRetryPolicy({ enabled: true, maxRetries: 0 }))).toBe(streamFn); + expect(withProviderRetryModels(models, resolveProviderRetryPolicy())).toBe(models); + expect(withProviderRetryModels(models, resolveProviderRetryPolicy({ enabled: true, maxRetries: 0 }))).toBe(models); + }); - expect(calls).toBe(2); - expect(events.some((event) => event.type === "error")).toBe(false); - expect(events.filter((event) => event.type === "text_delta").map((event) => event.delta)).toEqual(["done"]); + it("validates public policy values and the largest timer delay", () => { + for (const retry of [ + { enabled: "yes" }, + { maxRetries: -1 }, + { maxRetries: 1.5 }, + { maxRetries: Number.MAX_SAFE_INTEGER + 1 }, + { baseDelayMs: -1 }, + { baseDelayMs: Number.NaN }, + { baseDelayMs: Number.POSITIVE_INFINITY }, + { enabled: true, maxRetries: 2, baseDelayMs: 2_147_483_647 }, + ]) { + expect(() => resolveProviderRetryPolicy(retry as never)).toThrow(); + } + expect(() => resolveProviderRetryPolicy({ enabled: true, maxRetries: 0, baseDelayMs: Infinity })).toThrow(); + expect(resolveProviderRetryPolicy({ enabled: true, maxRetries: 100, baseDelayMs: 0 })).toMatchObject({ + enabled: true, + maxRetries: 100, + baseDelayMs: 0, + }); }); - it("uses exponential backoff and exposes the fourth error", async () => { + it("uses three retries with 2s/4s/8s defaults", async () => { vi.useFakeTimers(); let calls = 0; - const retrying = withProviderRetry(() => { + const retrying = withProviderRetry((() => { calls += 1; return errorStream("HTTP 429: rate limited"); - }); + }) as StreamFn, enabled); const result = retrying(model, context).result(); - await vi.advanceTimersByTimeAsync(2_000); - expect(calls).toBe(2); - await vi.advanceTimersByTimeAsync(4_000); - expect(calls).toBe(3); - await vi.advanceTimersByTimeAsync(8_000); - expect(calls).toBe(4); + for (const [delay, expectedCalls] of [[2_000, 2], [4_000, 3], [8_000, 4]] as const) { + await vi.advanceTimersByTimeAsync(delay - 1); + expect(calls).toBe(expectedCalls - 1); + await vi.advanceTimersByTimeAsync(1); + expect(calls).toBe(expectedCalls); + } expect((await result).stopReason).toBe("error"); }); - it.each([ - ["Please retry in 10.5s", 10_500], - ["retry delay of 2500 ms", 2_500], - ["Retry Delay: 3 s", 3_000], - ])("honors provider delay %s", async (message, delay) => { + it("uses custom retry count and delay without interpreting provider options", async () => { vi.useFakeTimers(); let calls = 0; - const retrying = withProviderRetry(() => (++calls === 1 ? errorStream(`HTTP 429: ${message}`) : textStream("ok"))); - const result = retrying(model, context).result(); - await vi.advanceTimersByTimeAsync(delay - 1); + const options = { maxRetries: 9, maxRetryDelayMs: 1 } satisfies SimpleStreamOptions; + const retrying = withProviderRetry(((_model, _context, received) => { + expect(received).toBe(options); + calls += 1; + return errorStream("HTTP 429: retry in 30s"); + }) as StreamFn, resolveProviderRetryPolicy({ enabled: true, maxRetries: 1, baseDelayMs: 25 })); + const result = retrying(model, context, options).result(); + await vi.advanceTimersByTimeAsync(24); expect(calls).toBe(1); await vi.advanceTimersByTimeAsync(1); - expect((await result).stopReason).toBe("stop"); + expect((await result).stopReason).toBe("error"); expect(calls).toBe(2); }); +}); + +describe("provider retry behavior", () => { + it("discards a failed attempt and replays one accepted snapshot sequence", async () => { + vi.useFakeTimers(); + let calls = 0; + const retrying = withProviderRetry((() => { + calls += 1; + return calls === 1 ? errorStream("HTTP 429", "discarded") : mutableTextStream(); + }) as StreamFn, enabled); + const eventsPromise = collect(retrying(model, context)); + await vi.advanceTimersByTimeAsync(2_000); + await vi.runAllTimersAsync(); + const events = await eventsPromise; + + expect(JSON.stringify(events)).not.toContain("discarded"); + expect(events.filter((event) => event.type === "text_delta").map((event) => event.partial.content)).toEqual([ + [{ type: "text", text: "a" }], + [{ type: "text", text: "ab" }], + ]); + expect(events.filter((event) => event.type === "done")).toHaveLength(1); + }); it.each([ - ["HTTP 429: retry in 61s", undefined], - ["HTTP 429: retry in 10s", 5_000], - ])("does not retry a provider delay above the cap", async (message, maxRetryDelayMs) => { + "invalid API key", + "insufficient_quota: quota exceeded", + "billing limit reached", + "ordinary model error", + ])("does not retry non-transient failure %s", async (error) => { let calls = 0; - const retrying = withProviderRetry(() => { + const retrying = withProviderRetry((() => { calls += 1; - return errorStream(message); - }); - expect((await retrying(model, context, { maxRetryDelayMs }).result()).stopReason).toBe("error"); + return errorStream(error); + }) as StreamFn, enabled); + expect((await retrying(model, context).result()).stopReason).toBe("error"); + expect(calls).toBe(1); + }); + + it("excludes context overflow before transient classification", async () => { + let calls = 0; + const retrying = withProviderRetry((() => { + calls += 1; + return errorStream("500: input exceeds the context window"); + }) as StreamFn, enabled); + await retrying(model, context).result(); expect(calls).toBe(1); }); - it.each(["retry in -3s", "retry delay: Infinity ms", "retry soon"])( - "falls back to exponential delay for malformed hint %s", - async (hint) => { + it.each(["HTTP 429", "503 service unavailable", "network error", "request timed out"])( + "retries classifier-accepted failure %s", + async (error) => { vi.useFakeTimers(); let calls = 0; - const retrying = withProviderRetry(() => (++calls === 1 ? errorStream(`HTTP 429: ${hint}`) : textStream("ok"))); + const retrying = withProviderRetry((() => (++calls === 1 ? errorStream(error) : textStream("ok"))) as StreamFn, enabled); const result = retrying(model, context).result(); - await vi.advanceTimersByTimeAsync(1_999); - expect(calls).toBe(1); - await vi.advanceTimersByTimeAsync(1); + await vi.advanceTimersByTimeAsync(2_000); expect((await result).stopReason).toBe("stop"); + expect(calls).toBe(2); }, ); - it("allows an uncapped provider delay when maxRetryDelayMs is zero", async () => { + it("normalizes retryable stream invocation and iteration failures", async () => { vi.useFakeTimers(); - let calls = 0; - const retrying = withProviderRetry(() => (++calls === 1 ? errorStream("HTTP 429: retry in 61s") : textStream("ok"))); - const result = retrying(model, context, { maxRetryDelayMs: 0 }).result(); - await vi.advanceTimersByTimeAsync(61_000); - expect((await result).stopReason).toBe("stop"); - expect(calls).toBe(2); + let invocationCalls = 0; + const invocation = withProviderRetry((() => { + if (++invocationCalls === 1) throw new Error("network error"); + return textStream("ok"); + }) as StreamFn, enabled); + const invocationResult = invocation(model, context).result(); + await vi.advanceTimersByTimeAsync(2_000); + expect((await invocationResult).stopReason).toBe("stop"); + + let rejectionCalls = 0; + const rejection = withProviderRetry((async () => { + if (++rejectionCalls === 1) throw new Error("fetch failed"); + return textStream("ok"); + }) as StreamFn, enabled); + const rejectionResult = rejection(model, context).result(); + await vi.advanceTimersByTimeAsync(2_000); + expect((await rejectionResult).stopReason).toBe("stop"); + + let iterationCalls = 0; + const iteration = withProviderRetry((() => { + iterationCalls += 1; + return iterationCalls === 1 ? throwingStream("socket hang up") : textStream("ok"); + }) as StreamFn, enabled); + const iterationResult = iteration(model, context).result(); + await vi.advanceTimersByTimeAsync(2_000); + expect((await iterationResult).stopReason).toBe("stop"); }); - it("waits the full uncapped delay when it exceeds the runtime timer limit", async () => { - vi.useFakeTimers(); + it("settles malformed and unsnapshotable streams without retrying", async () => { let calls = 0; - const delay = 3_000_000_000; - const retrying = withProviderRetry(() => - ++calls === 1 ? errorStream(`HTTP 429: retry in ${delay}ms`) : textStream("ok"), - ); - const result = retrying(model, context, { maxRetryDelayMs: 0 }).result(); + const malformed = withProviderRetry((() => { + calls += 1; + return emptyStream(); + }) as StreamFn, enabled); + const malformedEvents = await collect(malformed(model, context)); + expect(malformedEvents.at(-1)?.type).toBe("error"); + expect((await malformed(model, context).result()).errorMessage).toContain("without a terminal event"); - await vi.advanceTimersByTimeAsync(delay - 1); - expect(calls).toBe(1); - await vi.advanceTimersByTimeAsync(1); - expect((await result).stopReason).toBe("stop"); - expect(calls).toBe(2); + const unsnapshotable = withProviderRetry((() => { + calls += 1; + return uncloneableStream(); + }) as StreamFn, enabled); + const stream = unsnapshotable(model, context); + const [events, result] = await Promise.all([collect(stream), stream.result()]); + expect(events.at(-1)?.type).toBe("error"); + expect(result.stopReason).toBe("error"); + expect(calls).toBe(3); }); - it("does not retry non-retryable errors or aborted responses", async () => { - for (const message of [errorMessage("invalid API key"), errorMessage("aborted", "aborted")]) { - let calls = 0; - const retrying = withProviderRetry(() => { - calls += 1; - return message; - }); - await retrying(model, context).result(); - expect(calls).toBe(1); - } - }); + it("aborts before the first request and during backoff", async () => { + const alreadyAborted = new AbortController(); + alreadyAborted.abort(); + let calls = 0; + const retrying = withProviderRetry((() => { + calls += 1; + return errorStream("HTTP 429"); + }) as StreamFn, enabled); + const first = retrying(model, context, { signal: alreadyAborted.signal }); + const [firstEvents, firstResult] = await Promise.all([collect(first), first.result()]); + expect(firstEvents.filter((event) => event.type === "error")).toHaveLength(1); + expect(firstResult.stopReason).toBe("aborted"); + expect(calls).toBe(0); - it("emits only an aborted logical attempt when aborted during backoff", async () => { vi.useFakeTimers(); const controller = new AbortController(); - let calls = 0; - const retrying = withProviderRetry(() => { - calls += 1; - return errorStream("HTTP 429: retry in 10s", "discarded"); - }); - const eventsPromise = collect(retrying(model, context, { signal: controller.signal })); + const second = retrying(model, context, { signal: controller.signal }); await vi.advanceTimersByTimeAsync(1); controller.abort(); - const events = await eventsPromise; - + const [secondEvents, secondResult] = await Promise.all([collect(second), second.result()]); + expect(secondEvents.filter((event) => event.type === "error")).toHaveLength(1); + expect(secondResult.stopReason).toBe("aborted"); expect(calls).toBe(1); - expect(events.map((event) => event.type)).toEqual(["start", "error"]); - expect(events.at(-1)?.type === "error" && events.at(-1).error.stopReason).toBe("aborted"); - expect(JSON.stringify(events)).not.toContain("discarded"); }); - it("supports promise-returning streams and snapshots mutable partials", async () => { - const retrying = withProviderRetry(async () => mutableTextStream()); - const events = await collect(retrying(model, context)); - const deltas = events.filter((event) => event.type === "text_delta"); - expect(deltas[0].partial.content).toEqual([{ type: "text", text: "a" }]); - expect(deltas[1].partial.content).toEqual([{ type: "text", text: "ab" }]); + it("stops reading after the first terminal event", async () => { + let nextCalls = 0; + const retrying = withProviderRetry((() => terminalThenThrowStream(() => nextCalls++)) as StreamFn, enabled); + const stream = retrying(model, context); + const [events, result] = await Promise.all([collect(stream), stream.result()]); + expect(result.stopReason).toBe("stop"); + expect(events.filter((event) => event.type === "done")).toHaveLength(1); + expect(nextCalls).toBe(1); }); - it("applies retries to Models.completeSimple", async () => { - vi.useFakeTimers(); - let calls = 0; + it("wraps only Models simple paths and preserves generic delegates", async () => { + let simpleCalls = 0; + let genericCalls = 0; const models = { - streamSimple: () => (++calls === 1 ? errorStream("HTTP 429: rate limited") : textStream("done")), + streamSimple: () => (++simpleCalls === 1 ? errorStream("HTTP 429") : textStream("simple")), + stream: () => { + genericCalls += 1; + return textStream("generic"); + }, + complete: async () => { + genericCalls += 1; + return assistant("stop"); + }, } as Models; - const retryingModels = withProviderRetryModels(models); - const completion = retryingModels.completeSimple(model, context); + vi.useFakeTimers(); + const wrapped = withProviderRetryModels(models, enabled); + const simpleResult = wrapped.completeSimple(model, context); await vi.advanceTimersByTimeAsync(2_000); - - expect((await completion).content).toEqual([{ type: "text", text: "done" }]); - expect(calls).toBe(2); + expect((await simpleResult).content).toEqual([{ type: "text", text: "simple" }]); + await wrapped.stream(model, context).result(); + await wrapped.complete(model, context); + expect(simpleCalls).toBe(2); + expect(genericCalls).toBe(2); }); }); @@ -212,14 +299,6 @@ function errorStream(message: string, text?: string) { return stream; } -function errorMessage(message: string, reason: "error" | "aborted" = "error") { - const stream = createAssistantMessageEventStream(); - const value = assistant(reason, message); - stream.push({ type: "start", partial: value }); - stream.push({ type: "error", reason, error: value }); - return stream; -} - function textStream(text: string) { const stream = createAssistantMessageEventStream(); const value = assistant("stop"); @@ -245,6 +324,44 @@ function mutableTextStream() { return stream; } +function emptyStream() { + return { + async *[Symbol.asyncIterator]() {}, + } as ReturnType; +} + +function throwingStream(message: string) { + return { + async *[Symbol.asyncIterator]() { + throw new Error(message); + }, + } as ReturnType; +} + +function uncloneableStream() { + return { + async *[Symbol.asyncIterator]() { + yield { type: "start", partial: { ...assistant("stop"), uncloneable: () => {} } } as never; + }, + } as ReturnType; +} + +function terminalThenThrowStream(onNext: () => void) { + const value = assistant("stop"); + let index = 0; + return { + [Symbol.asyncIterator]() { + return { + async next() { + onNext(); + if (index++ === 0) return { done: false as const, value: { type: "done", reason: "stop", message: value } as const }; + throw new Error("must not read after terminal"); + }, + }; + }, + } as ReturnType; +} + async function collect(stream: AsyncIterable): Promise { const events: AssistantMessageEvent[] = []; for await (const event of stream) events.push(event); diff --git a/packages/cli/src/harness.ts b/packages/cli/src/harness.ts index 5dfaa97..3b1f83f 100644 --- a/packages/cli/src/harness.ts +++ b/packages/cli/src/harness.ts @@ -43,6 +43,8 @@ export interface BuildCuaHarnessOptions { extraTools?: CuaAgentHarnessOptions["extraTools"]; /** Override the pi `Models` collection requests stream through (mainly for tests). */ models?: Models; + /** Optional CUA-level retries around each provider request. Disabled by default. */ + retry?: CuaAgentHarnessOptions["retry"]; /** Override the catalog `baseUrl` on the resolved model (e.g. from `_BASE_URL`). */ modelBaseUrl?: string; } @@ -84,6 +86,7 @@ export function buildCuaHarness(opts: BuildCuaHarnessOptions): CuaAgentHarness { return composeSystemPrompt(runtime.defaultSystemPrompt, resources.skills ?? [], contextFiles); }, models: opts.models, + retry: opts.retry, }); return harness; } diff --git a/packages/cli/test/action-runner.test.ts b/packages/cli/test/action-runner.test.ts index ce5d3ba..333c87b 100644 --- a/packages/cli/test/action-runner.test.ts +++ b/packages/cli/test/action-runner.test.ts @@ -66,32 +66,34 @@ describe("action harness-runner", () => { try { fixture = await buildTestHarness({ turns: [ + { steps: [{ type: "tool_call", toolName: "click", args: { x: 9, y: 9 } }] }, { steps: [ { type: "text", text: "discarded" }, - { type: "tool_call", toolName: "click", args: { x: 9, y: 9 } }, { type: "error", message: "HTTP 429: Please retry in 10.367614288s" }, ], }, { steps: [{ type: "text", text: "done" }] }, ], + retry: { enabled: true }, }); const resultPromise = runAction( { action: "do", text: "recover" }, { harness: fixture.harness, browserHandle: handleFor(fixture), session: fixture.session, maxTurns: 3 }, ); - await vi.advanceTimersByTimeAsync(10_366); - expect(fixture.provider.callCount()).toBe(1); + await vi.advanceTimersByTimeAsync(1_999); + expect(fixture.provider.callCount()).toBe(2); await vi.advanceTimersByTimeAsync(1); const res = await resultPromise; - expect(fixture.provider.callCount()).toBe(2); + expect(fixture.provider.callCount()).toBe(3); expect(res.exitCode).toBe(0); expect(res.result.status).toBe("ok"); - expect(fixture.kernel.batchCalls).toHaveLength(0); + expect(fixture.kernel.batchCalls).toHaveLength(1); + expect(JSON.stringify(fixture.provider.lastContext())).toContain("toolResult"); const messages = (await fixture.session.getBranch()).filter((entry) => entry.type === "message"); - expect(messages).toHaveLength(2); + expect(messages).toHaveLength(4); expect(JSON.stringify(messages)).toContain("done"); expect(JSON.stringify(messages)).not.toContain("discarded"); } finally { diff --git a/packages/cli/test/fixtures/harness.ts b/packages/cli/test/fixtures/harness.ts index 83c0ef2..efc7a2e 100644 --- a/packages/cli/test/fixtures/harness.ts +++ b/packages/cli/test/fixtures/harness.ts @@ -25,6 +25,7 @@ export interface BuildTestHarnessOptions { skills?: Skill[]; /** CUA model ref. Defaults to "openai:gpt-5.5". */ modelRef?: string; + retry?: Parameters[0]["retry"]; } export async function buildTestHarness(opts: BuildTestHarnessOptions): Promise { @@ -46,6 +47,7 @@ export async function buildTestHarness(opts: BuildTestHarnessOptions): Promise Date: Thu, 9 Jul 2026 22:46:54 +0000 Subject: [PATCH 3/3] Preserve provider replay failures --- packages/agent/src/provider-retry.ts | 9 ++++----- packages/agent/test/provider-retry.test.ts | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/packages/agent/src/provider-retry.ts b/packages/agent/src/provider-retry.ts index bdc4880..3bfa848 100644 --- a/packages/agent/src/provider-retry.ts +++ b/packages/agent/src/provider-retry.ts @@ -109,6 +109,7 @@ async function runWithRetry( if (settled) return; settled = true; let terminalEmitted = false; + let result = message; try { if (events) { for (const event of events) { @@ -122,17 +123,15 @@ async function runWithRetry( } } catch (error) { if (!terminalEmitted) { - const failure = createFailureMessage(model, `Failed to replay provider response: ${errorMessage(error)}`); + result = createFailureMessage(model, `Failed to replay provider response: ${errorMessage(error)}`); try { - output.push({ type: "error", reason: "error", error: failure }); + output.push({ type: "error", reason: "error", error: result }); } catch { // end(result) still settles result() and any waiting iterator. } - output.end(failure); - return; } } finally { - output.end(message); + output.end(result); } }; diff --git a/packages/agent/test/provider-retry.test.ts b/packages/agent/test/provider-retry.test.ts index ed18180..edc89cf 100644 --- a/packages/agent/test/provider-retry.test.ts +++ b/packages/agent/test/provider-retry.test.ts @@ -205,6 +205,22 @@ describe("provider retry behavior", () => { expect(calls).toBe(3); }); + it("preserves a replay failure as the terminal result", async () => { + const retrying = withProviderRetry((() => textStream("ok")) as StreamFn, enabled); + const stream = retrying(model, context); + const push = stream.push.bind(stream); + let calls = 0; + stream.push = (event) => { + if (calls++ === 0) throw new Error("replay failed"); + return push(event); + }; + + const result = await stream.result(); + + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toContain("Failed to replay provider response: replay failed"); + }); + it("aborts before the first request and during backoff", async () => { const alreadyAborted = new AbortController(); alreadyAborted.abort();