From 162179af5a989c4b4fd62f109daee69b550428f2 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Tue, 21 Jul 2026 12:11:39 +0000 Subject: [PATCH] add branch/compact retries to agent-harness --- packages/agent/src/harness/agent-harness.ts | 26 ++- .../compaction/branch-summarization.ts | 24 ++- .../src/harness/compaction/compaction.ts | 50 ++++- packages/agent/src/harness/types.ts | 28 +++ .../agent/test/harness/agent-harness.test.ts | 174 ++++++++++++++++++ packages/ai/src/utils/retry.ts | 27 +-- packages/ai/test/retry.test.ts | 56 +++--- .../coding-agent/src/core/agent-session.ts | 12 +- .../src/modes/interactive/interactive-mode.ts | 4 +- ...tion-retries-transient-stream-drop.test.ts | 20 +- 10 files changed, 354 insertions(+), 67 deletions(-) diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 85af6b28..2273a9ed 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -4,6 +4,8 @@ import { type ImageContent, type Model, type Models, + type RetryCallbacks, + type RetryPolicy, type UserMessage, } from "@earendil-works/pi-ai"; import { runAgentLoop } from "../agent-loop.ts"; @@ -178,6 +180,7 @@ export class AgentHarness< private thinkingLevel: ThinkingLevel; private systemPrompt: AgentHarnessOptions["systemPrompt"]; private streamOptions: AgentHarnessStreamOptions; + private retry: RetryPolicy | undefined; private resources: AgentHarnessResources; private tools = new Map(); private activeToolNames: string[]; @@ -194,6 +197,7 @@ export class AgentHarness< this.models = options.models; this.resources = options.resources ?? {}; this.streamOptions = cloneStreamOptions(options.streamOptions); + this.retry = options.retry; this.systemPrompt = options.systemPrompt; this.validateUniqueNames( (options.tools ?? []).map((tool) => tool.name), @@ -256,6 +260,15 @@ export class AgentHarness< return lastResult; } + private retryCallbacks(operation: "compaction" | "branch_summary"): RetryCallbacks { + return { + onRetryScheduled: (attempt, maxAttempts, delayMs, errorMessage) => + this.emitOwn({ type: "retry_scheduled", operation, attempt, maxAttempts, delayMs, errorMessage }), + onRetryAttemptStart: () => this.emitOwn({ type: "retry_attempt_start", operation }), + onRetryFinished: () => this.emitOwn({ type: "retry_finished", operation }), + }; + } + private async emitBeforeProviderRequest( model: Model, sessionId: string, @@ -720,7 +733,16 @@ export class AgentHarness< const provided = hookResult?.compaction; const compactResult = provided ? { ok: true as const, value: provided } - : await compact(preparation, this.models, model, customInstructions, undefined, this.thinkingLevel); + : await compact( + preparation, + this.models, + model, + customInstructions, + undefined, + this.thinkingLevel, + this.retry, + this.retryCallbacks("compaction"), + ); if (!compactResult.ok) throw compactResult.error; const result = compactResult.value; const entryId = await this.session.appendCompaction( @@ -781,6 +803,8 @@ export class AgentHarness< signal: new AbortController().signal, customInstructions: hookResult?.customInstructions ?? options?.customInstructions, replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions, + retry: this.retry, + callbacks: this.retryCallbacks("branch_summary"), }); if (!branchSummary.ok) { if (branchSummary.error.code === "aborted") return { cancelled: true }; diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts index ab56f3ac..51683e09 100644 --- a/packages/agent/src/harness/compaction/branch-summarization.ts +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -1,4 +1,4 @@ -import { contentText, type Model, type Models } from "@earendil-works/pi-ai"; +import { contentText, type Model, type Models, type RetryCallbacks, type RetryPolicy } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; import { @@ -9,7 +9,7 @@ import { } from "../messages.ts"; import type { BranchSummaryResult, Session, SessionTreeEntry } from "../types.ts"; import { BranchSummaryError, err, ok, type Result, SessionError } from "../types.ts"; -import { estimateTokens, SUMMARIZATION_SYSTEM_PROMPT } from "./compaction.ts"; +import { completeSimpleWithRetries, estimateTokens, SUMMARIZATION_SYSTEM_PROMPT } from "./compaction.ts"; import { computeFileLists, createFileOps, @@ -61,6 +61,10 @@ export interface GenerateBranchSummaryOptions { replaceInstructions?: boolean; /** Tokens reserved for prompt and model output. Defaults to 16384. */ reserveTokens?: number; + /** Optional retry policy for transient summarization errors. */ + retry?: RetryPolicy; + /** Optional callbacks for retry reporting. */ + callbacks?: RetryCallbacks; } /** Collect entries that should be summarized before navigating to a different session tree entry. */ @@ -200,7 +204,16 @@ export async function generateBranchSummary( entries: SessionTreeEntry[], options: GenerateBranchSummaryOptions, ): Promise> { - const { models, model, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; + const { + models, + model, + signal, + customInstructions, + replaceInstructions, + reserveTokens = 16384, + retry, + callbacks, + } = options; const contextWindow = model.contextWindow || 128000; const tokenBudget = contextWindow - reserveTokens; @@ -228,10 +241,13 @@ export async function generateBranchSummary( timestamp: Date.now(), }, ]; - const response = await models.completeSimple( + const response = await completeSimpleWithRetries( + models, model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, { signal, maxTokens: 2048 }, + retry, + callbacks, ); if (response.stopReason === "aborted") { return err(new BranchSummaryError("aborted", response.errorMessage || "Branch summary aborted")); diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index 6c21a492..8ff0794b 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -1,9 +1,14 @@ import { type AssistantMessage, + type Context, contentText, type ImageContent, type Model, type Models, + type RetryCallbacks, + type RetryPolicy, + retryAssistantCall, + type SimpleStreamOptions, type TextContent, type Usage, } from "@earendil-works/pi-ai"; @@ -107,6 +112,17 @@ export interface CompactionResult { details?: T; } +export async function completeSimpleWithRetries( + models: Models, + model: Model, + context: Context, + options: SimpleStreamOptions, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, +): Promise { + return retryAssistantCall(() => models.completeSimple(model, context, options), retry, options.signal, callbacks); +} + function combineUsage(first: Usage, second: Usage): Usage { return { input: first.input + second.input, @@ -499,6 +515,8 @@ export async function generateSummary( customInstructions?: string, previousSummary?: string, thinkingLevel?: ThinkingLevel, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise> { const result = await generateSummaryWithUsage( currentMessages, @@ -509,6 +527,8 @@ export async function generateSummary( customInstructions, previousSummary, thinkingLevel, + retry, + callbacks, ); return result.ok ? ok(result.value.text) : err(result.error); } @@ -523,6 +543,8 @@ export async function generateSummaryWithUsage( customInstructions?: string, previousSummary?: string, thinkingLevel?: ThinkingLevel, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), @@ -553,10 +575,13 @@ export async function generateSummaryWithUsage( ? { maxTokens, signal, reasoning: thinkingLevel } : { maxTokens, signal }; - const response = await models.completeSimple( + const response = await completeSimpleWithRetries( + models, model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, completionOptions, + retry, + callbacks, ); if (response.stopReason === "aborted") { return err(new CompactionError("aborted", response.errorMessage || "Summarization aborted")); @@ -688,6 +713,8 @@ export async function compact( customInstructions?: string, signal?: AbortSignal, thinkingLevel?: ThinkingLevel, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise> { const { firstKeptEntryId, @@ -720,6 +747,8 @@ export async function compact( customInstructions, previousSummary, thinkingLevel, + retry, + callbacks, ); if (!historyResult.ok) return err(historyResult.error); historyText = historyResult.value.text; @@ -732,6 +761,8 @@ export async function compact( settings.reserveTokens, signal, thinkingLevel, + retry, + callbacks, ); if (!turnPrefixResult.ok) return err(turnPrefixResult.error); summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`; @@ -748,6 +779,8 @@ export async function compact( customInstructions, previousSummary, thinkingLevel, + retry, + callbacks, ); if (!summaryResult.ok) return err(summaryResult.error); summary = summaryResult.value.text; @@ -772,6 +805,8 @@ async function generateTurnPrefixSummary( reserveTokens: number, signal?: AbortSignal, thinkingLevel?: ThinkingLevel, + retry?: RetryPolicy, + callbacks?: RetryCallbacks, ): Promise> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), @@ -788,12 +823,17 @@ async function generateTurnPrefixSummary( }, ]; - const response = await models.completeSimple( - model, - { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + const completionOptions = model.reasoning && thinkingLevel && thinkingLevel !== "off" ? { maxTokens, signal, reasoning: thinkingLevel } - : { maxTokens, signal }, + : { maxTokens, signal }; + const response = await completeSimpleWithRetries( + models, + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + completionOptions, + retry, + callbacks, ); if (response.stopReason === "aborted") { return err(new CompactionError("aborted", response.errorMessage || "Turn prefix summarization aborted")); diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index cdc7a542..58d13451 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -2,6 +2,7 @@ import type { ImageContent, Model, Models, + RetryPolicy, SimpleStreamOptions, TextContent, Transport, @@ -613,6 +614,25 @@ export interface SessionTreeEvent { fromHook?: boolean; } +export interface RetryScheduledEvent { + type: "retry_scheduled"; + operation: "compaction" | "branch_summary"; + attempt: number; + maxAttempts: number; + delayMs: number; + errorMessage: string; +} + +export interface RetryAttemptStartEvent { + type: "retry_attempt_start"; + operation: "compaction" | "branch_summary"; +} + +export interface RetryFinishedEvent { + type: "retry_finished"; + operation: "compaction" | "branch_summary"; +} + export interface ModelUpdateEvent { type: "model_update"; model: Model; @@ -663,6 +683,9 @@ export type AgentHarnessOwnEvent< | SessionCompactEvent | SessionBeforeTreeEvent | SessionTreeEvent + | RetryScheduledEvent + | RetryAttemptStartEvent + | RetryFinishedEvent | ModelUpdateEvent | ThinkingLevelUpdateEvent | ResourcesUpdateEvent @@ -732,6 +755,9 @@ export type AgentHarnessEventResultMap = { session_compact: undefined; session_before_tree: SessionBeforeTreeResult | undefined; session_tree: undefined; + retry_scheduled: undefined; + retry_attempt_start: undefined; + retry_finished: undefined; model_update: undefined; thinking_level_update: undefined; resources_update: undefined; @@ -848,6 +874,8 @@ export interface AgentHarnessOptions< }) => string | Promise); /** Curated stream/provider request options. Snapshotted at turn start. */ streamOptions?: AgentHarnessStreamOptions; + /** Optional retry policy for generated compaction and branch-summary requests. */ + retry?: RetryPolicy; model: Model; thinkingLevel?: ThinkingLevel; activeToolNames?: string[]; diff --git a/packages/agent/test/harness/agent-harness.test.ts b/packages/agent/test/harness/agent-harness.test.ts index 48f18e24..48a3572a 100644 --- a/packages/agent/test/harness/agent-harness.test.ts +++ b/packages/agent/test/harness/agent-harness.test.ts @@ -550,6 +550,180 @@ describe("AgentHarness", () => { expect(compaction?.type === "compaction" ? compaction.usage : undefined).toEqual(usage); }); + describe("summarization retries", () => { + it("retries transient compaction errors and emits retry events", async () => { + const registration = newFaux(); + let calls = 0; + registration.setResponses([ + () => { + calls++; + return fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }); + }, + () => { + calls++; + return fauxAssistantMessage("## Goal\nRecovered summary"); + }, + ]); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + retry: { enabled: true, maxRetries: 1, baseDelayMs: 0 }, + }); + const retryEvents: string[] = []; + harness.subscribe((event) => { + if ( + event.type === "retry_scheduled" || + event.type === "retry_attempt_start" || + event.type === "retry_finished" + ) { + retryEvents.push(`${event.type}:${event.operation}`); + } + }); + + const result = await harness.compact(); + + expect(result.summary).toContain("Recovered summary"); + expect(calls).toBe(2); + expect(retryEvents).toEqual([ + "retry_scheduled:compaction", + "retry_attempt_start:compaction", + "retry_finished:compaction", + ]); + }); + + it("does not retry non-retryable compaction errors", async () => { + const registration = newFaux(); + let calls = 0; + registration.setResponses([ + () => { + calls++; + return fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }); + }, + ]); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + retry: { enabled: true, maxRetries: 1, baseDelayMs: 0 }, + }); + const retryEvents: string[] = []; + harness.subscribe((event) => { + if ( + event.type === "retry_scheduled" || + event.type === "retry_attempt_start" || + event.type === "retry_finished" + ) { + retryEvents.push(event.type); + } + }); + + await expect(harness.compact()).rejects.toThrow("insufficient_quota"); + + expect(calls).toBe(1); + expect(retryEvents).toEqual([]); + }); + + it("exhausts transient compaction retries after maxRetries failures", async () => { + const registration = newFaux(); + let calls = 0; + registration.setResponses( + Array.from({ length: 4 }, () => () => { + calls++; + return fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }); + }), + ); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + retry: { enabled: true, maxRetries: 3, baseDelayMs: 0 }, + }); + const retryEvents: string[] = []; + harness.subscribe((event) => { + if ( + event.type === "retry_scheduled" || + event.type === "retry_attempt_start" || + event.type === "retry_finished" + ) { + retryEvents.push(`${event.type}:${event.operation}`); + } + }); + + await expect(harness.compact()).rejects.toThrow("terminated"); + + expect(calls).toBe(4); + expect(retryEvents).toEqual([ + "retry_scheduled:compaction", + "retry_attempt_start:compaction", + "retry_scheduled:compaction", + "retry_attempt_start:compaction", + "retry_scheduled:compaction", + "retry_attempt_start:compaction", + "retry_finished:compaction", + ]); + }); + + it("retries transient branch summary errors and emits retry events", async () => { + const registration = newFaux(); + let calls = 0; + registration.setResponses([ + () => { + calls++; + return fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }); + }, + () => { + calls++; + return fauxAssistantMessage("## Goal\nRecovered branch summary"); + }, + ]); + const session = new Session(new InMemorySessionStorage()); + const targetId = await session.appendMessage(createUserMessage("first branch")); + await session.appendMessage(createAssistantMessage("first reply")); + await session.appendMessage(createUserMessage("abandoned work")); + await session.appendMessage(createAssistantMessage("abandoned reply")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + retry: { enabled: true, maxRetries: 1, baseDelayMs: 0 }, + }); + const retryEvents: string[] = []; + harness.subscribe((event) => { + if ( + event.type === "retry_scheduled" || + event.type === "retry_attempt_start" || + event.type === "retry_finished" + ) { + retryEvents.push(`${event.type}:${event.operation}`); + } + }); + + const result = await harness.navigateTree(targetId, { summarize: true }); + + expect(result.summaryEntry?.summary).toContain("Recovered branch summary"); + expect(calls).toBe(2); + expect(retryEvents).toEqual([ + "retry_scheduled:branch_summary", + "retry_attempt_start:branch_summary", + "retry_finished:branch_summary", + ]); + }); + }); + it("persists generated branch summary usage", async () => { const registration = newFaux(); registration.setResponses([fauxAssistantMessage("## Goal\nBranch summary")]); diff --git a/packages/ai/src/utils/retry.ts b/packages/ai/src/utils/retry.ts index 0a433439..441b40d7 100644 --- a/packages/ai/src/utils/retry.ts +++ b/packages/ai/src/utils/retry.ts @@ -102,11 +102,16 @@ export interface RetryPolicy { /** Optional callbacks emitted by {@link retryAssistantCall} around each retry. */ export interface RetryCallbacks { /** Emitted before the backoff sleep of each retry attempt (1-indexed). */ - onRetry?: (attempt: number, maxAttempts: number, delayMs: number, errorMessage: string) => void; + onRetryScheduled?: ( + attempt: number, + maxAttempts: number, + delayMs: number, + errorMessage: string, + ) => void | Promise; /** Emitted after the backoff sleep, immediately before the retried call starts. */ - onRetryAttemptStart?: () => void; + onRetryAttemptStart?: () => void | Promise; /** Emitted once when the loop ends: success if a later call returned a non-error message. */ - onRetryEnd?: (success: boolean, attempt: number, finalError?: string) => void; + onRetryFinished?: (success: boolean, attempt: number, finalError?: string) => void | Promise; } class RetrySleepAbortError extends Error { @@ -143,9 +148,9 @@ function sleep(ms: number, signal?: AbortSignal): Promise { * - A non-retryable error (per {@link isRetryableAssistantError}, including quota/ * billing exhaustion) is returned immediately so deterministic errors fail fast. * - Otherwise retries up to `maxRetries` times with exponential backoff, emitting - * `onRetry` before each sleep, `onRetryAttemptStart` after each sleep before the - * retried call starts, and `onRetryEnd` once at the end (whether the loop ends in - * success, exhausted retries, or an aborted backoff). + * `onRetryScheduled` before each sleep, `onRetryAttemptStart` after each sleep before + * the retried call starts, and `onRetryFinished` once at the end (whether the loop + * ends in success, exhausted retries, or an aborted backoff). * * When `policy` is undefined or disabled, the first response is returned unchanged * (equivalent to calling `produce()` directly). @@ -165,33 +170,33 @@ export async function retryAssistantCall( // Success or abort: never retry an aborted message; non-error returns as-is. if (response.stopReason !== "error") { - if (lastRetry) callbacks?.onRetryEnd?.(true, lastRetry.attempt); + if (lastRetry) await callbacks?.onRetryFinished?.(true, lastRetry.attempt); return response; } // Non-retryable, or budget exhausted: return the final error message. if (attempt >= maxAttempts || !isRetryableAssistantError(response)) { - if (lastRetry) callbacks?.onRetryEnd?.(false, lastRetry.attempt, response.errorMessage); + if (lastRetry) await callbacks?.onRetryFinished?.(false, lastRetry.attempt, response.errorMessage); return response; } attempt++; lastRetry = { attempt, errorMessage: response.errorMessage || "Unknown error" }; const delayMs = policy!.baseDelayMs * 2 ** (attempt - 1); - callbacks?.onRetry?.(attempt, maxAttempts, delayMs, lastRetry.errorMessage); + await callbacks?.onRetryScheduled?.(attempt, maxAttempts, delayMs, lastRetry.errorMessage); // Normalize aborts during retry backoff to the same AssistantMessage shape as // provider stream aborts, so callers do not need to care when cancellation happened. try { await sleep(delayMs, signal); } catch (error) { - callbacks?.onRetryEnd?.(false, attempt, lastRetry.errorMessage); + await callbacks?.onRetryFinished?.(false, attempt, lastRetry.errorMessage); if (error instanceof RetrySleepAbortError) { return { ...response, stopReason: "aborted", errorMessage: undefined }; } throw error; } - callbacks?.onRetryAttemptStart?.(); + await callbacks?.onRetryAttemptStart?.(); } } diff --git a/packages/ai/test/retry.test.ts b/packages/ai/test/retry.test.ts index 71743f67..eaae872e 100644 --- a/packages/ai/test/retry.test.ts +++ b/packages/ai/test/retry.test.ts @@ -80,35 +80,35 @@ describe("retryAssistantCall", () => { it("does not retry an aborted message", async () => { const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "aborted" })); - const onRetry = vi.fn(); - const res = await retryAssistantCall(produce, enabled, undefined, { onRetry }); + const onRetryScheduled = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled }); expect(res.stopReason).toBe("aborted"); expect(produce).toHaveBeenCalledTimes(1); - expect(onRetry).not.toHaveBeenCalled(); + expect(onRetryScheduled).not.toHaveBeenCalled(); }); it("does not retry a non-retryable error (quota/billing)", async () => { const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }), ); - const onRetry = vi.fn(); - const onRetryEnd = vi.fn(); - const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryEnd }); + const onRetryScheduled = vi.fn(); + const onRetryFinished = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled, onRetryFinished }); expect(res.stopReason).toBe("error"); expect(produce).toHaveBeenCalledTimes(1); - expect(onRetry).not.toHaveBeenCalled(); - expect(onRetryEnd).not.toHaveBeenCalled(); + expect(onRetryScheduled).not.toHaveBeenCalled(); + expect(onRetryFinished).not.toHaveBeenCalled(); }); it("retries a transient error up to maxRetries then returns the final error", async () => { const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); - const onRetry = vi.fn(); - const onRetryEnd = vi.fn(); - const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryEnd }); + const onRetryScheduled = vi.fn(); + const onRetryFinished = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled, onRetryFinished }); expect(res.stopReason).toBe("error"); expect(produce).toHaveBeenCalledTimes(4); // 1 initial + 3 retries - expect(onRetry).toHaveBeenCalledTimes(3); - expect(onRetryEnd).toHaveBeenCalledWith(false, 3, "terminated"); + expect(onRetryScheduled).toHaveBeenCalledTimes(3); + expect(onRetryFinished).toHaveBeenCalledWith(false, 3, "terminated"); }); it("stops retrying once a call succeeds", async () => { @@ -119,22 +119,22 @@ describe("retryAssistantCall", () => { ? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }) : fauxAssistantMessage("recovered"); }); - const onRetryEnd = vi.fn(); - const res = await retryAssistantCall(produce, enabled, undefined, { onRetryEnd }); + const onRetryFinished = vi.fn(); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetryFinished }); expect(res.content).toEqual([{ type: "text", text: "recovered" }]); expect(produce).toHaveBeenCalledTimes(3); - expect(onRetryEnd).toHaveBeenCalledWith(true, 2); + expect(onRetryFinished).toHaveBeenCalledWith(true, 2); }); it("does not retry when policy is disabled", async () => { const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); - const onRetry = vi.fn(); - const onRetryEnd = vi.fn(); - const res = await retryAssistantCall(produce, disabled, undefined, { onRetry, onRetryEnd }); + const onRetryScheduled = vi.fn(); + const onRetryFinished = vi.fn(); + const res = await retryAssistantCall(produce, disabled, undefined, { onRetryScheduled, onRetryFinished }); expect(res.stopReason).toBe("error"); expect(produce).toHaveBeenCalledTimes(1); - expect(onRetry).not.toHaveBeenCalled(); - expect(onRetryEnd).not.toHaveBeenCalled(); + expect(onRetryScheduled).not.toHaveBeenCalled(); + expect(onRetryFinished).not.toHaveBeenCalled(); }); it("emits onRetryAttemptStart after backoff before each retried call", async () => { @@ -147,15 +147,15 @@ describe("retryAssistantCall", () => { ? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }) : fauxAssistantMessage("recovered"); }); - const onRetry = vi.fn((attempt: number) => { + const onRetryScheduled = vi.fn((attempt: number) => { events.push(`retry:${attempt}`); }); const onRetryAttemptStart = vi.fn(() => { events.push("attempt-start"); }); - const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryAttemptStart }); + const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled, onRetryAttemptStart }); expect(res.content).toEqual([{ type: "text", text: "recovered" }]); - expect(onRetry).toHaveBeenCalledTimes(2); + expect(onRetryScheduled).toHaveBeenCalledTimes(2); expect(onRetryAttemptStart).toHaveBeenCalledTimes(2); expect(events).toEqual([ "produce:0", @@ -168,12 +168,12 @@ describe("retryAssistantCall", () => { ]); }); - it("aborts backoff sleep via signal, returns an aborted message, and emits onRetryEnd(false)", async () => { + it("aborts backoff sleep via signal, returns an aborted message, and emits onRetryFinished(false)", async () => { const controller = new AbortController(); const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); const policy: RetryPolicy = { enabled: true, maxRetries: 5, baseDelayMs: 10_000 }; - const onRetryEnd = vi.fn(); - const p = retryAssistantCall(produce, policy, controller.signal, { onRetryEnd }); + const onRetryFinished = vi.fn(); + const p = retryAssistantCall(produce, policy, controller.signal, { onRetryFinished }); // Let one error call resolve and the first backoff sleep start, then abort. await vi.waitFor(() => expect(produce).toHaveBeenCalled()); controller.abort(); @@ -181,6 +181,6 @@ describe("retryAssistantCall", () => { expect(res.stopReason).toBe("aborted"); expect(res.errorMessage).toBeUndefined(); expect(produce).toHaveBeenCalledTimes(1); - expect(onRetryEnd).toHaveBeenCalledWith(false, 1, "terminated"); + expect(onRetryFinished).toHaveBeenCalledWith(false, 1, "terminated"); }); }); diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 19adf6dd..4fdaaaa6 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -164,7 +164,7 @@ export type AgentSessionEvent = | { type: "auto_retry_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string } | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string } | { - type: "summarization_retry_start"; + type: "summarization_retry_scheduled"; attempt: number; maxAttempts: number; delayMs: number; @@ -176,7 +176,7 @@ export type AgentSessionEvent = source: "compaction"; reason: "manual" | "threshold" | "overflow"; } - | { type: "summarization_retry_end" }; + | { type: "summarization_retry_finished" }; /** Listener function for agent session events */ export type AgentSessionEventListener = (event: AgentSessionEvent) => void; @@ -2649,9 +2649,9 @@ export class AgentSession { source: { source: "branchSummary" } | { source: "compaction"; reason: "manual" | "threshold" | "overflow" }, ): RetryCallbacks { return { - onRetry: (attempt, maxAttempts, delayMs, errorMessage) => { + onRetryScheduled: (attempt, maxAttempts, delayMs, errorMessage) => { this._emit({ - type: "summarization_retry_start", + type: "summarization_retry_scheduled", attempt, maxAttempts, delayMs, @@ -2664,8 +2664,8 @@ export class AgentSession { ...source, }); }, - onRetryEnd: () => { - this._emit({ type: "summarization_retry_end" }); + onRetryFinished: () => { + this._emit({ type: "summarization_retry_finished" }); }, }; } diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 89e79c61..5a058c12 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -3111,7 +3111,7 @@ export class InteractiveMode { break; } - case "summarization_retry_start": { + case "summarization_retry_scheduled": { this.showError(event.errorMessage); this.showStatusIndicator( new RetryStatusIndicator(this.ui, event.attempt, event.maxAttempts, event.delayMs), @@ -3131,7 +3131,7 @@ export class InteractiveMode { break; } - case "summarization_retry_end": { + case "summarization_retry_finished": { this.clearStatusIndicator("retry"); this.ui.requestRender(); break; diff --git a/packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts b/packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts index 210b0a3b..88383687 100644 --- a/packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts +++ b/packages/coding-agent/test/suite/regressions/6647-compaction-retries-transient-stream-drop.test.ts @@ -100,13 +100,13 @@ describe("#6647 compaction retries transient summarization failures", () => { expect(result.summary).toContain("recovered summary"); expect(getCallCount()).toBe(3); // 1 initial + 2 retries - const starts = harness.eventsOfType("summarization_retry_start"); - const ends = harness.eventsOfType("summarization_retry_end"); + const starts = harness.eventsOfType("summarization_retry_scheduled"); + const ends = harness.eventsOfType("summarization_retry_finished"); expect(starts).toHaveLength(2); expect(ends).toHaveLength(1); expect(starts[0]).toMatchObject({ attempt: 1, maxAttempts: 3, errorMessage: "terminated" }); expect(starts[1]).toMatchObject({ attempt: 2, maxAttempts: 3 }); - expect(ends[0]).toMatchObject({ type: "summarization_retry_end" }); + expect(ends[0]).toMatchObject({ type: "summarization_retry_finished" }); // model.* referenced to keep imports honest expect(model.id).toBeTruthy(); }); @@ -125,7 +125,7 @@ describe("#6647 compaction retries transient summarization failures", () => { await expect(harness.session.compact()).rejects.toThrow("insufficient_quota"); expect(getCallCount()).toBe(1); - expect(harness.eventsOfType("summarization_retry_start")).toHaveLength(0); + expect(harness.eventsOfType("summarization_retry_scheduled")).toHaveLength(0); }); it("does not retry when retry is disabled", async () => { @@ -142,7 +142,7 @@ describe("#6647 compaction retries transient summarization failures", () => { await expect(harness.session.compact()).rejects.toThrow("terminated"); expect(getCallCount()).toBe(1); - expect(harness.eventsOfType("summarization_retry_start")).toHaveLength(0); + expect(harness.eventsOfType("summarization_retry_scheduled")).toHaveLength(0); }); it("stops retrying after maxRetries and reports failure", async () => { @@ -159,11 +159,11 @@ describe("#6647 compaction retries transient summarization failures", () => { await expect(harness.session.compact()).rejects.toThrow("terminated"); expect(getCallCount()).toBe(3); // 1 initial + 2 retries - const starts = harness.eventsOfType("summarization_retry_start"); - const ends = harness.eventsOfType("summarization_retry_end"); + const starts = harness.eventsOfType("summarization_retry_scheduled"); + const ends = harness.eventsOfType("summarization_retry_finished"); expect(starts).toHaveLength(2); expect(ends).toHaveLength(1); - expect(ends[0]).toMatchObject({ type: "summarization_retry_end" }); + expect(ends[0]).toMatchObject({ type: "summarization_retry_finished" }); }); it("aborts an in-flight retry backoff via abortCompaction", async () => { @@ -183,8 +183,8 @@ describe("#6647 compaction retries transient summarization failures", () => { await new Promise((resolve) => setTimeout(resolve, 0)); harness.session.abortCompaction(); - // The aborted retry backoff rejects with an AbortError (matching the real SDK - // abort path), which compaction classifies as aborted. + // The aborted retry backoff is normalized to an aborted assistant message, + // which compaction classifies as aborted. await expect(compactPromise).rejects.toThrow(); const compactionEnd = harness.eventsOfType("compaction_end").at(-1); expect(compactionEnd).toMatchObject({ aborted: true });