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

fix(coding-agent): count custom messages in compaction budget

Browse Source 
closes #6326
This commit is contained in:
Mario Zechner
2026-07-09 12:45:26 +02:00
parent e9fa5a68a1
commit a6f720e6ca
3 changed files with 97 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/core/compaction/compaction.ts
+61 -80
View File
@@ -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,
};
}