From faa4ec6accd8c7420a5bd8525c778c087bb81470 Mon Sep 17 00:00:00 2001 From: Micah Zoltu Date: Wed, 17 Jun 2026 18:35:38 +0000 Subject: [PATCH] Honor X-RateLimit-Reset header and extend retry budget to recover from AI provider 429s. Also extracts AI fetch retry logic into testable orchestration with injected leaf dependencies. --- source/ai.mts | 134 ++++++++++++++++++++++++++++++------ source/index.mts | 6 +- tests/ai.test.mts | 170 +++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 284 insertions(+), 26 deletions(-) diff --git a/source/ai.mts b/source/ai.mts index a0467d8..0cb8921 100644 --- a/source/ai.mts +++ b/source/ai.mts @@ -38,41 +38,131 @@ export function parseAiConfiguration(environment: Record Promise +export type Random = () => number +export type Now = () => number + +export interface FetchDependencies { + readonly httpFetch: Fetch + readonly sleep: Sleep + readonly random: Random + readonly now: Now +} + +export function createHttpFetch(configuration: AiConfiguration): Fetch { return async (signal, body, headers) => { const url = `${configuration.apiUrl.replace(/\/$/, "")}/chat/completions` const requestHeaders: Record = { ...headers } if (configuration.apiKey) requestHeaders["Authorization"] = `Bearer ${configuration.apiKey}` + return fetch(url, { method: "POST", headers: requestHeaders, body, signal }) + } +} - let lastResponse: Response | undefined - for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) { - if (attempt > 0) { - const backoff = Math.min(INITIAL_BACKOFF_MILLISECONDS * Math.pow(2, attempt - 1), MAX_BACKOFF_MILLISECONDS) - const jitter = backoff * (0.5 + Math.random() * 0.5) - await sleepWithSignal(jitter, signal) - } +export function createSleep(): Sleep { + return sleepWithSignal +} + +export function createRandom(): Random { + return () => Math.random() +} + +export function createNow(): Now { + return () => Date.now() +} + +interface RetryDelayInput { + readonly retryAfter: string | null + readonly rateLimitReset: string | null + readonly attempt: number + readonly deadlineRemainingMilliseconds: number + readonly now: number + readonly random: number +} - lastResponse = await fetch(url, { method: "POST", headers: requestHeaders, body, signal }) +// Precedence: Retry-After > X-RateLimit-Reset (epoch or duration, heuristic) > exponential backoff. +function computeRetryDelay(input: RetryDelayInput): number | null { + if (input.deadlineRemainingMilliseconds <= 0) return null - if (lastResponse.status === 429 || (lastResponse.status >= 500 && lastResponse.status <= 599)) { - if (attempt === MAX_RETRIES) return lastResponse - if (lastResponse.status === 429) { - const retryAfter = lastResponse.headers.get("Retry-After") - const delay = retryAfter ? Number.parseInt(retryAfter, 10) * 1000 : Math.min(INITIAL_BACKOFF_MILLISECONDS * Math.pow(2, attempt), MAX_BACKOFF_MILLISECONDS) - await sleepWithSignal(delay, signal) - } - continue + let computed: number | undefined + + if (input.retryAfter !== null) { + const seconds = Number.parseInt(input.retryAfter, 10) + if (Number.isFinite(seconds) && seconds >= 0) computed = seconds * 1000 + } + + if (computed === undefined && input.rateLimitReset !== null) { + const value = Number.parseFloat(input.rateLimitReset) + if (Number.isFinite(value) && value >= 0) { + const milliseconds = value < EPOCH_THRESHOLD + ? value * 1000 + : Math.max(0, value - input.now / 1000) * 1000 + computed = milliseconds + } + } + + if (computed === undefined) { + const backoff = INITIAL_BACKOFF_MILLISECONDS * Math.pow(2, input.attempt) + computed = Math.min(backoff, MAX_BACKOFF_MILLISECONDS) + } + + const capped = Math.min(computed, MAX_SINGLE_WAIT_MILLISECONDS, input.deadlineRemainingMilliseconds) + return capped * (0.5 + input.random * 0.5) +} + +function isRetryableStatus(status: number): boolean { + return status === 429 || (status >= 500 && status <= 599) +} + +async function fetchWithRetries(dependencies: FetchDependencies, signal: AbortSignal, body: string, headers?: Record): Promise { + const deadline = dependencies.now() + RETRY_DEADLINE_MILLISECONDS + let lastResponse: Response | undefined + for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) { + if (attempt > 0) { + const now = dependencies.now() + const delay = computeRetryDelay({ + retryAfter: lastResponse?.headers.get("Retry-After") ?? null, + rateLimitReset: lastResponse?.headers.get("X-RateLimit-Reset") ?? null, + attempt: attempt - 1, + deadlineRemainingMilliseconds: deadline - now, + now, + random: dependencies.random(), + }) + if (delay === null) { + if (lastResponse === undefined) throw new Error("Retry deadline exhausted before any response was received") + return lastResponse } + await dependencies.sleep(delay, signal) + } - return lastResponse + lastResponse = await dependencies.httpFetch(signal, body, headers) + + if (isRetryableStatus(lastResponse.status)) { + if (attempt === MAX_RETRIES) return lastResponse + if (dependencies.now() >= deadline) return lastResponse + continue } - return lastResponse! + return lastResponse } + + throw new Error("createFetch retry loop exited without returning a response") +} + +export function createFetch(dependencies: FetchDependencies): Fetch { + return (signal, body, headers) => fetchWithRetries(dependencies, signal, body, headers) } async function runAgent(dependencies: { fetch: Fetch; spawnGit: SpawnGit; logger: Logger; debugWriter: DebugWriter }, agent: Agent, baseCommit: string, baseCommitContext: BaseCommitContext, diffText: string, model: string, profile: ProviderProfile, agentInputs?: Map, outputValidator?: OutputValidator): Promise { diff --git a/source/index.mts b/source/index.mts index 417c677..f60a9db 100644 --- a/source/index.mts +++ b/source/index.mts @@ -1,8 +1,8 @@ -import { createFetch, parseAiConfiguration } from "./ai.mts" import { createReadAgentsFromDisk } from "./agents.mts" +import { createFetch, createHttpFetch, createNow, createRandom, createSleep, parseAiConfiguration } from "./ai.mts" import { getConfiguration } from "./configuration.mts" -import { createSpawnGit } from "./diff.mts" import { createDebugWriter } from "./debug.mts" +import { createSpawnGit } from "./diff.mts" import { createGithubFetch } from "./github.mts" import { createLogger } from "./logger.mts" import { runOnCommentTrigger, runOnLocalDiff, runOnPullRequest } from "./orchestrator.mts" @@ -21,7 +21,7 @@ async function main(): Promise { const readAgents = createReadAgentsFromDisk() const spawnGit = createSpawnGit(workspaceDirectory) const aiConfiguration = parseAiConfiguration(Bun.env) - const fetch = createFetch(aiConfiguration) + const fetch = createFetch({ httpFetch: createHttpFetch(aiConfiguration), sleep: createSleep(), random: createRandom(), now: createNow() }) const githubFetch = createGithubFetch(logger) const profile = selectProviderProfile(aiConfiguration.apiUrl, aiConfiguration.model) diff --git a/tests/ai.test.mts b/tests/ai.test.mts index eeb34fd..91986ff 100644 --- a/tests/ai.test.mts +++ b/tests/ai.test.mts @@ -1,5 +1,5 @@ import { describe, it, expect } from "bun:test" -import { analyze, parseAiConfiguration } from "../source/ai.mts" +import { analyze, createFetch, parseAiConfiguration, type Sleep, type Random, type Now } from "../source/ai.mts" import { IDENTITY_PROFILE } from "../source/provider-profiles.mts" import type { Agent } from "../source/agents.mts" import type { DebugWriter } from "../source/debug.mts" @@ -351,3 +351,171 @@ describe("analyze", () => { }) }) }) + +describe("createFetch", () => { + const BASE_TIME = 1_700_000_000_000 + + function response(status: number, headers: Record = {}): Response { + return new Response(null, { status, headers }) + } + + function createFakeHttpFetch(responses: readonly Response[]): { httpFetch: Fetch; callCount: () => number } { + let i = 0 + return { + httpFetch: async () => responses[Math.min(i++, responses.length - 1)]!, + callCount: () => i, + } + } + + function createFakeSleep(): { sleep: Sleep; delays: number[] } { + const delays: number[] = [] + return { + sleep: async (ms) => { delays.push(ms) }, + delays, + } + } + + const fixedRandom: Random = () => 1 + const constantNow = (time: number): Now => () => time + function scriptedNow(values: readonly number[]): Now { + let i = 0 + return () => values[Math.min(i++, values.length - 1)]! + } + + function buildFetch(responses: readonly Response[], now: Now = constantNow(BASE_TIME)): { fetch: Fetch; callCount: () => number; delays: () => number[] } { + const http = createFakeHttpFetch(responses) + const sleep = createFakeSleep() + const fetch = createFetch({ httpFetch: http.httpFetch, sleep: sleep.sleep, random: fixedRandom, now }) + return { fetch, callCount: http.callCount, delays: () => sleep.delays } + } + + it("returns a success response without retrying", async () => { + const { fetch, callCount, delays } = buildFetch([response(200)]) + const result = await fetch(new AbortController().signal, "body") + expect(result.status).toBe(200) + expect(callCount()).toBe(1) + expect(delays()).toEqual([]) + }) + + it("does not retry on a non-retryable 4xx status", async () => { + const { fetch, callCount, delays } = buildFetch([response(400)]) + const result = await fetch(new AbortController().signal, "body") + expect(result.status).toBe(400) + expect(callCount()).toBe(1) + expect(delays()).toEqual([]) + }) + + it("retries a 429 with exponential backoff when no headers are present", async () => { + const { fetch, callCount, delays } = buildFetch([response(429), response(200)]) + const result = await fetch(new AbortController().signal, "body") + expect(result.status).toBe(200) + expect(callCount()).toBe(2) + expect(delays()).toEqual([1_000]) + }) + + it("retries a 5xx status", async () => { + const { fetch, callCount, delays } = buildFetch([response(503), response(200)]) + const result = await fetch(new AbortController().signal, "body") + expect(result.status).toBe(200) + expect(callCount()).toBe(2) + expect(delays()).toEqual([1_000]) + }) + + it("honors the Retry-After header in seconds", async () => { + const { fetch, delays } = buildFetch([response(429, { "Retry-After": "5" }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([5_000]) + }) + + it("treats X-RateLimit-Reset below the epoch threshold as seconds-until-reset", async () => { + const { fetch, delays } = buildFetch([response(429, { "X-RateLimit-Reset": "30" }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([30_000]) + }) + + it("treats X-RateLimit-Reset at/above the epoch threshold as a Unix timestamp", async () => { + const resetEpochSeconds = (BASE_TIME + 10_000) / 1000 + const { fetch, delays } = buildFetch([response(429, { "X-RateLimit-Reset": String(resetEpochSeconds) }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([10_000]) + }) + + it("clamps a past Unix epoch reset to a zero base wait", async () => { + const resetEpochSeconds = (BASE_TIME - 5_000) / 1000 + const { fetch, delays } = buildFetch([response(429, { "X-RateLimit-Reset": String(resetEpochSeconds) }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([0]) + }) + + it("prefers Retry-After over X-RateLimit-Reset when both are present", async () => { + const { fetch, delays } = buildFetch([response(429, { "Retry-After": "2", "X-RateLimit-Reset": "999" }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([2_000]) + }) + + it("applies exponential backoff progression and caps at MAX_BACKOFF, then returns the final 429", async () => { + const elevenFailures = Array.from({ length: 11 }, () => response(429)) + const { fetch, callCount, delays } = buildFetch(elevenFailures) + const result = await fetch(new AbortController().signal, "body") + expect(result.status).toBe(429) + expect(callCount()).toBe(11) + expect(delays()).toEqual([1_000, 2_000, 4_000, 8_000, 16_000, 30_000, 30_000, 30_000, 30_000, 30_000]) + }) + + it("caps a far-future reset at MAX_SINGLE_WAIT_MILLISECONDS", async () => { + const farFutureEpochSeconds = (BASE_TIME + 10 * 60 * 1000) / 1000 + const { fetch, delays } = buildFetch([response(429, { "X-RateLimit-Reset": String(farFutureEpochSeconds) }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([240_000]) + }) + + it("returns the 429 without sleeping when the deadline is already exhausted after a fetch", async () => { + const now = scriptedNow([BASE_TIME, BASE_TIME + 301_000]) + const { fetch, callCount, delays } = buildFetch([response(429)], now) + const result = await fetch(new AbortController().signal, "body") + expect(result.status).toBe(429) + expect(callCount()).toBe(1) + expect(delays()).toEqual([]) + }) + + it("caps the wait by the remaining deadline when Retry-After exceeds it", async () => { + const now = scriptedNow([BASE_TIME, BASE_TIME, BASE_TIME + 290_000]) + const { fetch, callCount, delays } = buildFetch([response(429, { "Retry-After": "60" }), response(200)], now) + const result = await fetch(new AbortController().signal, "body") + expect(result.status).toBe(200) + expect(callCount()).toBe(2) + expect(delays()).toEqual([10_000]) + }) + + it("falls back to exponential backoff when Retry-After is non-numeric", async () => { + const { fetch, delays } = buildFetch([response(429, { "Retry-After": "soon" }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([1_000]) + }) + + it("falls back to exponential backoff when X-RateLimit-Reset is non-numeric", async () => { + const { fetch, delays } = buildFetch([response(429, { "X-RateLimit-Reset": "soon" }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([1_000]) + }) + + it("falls back to exponential backoff when X-RateLimit-Reset is negative", async () => { + const { fetch, delays } = buildFetch([response(429, { "X-RateLimit-Reset": "-5" }), response(200)]) + await fetch(new AbortController().signal, "body") + expect(delays()).toEqual([1_000]) + }) + + it("forwards the caller's abort signal to the in-flight httpFetch so requests can be cancelled mid-flight", async () => { + const controller = new AbortController() + let observedSignal: AbortSignal | undefined + const httpFetch: Fetch = async (signal) => { + observedSignal = signal + return response(200) + } + const fetch = createFetch({ httpFetch, sleep: createFakeSleep().sleep, random: fixedRandom, now: constantNow(BASE_TIME) }) + + await fetch(controller.signal, "body") + + expect(observedSignal).toBe(controller.signal) + }) +})