feat(FN-705): add rate-limit retry with exponential backoff for AI agents
- Add rate-limit retry utility with exponential backoff, jitter, and configurable limits - Integrate retry logic into executor, triage, and merger agent sessions - Export retryOnRateLimit from engine package index - Remove unused DirectoryPicker component, stale CSS, and dead dashboard routes - Update SetupWizardModal and usage tests to reflect cleanup
This commit is contained in:
@@ -51,6 +51,9 @@ vi.mock("node:child_process", () => ({
|
||||
vi.mock("node:fs", () => ({
|
||||
existsSync: vi.fn().mockReturnValue(true),
|
||||
}));
|
||||
vi.mock("./rate-limit-retry.js", () => ({
|
||||
withRateLimitRetry: (fn: () => Promise<any>) => fn(),
|
||||
}));
|
||||
|
||||
import { TaskExecutor, buildExecutionPrompt } from "./executor.js";
|
||||
import { createKbAgent } from "./pi.js";
|
||||
|
||||
@@ -14,6 +14,7 @@ import { AgentLogger } from "./agent-logger.js";
|
||||
import { executorLog, reviewerLog } from "./logger.js";
|
||||
import { isUsageLimitError, checkSessionError, type UsageLimitPauser } from "./usage-limit-detector.js";
|
||||
import { isTransientError } from "./transient-error-detector.js";
|
||||
import { withRateLimitRetry } from "./rate-limit-retry.js";
|
||||
import type { StuckTaskDetector } from "./stuck-task-detector.js";
|
||||
|
||||
// Re-export for backward compatibility (tests import from executor.ts)
|
||||
@@ -653,10 +654,18 @@ export class TaskExecutor {
|
||||
}
|
||||
};
|
||||
|
||||
const retryableWork = () => withRateLimitRetry(agentWork, {
|
||||
onRetry: (attempt, delayMs, error) => {
|
||||
const delaySec = Math.round(delayMs / 1000);
|
||||
executorLog.warn(`⏳ ${task.id} rate limited — retry ${attempt} in ${delaySec}s: ${error.message}`);
|
||||
this.store.logEntry(task.id, `Rate limited — retry ${attempt} in ${delaySec}s`).catch(() => {});
|
||||
},
|
||||
});
|
||||
|
||||
if (this.options.semaphore) {
|
||||
await this.options.semaphore.run(agentWork, PRIORITY_EXECUTE);
|
||||
await this.options.semaphore.run(retryableWork, PRIORITY_EXECUTE);
|
||||
} else {
|
||||
await agentWork();
|
||||
await retryableWork();
|
||||
}
|
||||
} catch (err: any) {
|
||||
if (this.depAborted.has(task.id)) {
|
||||
|
||||
@@ -9,6 +9,7 @@ export { createKbAgent, type AgentOptions, type AgentResult } from "./pi.js";
|
||||
export { WorktreePool, scanIdleWorktrees, cleanupOrphanedWorktrees } from "./worktree-pool.js";
|
||||
export { createLogger, type Logger } from "./logger.js";
|
||||
export { isUsageLimitError, UsageLimitPauser } from "./usage-limit-detector.js";
|
||||
export { withRateLimitRetry } from "./rate-limit-retry.js";
|
||||
export { PrMonitor, type PrComment, type TrackedPr, type OnNewCommentsCallback } from "./pr-monitor.js";
|
||||
export { PrCommentHandler } from "./pr-comment-handler.js";
|
||||
export { NtfyNotifier, type NtfyNotifierOptions } from "./notifier.js";
|
||||
|
||||
@@ -14,6 +14,10 @@ vi.mock("node:fs", () => ({
|
||||
readFileSync: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("./rate-limit-retry.js", () => ({
|
||||
withRateLimitRetry: (fn: () => Promise<any>) => fn(),
|
||||
}));
|
||||
|
||||
import {
|
||||
aiMergeTask,
|
||||
findWorktreeUser,
|
||||
|
||||
@@ -6,6 +6,7 @@ import type { WorktreePool } from "./worktree-pool.js";
|
||||
import { AgentLogger } from "./agent-logger.js";
|
||||
import { mergerLog } from "./logger.js";
|
||||
import { isUsageLimitError, checkSessionError, type UsageLimitPauser } from "./usage-limit-detector.js";
|
||||
import { withRateLimitRetry } from "./rate-limit-retry.js";
|
||||
import type { ToolDefinition } from "@mariozechner/pi-coding-agent";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
|
||||
@@ -1153,9 +1154,15 @@ async function runAiAgentForCommit(params: AiAgentParams): Promise<{ success: bo
|
||||
simplifiedContext,
|
||||
buildCommand,
|
||||
});
|
||||
await session.prompt(prompt);
|
||||
|
||||
checkSessionError(session);
|
||||
await withRateLimitRetry(async () => {
|
||||
await session.prompt(prompt);
|
||||
checkSessionError(session);
|
||||
}, {
|
||||
onRetry: (attempt, delayMs, error) => {
|
||||
const delaySec = Math.round(delayMs / 1000);
|
||||
mergerLog.warn(`⏳ ${taskId} rate limited — retry ${attempt} in ${delaySec}s: ${error.message}`);
|
||||
},
|
||||
});
|
||||
|
||||
// Check if build failed
|
||||
if (buildFailed) {
|
||||
|
||||
217
packages/engine/src/rate-limit-retry.test.ts
Normal file
217
packages/engine/src/rate-limit-retry.test.ts
Normal file
@@ -0,0 +1,217 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
import { withRateLimitRetry } from "./rate-limit-retry.js";
|
||||
|
||||
describe("withRateLimitRetry", () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
});
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("returns the result when fn succeeds on first call", async () => {
|
||||
const fn = vi.fn().mockResolvedValue("ok");
|
||||
const promise = withRateLimitRetry(fn);
|
||||
const result = await promise;
|
||||
expect(result).toBe("ok");
|
||||
expect(fn).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("retries on rate limit error and succeeds", async () => {
|
||||
const fn = vi
|
||||
.fn()
|
||||
.mockRejectedValueOnce(new Error("429 too many requests"))
|
||||
.mockResolvedValueOnce("recovered");
|
||||
|
||||
const onRetry = vi.fn();
|
||||
const promise = withRateLimitRetry(fn, {
|
||||
baseDelayMs: 1000,
|
||||
maxDelayMs: 10000,
|
||||
onRetry,
|
||||
});
|
||||
|
||||
// Advance past the first backoff delay (1000ms base + jitter)
|
||||
await vi.advanceTimersByTimeAsync(1500);
|
||||
|
||||
const result = await promise;
|
||||
expect(result).toBe("recovered");
|
||||
expect(fn).toHaveBeenCalledTimes(2);
|
||||
expect(onRetry).toHaveBeenCalledTimes(1);
|
||||
expect(onRetry).toHaveBeenCalledWith(1, expect.any(Number), expect.any(Error));
|
||||
});
|
||||
|
||||
it("throws after all retries are exhausted", async () => {
|
||||
const rateLimitErr = new Error("rate_limit exceeded");
|
||||
const fn = vi.fn().mockRejectedValue(rateLimitErr);
|
||||
const onRetry = vi.fn();
|
||||
|
||||
const promise = withRateLimitRetry(fn, {
|
||||
maxRetries: 2,
|
||||
baseDelayMs: 100,
|
||||
maxDelayMs: 1000,
|
||||
onRetry,
|
||||
});
|
||||
|
||||
// Advance enough to cover all backoff delays
|
||||
for (let i = 0; i < 10; i++) {
|
||||
await vi.advanceTimersByTimeAsync(500);
|
||||
}
|
||||
|
||||
await expect(promise).rejects.toThrow("rate_limit exceeded");
|
||||
expect(fn).toHaveBeenCalledTimes(3); // initial + 2 retries
|
||||
expect(onRetry).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("re-throws non-rate-limit errors immediately without retry", async () => {
|
||||
const fn = vi.fn().mockRejectedValue(new Error("ENOENT: file not found"));
|
||||
const onRetry = vi.fn();
|
||||
|
||||
await expect(
|
||||
withRateLimitRetry(fn, { baseDelayMs: 1000, onRetry }),
|
||||
).rejects.toThrow("ENOENT: file not found");
|
||||
|
||||
expect(fn).toHaveBeenCalledTimes(1);
|
||||
expect(onRetry).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("applies exponential backoff with increasing delays", async () => {
|
||||
const fn = vi
|
||||
.fn()
|
||||
.mockRejectedValueOnce(new Error("429"))
|
||||
.mockRejectedValueOnce(new Error("429"))
|
||||
.mockResolvedValueOnce("ok");
|
||||
|
||||
const delays: number[] = [];
|
||||
const onRetry = (_attempt: number, delayMs: number) => {
|
||||
delays.push(delayMs);
|
||||
};
|
||||
|
||||
// Use deterministic random for jitter
|
||||
vi.spyOn(Math, "random").mockReturnValue(0.5); // jitter = 0
|
||||
|
||||
const promise = withRateLimitRetry(fn, {
|
||||
baseDelayMs: 1000,
|
||||
maxDelayMs: 10000,
|
||||
onRetry,
|
||||
});
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1100); // 1st delay: 1000ms
|
||||
await vi.advanceTimersByTimeAsync(2100); // 2nd delay: 2000ms
|
||||
|
||||
await promise;
|
||||
|
||||
expect(delays).toEqual([1000, 2000]);
|
||||
expect(fn).toHaveBeenCalledTimes(3);
|
||||
|
||||
vi.spyOn(Math, "random").mockRestore();
|
||||
});
|
||||
|
||||
it("caps delay at maxDelayMs", async () => {
|
||||
const fn = vi
|
||||
.fn()
|
||||
.mockRejectedValueOnce(new Error("overloaded"))
|
||||
.mockResolvedValueOnce("ok");
|
||||
|
||||
const delays: number[] = [];
|
||||
vi.spyOn(Math, "random").mockReturnValue(0.5);
|
||||
|
||||
const promise = withRateLimitRetry(fn, {
|
||||
baseDelayMs: 100000, // would exceed maxDelayMs
|
||||
maxDelayMs: 5000,
|
||||
onRetry: (_a, d) => delays.push(d),
|
||||
});
|
||||
|
||||
await vi.advanceTimersByTimeAsync(6000);
|
||||
await promise;
|
||||
|
||||
// baseDelayMs * 2^0 = 100000, capped to 5000
|
||||
expect(delays[0]).toBe(5000);
|
||||
|
||||
vi.spyOn(Math, "random").mockRestore();
|
||||
});
|
||||
|
||||
it("cancels backoff sleep when abort signal fires", async () => {
|
||||
const fn = vi.fn().mockRejectedValue(new Error("429 rate limited"));
|
||||
const ac = new AbortController();
|
||||
|
||||
const promise = withRateLimitRetry(fn, {
|
||||
baseDelayMs: 60000,
|
||||
maxDelayMs: 120000,
|
||||
signal: ac.signal,
|
||||
});
|
||||
|
||||
// Let first call fail and start sleeping
|
||||
await vi.advanceTimersByTimeAsync(10);
|
||||
|
||||
// Abort during backoff
|
||||
ac.abort(new Error("Task paused"));
|
||||
|
||||
await expect(promise).rejects.toThrow("Task paused");
|
||||
expect(fn).toHaveBeenCalledTimes(1); // only initial call, no retry
|
||||
});
|
||||
|
||||
it("does not retry if abort signal is already aborted", async () => {
|
||||
const fn = vi.fn().mockRejectedValue(new Error("too many requests"));
|
||||
const ac = new AbortController();
|
||||
ac.abort(new Error("Already cancelled"));
|
||||
|
||||
await expect(
|
||||
withRateLimitRetry(fn, { signal: ac.signal }),
|
||||
).rejects.toThrow("too many requests");
|
||||
|
||||
// fn called once, then abort check triggers throw before sleep
|
||||
expect(fn).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("classifies various rate limit error patterns correctly", async () => {
|
||||
const patterns = [
|
||||
"overloaded",
|
||||
"rate limit exceeded",
|
||||
"429 Too Many Requests",
|
||||
"quota exceeded",
|
||||
"billing limit reached",
|
||||
"insufficient credit",
|
||||
];
|
||||
|
||||
for (const msg of patterns) {
|
||||
const fn = vi
|
||||
.fn()
|
||||
.mockRejectedValueOnce(new Error(msg))
|
||||
.mockResolvedValueOnce("ok");
|
||||
|
||||
const promise = withRateLimitRetry(fn, {
|
||||
baseDelayMs: 100,
|
||||
maxDelayMs: 100,
|
||||
});
|
||||
await vi.advanceTimersByTimeAsync(200);
|
||||
const result = await promise;
|
||||
expect(result).toBe("ok");
|
||||
expect(fn).toHaveBeenCalledTimes(2);
|
||||
}
|
||||
});
|
||||
|
||||
it("handles non-Error thrown values", async () => {
|
||||
const fn = vi.fn().mockRejectedValue("string error");
|
||||
|
||||
await expect(
|
||||
withRateLimitRetry(fn, { baseDelayMs: 100 }),
|
||||
).rejects.toThrow("string error");
|
||||
|
||||
expect(fn).toHaveBeenCalledTimes(1); // not a rate limit error string
|
||||
});
|
||||
|
||||
it("uses default options when none provided", async () => {
|
||||
const fn = vi
|
||||
.fn()
|
||||
.mockRejectedValueOnce(new Error("429"))
|
||||
.mockResolvedValueOnce("ok");
|
||||
|
||||
const promise = withRateLimitRetry(fn);
|
||||
|
||||
// Default baseDelayMs is 30000
|
||||
await vi.advanceTimersByTimeAsync(35000);
|
||||
|
||||
const result = await promise;
|
||||
expect(result).toBe("ok");
|
||||
});
|
||||
});
|
||||
142
packages/engine/src/rate-limit-retry.ts
Normal file
142
packages/engine/src/rate-limit-retry.ts
Normal file
@@ -0,0 +1,142 @@
|
||||
/**
|
||||
* Rate Limit Retry — wraps async agent work with exponential backoff
|
||||
* specifically for rate-limit / usage-limit errors.
|
||||
*
|
||||
* When an AI model returns a rate limit error (429, overloaded, quota, etc.),
|
||||
* this utility retries the operation with exponential backoff before letting
|
||||
* the error propagate to the caller's catch block, which triggers a global
|
||||
* pause via `UsageLimitPauser`.
|
||||
*
|
||||
* **Backoff strategy:** `delay = min(baseDelayMs × 2^attempt, maxDelayMs)` with
|
||||
* ±10 % jitter to avoid thundering-herd effects across concurrent agents.
|
||||
*
|
||||
* **Abort support:** An optional `AbortSignal` allows the engine to cancel
|
||||
* pending retries when a task is paused, cancelled, or the engine is shutting
|
||||
* down — so agents don't sit in a 2-minute sleep unnecessarily.
|
||||
*
|
||||
* **Scope:** Only rate-limit errors (as classified by `isUsageLimitError`) are
|
||||
* retried. All other error types are re-thrown immediately so existing error
|
||||
* handling (transient-error retry, failure marking, etc.) is unaffected.
|
||||
*/
|
||||
|
||||
import { isUsageLimitError } from "./usage-limit-detector.js";
|
||||
|
||||
export interface RateLimitRetryOptions {
|
||||
/** Maximum number of retry attempts before re-throwing (default: 3). */
|
||||
maxRetries?: number;
|
||||
/** Initial backoff delay in milliseconds (default: 30 000 — 30 s). */
|
||||
baseDelayMs?: number;
|
||||
/** Upper bound on backoff delay in milliseconds (default: 120 000 — 2 min). */
|
||||
maxDelayMs?: number;
|
||||
/**
|
||||
* Called before each retry with the attempt number (1-based) and the
|
||||
* computed delay. Use this to log retry activity to the task and agent logs.
|
||||
*/
|
||||
onRetry?: (attempt: number, delayMs: number, error: Error) => void;
|
||||
/**
|
||||
* Abort signal that, when triggered, cancels any pending backoff sleep and
|
||||
* re-throws the last error immediately. Essential for paused / cancelled tasks.
|
||||
*/
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
|
||||
/**
|
||||
* Wrap an async function with rate-limit-aware exponential backoff.
|
||||
*
|
||||
* The wrapper calls `fn()`. If it throws a rate-limit error (detected via
|
||||
* `isUsageLimitError`), it sleeps with exponential backoff and retries up to
|
||||
* `maxRetries` times. Non-rate-limit errors are re-thrown immediately.
|
||||
*
|
||||
* After all retries are exhausted, the **original** error is thrown so the
|
||||
* caller's existing catch block can trigger the global pause via
|
||||
* `UsageLimitPauser`.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* await withRateLimitRetry(() => agentWork(), {
|
||||
* onRetry: (attempt, delayMs) =>
|
||||
* store.logEntry(taskId, `Rate limited — retry ${attempt} in ${delayMs}ms`),
|
||||
* signal: abortController.signal,
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
export async function withRateLimitRetry<T>(
|
||||
fn: () => Promise<T>,
|
||||
options: RateLimitRetryOptions = {},
|
||||
): Promise<T> {
|
||||
const {
|
||||
maxRetries = 3,
|
||||
baseDelayMs = 30_000,
|
||||
maxDelayMs = 120_000,
|
||||
onRetry,
|
||||
signal,
|
||||
} = options;
|
||||
|
||||
let lastError: Error | undefined;
|
||||
|
||||
for (let attempt = 0; attempt <= maxRetries; attempt++) {
|
||||
try {
|
||||
return await fn();
|
||||
} catch (err: any) {
|
||||
const error = err instanceof Error ? err : new Error(String(err));
|
||||
|
||||
// Non-rate-limit errors: re-throw immediately — no retry
|
||||
if (!isUsageLimitError(error.message)) {
|
||||
throw error;
|
||||
}
|
||||
|
||||
lastError = error;
|
||||
|
||||
// All retries exhausted — throw so caller can trigger global pause
|
||||
if (attempt >= maxRetries) {
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
// Check abort before sleeping
|
||||
if (signal?.aborted) {
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
// Exponential backoff with ±10 % jitter
|
||||
const rawDelay = Math.min(baseDelayMs * 2 ** attempt, maxDelayMs);
|
||||
const jitter = rawDelay * 0.1 * (2 * Math.random() - 1); // ±10 %
|
||||
const delay = Math.max(0, Math.round(rawDelay + jitter));
|
||||
|
||||
onRetry?.(attempt + 1, delay, error);
|
||||
|
||||
await sleep(delay, signal);
|
||||
}
|
||||
}
|
||||
|
||||
// Unreachable, but satisfies TypeScript
|
||||
throw lastError ?? new Error("withRateLimitRetry: unexpected state");
|
||||
}
|
||||
|
||||
/**
|
||||
* Sleep for `ms` milliseconds, cancellable via an `AbortSignal`.
|
||||
* @internal exported for testing only
|
||||
*/
|
||||
export function sleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
if (signal?.aborted) {
|
||||
reject(signal.reason ?? new Error("Aborted"));
|
||||
return;
|
||||
}
|
||||
|
||||
const timer = setTimeout(resolve, ms);
|
||||
|
||||
if (signal) {
|
||||
const onAbort = () => {
|
||||
clearTimeout(timer);
|
||||
reject(signal.reason ?? new Error("Aborted"));
|
||||
};
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
// Clean up listener when timer fires normally
|
||||
const origResolve = resolve;
|
||||
resolve = () => {
|
||||
signal.removeEventListener("abort", onAbort);
|
||||
origResolve();
|
||||
};
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -22,6 +22,7 @@ import {
|
||||
type UsageLimitPauser,
|
||||
} from "./usage-limit-detector.js";
|
||||
import { isTransientError } from "./transient-error-detector.js";
|
||||
import { withRateLimitRetry } from "./rate-limit-retry.js";
|
||||
|
||||
export const TRIAGE_SYSTEM_PROMPT = `You are a task specification agent for "kb", an AI-orchestrated task board.
|
||||
|
||||
@@ -640,10 +641,18 @@ export class TriageProcessor {
|
||||
}
|
||||
};
|
||||
|
||||
const retryableWork = () => withRateLimitRetry(agentWork, {
|
||||
onRetry: (attempt, delayMs, error) => {
|
||||
const delaySec = Math.round(delayMs / 1000);
|
||||
triageLog.warn(`⏳ ${task.id} rate limited — retry ${attempt} in ${delaySec}s: ${error.message}`);
|
||||
this.store.logEntry(task.id, `Rate limited — retry ${attempt} in ${delaySec}s`).catch(() => {});
|
||||
},
|
||||
});
|
||||
|
||||
if (this.options.semaphore) {
|
||||
await this.options.semaphore.run(agentWork, PRIORITY_SPECIFY);
|
||||
await this.options.semaphore.run(retryableWork, PRIORITY_SPECIFY);
|
||||
} else {
|
||||
await agentWork();
|
||||
await retryableWork();
|
||||
}
|
||||
} catch (err: any) {
|
||||
// Race condition: task was deleted (e.g. as a duplicate) between listTasks()
|
||||
|
||||
Reference in New Issue
Block a user