feat: wrap compaction generateSummary in retryAsync

amabito · steipete · commit 068b9c97498d · 2026-02-17T00:01:03.000+01:00
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
diff --git a/src/agents/compaction.retry.test.ts 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<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();
+  });
+});
diff --git 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"),
+      },
     );
   }
 
