feat: wrap compaction generateSummary in retryAsync

Integrate retry logic with abort-classifier for /compact endpoint:
- Wrap generateSummary calls in retryAsync with exponential backoff
- Auto-skip retry on user cancellation and gateway restart (AbortError)
- Config: 3 attempts, 500ms-5s delay, 20% jitter
- Add comprehensive Vitest tests (5/5 passed)

Related: #16809, #5744, #17143
This commit is contained in:
amabito
2026-02-16 20:12:44 +09:00
committed by Peter Steinberger
parent 990cf2d226
commit 068b9c9749
2 changed files with 202 additions and 8 deletions

View File

@@ -0,0 +1,182 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
import * as piCodingAgent from "@mariozechner/pi-coding-agent";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { retryAsync } from "../infra/retry.js";
// Mock the external generateSummary function
vi.mock("@mariozechner/pi-coding-agent", async (importOriginal) => {
const actual = await importOriginal<typeof piCodingAgent>();
return {
...actual,
generateSummary: vi.fn(),
};
});
const mockGenerateSummary = vi.mocked(piCodingAgent.generateSummary);
describe("compaction retry integration", () => {
beforeEach(() => {
mockGenerateSummary.mockClear();
});
afterEach(() => {
vi.clearAllTimers();
vi.useRealTimers();
});
const testMessages: AgentMessage[] = [
{ role: "user", content: "Test message" },
{ role: "assistant", content: "Test response" },
];
const testModel: NonNullable<ExtensionContext["model"]> = {
provider: "anthropic",
model: "claude-3-opus",
};
it("should successfully call generateSummary with retry wrapper", async () => {
mockGenerateSummary.mockResolvedValueOnce("Test summary");
const result = await retryAsync(
() =>
mockGenerateSummary(
testMessages,
testModel,
1000,
"test-api-key",
new AbortController().signal,
),
{
attempts: 3,
minDelayMs: 500,
maxDelayMs: 5000,
jitter: 0.2,
label: "compaction/generateSummary",
},
);
expect(result).toBe("Test summary");
expect(mockGenerateSummary).toHaveBeenCalledTimes(1);
});
it("should retry on transient error and succeed", async () => {
mockGenerateSummary
.mockRejectedValueOnce(new Error("Network timeout"))
.mockResolvedValueOnce("Success after retry");
const result = await retryAsync(
() =>
mockGenerateSummary(
testMessages,
testModel,
1000,
"test-api-key",
new AbortController().signal,
),
{
attempts: 3,
minDelayMs: 0,
maxDelayMs: 0,
label: "compaction/generateSummary",
},
);
expect(result).toBe("Success after retry");
expect(mockGenerateSummary).toHaveBeenCalledTimes(2);
});
it("should NOT retry on user abort", async () => {
const abortErr = new Error("aborted");
abortErr.name = "AbortError";
(abortErr as { cause?: unknown }).cause = { source: "user" };
mockGenerateSummary.mockRejectedValueOnce(abortErr);
await expect(
retryAsync(
() =>
mockGenerateSummary(
testMessages,
testModel,
1000,
"test-api-key",
new AbortController().signal,
),
{
attempts: 3,
minDelayMs: 0,
label: "compaction/generateSummary",
shouldRetry: (err: unknown) => !(err instanceof Error && err.name === "AbortError"),
},
),
).rejects.toThrow("aborted");
// Should NOT retry on user cancellation (AbortError filtered by shouldRetry)
expect(mockGenerateSummary).toHaveBeenCalledTimes(1);
});
it("should retry up to 3 times and then fail", async () => {
mockGenerateSummary.mockRejectedValue(new Error("Persistent API error"));
await expect(
retryAsync(
() =>
mockGenerateSummary(
testMessages,
testModel,
1000,
"test-api-key",
new AbortController().signal,
),
{
attempts: 3,
minDelayMs: 0,
maxDelayMs: 0,
label: "compaction/generateSummary",
},
),
).rejects.toThrow("Persistent API error");
expect(mockGenerateSummary).toHaveBeenCalledTimes(3);
});
it("should apply exponential backoff", async () => {
vi.useFakeTimers();
mockGenerateSummary
.mockRejectedValueOnce(new Error("Error 1"))
.mockRejectedValueOnce(new Error("Error 2"))
.mockResolvedValueOnce("Success on 3rd attempt");
const delays: number[] = [];
const promise = retryAsync(
() =>
mockGenerateSummary(
testMessages,
testModel,
1000,
"test-api-key",
new AbortController().signal,
),
{
attempts: 3,
minDelayMs: 500,
maxDelayMs: 5000,
jitter: 0,
label: "compaction/generateSummary",
onRetry: (info) => delays.push(info.delayMs),
},
);
await vi.runAllTimersAsync();
const result = await promise;
expect(result).toBe("Success on 3rd attempt");
expect(mockGenerateSummary).toHaveBeenCalledTimes(3);
// First retry: 500ms, second retry: 1000ms
expect(delays[0]).toBe(500);
expect(delays[1]).toBe(1000);
vi.useRealTimers();
});
});

View File

@@ -1,6 +1,7 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
import { estimateTokens, generateSummary } from "@mariozechner/pi-coding-agent";
import { retryAsync } from "../infra/retry.js";
import { DEFAULT_CONTEXT_TOKENS } from "./defaults.js";
import { repairToolUseResultPairing, stripToolResultDetails } from "./session-transcript-repair.js";
@@ -159,14 +160,25 @@ async function summarizeChunks(params: {
let summary = params.previousSummary;
for (const chunk of chunks) {
summary = await generateSummary(
chunk,
params.model,
params.reserveTokens,
params.apiKey,
params.signal,
params.customInstructions,
summary,
summary = await retryAsync(
() =>
generateSummary(
chunk,
params.model,
params.reserveTokens,
params.apiKey,
params.signal,
params.customInstructions,
summary,
),
{
attempts: 3,
minDelayMs: 500,
maxDelayMs: 5000,
jitter: 0.2,
label: "compaction/generateSummary",
shouldRetry: (err) => !(err instanceof Error && err.name === "AbortError"),
},
);
}