Merge pull request #6901 from earendil-works/fix/issue-6647-retry-summary-requests-2
compaction & branch summarization follow retry policy
This commit is contained in:
@@ -175,6 +175,16 @@ Summary:
|
||||
|
||||
Event payloads describe what is happening. Harness getters describe latest config for future snapshots. Hook and listener settlement should be awaited in lifecycle order where possible; transport backpressure is handled below the harness by `AssistantMessageStream`, so the harness does not need a separate async event queue merely to keep SSE or websocket reads flowing.
|
||||
|
||||
### Summarization retry events
|
||||
|
||||
When the harness is configured with a retry policy, generated compaction and branch-summary requests emit retry lifecycle events for transient provider errors:
|
||||
|
||||
- `retry_scheduled`: a retry was scheduled. Includes `operation: "compaction" | "branch_summary"`, `attempt`, `maxAttempts`, `delayMs`, and `errorMessage`.
|
||||
- `retry_attempt_start`: the backoff delay completed and the retried summarization request is starting. Includes `operation`.
|
||||
- `retry_finished`: the retry loop finished after success, exhaustion, or abort. Includes `operation`.
|
||||
|
||||
These events are observational and do not accept hook results.
|
||||
|
||||
## Planned session facade
|
||||
|
||||
Extensions should eventually interact with a harness-scoped `HarnessSession` facade rather than the raw session. The facade should wrap the internal session and enforce harness pending-write ordering semantics. Once this exists, hooks and event listeners can receive a context that exposes the full `AgentHarness` plus the session facade without giving direct access to unordered raw session writes.
|
||||
|
||||
@@ -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(
|
||||
@@ -782,6 +804,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";
|
||||
@@ -109,6 +114,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,
|
||||
@@ -501,6 +517,8 @@ export async function generateSummary(
|
||||
customInstructions?: string,
|
||||
previousSummary?: string,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<Result<string, CompactionError>> {
|
||||
const result = await generateSummaryWithUsage(
|
||||
currentMessages,
|
||||
@@ -511,6 +529,8 @@ export async function generateSummary(
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
return result.ok ? ok(result.value.text) : err(result.error);
|
||||
}
|
||||
@@ -525,6 +545,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),
|
||||
@@ -555,10 +577,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"));
|
||||
@@ -700,6 +725,8 @@ export async function compact(
|
||||
customInstructions?: string,
|
||||
signal?: AbortSignal,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<Result<CompactionResult, CompactionError>> {
|
||||
const {
|
||||
firstKeptEntryId,
|
||||
@@ -733,6 +760,8 @@ export async function compact(
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
if (!historyResult.ok) return err(historyResult.error);
|
||||
historyText = historyResult.value.text;
|
||||
@@ -745,6 +774,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}`;
|
||||
@@ -761,6 +792,8 @@ export async function compact(
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
if (!summaryResult.ok) return err(summaryResult.error);
|
||||
summary = summaryResult.value.text;
|
||||
@@ -786,6 +819,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),
|
||||
@@ -802,12 +837,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"));
|
||||
|
||||
@@ -2,6 +2,7 @@ import type {
|
||||
ImageContent,
|
||||
Model,
|
||||
Models,
|
||||
RetryPolicy,
|
||||
SimpleStreamOptions,
|
||||
TextContent,
|
||||
Transport,
|
||||
@@ -629,6 +630,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>;
|
||||
@@ -679,6 +699,9 @@ export type AgentHarnessOwnEvent<
|
||||
| SessionCompactEvent
|
||||
| SessionBeforeTreeEvent
|
||||
| SessionTreeEvent
|
||||
| RetryScheduledEvent
|
||||
| RetryAttemptStartEvent
|
||||
| RetryFinishedEvent
|
||||
| ModelUpdateEvent
|
||||
| ThinkingLevelUpdateEvent
|
||||
| ResourcesUpdateEvent<TSkill, TPromptTemplate>
|
||||
@@ -748,6 +771,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;
|
||||
@@ -866,6 +892,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")]);
|
||||
|
||||
@@ -85,6 +85,128 @@ const RETRYABLE_PROVIDER_ERROR_PATTERN = buildProviderErrorPattern([
|
||||
"ResourceExhausted",
|
||||
]);
|
||||
|
||||
/**
|
||||
* Retry policy: bounded attempts with exponential backoff (`baseDelayMs * 2^(attempt-1)`).
|
||||
* Matches `settings.retry` (`enabled`, `maxRetries`, `baseDelayMs`) in coding-agent; kept
|
||||
* here so the classifier and the policy-driven retry loop live together and stay reusable
|
||||
* by the SDK and other callers.
|
||||
*/
|
||||
export interface RetryPolicy {
|
||||
enabled: boolean;
|
||||
/** Max retry attempts (0 = no retries). The initial call never counts as a retry. */
|
||||
maxRetries: number;
|
||||
/** Base delay in ms. Per-attempt delay is `baseDelayMs * 2^(attempt-1)` before jitter. */
|
||||
baseDelayMs: number;
|
||||
}
|
||||
|
||||
/** Optional callbacks emitted by {@link retryAssistantCall} around each retry. */
|
||||
export interface RetryCallbacks {
|
||||
/** Emitted before the backoff sleep of each retry attempt (1-indexed). */
|
||||
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 | Promise<void>;
|
||||
/** Emitted once when the loop ends: success if a later call completed normally. */
|
||||
onRetryFinished?: (success: boolean, attempt: number, finalError?: string) => void | Promise<void>;
|
||||
}
|
||||
|
||||
class RetrySleepAbortError extends Error {
|
||||
constructor() {
|
||||
super("Aborted");
|
||||
}
|
||||
}
|
||||
|
||||
function sleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (signal?.aborted) {
|
||||
reject(new RetrySleepAbortError());
|
||||
return;
|
||||
}
|
||||
const timeout = setTimeout(resolve, ms);
|
||||
signal?.addEventListener(
|
||||
"abort",
|
||||
() => {
|
||||
clearTimeout(timeout);
|
||||
reject(new RetrySleepAbortError());
|
||||
},
|
||||
{ once: true },
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Run a single assistant-producing call with bounded retry on transient errors.
|
||||
*
|
||||
* Behavior:
|
||||
* - A successful response is returned immediately. Aborts are terminal and never
|
||||
* retried, but reported as unsuccessful if they happen after a retry was scheduled.
|
||||
* Aborts during the backoff sleep are normalized to an aborted `AssistantMessage`
|
||||
* too, so callers do not need to care when cancellation happened.
|
||||
* - 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
|
||||
* `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).
|
||||
*/
|
||||
export async function retryAssistantCall(
|
||||
produce: () => Promise<AssistantMessage>,
|
||||
policy: RetryPolicy | undefined,
|
||||
signal: AbortSignal | undefined,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<AssistantMessage> {
|
||||
const maxAttempts = policy?.enabled ? policy.maxRetries : 0;
|
||||
|
||||
let attempt = 0;
|
||||
let lastRetry: { attempt: number; errorMessage: string } | undefined;
|
||||
for (;;) {
|
||||
const response = await produce();
|
||||
|
||||
// Abort: terminal but not successful. Never retry an aborted message.
|
||||
if (response.stopReason === "aborted") {
|
||||
if (lastRetry) await callbacks?.onRetryFinished?.(false, lastRetry.attempt);
|
||||
return response;
|
||||
}
|
||||
|
||||
// Success: non-error, non-abort responses return as-is.
|
||||
if (response.stopReason !== "error") {
|
||||
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) 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);
|
||||
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) {
|
||||
await callbacks?.onRetryFinished?.(false, attempt, lastRetry.errorMessage);
|
||||
if (error instanceof RetrySleepAbortError) {
|
||||
return { ...response, stopReason: "aborted", errorMessage: undefined };
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
await callbacks?.onRetryAttemptStart?.();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Classifies whether a failed assistant message looks like a transient provider
|
||||
* or transport error, so callers can decide if the last assistant turn should be
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { fauxAssistantMessage } from "../src/providers/faux.ts";
|
||||
import { isRetryableAssistantError } from "../src/utils/retry.ts";
|
||||
import { isRetryableAssistantError, type RetryPolicy, retryAssistantCall } from "../src/utils/retry.ts";
|
||||
|
||||
const openAIExplicitRetryMessage =
|
||||
"An error occurred while processing your request. You can retry your request, or contact us through our help center at help.openai.com if the error persists. Please include the request ID req_******** in your message.";
|
||||
@@ -66,3 +66,136 @@ describe("provider retry classification", () => {
|
||||
expect(isRetryableAssistantError(fauxAssistantMessage("not an error"))).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("retryAssistantCall", () => {
|
||||
const disabled: RetryPolicy = { enabled: false, maxRetries: 3, baseDelayMs: 0 };
|
||||
const enabled: RetryPolicy = { enabled: true, maxRetries: 3, baseDelayMs: 0 };
|
||||
|
||||
it("returns a successful response immediately without retrying", async () => {
|
||||
const produce = vi.fn(async () => fauxAssistantMessage("ok"));
|
||||
const res = await retryAssistantCall(produce, enabled, undefined);
|
||||
expect(res.content).toEqual([{ type: "text", text: "ok" }]);
|
||||
expect(produce).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("does not retry an aborted message", async () => {
|
||||
const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "aborted" }));
|
||||
const onRetryScheduled = vi.fn();
|
||||
const res = await retryAssistantCall(produce, enabled, undefined, { onRetryScheduled });
|
||||
expect(res.stopReason).toBe("aborted");
|
||||
expect(produce).toHaveBeenCalledTimes(1);
|
||||
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 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(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 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(onRetryScheduled).toHaveBeenCalledTimes(3);
|
||||
expect(onRetryFinished).toHaveBeenCalledWith(false, 3, "terminated");
|
||||
});
|
||||
|
||||
it("stops retrying once a call succeeds", async () => {
|
||||
let n = 0;
|
||||
const produce = vi.fn(async () => {
|
||||
n++;
|
||||
return n < 3
|
||||
? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })
|
||||
: fauxAssistantMessage("recovered");
|
||||
});
|
||||
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(onRetryFinished).toHaveBeenCalledWith(true, 2);
|
||||
});
|
||||
|
||||
it("reports an aborted retried call as unsuccessful", async () => {
|
||||
let n = 0;
|
||||
const produce = vi.fn(async () => {
|
||||
n++;
|
||||
return n === 1
|
||||
? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })
|
||||
: fauxAssistantMessage("", { stopReason: "aborted" });
|
||||
});
|
||||
const onRetryFinished = vi.fn();
|
||||
const res = await retryAssistantCall(produce, enabled, undefined, { onRetryFinished });
|
||||
expect(res.stopReason).toBe("aborted");
|
||||
expect(produce).toHaveBeenCalledTimes(2);
|
||||
expect(onRetryFinished).toHaveBeenCalledWith(false, 1);
|
||||
});
|
||||
|
||||
it("does not retry when policy is disabled", async () => {
|
||||
const produce = vi.fn(async () => fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }));
|
||||
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(onRetryScheduled).not.toHaveBeenCalled();
|
||||
expect(onRetryFinished).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("emits onRetryAttemptStart after backoff before each retried call", async () => {
|
||||
const events: string[] = [];
|
||||
let n = 0;
|
||||
const produce = vi.fn(async () => {
|
||||
events.push(`produce:${n}`);
|
||||
n++;
|
||||
return n < 3
|
||||
? fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" })
|
||||
: fauxAssistantMessage("recovered");
|
||||
});
|
||||
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, { onRetryScheduled, onRetryAttemptStart });
|
||||
expect(res.content).toEqual([{ type: "text", text: "recovered" }]);
|
||||
expect(onRetryScheduled).toHaveBeenCalledTimes(2);
|
||||
expect(onRetryAttemptStart).toHaveBeenCalledTimes(2);
|
||||
expect(events).toEqual([
|
||||
"produce:0",
|
||||
"retry:1",
|
||||
"attempt-start",
|
||||
"produce:1",
|
||||
"retry:2",
|
||||
"attempt-start",
|
||||
"produce:2",
|
||||
]);
|
||||
});
|
||||
|
||||
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 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();
|
||||
const res = await p;
|
||||
expect(res.stopReason).toBe("aborted");
|
||||
expect(res.errorMessage).toBeUndefined();
|
||||
expect(produce).toHaveBeenCalledTimes(1);
|
||||
expect(onRetryFinished).toHaveBeenCalledWith(false, 1, "terminated");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -17,7 +17,11 @@ type AgentSessionEvent =
|
||||
| { type: "compaction_start"; reason: "manual" | "threshold" | "overflow" }
|
||||
| { type: "compaction_end"; reason: "manual" | "threshold" | "overflow"; result: CompactionResult | undefined; aborted: boolean; willRetry: boolean; 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_scheduled"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string }
|
||||
| { type: "summarization_retry_attempt_start"; source: "branchSummary" }
|
||||
| { type: "summarization_retry_attempt_start"; source: "compaction"; reason: "manual" | "threshold" | "overflow" }
|
||||
| { type: "summarization_retry_finished" };
|
||||
```
|
||||
|
||||
`queue_update` emits the full pending steering and follow-up queues whenever they change. `compaction_start` and `compaction_end` cover both manual and automatic compaction.
|
||||
|
||||
@@ -851,6 +851,9 @@ Events are streamed to stdout as JSON lines during agent operation. Events do NO
|
||||
| `compaction_end` | Compaction completes |
|
||||
| `auto_retry_start` | Auto-retry begins (after transient error) |
|
||||
| `auto_retry_end` | Auto-retry completes (success or final failure) |
|
||||
| `summarization_retry_scheduled` | Retry scheduled for a transient compaction or branch-summary summarization error |
|
||||
| `summarization_retry_attempt_start` | Retried summarization request starts |
|
||||
| `summarization_retry_finished` | Summarization retry loop completes |
|
||||
| `extension_error` | Extension threw an error |
|
||||
|
||||
### agent_start
|
||||
@@ -1077,6 +1080,36 @@ On final failure (max retries exceeded):
|
||||
}
|
||||
```
|
||||
|
||||
### summarization_retry_scheduled / summarization_retry_attempt_start / summarization_retry_finished
|
||||
|
||||
Emitted when compaction or branch-summary summarization retries after a transient provider error. These events use the same retry settings as automatic assistant-turn retries.
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "summarization_retry_scheduled",
|
||||
"attempt": 1,
|
||||
"maxAttempts": 3,
|
||||
"delayMs": 2000,
|
||||
"errorMessage": "terminated"
|
||||
}
|
||||
```
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "summarization_retry_attempt_start",
|
||||
"source": "compaction",
|
||||
"reason": "threshold"
|
||||
}
|
||||
```
|
||||
|
||||
For branch summaries, `source` is `"branchSummary"` and no `reason` is present.
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "summarization_retry_finished"
|
||||
}
|
||||
```
|
||||
|
||||
### extension_error
|
||||
|
||||
Emitted when an extension throws an error.
|
||||
|
||||
@@ -319,6 +319,9 @@ session.subscribe((event) => {
|
||||
case "compaction_end":
|
||||
case "auto_retry_start":
|
||||
case "auto_retry_end":
|
||||
case "summarization_retry_scheduled":
|
||||
case "summarization_retry_attempt_start":
|
||||
case "summarization_retry_finished":
|
||||
break;
|
||||
}
|
||||
});
|
||||
|
||||
@@ -41,6 +41,7 @@ import {
|
||||
isContextOverflow,
|
||||
isRetryableAssistantError,
|
||||
modelsAreEqual,
|
||||
type RetryCallbacks,
|
||||
resetApiProviders,
|
||||
streamSimple,
|
||||
} from "@earendil-works/pi-ai/compat";
|
||||
@@ -161,7 +162,21 @@ export type AgentSessionEvent =
|
||||
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_scheduled";
|
||||
attempt: number;
|
||||
maxAttempts: number;
|
||||
delayMs: number;
|
||||
errorMessage: string;
|
||||
}
|
||||
| { type: "summarization_retry_attempt_start"; source: "branchSummary" }
|
||||
| {
|
||||
type: "summarization_retry_attempt_start";
|
||||
source: "compaction";
|
||||
reason: "manual" | "threshold" | "overflow";
|
||||
}
|
||||
| { type: "summarization_retry_finished" };
|
||||
|
||||
/** Listener function for agent session events */
|
||||
export type AgentSessionEventListener = (event: AgentSessionEvent) => void;
|
||||
@@ -1838,6 +1853,8 @@ export class AgentSession {
|
||||
this.thinkingLevel,
|
||||
this.agent.streamFunction,
|
||||
env,
|
||||
this.settingsManager.getRetrySettings(),
|
||||
this._summarizationRetryCallbacks({ source: "compaction", reason: "manual" }),
|
||||
);
|
||||
summary = result.summary;
|
||||
firstKeptEntryId = result.firstKeptEntryId;
|
||||
@@ -2114,6 +2131,8 @@ export class AgentSession {
|
||||
this.thinkingLevel,
|
||||
this.agent.streamFunction,
|
||||
env,
|
||||
this.settingsManager.getRetrySettings(),
|
||||
this._summarizationRetryCallbacks({ source: "compaction", reason }),
|
||||
);
|
||||
summary = compactResult.summary;
|
||||
firstKeptEntryId = compactResult.firstKeptEntryId;
|
||||
@@ -2620,6 +2639,37 @@ export class AgentSession {
|
||||
return isRetryableAssistantError(message);
|
||||
}
|
||||
|
||||
/**
|
||||
* Retry policy + callbacks shared by compaction and branch-summary summarization calls.
|
||||
* Uses the same `settings.retry` budget/backoff as agent-turn retries so a single transient
|
||||
* stream drop no longer fails the whole operation. `source` carries the context
|
||||
* the TUI needs to render the retry and recreate the underlying indicator.
|
||||
*/
|
||||
private _summarizationRetryCallbacks(
|
||||
source: { source: "branchSummary" } | { source: "compaction"; reason: "manual" | "threshold" | "overflow" },
|
||||
): RetryCallbacks {
|
||||
return {
|
||||
onRetryScheduled: (attempt, maxAttempts, delayMs, errorMessage) => {
|
||||
this._emit({
|
||||
type: "summarization_retry_scheduled",
|
||||
attempt,
|
||||
maxAttempts,
|
||||
delayMs,
|
||||
errorMessage,
|
||||
});
|
||||
},
|
||||
onRetryAttemptStart: () => {
|
||||
this._emit({
|
||||
type: "summarization_retry_attempt_start",
|
||||
...source,
|
||||
});
|
||||
},
|
||||
onRetryFinished: () => {
|
||||
this._emit({ type: "summarization_retry_finished" });
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepare a retryable error for continuation with exponential backoff.
|
||||
* @returns true if the caller should continue the agent, false otherwise
|
||||
@@ -2934,6 +2984,8 @@ export class AgentSession {
|
||||
replaceInstructions,
|
||||
reserveTokens: branchSummarySettings.reserveTokens,
|
||||
streamFn: this.agent.streamFunction,
|
||||
retry: this.settingsManager.getRetrySettings(),
|
||||
callbacks: this._summarizationRetryCallbacks({ source: "branchSummary" }),
|
||||
});
|
||||
if (result.aborted) {
|
||||
return { cancelled: true, aborted: true };
|
||||
|
||||
@@ -6,9 +6,9 @@
|
||||
*/
|
||||
|
||||
import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core";
|
||||
import type { RetryCallbacks, RetryPolicy } from "@earendil-works/pi-ai";
|
||||
import { contentText } from "@earendil-works/pi-ai";
|
||||
import type { Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat";
|
||||
import { completeSimple } from "@earendil-works/pi-ai/compat";
|
||||
import {
|
||||
convertToLlm,
|
||||
createBranchSummaryMessage,
|
||||
@@ -16,7 +16,7 @@ import {
|
||||
createCustomMessage,
|
||||
} from "../messages.ts";
|
||||
import type { ReadonlySessionManager, SessionEntry } from "../session-manager.ts";
|
||||
import { estimateTokens } from "./compaction.ts";
|
||||
import { completeSummarization, estimateTokens } from "./compaction.ts";
|
||||
import {
|
||||
computeFileLists,
|
||||
createFileOps,
|
||||
@@ -83,6 +83,10 @@ export interface GenerateBranchSummaryOptions {
|
||||
reserveTokens?: number;
|
||||
/** Optional session stream function. Used to preserve SDK request behavior without mutating agent state. */
|
||||
streamFn?: StreamFn;
|
||||
/** Retry policy for transient summarization errors. Reuses coding-agent's `settings.retry`. */
|
||||
retry?: RetryPolicy;
|
||||
/** Optional callbacks for retry reporting (e.g. TUI retry indicators). */
|
||||
callbacks?: RetryCallbacks;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
@@ -300,6 +304,8 @@ export async function generateBranchSummary(
|
||||
replaceInstructions,
|
||||
reserveTokens = 16384,
|
||||
streamFn,
|
||||
retry,
|
||||
callbacks,
|
||||
} = options;
|
||||
|
||||
// Token budget = context window minus reserved space for prompt + response
|
||||
@@ -338,12 +344,11 @@ export async function generateBranchSummary(
|
||||
|
||||
// Call LLM for summarization. Prefer the session stream function so SDK
|
||||
// request behavior (timeouts, retries, attribution headers) stays consistent
|
||||
// without running through agent state/events.
|
||||
// without running through agent state/events. Retried via completeSummarization
|
||||
// so transient stream drops reuse the configured retry policy.
|
||||
const context = { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages };
|
||||
const requestOptions: SimpleStreamOptions = { apiKey, headers, env, signal, maxTokens: 2048 };
|
||||
const response = streamFn
|
||||
? await (await streamFn(model, context, requestOptions)).result()
|
||||
: await completeSimple(model, context, requestOptions);
|
||||
const response = await completeSummarization(model, context, requestOptions, streamFn, retry, callbacks);
|
||||
|
||||
// Check if aborted or errored
|
||||
if (response.stopReason === "aborted") {
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
*/
|
||||
|
||||
import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core";
|
||||
import { contentText } from "@earendil-works/pi-ai";
|
||||
import { contentText, type RetryCallbacks, type RetryPolicy, retryAssistantCall } from "@earendil-works/pi-ai";
|
||||
import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat";
|
||||
import { completeSimple } from "@earendil-works/pi-ai/compat";
|
||||
import { convertToLlm } from "../messages.ts";
|
||||
@@ -552,17 +552,24 @@ function createSummarizationOptions(
|
||||
return options;
|
||||
}
|
||||
|
||||
async function completeSummarization(
|
||||
/**
|
||||
* Shared choke point for every compaction/branch-summary summarization call. Wraps the
|
||||
* single LLM call in {@link retryAssistantCall} so transient stream drops (e.g.
|
||||
* `terminated`, socket close) honor the configured retry policy instead of failing
|
||||
* the whole compaction on the first attempt. Deterministic errors and aborts return
|
||||
* immediately (see {@link retryAssistantCall}).
|
||||
*/
|
||||
export async function completeSummarization(
|
||||
model: Model<any>,
|
||||
context: Context,
|
||||
options: SimpleStreamOptions,
|
||||
streamFn?: StreamFn,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<AssistantMessage> {
|
||||
if (!streamFn) {
|
||||
return completeSimple(model, context, options);
|
||||
}
|
||||
const stream = await streamFn(model, context, options);
|
||||
return stream.result();
|
||||
const produce = async (): Promise<AssistantMessage> =>
|
||||
streamFn ? (await streamFn(model, context, options)).result() : completeSimple(model, context, options);
|
||||
return retryAssistantCall(produce, retry, options.signal, callbacks);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -581,6 +588,8 @@ export async function generateSummary(
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
streamFn?: StreamFn,
|
||||
env?: Record<string, string>,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<string> {
|
||||
return (
|
||||
await generateSummaryWithUsage(
|
||||
@@ -595,6 +604,8 @@ export async function generateSummary(
|
||||
thinkingLevel,
|
||||
streamFn,
|
||||
env,
|
||||
retry,
|
||||
callbacks,
|
||||
)
|
||||
).text;
|
||||
}
|
||||
@@ -612,6 +623,8 @@ export async function generateSummaryWithUsage(
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
streamFn?: StreamFn,
|
||||
env?: Record<string, string>,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<{ text: string; usage: Usage }> {
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.8 * reserveTokens),
|
||||
@@ -651,6 +664,8 @@ export async function generateSummaryWithUsage(
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
completionOptions,
|
||||
streamFn,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
|
||||
if (response.stopReason === "error") {
|
||||
@@ -801,6 +816,8 @@ export async function compact(
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
streamFn?: StreamFn,
|
||||
env?: Record<string, string>,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<CompactionResult> {
|
||||
const {
|
||||
firstKeptEntryId,
|
||||
@@ -833,6 +850,8 @@ export async function compact(
|
||||
thinkingLevel,
|
||||
streamFn,
|
||||
env,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
historyText = historyResult.text;
|
||||
historyUsage = historyResult.usage;
|
||||
@@ -847,6 +866,8 @@ export async function compact(
|
||||
signal,
|
||||
thinkingLevel,
|
||||
streamFn,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
// Merge into single summary
|
||||
summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.text}`;
|
||||
@@ -865,6 +886,8 @@ export async function compact(
|
||||
thinkingLevel,
|
||||
streamFn,
|
||||
env,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
summary = result.text;
|
||||
summaryUsage = result.usage;
|
||||
@@ -900,6 +923,8 @@ async function generateTurnPrefixSummary(
|
||||
signal?: AbortSignal,
|
||||
thinkingLevel?: ThinkingLevel,
|
||||
streamFn?: StreamFn,
|
||||
retry?: RetryPolicy,
|
||||
callbacks?: RetryCallbacks,
|
||||
): Promise<{ text: string; usage: Usage }> {
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.5 * reserveTokens),
|
||||
@@ -921,6 +946,8 @@ async function generateTurnPrefixSummary(
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel),
|
||||
streamFn,
|
||||
retry,
|
||||
callbacks,
|
||||
);
|
||||
|
||||
if (response.stopReason === "error") {
|
||||
|
||||
@@ -3110,6 +3110,32 @@ export class InteractiveMode {
|
||||
this.ui.requestRender();
|
||||
break;
|
||||
}
|
||||
|
||||
case "summarization_retry_scheduled": {
|
||||
this.showError(event.errorMessage);
|
||||
this.showStatusIndicator(
|
||||
new RetryStatusIndicator(this.ui, event.attempt, event.maxAttempts, event.delayMs),
|
||||
);
|
||||
this.ui.requestRender();
|
||||
break;
|
||||
}
|
||||
|
||||
case "summarization_retry_attempt_start": {
|
||||
this.clearStatusIndicator("retry");
|
||||
if (event.source === "branchSummary") {
|
||||
this.showStatusIndicator(new BranchSummaryStatusIndicator(this.ui));
|
||||
} else {
|
||||
this.showStatusIndicator(new CompactionStatusIndicator(this.ui, event.reason));
|
||||
}
|
||||
this.ui.requestRender();
|
||||
break;
|
||||
}
|
||||
|
||||
case "summarization_retry_finished": {
|
||||
this.clearStatusIndicator("retry");
|
||||
this.ui.requestRender();
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+192
@@ -0,0 +1,192 @@
|
||||
import type { StreamFn } from "@earendil-works/pi-agent-core";
|
||||
import { type AssistantMessage, createAssistantMessageEventStream, fauxAssistantMessage } from "@earendil-works/pi-ai";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { createHarness, type Harness } from "../harness.ts";
|
||||
|
||||
/**
|
||||
* Regression for #6647: compaction runs a single non-retried summarization call, so a
|
||||
* transient mid-stream socket death (`terminated`) failed the whole compaction.
|
||||
* Verifies that summarization now reuses `settings.retry` (bounded retries with
|
||||
* exponential backoff gated on isRetryableAssistantError), emits
|
||||
* `summarization_retry_*` events, and that aborts / non-retryable errors are not retried.
|
||||
*/
|
||||
describe("#6647 compaction retries transient summarization failures", () => {
|
||||
const harnesses: Harness[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (harnesses.length > 0) {
|
||||
harnesses.pop()?.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
function createUsage(totalTokens: number) {
|
||||
return {
|
||||
input: totalTokens,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
function seedCompactableSession(harness: Harness): void {
|
||||
harness.settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } });
|
||||
const now = Date.now();
|
||||
harness.sessionManager.appendMessage({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "message to compact" }],
|
||||
timestamp: now - 1000,
|
||||
});
|
||||
const model = harness.getModel();
|
||||
const assistant: AssistantMessage = {
|
||||
...fauxAssistantMessage("", { stopReason: "stop", timestamp: now - 500 }),
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: createUsage(100),
|
||||
};
|
||||
assistant.content = [{ type: "text", text: "assistant response to compact" }];
|
||||
harness.sessionManager.appendMessage(assistant);
|
||||
harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages;
|
||||
}
|
||||
|
||||
/** streamFn that responds with the given sequence of assistant messages across calls. */
|
||||
function useScriptedStreamFn(harness: Harness, script: AssistantMessage[]): () => number {
|
||||
let callCount = 0;
|
||||
const streamFunction: StreamFn = (model) => {
|
||||
const message = script[callCount] ?? script[script.length - 1]!;
|
||||
callCount++;
|
||||
const stream = createAssistantMessageEventStream();
|
||||
queueMicrotask(() => {
|
||||
if (message.stopReason === "error" || message.stopReason === "aborted") {
|
||||
stream.push({
|
||||
type: "error",
|
||||
reason: message.stopReason,
|
||||
error: { ...message, api: model.api, provider: model.provider, model: model.id },
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "done",
|
||||
reason: message.stopReason,
|
||||
message: { ...message, api: model.api, provider: model.provider, model: model.id },
|
||||
});
|
||||
}
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
harness.session.agent.streamFunction = streamFunction;
|
||||
return () => callCount;
|
||||
}
|
||||
|
||||
it("retries a transient `terminated` summarization error and compacts successfully", async () => {
|
||||
const harness = await createHarness({ withConfiguredAuth: false });
|
||||
harnesses.push(harness);
|
||||
seedCompactableSession(harness);
|
||||
harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 3, baseDelayMs: 0 } });
|
||||
|
||||
const model = harness.getModel();
|
||||
const error = (errorMessage: string): AssistantMessage => ({
|
||||
...fauxAssistantMessage("", { stopReason: "error", errorMessage }),
|
||||
usage: createUsage(10),
|
||||
});
|
||||
const success: AssistantMessage = {
|
||||
...fauxAssistantMessage("recovered summary"),
|
||||
usage: createUsage(10),
|
||||
};
|
||||
const getCallCount = useScriptedStreamFn(harness, [error("terminated"), error("terminated"), success]);
|
||||
|
||||
const result = await harness.session.compact();
|
||||
|
||||
expect(result.summary).toContain("recovered summary");
|
||||
expect(getCallCount()).toBe(3); // 1 initial + 2 retries
|
||||
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_finished" });
|
||||
// model.* referenced to keep imports honest
|
||||
expect(model.id).toBeTruthy();
|
||||
});
|
||||
|
||||
it("does not retry a non-retryable error (insufficient_quota)", async () => {
|
||||
const harness = await createHarness({ withConfiguredAuth: false });
|
||||
harnesses.push(harness);
|
||||
seedCompactableSession(harness);
|
||||
harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 3, baseDelayMs: 0 } });
|
||||
|
||||
const error: AssistantMessage = {
|
||||
...fauxAssistantMessage("", { stopReason: "error", errorMessage: "insufficient_quota" }),
|
||||
usage: createUsage(10),
|
||||
};
|
||||
const getCallCount = useScriptedStreamFn(harness, [error]);
|
||||
|
||||
await expect(harness.session.compact()).rejects.toThrow("insufficient_quota");
|
||||
expect(getCallCount()).toBe(1);
|
||||
expect(harness.eventsOfType("summarization_retry_scheduled")).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("does not retry when retry is disabled", async () => {
|
||||
const harness = await createHarness({ withConfiguredAuth: false });
|
||||
harnesses.push(harness);
|
||||
seedCompactableSession(harness);
|
||||
harness.settingsManager.applyOverrides({ retry: { enabled: false, maxRetries: 3, baseDelayMs: 0 } });
|
||||
|
||||
const error: AssistantMessage = {
|
||||
...fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }),
|
||||
usage: createUsage(10),
|
||||
};
|
||||
const getCallCount = useScriptedStreamFn(harness, [error]);
|
||||
|
||||
await expect(harness.session.compact()).rejects.toThrow("terminated");
|
||||
expect(getCallCount()).toBe(1);
|
||||
expect(harness.eventsOfType("summarization_retry_scheduled")).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("stops retrying after maxRetries and reports failure", async () => {
|
||||
const harness = await createHarness({ withConfiguredAuth: false });
|
||||
harnesses.push(harness);
|
||||
seedCompactableSession(harness);
|
||||
harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 2, baseDelayMs: 0 } });
|
||||
|
||||
const error: AssistantMessage = {
|
||||
...fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }),
|
||||
usage: createUsage(10),
|
||||
};
|
||||
const getCallCount = useScriptedStreamFn(harness, [error, error, error]);
|
||||
|
||||
await expect(harness.session.compact()).rejects.toThrow("terminated");
|
||||
expect(getCallCount()).toBe(3); // 1 initial + 2 retries
|
||||
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_finished" });
|
||||
});
|
||||
|
||||
it("aborts an in-flight retry backoff via abortCompaction", async () => {
|
||||
const harness = await createHarness({ withConfiguredAuth: false });
|
||||
harnesses.push(harness);
|
||||
seedCompactableSession(harness);
|
||||
harness.settingsManager.applyOverrides({ retry: { enabled: true, maxRetries: 5, baseDelayMs: 30_000 } });
|
||||
|
||||
const error: AssistantMessage = {
|
||||
...fauxAssistantMessage("", { stopReason: "error", errorMessage: "terminated" }),
|
||||
usage: createUsage(10),
|
||||
};
|
||||
useScriptedStreamFn(harness, [error, error, error]);
|
||||
|
||||
const compactPromise = harness.session.compact();
|
||||
// Let the first error resolve and the retry backoff sleep start.
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
harness.session.abortCompaction();
|
||||
|
||||
// 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 });
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user