fix(coding-agent): preserve kept messages across repeated compaction closes #2608

packages/coding-agent/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixed
 
+- Fixed repeated compactions dropping messages that were kept by an earlier compaction by re-summarizing from the previous kept boundary and recalculating `tokensBefore` from the rebuilt session context ([#2608](https://github.com/badlogic/pi-mono/issues/2608))
 - Fixed interactive compaction UI updates so `ctx.compact()` rebuilds the chat through unified compaction events, manual compaction no longer duplicates the summary block, and the `trigger-compact` example only fires when context usage crosses its threshold ([#2617](https://github.com/badlogic/pi-mono/issues/2617))
 - Fixed auto-compaction overflow recovery for Ollama models when the backend returns explicit `prompt too long; exceeded max context length ...` errors instead of silently truncating input ([#2626](https://github.com/badlogic/pi-mono/issues/2626))
 - Fixed built-in tool overrides that reuse built-in parameter schemas to still honor custom `renderCall` and `renderResult` renderers in the interactive TUI, restoring the `minimal-mode` example ([#2595](https://github.com/badlogic/pi-mono/issues/2595))

packages/coding-agent/src/core/compaction/compaction.ts

@@ -14,7 +14,7 @@ import {
 	createCompactionSummaryMessage,
 	createCustomMessage,
 } from "../messages.js";
-import type { CompactionEntry, SessionEntry } from "../session-manager.js";
+import { buildSessionContext, type CompactionEntry, type SessionEntry } from "../session-manager.js";
 import {
 	computeFileLists,
 	createFileOps,
@@ -92,6 +92,13 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
 	return undefined;
 }
 
+function getMessageFromEntryForCompaction(entry: SessionEntry): AgentMessage | undefined {
+	if (entry.type === "compaction") {
+		return undefined;
+	}
+	return getMessageFromEntry(entry);
+}
+
 /** Result from compact() - SessionManager adds uuid/parentUuid when saving */
 export interface CompactionResult<T = unknown> {
 	summary: string;
@@ -617,16 +624,18 @@ export function prepareCompaction(
 			break;
 		}
 	}
-	const boundaryStart = prevCompactionIndex + 1;
+
+	let previousSummary: string | undefined;
+	let boundaryStart = 0;
+	if (prevCompactionIndex >= 0) {
+		const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry;
+		previousSummary = prevCompaction.summary;
+		const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId);
+		boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1;
+	}
 	const boundaryEnd = pathEntries.length;
 
-	const usageStart = prevCompactionIndex >= 0 ? prevCompactionIndex : 0;
+	const tokensBefore = estimateContextTokens(buildSessionContext(pathEntries).messages).tokens;
-	const usageMessages: AgentMessage[] = [];
-	for (let i = usageStart; i < boundaryEnd; i++) {
-		const msg = getMessageFromEntry(pathEntries[i]);
-		if (msg) usageMessages.push(msg);
-	}
-	const tokensBefore = estimateContextTokens(usageMessages).tokens;
 
 	const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
 
@@ -642,7 +651,7 @@ export function prepareCompaction(
 	// Messages to summarize (will be discarded after summary)
 	const messagesToSummarize: AgentMessage[] = [];
 	for (let i = boundaryStart; i < historyEnd; i++) {
-		const msg = getMessageFromEntry(pathEntries[i]);
+		const msg = getMessageFromEntryForCompaction(pathEntries[i]);
 		if (msg) messagesToSummarize.push(msg);
 	}
 
@@ -650,18 +659,11 @@ export function prepareCompaction(
 	const turnPrefixMessages: AgentMessage[] = [];
 	if (cutPoint.isSplitTurn) {
 		for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) {
-			const msg = getMessageFromEntry(pathEntries[i]);
+			const msg = getMessageFromEntryForCompaction(pathEntries[i]);
 			if (msg) turnPrefixMessages.push(msg);
 		}
 	}
 
-	// Get previous summary for iterative update
-	let previousSummary: string | undefined;
-	if (prevCompactionIndex >= 0) {
-		const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry;
-		previousSummary = prevCompaction.summary;
-	}
-
 	// Extract file operations from messages and previous compaction
 	const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex);
 

packages/coding-agent/test/compaction.test.ts

@@ -9,6 +9,7 @@ import {
 	calculateContextTokens,
 	compact,
 	DEFAULT_COMPACTION_SETTINGS,
+	estimateContextTokens,
 	findCutPoint,
 	getLastAssistantUsage,
 	prepareCompaction,
@@ -133,6 +134,42 @@ function createThinkingLevelEntry(thinkingLevel: string): ThinkingLevelChangeEnt
 	return entry;
 }
 
+function extractText(messages: AgentMessage[]): string {
+	return messages
+		.map((message) => {
+			switch (message.role) {
+				case "user":
+					return typeof message.content === "string"
+						? message.content
+						: message.content
+								.filter((block): block is { type: "text"; text: string } => block.type === "text")
+								.map((block) => block.text)
+								.join(" ");
+				case "assistant":
+					return message.content
+						.filter((block): block is { type: "text"; text: string } => block.type === "text")
+						.map((block) => block.text)
+						.join(" ");
+				case "branchSummary":
+				case "compactionSummary":
+					return message.summary;
+				case "custom":
+				case "toolResult":
+					return typeof message.content === "string"
+						? message.content
+						: message.content
+								.filter((block): block is { type: "text"; text: string } => block.type === "text")
+								.map((block) => block.text)
+								.join(" ");
+				case "bashExecution":
+					return `${message.command}\n${message.output}`;
+				default:
+					return "";
+			}
+		})
+		.join("\n");
+}
+
 // ============================================================================
 // Unit tests
 // ============================================================================
@@ -358,6 +395,70 @@ describe("buildSessionContext", () => {
 	});
 });
 
+describe("prepareCompaction with previous compaction", () => {
+	it("should preserve kept messages across repeated compactions when they still fit", () => {
+		const u1 = createMessageEntry(createUserMessage("user msg 1 (summarized by compaction1)"));
+		const a1 = createMessageEntry(createAssistantMessage("assistant msg 1"));
+		const u2 = createMessageEntry(createUserMessage("user msg 2 - kept by compaction1"));
+		const a2 = createMessageEntry(createAssistantMessage("assistant msg 2"));
+		const u3 = createMessageEntry(createUserMessage("user msg 3 - kept by compaction1"));
+		const a3 = createMessageEntry(createAssistantMessage("assistant msg 3", createMockUsage(5000, 1000)));
+		const compaction1 = createCompactionEntry("First summary", u2.id);
+		const u4 = createMessageEntry(createUserMessage("user msg 4 (new after compaction1)"));
+		const a4 = createMessageEntry(createAssistantMessage("assistant msg 4", createMockUsage(8000, 2000)));
+
+		const pathEntries = [u1, a1, u2, a2, u3, a3, compaction1, u4, a4];
+		const contextBefore = buildSessionContext(pathEntries);
+		const preparation = prepareCompaction(pathEntries, DEFAULT_COMPACTION_SETTINGS);
+
+		expect(preparation).toBeDefined();
+		expect(preparation!.firstKeptEntryId).toBe(u2.id);
+		expect(preparation!.previousSummary).toBe("First summary");
+		expect(extractText(preparation!.messagesToSummarize)).not.toContain("First summary");
+		expect(preparation!.tokensBefore).toBe(estimateContextTokens(contextBefore.messages).tokens);
+
+		const compaction2: CompactionEntry = {
+			type: "compaction",
+			id: "compaction2-id",
+			parentId: a4.id,
+			timestamp: new Date().toISOString(),
+			summary: "Second summary",
+			firstKeptEntryId: preparation!.firstKeptEntryId,
+			tokensBefore: preparation!.tokensBefore,
+		};
+		const contextAfter = buildSessionContext([...pathEntries, compaction2]);
+		const contextAfterText = extractText(contextAfter.messages);
+
+		expect(contextAfterText).toContain("user msg 2 - kept by compaction1");
+		expect(contextAfterText).toContain("user msg 3 - kept by compaction1");
+	});
+
+	it("should re-summarize previously kept messages when the recent window moves past them", () => {
+		const u1 = createMessageEntry(createUserMessage("user msg 1 (summarized by compaction1)".repeat(4)));
+		const a1 = createMessageEntry(createAssistantMessage("assistant msg 1".repeat(4)));
+		const u2 = createMessageEntry(createUserMessage("user msg 2 - kept by compaction1 ".repeat(12)));
+		const a2 = createMessageEntry(createAssistantMessage("assistant msg 2 ".repeat(12)));
+		const u3 = createMessageEntry(createUserMessage("user msg 3 - kept by compaction1 ".repeat(12)));
+		const a3 = createMessageEntry(createAssistantMessage("assistant msg 3 ".repeat(12), createMockUsage(5000, 1000)));
+		const compaction1 = createCompactionEntry("First summary", u2.id);
+		const u4 = createMessageEntry(createUserMessage("user msg 4 (new after compaction1) ".repeat(12)));
+		const a4 = createMessageEntry(createAssistantMessage("assistant msg 4 ".repeat(12), createMockUsage(8000, 2000)));
+
+		const settings: CompactionSettings = {
+			...DEFAULT_COMPACTION_SETTINGS,
+			keepRecentTokens: 100,
+		};
+		const preparation = prepareCompaction([u1, a1, u2, a2, u3, a3, compaction1, u4, a4], settings);
+
+		expect(preparation).toBeDefined();
+		const summarizedText = extractText(preparation!.messagesToSummarize);
+		expect(summarizedText).toContain("user msg 2 - kept by compaction1");
+		expect(summarizedText).toContain("user msg 3 - kept by compaction1");
+		expect(summarizedText).not.toContain("First summary");
+		expect(preparation!.previousSummary).toBe("First summary");
+	});
+});
+
 // ============================================================================
 // Integration tests with real session data
 // ============================================================================