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 ImageContent,
type Model, type Model,
type Models, type Models,
type RetryCallbacks,
type RetryPolicy,
type UserMessage, type UserMessage,
} from "@earendil-works/pi-ai"; } from "@earendil-works/pi-ai";
import { runAgentLoop } from "../agent-loop.ts"; import { runAgentLoop } from "../agent-loop.ts";
@@ -178,6 +180,7 @@ export class AgentHarness<
private thinkingLevel: ThinkingLevel; private thinkingLevel: ThinkingLevel;
private systemPrompt: AgentHarnessOptions<TSkill, TPromptTemplate, TTool>["systemPrompt"]; private systemPrompt: AgentHarnessOptions<TSkill, TPromptTemplate, TTool>["systemPrompt"];
private streamOptions: AgentHarnessStreamOptions; private streamOptions: AgentHarnessStreamOptions;
private retry: RetryPolicy | undefined;
private resources: AgentHarnessResources<TSkill, TPromptTemplate>; private resources: AgentHarnessResources<TSkill, TPromptTemplate>;
private tools = new Map<string, TTool>(); private tools = new Map<string, TTool>();
private activeToolNames: string[]; private activeToolNames: string[];
@@ -194,6 +197,7 @@ export class AgentHarness<
this.models = options.models; this.models = options.models;
this.resources = options.resources ?? {}; this.resources = options.resources ?? {};
this.streamOptions = cloneStreamOptions(options.streamOptions); this.streamOptions = cloneStreamOptions(options.streamOptions);
this.retry = options.retry;
this.systemPrompt = options.systemPrompt; this.systemPrompt = options.systemPrompt;
this.validateUniqueNames( this.validateUniqueNames(
(options.tools ?? []).map((tool) => tool.name), (options.tools ?? []).map((tool) => tool.name),
@@ -256,6 +260,15 @@ export class AgentHarness<
return lastResult; 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( private async emitBeforeProviderRequest(
model: Model<any>, model: Model<any>,
sessionId: string, sessionId: string,
@@ -720,7 +733,16 @@ export class AgentHarness<
const provided = hookResult?.compaction; const provided = hookResult?.compaction;
const compactResult = provided const compactResult = provided
? { ok: true as const, value: 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; if (!compactResult.ok) throw compactResult.error;
const result = compactResult.value; const result = compactResult.value;
const entryId = await this.session.appendCompaction( const entryId = await this.session.appendCompaction(
@@ -781,6 +803,8 @@ export class AgentHarness<
signal: new AbortController().signal, signal: new AbortController().signal,
customInstructions: hookResult?.customInstructions ?? options?.customInstructions, customInstructions: hookResult?.customInstructions ?? options?.customInstructions,
replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions, replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions,
retry: this.retry,
callbacks: this.retryCallbacks("branch_summary"),
}); });
if (!branchSummary.ok) { if (!branchSummary.ok) {
if (branchSummary.error.code === "aborted") return { cancelled: true }; 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 type { AgentMessage } from "../../types.ts";
import { import {
@@ -9,7 +9,7 @@ import {
} from "../messages.ts"; } from "../messages.ts";
import type { BranchSummaryResult, Session, SessionTreeEntry } from "../types.ts"; import type { BranchSummaryResult, Session, SessionTreeEntry } from "../types.ts";
import { BranchSummaryError, err, ok, type Result, SessionError } 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 { import {
computeFileLists, computeFileLists,
createFileOps, createFileOps,
@@ -61,6 +61,10 @@ export interface GenerateBranchSummaryOptions {
replaceInstructions?: boolean; replaceInstructions?: boolean;
/** Tokens reserved for prompt and model output. Defaults to 16384. */ /** Tokens reserved for prompt and model output. Defaults to 16384. */
reserveTokens?: number; 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. */ /** Collect entries that should be summarized before navigating to a different session tree entry. */
@@ -200,7 +204,16 @@ export async function generateBranchSummary(
entries: SessionTreeEntry[], entries: SessionTreeEntry[],
options: GenerateBranchSummaryOptions, options: GenerateBranchSummaryOptions,
): Promise<Result<BranchSummaryResult, BranchSummaryError>> { ): 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 contextWindow = model.contextWindow || 128000;
const tokenBudget = contextWindow - reserveTokens; const tokenBudget = contextWindow - reserveTokens;
@@ -228,10 +241,13 @@ export async function generateBranchSummary(
timestamp: Date.now(), timestamp: Date.now(),
}, },
]; ];
const response = await models.completeSimple( const response = await completeSimpleWithRetries(
models,
model, model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
{ signal, maxTokens: 2048 }, { signal, maxTokens: 2048 },
retry,
callbacks,
); );
if (response.stopReason === "aborted") { if (response.stopReason === "aborted") {
return err(new BranchSummaryError("aborted", response.errorMessage || "Branch summary aborted")); return err(new BranchSummaryError("aborted", response.errorMessage || "Branch summary aborted"));
@@ -1,9 +1,14 @@
import { import {
type AssistantMessage, type AssistantMessage,
type Context,
contentText, contentText,
type ImageContent, type ImageContent,
type Model, type Model,
type Models, type Models,
type RetryCallbacks,
type RetryPolicy,
retryAssistantCall,
type SimpleStreamOptions,
type TextContent, type TextContent,
type Usage, type Usage,
} from "@earendil-works/pi-ai"; } from "@earendil-works/pi-ai";
@@ -107,6 +112,17 @@ export interface CompactionResult<T = unknown> {
details?: T; 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 { function combineUsage(first: Usage, second: Usage): Usage {
return { return {
input: first.input + second.input, input: first.input + second.input,
@@ -499,6 +515,8 @@ export async function generateSummary(
customInstructions?: string, customInstructions?: string,
previousSummary?: string, previousSummary?: string,
thinkingLevel?: ThinkingLevel, thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<string, CompactionError>> { ): Promise<Result<string, CompactionError>> {
const result = await generateSummaryWithUsage( const result = await generateSummaryWithUsage(
currentMessages, currentMessages,
@@ -509,6 +527,8 @@ export async function generateSummary(
customInstructions, customInstructions,
previousSummary, previousSummary,
thinkingLevel, thinkingLevel,
retry,
callbacks,
); );
return result.ok ? ok(result.value.text) : err(result.error); return result.ok ? ok(result.value.text) : err(result.error);
} }
@@ -523,6 +543,8 @@ export async function generateSummaryWithUsage(
customInstructions?: string, customInstructions?: string,
previousSummary?: string, previousSummary?: string,
thinkingLevel?: ThinkingLevel, thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> { ): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
const maxTokens = Math.min( const maxTokens = Math.min(
Math.floor(0.8 * reserveTokens), Math.floor(0.8 * reserveTokens),
@@ -553,10 +575,13 @@ export async function generateSummaryWithUsage(
? { maxTokens, signal, reasoning: thinkingLevel } ? { maxTokens, signal, reasoning: thinkingLevel }
: { maxTokens, signal }; : { maxTokens, signal };
const response = await models.completeSimple( const response = await completeSimpleWithRetries(
models,
model, model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
completionOptions, completionOptions,
retry,
callbacks,
); );
if (response.stopReason === "aborted") { if (response.stopReason === "aborted") {
return err(new CompactionError("aborted", response.errorMessage || "Summarization aborted")); return err(new CompactionError("aborted", response.errorMessage || "Summarization aborted"));
@@ -688,6 +713,8 @@ export async function compact(
customInstructions?: string, customInstructions?: string,
signal?: AbortSignal, signal?: AbortSignal,
thinkingLevel?: ThinkingLevel, thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<CompactionResult, CompactionError>> { ): Promise<Result<CompactionResult, CompactionError>> {
const { const {
firstKeptEntryId, firstKeptEntryId,
@@ -720,6 +747,8 @@ export async function compact(
customInstructions, customInstructions,
previousSummary, previousSummary,
thinkingLevel, thinkingLevel,
retry,
callbacks,
); );
if (!historyResult.ok) return err(historyResult.error); if (!historyResult.ok) return err(historyResult.error);
historyText = historyResult.value.text; historyText = historyResult.value.text;
@@ -732,6 +761,8 @@ export async function compact(
settings.reserveTokens, settings.reserveTokens,
signal, signal,
thinkingLevel, thinkingLevel,
retry,
callbacks,
); );
if (!turnPrefixResult.ok) return err(turnPrefixResult.error); if (!turnPrefixResult.ok) return err(turnPrefixResult.error);
summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`; summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`;
@@ -748,6 +779,8 @@ export async function compact(
customInstructions, customInstructions,
previousSummary, previousSummary,
thinkingLevel, thinkingLevel,
retry,
callbacks,
); );
if (!summaryResult.ok) return err(summaryResult.error); if (!summaryResult.ok) return err(summaryResult.error);
summary = summaryResult.value.text; summary = summaryResult.value.text;
@@ -772,6 +805,8 @@ async function generateTurnPrefixSummary(
reserveTokens: number, reserveTokens: number,
signal?: AbortSignal, signal?: AbortSignal,
thinkingLevel?: ThinkingLevel, thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> { ): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
const maxTokens = Math.min( const maxTokens = Math.min(
Math.floor(0.5 * reserveTokens), Math.floor(0.5 * reserveTokens),
@@ -788,12 +823,17 @@ async function generateTurnPrefixSummary(
}, },
]; ];
const response = await models.completeSimple( const completionOptions =
model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
model.reasoning && thinkingLevel && thinkingLevel !== "off" model.reasoning && thinkingLevel && thinkingLevel !== "off"
? { maxTokens, signal, reasoning: thinkingLevel } ? { 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") { if (response.stopReason === "aborted") {
return err(new CompactionError("aborted", response.errorMessage || "Turn prefix summarization aborted")); return err(new CompactionError("aborted", response.errorMessage || "Turn prefix summarization aborted"));
+28
View File
@@ -2,6 +2,7 @@ import type {
ImageContent, ImageContent,
Model, Model,
Models, Models,
RetryPolicy,
SimpleStreamOptions, SimpleStreamOptions,
TextContent, TextContent,
Transport, Transport,
@@ -613,6 +614,25 @@ export interface SessionTreeEvent {
fromHook?: boolean; 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 { export interface ModelUpdateEvent {
type: "model_update"; type: "model_update";
model: Model<any>; model: Model<any>;
@@ -663,6 +683,9 @@ export type AgentHarnessOwnEvent<
| SessionCompactEvent | SessionCompactEvent
| SessionBeforeTreeEvent | SessionBeforeTreeEvent
| SessionTreeEvent | SessionTreeEvent
| RetryScheduledEvent
| RetryAttemptStartEvent
| RetryFinishedEvent
| ModelUpdateEvent | ModelUpdateEvent
| ThinkingLevelUpdateEvent | ThinkingLevelUpdateEvent
| ResourcesUpdateEvent<TSkill, TPromptTemplate> | ResourcesUpdateEvent<TSkill, TPromptTemplate>
@@ -732,6 +755,9 @@ export type AgentHarnessEventResultMap = {
session_compact: undefined; session_compact: undefined;
session_before_tree: SessionBeforeTreeResult | undefined; session_before_tree: SessionBeforeTreeResult | undefined;
session_tree: undefined; session_tree: undefined;
retry_scheduled: undefined;
retry_attempt_start: undefined;
retry_finished: undefined;
model_update: undefined; model_update: undefined;
thinking_level_update: undefined; thinking_level_update: undefined;
resources_update: undefined; resources_update: undefined;
@@ -848,6 +874,8 @@ export interface AgentHarnessOptions<
}) => string | Promise<string>); }) => string | Promise<string>);
/** Curated stream/provider request options. Snapshotted at turn start. */ /** Curated stream/provider request options. Snapshotted at turn start. */
streamOptions?: AgentHarnessStreamOptions; streamOptions?: AgentHarnessStreamOptions;
/** Optional retry policy for generated compaction and branch-summary requests. */
retry?: RetryPolicy;
model: Model<any>; model: Model<any>;
thinkingLevel?: ThinkingLevel; thinkingLevel?: ThinkingLevel;
activeToolNames?: string[]; activeToolNames?: string[];
@@ -550,6 +550,180 @@ describe("AgentHarness", () => {
expect(compaction?.type === "compaction" ? compaction.usage : undefined).toEqual(usage); 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 () => { it("persists generated branch summary usage", async () => {
const registration = newFaux(); const registration = newFaux();
registration.setResponses([fauxAssistantMessage("## Goal\nBranch summary")]); 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. */ /** Optional callbacks emitted by {@link retryAssistantCall} around each retry. */
export interface RetryCallbacks { export interface RetryCallbacks {
/** Emitted before the backoff sleep of each retry attempt (1-indexed). */ /** 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. */ /** 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. */ /** 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 { 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/ * - A non-retryable error (per {@link isRetryableAssistantError}, including quota/
* billing exhaustion) is returned immediately so deterministic errors fail fast. * billing exhaustion) is returned immediately so deterministic errors fail fast.
* - Otherwise retries up to `maxRetries` times with exponential backoff, emitting * - Otherwise retries up to `maxRetries` times with exponential backoff, emitting
* `onRetry` before each sleep, `onRetryAttemptStart` after each sleep before the * `onRetryScheduled` before each sleep, `onRetryAttemptStart` after each sleep before
* retried call starts, and `onRetryEnd` once at the end (whether the loop ends in * the retried call starts, and `onRetryFinished` once at the end (whether the loop
* success, exhausted retries, or an aborted backoff). * ends in success, exhausted retries, or an aborted backoff).
* *
* When `policy` is undefined or disabled, the first response is returned unchanged * When `policy` is undefined or disabled, the first response is returned unchanged
* (equivalent to calling `produce()` directly). * (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. // Success or abort: never retry an aborted message; non-error returns as-is.
if (response.stopReason !== "error") { if (response.stopReason !== "error") {
if (lastRetry) callbacks?.onRetryEnd?.(true, lastRetry.attempt); if (lastRetry) await callbacks?.onRetryFinished?.(true, lastRetry.attempt);
return response; return response;
} }
// Non-retryable, or budget exhausted: return the final error message. // Non-retryable, or budget exhausted: return the final error message.
if (attempt >= maxAttempts || !isRetryableAssistantError(response)) { 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; return response;
} }
attempt++; attempt++;
lastRetry = { attempt, errorMessage: response.errorMessage || "Unknown error" }; lastRetry = { attempt, errorMessage: response.errorMessage || "Unknown error" };
const delayMs = policy!.baseDelayMs * 2 ** (attempt - 1); 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 // Normalize aborts during retry backoff to the same AssistantMessage shape as
// provider stream aborts, so callers do not need to care when cancellation happened. // provider stream aborts, so callers do not need to care when cancellation happened.
try { try {
await sleep(delayMs, signal); await sleep(delayMs, signal);
} catch (error) { } catch (error) {
callbacks?.onRetryEnd?.(false, attempt, lastRetry.errorMessage); await callbacks?.onRetryFinished?.(false, attempt, lastRetry.errorMessage);
if (error instanceof RetrySleepAbortError) { if (error instanceof RetrySleepAbortError) {
return { ...response, stopReason: "aborted", errorMessage: undefined }; return { ...response, stopReason: "aborted", errorMessage: undefined };
} }
throw error; 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 () => { it("does not retry an aborted message", async () => {
const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "aborted" })); const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "aborted" }));
const onRetry = vi.fn(); const onRetryScheduled = vi.fn();
const res = await retryAssistantCall(produce, enabled, undefined, { onRetry }); const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled });
expect(res.stopReason).toBe("aborted"); expect(res.stopReason).toBe("aborted");
expect(produce).toHaveBeenCalledTimes(1); expect(produce).toHaveBeenCalledTimes(1);
expect(onRetry).not.toHaveBeenCalled(); expect(onRetryScheduled).not.toHaveBeenCalled();
}); });
it("does not retry a non-retryable error (quota/billing)", async () => { it("does not retry a non-retryable error (quota/billing)", async () => {
const produce = vi.fn(async () => const produce = vi.fn(async () =>
fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }), fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }),
); );
const onRetry = vi.fn(); const onRetryScheduled = vi.fn();
const onRetryEnd = vi.fn(); const onRetryFinished = vi.fn();
const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryEnd }); const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled, onRetryFinished });
expect(res.stopReason).toBe("error"); expect(res.stopReason).toBe("error");
expect(produce).toHaveBeenCalledTimes(1); expect(produce).toHaveBeenCalledTimes(1);
expect(onRetry).not.toHaveBeenCalled(); expect(onRetryScheduled).not.toHaveBeenCalled();
expect(onRetryEnd).not.toHaveBeenCalled(); expect(onRetryFinished).not.toHaveBeenCalled();
}); });
it("retries a transient error up to maxRetries then returns the final error", async () => { 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 produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }));
const onRetry = vi.fn(); const onRetryScheduled = vi.fn();
const onRetryEnd = vi.fn(); const onRetryFinished = vi.fn();
const res = await retryAssistantCall(produce, enabled, undefined, { onRetry, onRetryEnd }); const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled, onRetryFinished });
expect(res.stopReason).toBe("error"); expect(res.stopReason).toBe("error");
expect(produce).toHaveBeenCalledTimes(4); // 1 initial + 3 retries expect(produce).toHaveBeenCalledTimes(4); // 1 initial + 3 retries
expect(onRetry).toHaveBeenCalledTimes(3); expect(onRetryScheduled).toHaveBeenCalledTimes(3);
expect(onRetryEnd).toHaveBeenCalledWith(false, 3, "terminated"); expect(onRetryFinished).toHaveBeenCalledWith(false, 3, "terminated");
}); });
it("stops retrying once a call succeeds", async () => { it("stops retrying once a call succeeds", async () => {
@@ -119,22 +119,22 @@ describe("retryAssistantCall", () => {
? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }) ? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })
: fauxAssistantMessage("recovered"); : fauxAssistantMessage("recovered");
}); });
const onRetryEnd = vi.fn(); const onRetryFinished = vi.fn();
const res = await retryAssistantCall(produce, enabled, undefined, { onRetryEnd }); const res = await retryAssistantCall(produce, enabled, undefined, { onRetryFinished });
expect(res.content).toEqual([{ type: "text", text: "recovered" }]); expect(res.content).toEqual([{ type: "text", text: "recovered" }]);
expect(produce).toHaveBeenCalledTimes(3); expect(produce).toHaveBeenCalledTimes(3);
expect(onRetryEnd).toHaveBeenCalledWith(true, 2); expect(onRetryFinished).toHaveBeenCalledWith(true, 2);
}); });
it("does not retry when policy is disabled", async () => { it("does not retry when policy is disabled", async () => {
const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }));
const onRetry = vi.fn(); const onRetryScheduled = vi.fn();
const onRetryEnd = vi.fn(); const onRetryFinished = vi.fn();
const res = await retryAssistantCall(produce, disabled, undefined, { onRetry, onRetryEnd }); const res = await retryAssistantCall(produce, disabled, undefined, { onRetryScheduled, onRetryFinished });
expect(res.stopReason).toBe("error"); expect(res.stopReason).toBe("error");
expect(produce).toHaveBeenCalledTimes(1); expect(produce).toHaveBeenCalledTimes(1);
expect(onRetry).not.toHaveBeenCalled(); expect(onRetryScheduled).not.toHaveBeenCalled();
expect(onRetryEnd).not.toHaveBeenCalled(); expect(onRetryFinished).not.toHaveBeenCalled();
}); });
it("emits onRetryAttemptStart after backoff before each retried call", async () => { it("emits onRetryAttemptStart after backoff before each retried call", async () => {
@@ -147,15 +147,15 @@ describe("retryAssistantCall", () => {
? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }) ? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })
: fauxAssistantMessage("recovered"); : fauxAssistantMessage("recovered");
}); });
const onRetry = vi.fn((attempt: number) => { const onRetryScheduled = vi.fn((attempt: number) => {
events.push(`retry:${attempt}`); events.push(`retry:${attempt}`);
}); });
const onRetryAttemptStart = vi.fn(() => { const onRetryAttemptStart = vi.fn(() => {
events.push("attempt-start"); 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(res.content).toEqual([{ type: "text", text: "recovered" }]);
expect(onRetry).toHaveBeenCalledTimes(2); expect(onRetryScheduled).toHaveBeenCalledTimes(2);
expect(onRetryAttemptStart).toHaveBeenCalledTimes(2); expect(onRetryAttemptStart).toHaveBeenCalledTimes(2);
expect(events).toEqual([ expect(events).toEqual([
"produce:0", "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 controller = new AbortController();
const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })); const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }));
const policy: RetryPolicy = { enabled: true, maxRetries: 5, baseDelayMs: 10_000 }; const policy: RetryPolicy = { enabled: true, maxRetries: 5, baseDelayMs: 10_000 };
const onRetryEnd = vi.fn(); const onRetryFinished = vi.fn();
const p = retryAssistantCall(produce, policy, controller.signal, { onRetryEnd }); const p = retryAssistantCall(produce, policy, controller.signal, { onRetryFinished });
// Let one error call resolve and the first backoff sleep start, then abort. // Let one error call resolve and the first backoff sleep start, then abort.
await vi.waitFor(() => expect(produce).toHaveBeenCalled()); await vi.waitFor(() => expect(produce).toHaveBeenCalled());
controller.abort(); controller.abort();
@@ -181,6 +181,6 @@ describe("retryAssistantCall", () => {
expect(res.stopReason).toBe("aborted"); expect(res.stopReason).toBe("aborted");
expect(res.errorMessage).toBeUndefined(); expect(res.errorMessage).toBeUndefined();
expect(produce).toHaveBeenCalledTimes(1); 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_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string }
| { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string } | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string }
| { | {
type: "summarization_retry_start"; type: "summarization_retry_scheduled";
attempt: number; attempt: number;
maxAttempts: number; maxAttempts: number;
delayMs: number; delayMs: number;
@@ -176,7 +176,7 @@ export type AgentSessionEvent =
source: "compaction"; source: "compaction";
reason: "manual" | "threshold" | "overflow"; reason: "manual" | "threshold" | "overflow";
} }
| { type: "summarization_retry_end" }; | { type: "summarization_retry_finished" };
/** Listener function for agent session events */ /** Listener function for agent session events */
export type AgentSessionEventListener = (event: AgentSessionEvent) => void; export type AgentSessionEventListener = (event: AgentSessionEvent) => void;
@@ -2649,9 +2649,9 @@ export class AgentSession {
source: { source: "branchSummary" } | { source: "compaction"; reason: "manual" | "threshold" | "overflow" }, source: { source: "branchSummary" } | { source: "compaction"; reason: "manual" | "threshold" | "overflow" },
): RetryCallbacks { ): RetryCallbacks {
return { return {
onRetry: (attempt, maxAttempts, delayMs, errorMessage) => { onRetryScheduled: (attempt, maxAttempts, delayMs, errorMessage) => {
this._emit({ this._emit({
type: "summarization_retry_start", type: "summarization_retry_scheduled",
attempt, attempt,
maxAttempts, maxAttempts,
delayMs, delayMs,
@@ -2664,8 +2664,8 @@ export class AgentSession {
...source, ...source,
}); });
}, },
onRetryEnd: () => { onRetryFinished: () => {
this._emit({ type: "summarization_retry_end" }); this._emit({ type: "summarization_retry_finished" });
}, },
}; };
} }
@@ -3111,7 +3111,7 @@ export class InteractiveMode {
break; break;
} }
case "summarization_retry_start": { case "summarization_retry_scheduled": {
this.showError(event.errorMessage); this.showError(event.errorMessage);
this.showStatusIndicator( this.showStatusIndicator(
new RetryStatusIndicator(this.ui, event.attempt, event.maxAttempts, event.delayMs), new RetryStatusIndicator(this.ui, event.attempt, event.maxAttempts, event.delayMs),
@@ -3131,7 +3131,7 @@ export class InteractiveMode {
break; break;
} }
case "summarization_retry_end": { case "summarization_retry_finished": {
this.clearStatusIndicator("retry"); this.clearStatusIndicator("retry");
this.ui.requestRender(); this.ui.requestRender();
break; break;
@@ -100,13 +100,13 @@ describe("#6647 compaction retries transient summarization failures", () => {
expect(result.summary).toContain("recovered summary"); expect(result.summary).toContain("recovered summary");
expect(getCallCount()).toBe(3); // 1 initial + 2 retries expect(getCallCount()).toBe(3); // 1 initial + 2 retries
const starts = harness.eventsOfType("summarization_retry_start"); const starts = harness.eventsOfType("summarization_retry_scheduled");
const ends = harness.eventsOfType("summarization_retry_end"); const ends = harness.eventsOfType("summarization_retry_finished");
expect(starts).toHaveLength(2); expect(starts).toHaveLength(2);
expect(ends).toHaveLength(1); expect(ends).toHaveLength(1);
expect(starts[0]).toMatchObject({ attempt: 1, maxAttempts: 3, errorMessage: "terminated" }); expect(starts[0]).toMatchObject({ attempt: 1, maxAttempts: 3, errorMessage: "terminated" });
expect(starts[1]).toMatchObject({ attempt: 2, maxAttempts: 3 }); 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 // model.* referenced to keep imports honest
expect(model.id).toBeTruthy(); expect(model.id).toBeTruthy();
}); });
@@ -125,7 +125,7 @@ describe("#6647 compaction retries transient summarization failures", () => {
await expect(harness.session.compact()).rejects.toThrow("insufficient_quota"); await expect(harness.session.compact()).rejects.toThrow("insufficient_quota");
expect(getCallCount()).toBe(1); 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 () => { 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"); await expect(harness.session.compact()).rejects.toThrow("terminated");
expect(getCallCount()).toBe(1); 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 () => { 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"); await expect(harness.session.compact()).rejects.toThrow("terminated");
expect(getCallCount()).toBe(3); // 1 initial + 2 retries expect(getCallCount()).toBe(3); // 1 initial + 2 retries
const starts = harness.eventsOfType("summarization_retry_start"); const starts = harness.eventsOfType("summarization_retry_scheduled");
const ends = harness.eventsOfType("summarization_retry_end"); const ends = harness.eventsOfType("summarization_retry_finished");
expect(starts).toHaveLength(2); expect(starts).toHaveLength(2);
expect(ends).toHaveLength(1); 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 () => { 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)); await new Promise((resolve) => setTimeout(resolve, 0));
harness.session.abortCompaction(); harness.session.abortCompaction();
// The aborted retry backoff rejects with an AbortError (matching the real SDK // The aborted retry backoff is normalized to an aborted assistant message,
// abort path), which compaction classifies as aborted. // which compaction classifies as aborted.
await expect(compactPromise).rejects.toThrow(); await expect(compactPromise).rejects.toThrow();
const compactionEnd = harness.eventsOfType("compaction_end").at(-1); const compactionEnd = harness.eventsOfType("compaction_end").at(-1);
expect(compactionEnd).toMatchObject({ aborted: true }); expect(compactionEnd).toMatchObject({ aborted: true });