From a6f720e6caf1cf429e382011156c015fa204c512 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 9 Jul 2026 12:45:26 +0200 Subject: [PATCH] fix(coding-agent): count custom messages in compaction budget closes #6326 --- packages/coding-agent/CHANGELOG.md | 1 + .../src/core/compaction/compaction.ts | 141 ++++++++---------- packages/coding-agent/test/compaction.test.ts | 35 +++++ 3 files changed, 97 insertions(+), 80 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 1aee83ad3..3df856298 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -16,6 +16,7 @@ - Fixed the question extension example to run question tool calls sequentially so multiple questions in one assistant turn remain answerable ([#6189](https://github.com/earendil-works/pi/issues/6189)). - Fixed `/login` to report auth storage persistence failures instead of claiming credentials were saved when `auth.json` is locked ([#6223](https://github.com/earendil-works/pi/issues/6223)). - Fixed split-turn compaction to serialize summary requests so single-concurrency local providers do not fail with 429 errors ([#5536](https://github.com/earendil-works/pi/issues/5536)). +- Fixed compaction retained-token budgeting to count context-visible custom messages ([#6326](https://github.com/earendil-works/pi/issues/6326)). - Fixed custom session entries appended during assistant streaming to render before the live assistant message, matching persisted session order. - Fixed oversized bash tool timeouts to fail with a clear validation error instead of being clamped to an immediate timeout ([#6181](https://github.com/earendil-works/pi/issues/6181)). - Fixed the edit tool schema to allow model-invented extra replacement fields instead of rejecting otherwise valid edits ([#6278](https://github.com/earendil-works/pi/issues/6278)). diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 46da21c4a..39b1042d3 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -8,13 +8,13 @@ import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; 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"; import { - convertToLlm, - createBranchSummaryMessage, - createCompactionSummaryMessage, - createCustomMessage, -} from "../messages.ts"; -import { buildSessionContext, type CompactionEntry, type SessionEntry } from "../session-manager.ts"; + buildSessionContext, + type CompactionEntry, + type SessionEntry, + sessionEntryToContextMessages, +} from "../session-manager.ts"; import { computeFileLists, createFileOps, @@ -76,27 +76,11 @@ function extractFileOperations( * Extract AgentMessage from an entry if it produces one. * Returns undefined for entries that don't contribute to LLM context. */ -function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { - if (entry.type === "message") { - return entry.message; - } - if (entry.type === "custom_message") { - return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); - } - if (entry.type === "branch_summary") { - return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); - } - if (entry.type === "compaction") { - return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp); - } - return undefined; -} - function getMessageFromEntryForCompaction(entry: SessionEntry): AgentMessage | undefined { if (entry.type === "compaction") { return undefined; } - return getMessageFromEntry(entry); + return sessionEntryToContextMessages(entry)[0]; } /** Result from compact() - SessionManager adds uuid/parentUuid when saving */ @@ -295,47 +279,57 @@ export function estimateTokens(message: AgentMessage): number { return 0; } +function isCutPointMessage(message: AgentMessage): boolean { + switch (message.role) { + case "user": + case "assistant": + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + return true; + case "toolResult": + return false; + } + return false; +} + +function isTurnStartMessage(message: AgentMessage): boolean { + switch (message.role) { + case "user": + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + return true; + case "assistant": + case "toolResult": + return false; + } + return false; +} + +function isTurnStartEntry(entry: SessionEntry): boolean { + if (entry.type === "compaction") { + return false; + } + return sessionEntryToContextMessages(entry).some(isTurnStartMessage); +} + /** - * Find valid cut points: indices of user, assistant, custom, or bashExecution messages. + * Find valid cut points: indices of context-visible user-like or assistant messages. * Never cut at tool results (they must follow their tool call). * When we cut at an assistant message with tool calls, its tool results follow it * and will be kept. - * BashExecutionMessage is treated like a user message (user-initiated context). */ function findValidCutPoints(entries: SessionEntry[], startIndex: number, endIndex: number): number[] { const cutPoints: number[] = []; for (let i = startIndex; i < endIndex; i++) { const entry = entries[i]; - switch (entry.type) { - case "message": { - const role = entry.message.role; - switch (role) { - case "bashExecution": - case "custom": - case "branchSummary": - case "compactionSummary": - case "user": - case "assistant": - cutPoints.push(i); - break; - case "toolResult": - break; - } - break; - } - case "thinking_level_change": - case "model_change": - case "compaction": - case "branch_summary": - case "custom": - case "custom_message": - case "label": - case "session_info": - break; + if (entry.type === "compaction") { + continue; } - - // branch_summary and custom_message are user-role messages, valid cut points - if (entry.type === "branch_summary" || entry.type === "custom_message") { + if (sessionEntryToContextMessages(entry).some(isCutPointMessage)) { cutPoints.push(i); } } @@ -343,23 +337,14 @@ function findValidCutPoints(entries: SessionEntry[], startIndex: number, endInde } /** - * Find the user message (or bashExecution) that starts the turn containing the given entry index. + * Find the context-visible user-role message that starts the turn containing the given entry index. * Returns -1 if no turn start found before the index. - * BashExecutionMessage is treated like a user message for turn boundaries. */ export function findTurnStartIndex(entries: SessionEntry[], entryIndex: number, startIndex: number): number { for (let i = entryIndex; i >= startIndex; i--) { - const entry = entries[i]; - // branch_summary and custom_message are user-role messages, can start a turn - if (entry.type === "branch_summary" || entry.type === "custom_message") { + if (isTurnStartEntry(entries[i])) { return i; } - if (entry.type === "message") { - const role = entry.message.role; - if (role === "user" || role === "bashExecution") { - return i; - } - } } return -1; } @@ -407,10 +392,11 @@ export function findCutPoint( for (let i = endIndex - 1; i >= startIndex; i--) { const entry = entries[i]; - if (entry.type !== "message") continue; - - // Estimate this message's size - const messageTokens = estimateTokens(entry.message); + const messageTokens = sessionEntryToContextMessages(entry).reduce( + (sum, message) => sum + estimateTokens(message), + 0, + ); + if (messageTokens === 0) continue; accumulatedTokens += messageTokens; // Check if we've exceeded the budget @@ -426,30 +412,25 @@ export function findCutPoint( } } - // Scan backwards from cutIndex to include any non-message entries (bash, settings, etc.) + // Scan backwards from cutIndex to include adjacent metadata entries that do not affect context. while (cutIndex > startIndex) { const prevEntry = entries[cutIndex - 1]; - // Stop at session header or compaction boundaries - if (prevEntry.type === "compaction") { + // Stop at compaction boundaries or context-visible entries. + if (prevEntry.type === "compaction" || sessionEntryToContextMessages(prevEntry).length > 0) { break; } - if (prevEntry.type === "message") { - // Stop if we hit any message - break; - } - // Include this non-message entry (bash, settings change, etc.) cutIndex--; } // Determine if this is a split turn const cutEntry = entries[cutIndex]; - const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user"; - const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); + const startsTurn = isTurnStartEntry(cutEntry); + const turnStartIndex = startsTurn ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); return { firstKeptEntryIndex: cutIndex, turnStartIndex, - isSplitTurn: !isUserMessage && turnStartIndex !== -1, + isSplitTurn: !startsTurn && turnStartIndex !== -1, }; } diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index ab09c3bf6..12c9d0be2 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -18,6 +18,7 @@ import { import { buildSessionContext, type CompactionEntry, + type CustomMessageEntry, type ModelChangeEntry, migrateSessionEntries, parseSessionEntries, @@ -134,6 +135,21 @@ function createThinkingLevelEntry(thinkingLevel: string): ThinkingLevelChangeEnt return entry; } +function createCustomMessageEntry(content: string): CustomMessageEntry { + const id = `test-id-${entryCounter++}`; + const entry: CustomMessageEntry = { + type: "custom_message", + id, + parentId: lastId, + timestamp: new Date().toISOString(), + customType: "test", + content, + display: true, + }; + lastId = id; + return entry; +} + function extractText(messages: AgentMessage[]): string { return messages .map((message) => { @@ -338,6 +354,25 @@ describe("findCutPoint", () => { expect(result.turnStartIndex).toBe(2); // Turn 2 starts at index 2 } }); + + it("should budget context-visible custom message entries", () => { + const entries: SessionEntry[] = [ + createMessageEntry(createUserMessage("hi")), + createMessageEntry(createAssistantMessage("hello")), + createCustomMessageEntry("x".repeat(4000)), + createMessageEntry(createAssistantMessage("ok")), + ]; + + const tinyBudget = findCutPoint(entries, 0, entries.length, 1); + expect(tinyBudget.firstKeptEntryIndex).toBe(3); + expect(tinyBudget.isSplitTurn).toBe(true); + expect(tinyBudget.turnStartIndex).toBe(2); + + const customFitsBudget = findCutPoint(entries, 0, entries.length, 2); + expect(customFitsBudget.firstKeptEntryIndex).toBe(2); + expect(customFitsBudget.isSplitTurn).toBe(false); + expect(customFitsBudget.turnStartIndex).toBe(-1); + }); }); describe("buildSessionContext", () => {