diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 38ca37b..cde7de4 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -26,6 +26,12 @@ import { type SimpleStreamOptions, } from "@onkernel/cua-ai"; import type Kernel from "@onkernel/sdk"; +import { + type CuaRetryOptions, + resolveProviderRetryPolicy, + withProviderRetry, + withProviderRetryModels, +} from "./provider-retry"; import { buildCuaComputerTools } from "./tools"; import { InternalComputerTranslator, type KernelBrowser } from "./translator/translator"; @@ -78,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; }; /** @@ -116,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; }; /** @@ -296,6 +306,7 @@ export class CuaAgent extends Agent { mode, nativeTool, playwright, + retry, ...agentOptions } = options; const runtime = new CuaRuntimeController({ @@ -308,13 +319,17 @@ export class CuaAgent extends Agent { playwright, onPayload, }); + const retryingStream = withProviderRetry( + streamFn ?? defaultCuaStream, + resolveProviderRetryPolicy(retry), + ); 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({ @@ -446,6 +461,7 @@ export class CuaAgentHarness< systemPrompt, onPayload, activeToolNames, + retry, ...harnessOptions } = options; const runtime = new CuaRuntimeController({ @@ -459,11 +475,15 @@ export class CuaAgentHarness< onPayload, }); const resolvedTools = runtime.tools(); + const retryingModels = withProviderRetryModels( + models ?? cuaModels(), + resolveProviderRetryPolicy(retry), + ); 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/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 new file mode 100644 index 0000000..3bfa848 --- /dev/null +++ b/packages/agent/src/provider-retry.ts @@ -0,0 +1,259 @@ +import type { StreamFn } from "@earendil-works/pi-agent-core"; +import { + createAssistantMessageEventStream, + isContextOverflow, + isRetryableAssistantError, + type Api, + type AssistantMessage, + type AssistantMessageEvent, + type Context, + type Model, + type Models, + type SimpleStreamOptions, +} from "@earendil-works/pi-ai"; + +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, +) => ReturnType; + +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; + + 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`); + } + } + + 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, 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), + getModels: (provider) => models.getModels(provider), + getModel: (provider, id) => models.getModel(provider, id), + refresh: (provider) => models.refresh(provider), + getAuth: (model) => models.getAuth(model), + 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(), + }; +} + +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; + let result = message; + 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) { + result = createFailureMessage(model, `Failed to replay provider response: ${errorMessage(error)}`); + try { + output.push({ type: "error", reason: "error", error: result }); + } catch { + // end(result) still settles result() and any waiting iterator. + } + } + } finally { + output.end(result); + } + }; + + 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 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 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, + 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: message, + timestamp: Date.now(), + }; +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +class SnapshotError extends Error {} + +function wait(delayMs: number, signal: AbortSignal | undefined): Promise { + if (signal?.aborted) return Promise.resolve(false); + return new Promise((resolve) => { + 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 d7d2273..b699f7c 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -231,6 +231,66 @@ 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 { + 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, + retry: { enabled: true }, + 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..edc89cf --- /dev/null +++ b/packages/agent/test/provider-retry.test.ts @@ -0,0 +1,385 @@ +import { + createAssistantMessageEventStream, + type Api, + type AssistantMessage, + 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 { + 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: [] }; + +const enabled = resolveProviderRetryPolicy({ enabled: true }); + +afterEach(() => vi.useRealTimers()); + +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); + }); + + 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 three retries with 2s/4s/8s defaults", async () => { + vi.useFakeTimers(); + let calls = 0; + const retrying = withProviderRetry((() => { + calls += 1; + return errorStream("HTTP 429: rate limited"); + }) as StreamFn, enabled); + const result = retrying(model, context).result(); + + 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("uses custom retry count and delay without interpreting provider options", async () => { + vi.useFakeTimers(); + let calls = 0; + 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("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([ + "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((() => { + calls += 1; + 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(["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(error) : textStream("ok"))) as StreamFn, enabled); + const result = retrying(model, context).result(); + await vi.advanceTimersByTimeAsync(2_000); + expect((await result).stopReason).toBe("stop"); + expect(calls).toBe(2); + }, + ); + + it("normalizes retryable stream invocation and iteration failures", async () => { + vi.useFakeTimers(); + 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("settles malformed and unsnapshotable streams without retrying", async () => { + let calls = 0; + 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"); + + 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("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(); + 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); + + vi.useFakeTimers(); + const controller = new AbortController(); + const second = retrying(model, context, { signal: controller.signal }); + await vi.advanceTimersByTimeAsync(1); + controller.abort(); + 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); + }); + + 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("wraps only Models simple paths and preserves generic delegates", async () => { + let simpleCalls = 0; + let genericCalls = 0; + const models = { + streamSimple: () => (++simpleCalls === 1 ? errorStream("HTTP 429") : textStream("simple")), + stream: () => { + genericCalls += 1; + return textStream("generic"); + }, + complete: async () => { + genericCalls += 1; + return assistant("stop"); + }, + } as Models; + vi.useFakeTimers(); + const wrapped = withProviderRetryModels(models, enabled); + const simpleResult = wrapped.completeSimple(model, context); + await vi.advanceTimersByTimeAsync(2_000); + 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); + }); +}); + +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 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; +} + +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); + return events; +} 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 6833e79..333c87b 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,46 @@ 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: "tool_call", toolName: "click", args: { x: 9, y: 9 } }] }, + { + steps: [ + { type: "text", text: "discarded" }, + { 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(1_999); + expect(fixture.provider.callCount()).toBe(2); + await vi.advanceTimersByTimeAsync(1); + const res = await resultPromise; + + expect(fixture.provider.callCount()).toBe(3); + expect(res.exitCode).toBe(0); + expect(res.result.status).toBe("ok"); + 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(4); + 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: [ 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