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:
gsxdsm
2026-04-02 13:39:04 -07:00
parent a4896f034e
commit 657c7bf238
8 changed files with 399 additions and 7 deletions

View File

@@ -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";

View File

@@ -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)) {

View File

@@ -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";

View File

@@ -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,

View File

@@ -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) {

View 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");
});
});

View 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();
};
}
});
}

View File

@@ -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()