From 068b9c97498d36e5d155404ff4aa7cd51a57afde Mon Sep 17 00:00:00 2001 From: amabito Date: Mon, 16 Feb 2026 20:12:44 +0900 Subject: [PATCH] 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 --- src/agents/compaction.retry.test.ts | 182 ++++++++++++++++++++++++++++ src/agents/compaction.ts | 28 +++-- 2 files changed, 202 insertions(+), 8 deletions(-) create mode 100644 src/agents/compaction.retry.test.ts diff --git a/src/agents/compaction.retry.test.ts b/src/agents/compaction.retry.test.ts new file mode 100644 index 0000000000..0513d64d89 --- /dev/null +++ b/src/agents/compaction.retry.test.ts @@ -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(); + 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 = { + 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(); + }); +}); diff --git a/src/agents/compaction.ts b/src/agents/compaction.ts index 7c9798dd26..d60d1af2ad 100644 --- a/src/agents/compaction.ts +++ b/src/agents/compaction.ts @@ -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"), + }, ); }