add branch/compact retries to agent-harness

This commit is contained in:
David Brailovsky
2026-07-21 12:11:39 +00:00
parent 8e53e0e49c
commit 162179af5a
10 changed files with 354 additions and 67 deletions
+25 -1
View File
@@ -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<TSkill, TPromptTemplate, TTool>["systemPrompt"];
private streamOptions: AgentHarnessStreamOptions;
private retry: RetryPolicy | undefined;
private resources: AgentHarnessResources<TSkill, TPromptTemplate>;
private tools = new Map<string, TTool>();
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<any>,
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 };
@@ -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<Result<BranchSummaryResult, BranchSummaryError>> {
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"));
@@ -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<T = unknown> {
details?: T;
}
export async function completeSimpleWithRetries(
models: Models,
model: Model<any>,
context: Context,
options: SimpleStreamOptions,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<AssistantMessage> {
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<Result<string, CompactionError>> {
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<Result<{ text: string; usage: Usage }, CompactionError>> {
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<Result<CompactionResult, CompactionError>> {
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<Result<{ text: string; usage: Usage }, CompactionError>> {
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"));
+28
View File
@@ -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<any>;
@@ -663,6 +683,9 @@ export type AgentHarnessOwnEvent<
| SessionCompactEvent
| SessionBeforeTreeEvent
| SessionTreeEvent
| RetryScheduledEvent
| RetryAttemptStartEvent
| RetryFinishedEvent
| ModelUpdateEvent
| ThinkingLevelUpdateEvent
| ResourcesUpdateEvent<TSkill, TPromptTemplate>
@@ -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<string>);
/** 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<any>;
thinkingLevel?: ThinkingLevel;
activeToolNames?: string[];
@@ -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")]);
+16 -11
View File
@@ -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<void>;
/** Emitted after the backoff sleep, immediately before the retried call starts. */
onRetryAttemptStart?: () => void;
onRetryAttemptStart?: () => void | Promise<void>;
/** 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<void>;
}
class RetrySleepAbortError extends Error {
@@ -143,9 +148,9 @@ function sleep(ms: number, signal?: AbortSignal): Promise<void> {
* - 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?.();
}
}
+28 -28
View File
@@ -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");
});
});
@@ -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" });
},
};
}
@@ -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;
@@ -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 });