From c925a4029625567811c8d173c2696eba260443d1 Mon Sep 17 00:00:00 2001 From: Rodrigo Uroz Date: Tue, 24 Feb 2026 16:32:29 +0000 Subject: [PATCH] Compaction/Safeguard: preserve recent turns verbatim --- .../compaction-safeguard-runtime.ts | 1 + .../compaction-safeguard.test.ts | 39 ++++++ .../pi-extensions/compaction-safeguard.ts | 111 ++++++++++++++++++ 3 files changed, 151 insertions(+) diff --git a/src/agents/pi-extensions/compaction-safeguard-runtime.ts b/src/agents/pi-extensions/compaction-safeguard-runtime.ts index 7391e3c1cba..b33bc82ffe9 100644 --- a/src/agents/pi-extensions/compaction-safeguard-runtime.ts +++ b/src/agents/pi-extensions/compaction-safeguard-runtime.ts @@ -10,6 +10,7 @@ export type CompactionSafeguardRuntimeValue = { * (extensionRunner.initialize() is never called in that path). */ model?: Model; + recentTurnsPreserve?: number; }; const registry = createSessionManagerRuntimeRegistry(); diff --git a/src/agents/pi-extensions/compaction-safeguard.test.ts b/src/agents/pi-extensions/compaction-safeguard.test.ts index 60d3858c5d0..3b6ba42dcad 100644 --- a/src/agents/pi-extensions/compaction-safeguard.test.ts +++ b/src/agents/pi-extensions/compaction-safeguard.test.ts @@ -11,6 +11,9 @@ import compactionSafeguardExtension, { __testing } from "./compaction-safeguard. const { collectToolFailures, formatToolFailuresSection, + splitPreservedRecentTurns, + formatPreservedTurnsSection, + resolveRecentTurnsPreserve, computeAdaptiveChunkRatio, isOversizedForSummary, BASE_CHUNK_RATIO, @@ -363,6 +366,42 @@ describe("compaction-safeguard runtime registry", () => { }); }); +describe("compaction-safeguard recent-turn preservation", () => { + it("preserves the most recent user/assistant messages", () => { + const messages: AgentMessage[] = [ + { role: "user", content: "older ask", timestamp: 1 }, + { + role: "assistant", + content: [{ type: "text", text: "older answer" }], + timestamp: 2, + } as unknown as AgentMessage, + { role: "user", content: "recent ask", timestamp: 3 }, + { + role: "assistant", + content: [{ type: "text", text: "recent answer" }], + timestamp: 4, + } as unknown as AgentMessage, + ]; + + const split = splitPreservedRecentTurns({ + messages, + recentTurnsPreserve: 1, + }); + + expect(split.preservedMessages).toHaveLength(2); + expect(split.summarizableMessages).toHaveLength(2); + expect(formatPreservedTurnsSection(split.preservedMessages)).toContain( + "## Recent turns preserved verbatim", + ); + }); + + it("clamps preserve count into a safe range", () => { + expect(resolveRecentTurnsPreserve(undefined)).toBe(3); + expect(resolveRecentTurnsPreserve(-1)).toBe(0); + expect(resolveRecentTurnsPreserve(99)).toBe(12); + }); +}); + describe("compaction-safeguard extension model fallback", () => { it("uses runtime.model when ctx.model is undefined (compact.ts workflow)", async () => { // This test verifies the root-cause fix: when extensionRunner.initialize() is not called diff --git a/src/agents/pi-extensions/compaction-safeguard.ts b/src/agents/pi-extensions/compaction-safeguard.ts index fbcf82b2003..8eb613da42e 100644 --- a/src/agents/pi-extensions/compaction-safeguard.ts +++ b/src/agents/pi-extensions/compaction-safeguard.ts @@ -28,6 +28,9 @@ const TURN_PREFIX_INSTRUCTIONS = " early progress, and any details needed to understand the retained suffix."; const MAX_TOOL_FAILURES = 8; const MAX_TOOL_FAILURE_CHARS = 240; +const DEFAULT_RECENT_TURNS_PRESERVE = 3; +const MAX_RECENT_TURNS_PRESERVE = 12; +const MAX_RECENT_TURN_TEXT_CHARS = 600; type ToolFailure = { toolCallId: string; @@ -36,6 +39,18 @@ type ToolFailure = { meta?: string; }; +function clampNonNegativeInt(value: unknown, fallback: number): number { + const normalized = typeof value === "number" && Number.isFinite(value) ? value : fallback; + return Math.max(0, Math.floor(normalized)); +} + +function resolveRecentTurnsPreserve(value: unknown): number { + return Math.min( + MAX_RECENT_TURNS_PRESERVE, + clampNonNegativeInt(value, DEFAULT_RECENT_TURNS_PRESERVE), + ); +} + function normalizeFailureText(text: string): string { return text.replace(/\s+/g, " ").trim(); } @@ -158,6 +173,87 @@ function formatFileOperations(readFiles: string[], modifiedFiles: string[]): str return `\n\n${sections.join("\n\n")}`; } +function extractMessageText(message: AgentMessage): string { + const content = (message as { content?: unknown }).content; + if (typeof content === "string") { + return content.trim(); + } + if (!Array.isArray(content)) { + return ""; + } + const parts: string[] = []; + for (const block of content) { + if (!block || typeof block !== "object") { + continue; + } + const text = (block as { text?: unknown }).text; + if (typeof text === "string" && text.trim().length > 0) { + parts.push(text.trim()); + } + } + return parts.join("\n").trim(); +} + +function splitPreservedRecentTurns(params: { + messages: AgentMessage[]; + recentTurnsPreserve: number; +}): { summarizableMessages: AgentMessage[]; preservedMessages: AgentMessage[] } { + const preserveTurns = Math.min( + MAX_RECENT_TURNS_PRESERVE, + clampNonNegativeInt(params.recentTurnsPreserve, 0), + ); + if (preserveTurns <= 0) { + return { summarizableMessages: params.messages, preservedMessages: [] }; + } + const preserveMessages = preserveTurns * 2; + const candidateIndexes: number[] = []; + for (let i = params.messages.length - 1; i >= 0; i -= 1) { + const role = (params.messages[i] as { role?: unknown }).role; + if (role === "user" || role === "assistant") { + candidateIndexes.push(i); + } + if (candidateIndexes.length >= preserveMessages) { + break; + } + } + if (candidateIndexes.length === 0) { + return { summarizableMessages: params.messages, preservedMessages: [] }; + } + const preservedIndexSet = new Set(candidateIndexes); + const summarizableMessages = params.messages.filter((_, idx) => !preservedIndexSet.has(idx)); + const preservedMessages = params.messages + .filter((_, idx) => preservedIndexSet.has(idx)) + .filter((msg) => { + const role = (msg as { role?: unknown }).role; + return role === "user" || role === "assistant"; + }); + return { summarizableMessages, preservedMessages }; +} + +function formatPreservedTurnsSection(messages: AgentMessage[]): string { + if (messages.length === 0) { + return ""; + } + const lines = messages + .map((message) => { + const role = message.role === "assistant" ? "Assistant" : "User"; + const text = extractMessageText(message); + if (!text) { + return null; + } + const trimmed = + text.length > MAX_RECENT_TURN_TEXT_CHARS + ? `${text.slice(0, MAX_RECENT_TURN_TEXT_CHARS)}...` + : text; + return `- ${role}: ${trimmed}`; + }) + .filter((line): line is string => Boolean(line)); + if (lines.length === 0) { + return ""; + } + return `\n\n## Recent turns preserved verbatim\n${lines.join("\n")}`; +} + /** * Read and format critical workspace context for compaction summary. * Extracts "Session Startup" and "Red Lines" from AGENTS.md. @@ -240,6 +336,7 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { const contextWindowTokens = runtime?.contextWindowTokens ?? modelContextWindow; const turnPrefixMessages = preparation.turnPrefixMessages ?? []; let messagesToSummarize = preparation.messagesToSummarize; + const recentTurnsPreserve = resolveRecentTurnsPreserve(runtime?.recentTurnsPreserve); const maxHistoryShare = runtime?.maxHistoryShare ?? 0.5; @@ -309,6 +406,16 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { } } + const { + summarizableMessages: summaryTargetMessages, + preservedMessages: preservedRecentMessages, + } = splitPreservedRecentTurns({ + messages: messagesToSummarize, + recentTurnsPreserve, + }); + messagesToSummarize = summaryTargetMessages; + const preservedTurnsSection = formatPreservedTurnsSection(preservedRecentMessages); + // Use adaptive chunk ratio based on message sizes, reserving headroom for // the summarization prompt, system prompt, previous summary, and reasoning budget // that generateSummary adds on top of the serialized conversation chunk. @@ -351,6 +458,7 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { }); summary = `${historySummary}\n\n---\n\n**Turn Context (split turn):**\n\n${prefixSummary}`; } + summary += preservedTurnsSection; summary += toolFailureSection; summary += fileOpsSummary; @@ -383,6 +491,9 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { export const __testing = { collectToolFailures, formatToolFailuresSection, + splitPreservedRecentTurns, + formatPreservedTurnsSection, + resolveRecentTurnsPreserve, computeAdaptiveChunkRatio, isOversizedForSummary, BASE_CHUNK_RATIO,