ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

compaction & branch summarization follow retry policy

Browse Source 
fixes #6647

compaction (auto & manual) and branch summarization retry on transient failures.
use the same retry policy from settings.
emit events for the tui to show indication of retries
This commit is contained in:
David Brailovsky
2026-07-21 10:07:20 +00:00
parent 890b3547af
commit 8e53e0e49c
7 changed files with 546 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/core/agent-session.ts
+53 -1
View File
@@ -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_start";
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_end" };
/** 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 {
onRetry: (attempt, maxAttempts, delayMs, errorMessage) => {
this._emit({
type: "summarization_retry_start",
attempt,
maxAttempts,
delayMs,
errorMessage,
});
},
onRetryAttemptStart: () => {
this._emit({
type: "summarization_retry_attempt_start",
...source,
});
},
onRetryEnd: () => {
this._emit({ type: "summarization_retry_end" });
},
};
}
/**
* 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 };
packages/coding-agent/src/core/compaction/branch-summarization.ts
+11 -6
View File
@@ -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") {
packages/coding-agent/src/core/compaction/compaction.ts
+34 -7
View File
@@ -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") {