feat(core): add retry for transient HTTP and stream-level errors

Add retry mechanism driven by pi settings.json retry.provider config (timeoutMs, maxRetries, maxRetryDelayMs). HTTP-level retries handle 429/5xx with exponential backoff and jitter, respecting Retry-After headers (seconds and HTTP-date formats). Stream-level retries handle cases where the API returns 200 OK but sends an error event in the stream body. Retries only when no content has been emitted yet. Per-attempt timeout via AbortController with automatic retry. Clean abort propagation through the retry loop.
2026-06-02 11:38:10 +08:00
parent f48291f252
commit c3107dbc1c
7 changed files with 610 additions and 1995 deletions
@@ -112,6 +112,7 @@ export function createTestDeps(overrides: Partial<CoreDependencies> = {}): TestD
    now: () => new Date("2026-05-05T12:00:00Z").getTime(),
    uuid: () => "00000000-0000-4000-8000-000000000000",
    cwd: () => "/repo",
+    delay: async () => {},
    ...overrides,
  })
  return { streamCommandCode, calculatedUsages }
@@ -124,12 +125,15 @@ type SuccessPlan = {
  chunks?: string[]
  delays?: number[]
  hangAfterLast?: boolean
+  /** Delay in ms before the server starts sending the response. */
+  responseDelay?: number
 }

 type ErrorPlan = {
  type: "error"
  status: number
  body: string
+  headers?: Record<string, string>
 }

 export type ResponsePlan = SuccessPlan | ErrorPlan
@@ -146,6 +150,7 @@ function headersToRecord(headers: IncomingHttpHeaders): Record<string, string> {
 export interface MockCommandCodeServer {
  baseUrl(): string
  mockResponse(plan: ResponsePlan): void
+  mockResponseQueue(plans: ResponsePlan[]): void
  reset(): void
  close(): Promise<void>
  lastRequestBody(): unknown
@@ -155,7 +160,7 @@ export interface MockCommandCodeServer {
 }

 export async function startMockCommandCodeServer(): Promise<MockCommandCodeServer> {
-  let nextPlan: ResponsePlan = { type: "success", events: [] }
+  let planQueue: ResponsePlan[] = [{ type: "success", events: [] }]
  let lastBody: unknown
  let lastHeaders: Record<string, string> = {}
  let requests = 0
@@ -183,9 +188,12 @@ export async function startMockCommandCodeServer(): Promise<MockCommandCodeServe
        lastBody = undefined
      }

-      const plan = nextPlan
+      // Pop the first plan from the queue; keep the last one as fallback.
+      const plan = planQueue.length > 1 ? planQueue.shift()! : planQueue[0]
+
      if (plan.type === "error") {
-        res.writeHead(plan.status, { "Content-Type": "text/plain" })
+        const headers: Record<string, string> = { "Content-Type": "text/plain", ...plan.headers }
+        res.writeHead(plan.status, headers)
        res.end(plan.body)
        return
      }
@@ -223,7 +231,11 @@ export async function startMockCommandCodeServer(): Promise<MockCommandCodeServe
        }
      }

-      sendNext()
+      if (plan.responseDelay) {
+        setTimeout(sendNext, plan.responseDelay)
+      } else {
+        sendNext()
+      }
    })
  })

@@ -238,10 +250,13 @@ export async function startMockCommandCodeServer(): Promise<MockCommandCodeServe
  return {
    baseUrl: () => `http://127.0.0.1:${port}`,
    mockResponse(plan: ResponsePlan) {
-      nextPlan = plan
+      planQueue = [plan]
+    },
+    mockResponseQueue(plans: ResponsePlan[]) {
+      planQueue = [...plans]
    },
    reset() {
-      nextPlan = { type: "success", events: [] }
+      planQueue = [{ type: "success", events: [] }]
      lastBody = undefined
      lastHeaders = {}
      requests = 0
@@ -0,0 +1,367 @@
+/**
+ * Tests for retry and timeout behaviour driven by pi settings.json
+ * retry config (timeoutMs, maxRetries, maxRetryDelayMs).
+ */
+
+import assert from "node:assert/strict"
+import { after, before, beforeEach, describe, it } from "node:test"
+
+import type { AssistantMessageEvent } from "../src/core.ts"
+import {
+  collectEvents,
+  createTestDeps,
+  makeContext,
+  makeModel,
+  startMockCommandCodeServer,
+  type MockCommandCodeServer,
+} from "./helpers.ts"
+
+let server: MockCommandCodeServer
+
+before(async () => {
+  server = await startMockCommandCodeServer()
+})
+
+after(async () => {
+  await server.close()
+})
+
+beforeEach(() => {
+  server.reset()
+})
+
+function eventTypes(events: readonly AssistantMessageEvent[]): string[] {
+  return events.map((event) => event.type)
+}
+
+describe("streamCommandCode — retry on transient errors", () => {
+  it("retries on 429 and succeeds on the second attempt", async () => {
+    server.mockResponseQueue([
+      { type: "error", status: 429, body: "rate limited" },
+      {
+        type: "success",
+        events: [
+          JSON.stringify({ type: "text-delta", text: "ok" }),
+          JSON.stringify({ type: "finish", finishReason: "stop" }),
+        ],
+      },
+    ])
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    assert.equal(server.requestCount(), 2)
+    assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "text_end", "done"])
+    const done = events.at(-1)
+    if (done?.type !== "done") throw new Error("expected done")
+    assert.equal(done.reason, "stop")
+  })
+
+  it("retries on 500 and succeeds on the second attempt", async () => {
+    server.mockResponseQueue([
+      { type: "error", status: 500, body: "internal server error" },
+      {
+        type: "success",
+        events: [JSON.stringify({ type: "finish", finishReason: "stop" })],
+      },
+    ])
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    assert.equal(server.requestCount(), 2)
+    assert.equal(events.at(-1)?.type, "done")
+  })
+
+  it("does NOT retry on 400 (non-retryable client error)", async () => {
+    server.mockResponse({ type: "error", status: 400, body: "bad request" })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    assert.equal(server.requestCount(), 1)
+    assert.deepEqual(eventTypes(events), ["start", "error"])
+    const last = events.at(-1)
+    if (last?.type !== "error") throw new Error("expected error")
+    assert.match(last.error.errorMessage ?? "", /400/)
+  })
+
+  it("exhausts maxRetries and emits an error", async () => {
+    server.mockResponse({ type: "error", status: 503, body: "unavailable" })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        maxRetries: 3,
+      }),
+    )
+
+    // initial attempt + 3 retries = 4 total
+    assert.equal(server.requestCount(), 4)
+    assert.deepEqual(eventTypes(events), ["start", "error"])
+    const last503 = events.at(-1)
+    if (last503?.type !== "error") throw new Error("expected error")
+    assert.match(last503.error.errorMessage ?? "", /503/)
+  })
+})
+
+describe("streamCommandCode — Retry-After header", () => {
+  it("respects Retry-After delay in seconds", async () => {
+    let delayCalled = false
+    server.mockResponseQueue([
+      {
+        type: "error",
+        status: 429,
+        body: "rate limited",
+        headers: { "retry-after": "2" },
+      },
+      {
+        type: "success",
+        events: [JSON.stringify({ type: "finish", finishReason: "stop" })],
+      },
+    ])
+    const { streamCommandCode } = createTestDeps({
+      apiBase: server.baseUrl(),
+      delay: async (ms: number) => {
+        delayCalled = true
+        assert.equal(ms, 2000)
+      },
+    })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    assert.equal(server.requestCount(), 2)
+    assert.equal(events.at(-1)?.type, "done")
+    assert.ok(delayCalled, "delay should have been called with Retry-After value")
+  })
+
+  it("fails immediately when Retry-After exceeds maxRetryDelayMs", async () => {
+    server.mockResponse({
+      type: "error",
+      status: 429,
+      body: "rate limited",
+      headers: { "retry-after": "300" },
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        maxRetryDelayMs: 10_000,
+      }),
+    )
+
+    assert.equal(server.requestCount(), 1)
+    assert.deepEqual(eventTypes(events), ["start", "error"])
+    const lastMax = events.at(-1)
+    if (lastMax?.type !== "error") throw new Error("expected error")
+    assert.match(lastMax.error.errorMessage ?? "", /exceeds max/)
+  })
+})
+
+describe("streamCommandCode — timeout", () => {
+  it("retries on per-attempt timeout and succeeds", async () => {
+    server.mockResponseQueue([
+      {
+        type: "success",
+        events: [JSON.stringify({ type: "finish", finishReason: "stop" })],
+        hangAfterLast: true,
+        responseDelay: 200,
+      },
+      {
+        type: "success",
+        events: [
+          JSON.stringify({ type: "text-delta", text: "fast" }),
+          JSON.stringify({ type: "finish", finishReason: "stop" }),
+        ],
+      },
+    ])
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        timeoutMs: 50,
+      }),
+      5_000,
+    )
+
+    assert.equal(server.requestCount(), 2)
+    assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "text_end", "done"])
+  })
+
+  it("emits error when all retry attempts time out", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [JSON.stringify({ type: "finish", finishReason: "stop" })],
+      hangAfterLast: true,
+      responseDelay: 200,
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        timeoutMs: 50,
+        maxRetries: 1,
+      }),
+      5_000,
+    )
+
+    // initial + 1 retry = 2
+    assert.equal(server.requestCount(), 2)
+    assert.deepEqual(eventTypes(events), ["start", "error"])
+  })
+})
+
+describe("streamCommandCode — abort cancels retry loop", () => {
+  it("user abort stops retries immediately", async () => {
+    server.mockResponse({ type: "error", status: 500, body: "error" })
+    const controller = new AbortController()
+    const { streamCommandCode } = createTestDeps({
+      apiBase: server.baseUrl(),
+      delay: async (_ms: number, signal: AbortSignal) => {
+        // Abort during the retry delay
+        controller.abort()
+        // Simulate the real delay which rejects on abort
+        return new Promise<void>((_, reject) => {
+          if (signal.aborted) reject(new DOMException("Aborted", "AbortError"))
+          signal.addEventListener("abort", () => reject(new DOMException("Aborted", "AbortError")))
+        })
+      },
+    })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        signal: controller.signal,
+        maxRetries: 10,
+      }),
+    )
+
+    // Should only have made 1 request (the initial one), then aborted during delay
+    assert.equal(server.requestCount(), 1)
+    assert.deepEqual(eventTypes(events), ["start", "error"])
+    const error = events.at(-1)
+    if (error?.type !== "error") throw new Error("expected error")
+    assert.equal(error.reason, "aborted")
+  })
+})
+
+describe("streamCommandCode — retry defaults", () => {
+  it("uses default maxRetries of 2 when not specified", async () => {
+    server.mockResponse({ type: "error", status: 500, body: "error" })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    await collectEvents(streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }))
+
+    // initial + 2 retries = 3
+    assert.equal(server.requestCount(), 3)
+  })
+
+  it("respects maxRetries: 0 (no retries)", async () => {
+    server.mockResponse({ type: "error", status: 500, body: "error" })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        maxRetries: 0,
+      }),
+    )
+
+    assert.equal(server.requestCount(), 1)
+    assert.deepEqual(eventTypes(events), ["start", "error"])
+  })
+})
+
+describe("streamCommandCode — stream-level error retry", () => {
+  it("retries when API returns 200 OK but stream contains an error event", async () => {
+    server.mockResponseQueue([
+      {
+        type: "success",
+        events: [
+          JSON.stringify({
+            type: "error",
+            error: "Service temporarily unavailable. Please try again shortly.",
+          }),
+        ],
+      },
+      {
+        type: "success",
+        events: [
+          JSON.stringify({ type: "text-delta", text: "ok" }),
+          JSON.stringify({ type: "finish", finishReason: "stop" }),
+        ],
+      },
+    ])
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    assert.equal(server.requestCount(), 2)
+    assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "text_end", "done"])
+  })
+
+  it("exhausts retries on persistent stream-level errors", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [
+        JSON.stringify({
+          type: "error",
+          error: "Service temporarily unavailable. Please try again shortly.",
+        }),
+      ],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), {
+        apiKey: "mock-key",
+        maxRetries: 3,
+      }),
+    )
+
+    // initial + 3 retries = 4
+    assert.equal(server.requestCount(), 4)
+    assert.deepEqual(eventTypes(events), ["start", "error"])
+    const last = events.at(-1)
+    if (last?.type !== "error") throw new Error("expected error")
+    assert.match(last.error.errorMessage ?? "", /temporarily unavailable/)
+  })
+
+  it("does NOT retry stream error when content was already emitted", async () => {
+    server.mockResponseQueue([
+      {
+        type: "success",
+        events: [
+          JSON.stringify({ type: "text-delta", text: "partial" }),
+          JSON.stringify({
+            type: "error",
+            error: "Service temporarily unavailable",
+          }),
+        ],
+      },
+    ])
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    // Only 1 request — no retry because content was already emitted.
+    assert.equal(server.requestCount(), 1)
+    assert.deepEqual(eventTypes(events), ["start", "text_start", "text_delta", "error"])
+  })
+})