feat: implemented token tracking and metadata options for compaction summaries

- Added `CompactionSummaryMessageOptions` interface to support optional fields and metadata in compaction summary messages.
- Included compaction method labels and before-after amount badge rendering in compaction views.
- Updated session maintenance and compaction methods to track and compute post-compaction token counts.
- Updated calls across agent and coding-agent packages to use the new compaction options object.
This commit is contained in:
can1357
2026-08-20 07:18:34 +02:00
parent 6740f146b0
commit d820cf75bb
16 changed files with 205 additions and 83 deletions
+1
View File
@@ -16,6 +16,7 @@
### Changed
- Expanded native tokenizer support across catalog models, adding exact embedded token counting for Claude, Qwen 3.5+, DeepSeek V3/V4/R1, Kimi K2/K3, and GLM-5+ models. `Tokenizer` now constructs from a resolved catalog `Model`.
- `createCompactionSummaryMessage` takes an options object after `(summary, tokensBefore, timestamp)`; `CompactionSummaryMessage` gained optional `method` and `tokensAfter` display metadata.
## [17.3.8] - 2026-08-19
@@ -191,7 +191,9 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
case "compaction":
return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp, entry.shortSummary);
return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp, {
shortSummary: entry.shortSummary,
});
// These don't contribute to conversation content
case "thinking_level_change":
+21 -5
View File
@@ -49,6 +49,10 @@ export interface CompactionSummaryMessage {
summary: string;
shortSummary?: string;
tokensBefore: number;
/** Estimated context tokens after the rewrite (display metadata). */
tokensAfter?: number;
/** Harness compaction method that produced this summary (display metadata). */
method?: string;
providerPayload?: ProviderPayload;
/** Runtime-only ordered archive blocks for snapcompact: old text region,
* imaged middle, then new text region. When present, `summary` is already
@@ -99,16 +103,26 @@ export function createBranchSummaryMessage(summary: string, fromId: string, time
};
}
/** Optional metadata for {@link createCompactionSummaryMessage}. */
export interface CompactionSummaryMessageOptions {
shortSummary?: string;
providerPayload?: ProviderPayload;
images?: ImageContent[];
blocks?: (TextContent | ImageContent)[];
warning?: string;
/** Harness compaction method that produced this summary (e.g. "remote", "soft", "handoff"). */
method?: string;
/** Estimated context tokens after the rewrite, for display alongside `tokensBefore`. */
tokensAfter?: number;
}
export function createCompactionSummaryMessage(
summary: string,
tokensBefore: number,
timestamp: string,
shortSummary?: string,
providerPayload?: ProviderPayload,
images?: ImageContent[],
blocks?: (TextContent | ImageContent)[],
warning?: string,
options: CompactionSummaryMessageOptions = {},
): CompactionSummaryMessage {
const { shortSummary, providerPayload, images, blocks, warning, method, tokensAfter } = options;
const imageBlocks =
blocks?.filter((block): block is ImageContent => block.type === "image") ??
(images && images.length > 0 ? images : undefined);
@@ -117,6 +131,8 @@ export function createCompactionSummaryMessage(
summary,
shortSummary,
tokensBefore,
tokensAfter,
method,
providerPayload,
blocks: blocks && blocks.length > 0 ? blocks : undefined,
images: imageBlocks && imageBlocks.length > 0 ? imageBlocks : undefined,
+3 -15
View File
@@ -14,28 +14,16 @@ describe("compaction summary message with snapcompact frames", () => {
it("countMessage charges per attached frame", () => {
const bare = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString());
const withFrames = createCompactionSummaryMessage(
"summary text",
1000,
new Date().toISOString(),
undefined,
undefined,
images,
);
const withFrames = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString(), { images });
expect(tokenizer.countMessage(withFrames) - tokenizer.countMessage(bare)).toBe(
2 * snapcompact.FRAME_TOKEN_ESTIMATE,
);
});
it("defaultConvertToLlm appends frames as image blocks after the summary text", () => {
const message = createCompactionSummaryMessage(
"the snapcompact archive",
1000,
new Date().toISOString(),
undefined,
undefined,
const message = createCompactionSummaryMessage("the snapcompact archive", 1000, new Date().toISOString(), {
images,
);
});
const [converted] = defaultConvertToLlm([message]);
expect(converted.role).toBe("user");
const content = converted.content as Array<{ type: string; text?: string; data?: string }>;
+1
View File
@@ -20,6 +20,7 @@
### Changed
- Revamped the todo HUD — overall progress renders along the tree-spine connector with smooth completion transitions.
- Compaction divider now names the maintenance method that fired (`remote-compacted`, `soft-compacted`, `handed-off`, `snap-compacted`) and shows the before → after context size (e.g. `256K→20K`).
- `/handoff` (and automatic handoff compaction) now compacts in place, replacing the session context instead of forking a new session.
- Compaction method priorities — `compaction.methodOrder` takes an ordered preference list (e.g. `[remote, snap]` uses remote compaction where the provider supports it, such as OpenAI, and snap everywhere else), replacing `compaction.strategy`/`compaction.remoteEnabled`.
- Unified inline overlays and selectors (model picker, settings, `/cleanse`) into one titled rounded-box panel style.
@@ -1,7 +1,23 @@
import { Box, type Component, Markdown } from "@oh-my-pi/pi-tui";
import { formatNumber } from "@oh-my-pi/pi-utils";
import { getMarkdownTheme, theme } from "../../modes/theme/theme";
import type { BranchSummaryMessage, CompactionSummaryMessage, CustomMessage } from "../../session/messages";
/** Divider labels per compaction method; unknown/legacy methods fall back to "compacted". */
const COMPACTION_METHOD_LABELS: Record<string, string> = {
remote: "remote-compacted",
soft: "soft-compacted",
handoff: "handed-off",
snapcompact: "snap-compacted",
shake: "shaken",
};
/** `256K→20K` amount badge, or undefined when the entry predates `tokensAfter`. */
function compactionAmount(message: CompactionSummaryMessage): string | undefined {
if (message.tokensAfter === undefined || message.tokensBefore <= 0) return undefined;
return `${formatNumber(message.tokensBefore)}→${formatNumber(message.tokensAfter)}`;
}
interface SummaryDividerOptions {
label: () => string;
detailMarkdown: () => string;
@@ -76,9 +92,12 @@ class SummaryDividerComponent implements Component {
/**
* Compaction point in the transcript, rendered as a slim horizontal divider:
*
* ──────── 📷 compacted · ctrl+o ────────
* ──────── 📷 remote-compacted · 256K→20K · ctrl+o ────────
*
* The conversation above the divider stays visible (display transcript keeps
* The label names the maintenance method that fired (remote/soft/handoff/
* snapcompact; "compacted" for legacy or extension-provided entries) and the
* before → after context amounts when the entry recorded them. The
* conversation above the divider stays visible (display transcript keeps
* full history); only the LLM context was reset. Expanding (ctrl+o) reveals
* the compaction summary below the divider.
*/
@@ -89,14 +108,20 @@ export class CompactionSummaryMessageComponent implements Component {
this.#divider = new SummaryDividerComponent({
// A dead-end warning stamped by the progress guard badges the bar;
// the full text lives in the ctrl+o detail block below.
label: () =>
this.message.warning
? `${theme.icon.camera} compacted ${theme.fg("warning", theme.icon.warning)}`
: `${theme.icon.camera} compacted`,
label: () => this.#label(),
detailMarkdown: () => this.#detailMarkdown(),
});
}
#label(): string {
const name = (this.message.method && COMPACTION_METHOD_LABELS[this.message.method]) || "compacted";
let label = `${theme.icon.camera} ${name}`;
const amount = compactionAmount(this.message);
if (amount) label += `${theme.sep.dot}${amount}`;
if (this.message.warning) label += ` ${theme.fg("warning", theme.icon.warning)}`;
return label;
}
setExpanded(expanded: boolean): void {
this.#divider.setExpanded(expanded);
}
@@ -110,7 +135,10 @@ export class CompactionSummaryMessageComponent implements Component {
}
#detailMarkdown(): string {
const tokenStr = this.message.tokensBefore.toLocaleString();
const tokenStr =
this.message.tokensAfter !== undefined
? `${this.message.tokensBefore.toLocaleString()} to ${this.message.tokensAfter.toLocaleString()}`
: this.message.tokensBefore.toLocaleString();
const frameCount = this.message.images?.length ?? 0;
const frameNote =
frameCount > 0 ? `\n\n_${frameCount} snapcompact frame${frameCount === 1 ? "" : "s"} attached_` : "";
@@ -129,7 +157,7 @@ export class HandoffSummaryMessageComponent implements Component {
constructor(private readonly message: CustomMessage<unknown>) {
this.#divider = new SummaryDividerComponent({
label: () => `${theme.icon.context} handoff`,
label: () => `${theme.icon.context} handed-off`,
detailMarkdown: () => this.#detailMarkdown(),
});
}
@@ -1501,7 +1501,7 @@ export class SessionAdvisors {
// only assistants appended afterward can become the next usage anchor.
const advisorUsageAnchorStartIndex = preparation.recentMessages.length + 1;
const summaryMessage = {
...createCompactionSummaryMessage(summary, tokensBefore, new Date().toISOString(), shortSummary),
...createCompactionSummaryMessage(summary, tokensBefore, new Date().toISOString(), { shortSummary }),
firstKeptEntryId,
advisorUsageAnchorStartIndex,
} satisfies AdvisorCompactionSummaryMessage;
@@ -372,11 +372,13 @@ export function buildSessionContext(
active ? entry.summary : SUPERSEDED_COMPACTION_SUMMARY,
entry.tokensBefore,
entry.timestamp,
active ? entry.shortSummary : SUPERSEDED_COMPACTION_SHORT_SUMMARY,
undefined,
undefined,
snapcompactHistoryBlocksForContext(snapcompactArchive, options),
entry.warning,
{
shortSummary: active ? entry.shortSummary : SUPERSEDED_COMPACTION_SHORT_SUMMARY,
blocks: snapcompactHistoryBlocksForContext(snapcompactArchive, options),
warning: entry.warning,
method: entry.method,
tokensAfter: entry.tokensAfter,
},
),
);
} else {
@@ -414,11 +416,14 @@ export function buildSessionContext(
compaction.summary,
compaction.tokensBefore,
compaction.timestamp,
compaction.shortSummary,
providerPayload,
undefined,
snapcompactHistoryBlocksForContext(snapcompactArchive, options),
compaction.warning,
{
shortSummary: compaction.shortSummary,
providerPayload,
blocks: snapcompactHistoryBlocksForContext(snapcompactArchive, options),
warning: compaction.warning,
method: compaction.method,
tokensAfter: compaction.tokensAfter,
},
);
// Agent context (non-transcript): summary first so the LLM sees the
// compacted context before recent messages.
@@ -1,6 +1,7 @@
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
import type { ImageContent, MessageAttribution, ServiceTierByFamily, TextContent } from "@oh-my-pi/pi-ai";
import type { StructuredSubagentSchemaMode } from "../task/types";
import type { CompactionMethod } from "./compaction-methods";
export const CURRENT_SESSION_VERSION = 3;
@@ -99,6 +100,10 @@ export interface CompactionEntry<T = unknown> extends SessionEntryBase {
shortSummary?: string;
firstKeptEntryId: string;
tokensBefore: number;
/** Estimated context tokens after the rewrite (display metadata). */
tokensAfter?: number;
/** Method that produced this entry; absent on legacy sessions and extension-provided compactions. */
method?: CompactionMethod;
/** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */
details?: T;
/** Hook-provided data to persist across compaction */
@@ -971,6 +971,7 @@ export class SessionMaintenance {
details,
fromExtension,
preserveData,
method: fromExtension ? undefined : selectedMethod,
codexCompaction,
advisorResetReason: "compact",
});
@@ -1101,6 +1102,7 @@ export class SessionMaintenance {
details,
fromExtension: false,
preserveData: undefined,
method: "handoff",
codexCompaction: undefined,
advisorResetReason: "handoff",
});
@@ -1310,6 +1312,7 @@ export class SessionMaintenance {
details: unknown;
fromExtension: boolean;
preserveData: Record<string, unknown> | undefined;
method: CompactionMethod | undefined;
codexCompaction: CodexCompactionContext | undefined;
advisorResetReason: string;
detachExtensionEmit?: boolean;
@@ -1319,9 +1322,13 @@ export class SessionMaintenance {
args.shortSummary,
args.firstKeptEntryId,
args.tokensBefore,
args.details,
args.fromExtension,
args.preserveData,
{
details: args.details,
fromExtension: args.fromExtension,
preserveData: args.preserveData,
method: args.method,
tokensAfter: this.#projectCompactedContextTokens(args),
},
);
const newEntries = this.#host.sessionManager.getEntries();
const sessionContext = this.#host.buildDisplaySessionContext();
@@ -2164,10 +2171,10 @@ export class SessionMaintenance {
result.summary,
result.tokensBefore,
new Date().toISOString(),
result.shortSummary,
undefined,
undefined,
blocks,
{
shortSummary: result.shortSummary,
blocks,
},
);
let tokens =
computeNonMessageTokens(this.#host.nonMessageTokenSource(), this.#tokenizer) +
@@ -2176,6 +2183,40 @@ export class SessionMaintenance {
return tokens;
}
/**
* Estimated context tokens after a compaction commit: fixed non-message
* overhead + the summary message (with any snapcompact frames re-attached)
* + every message from `firstKeptEntryId` to the branch leaf. Mirrors the
* post-commit context rebuild; persisted as `tokensAfter` on the entry so
* the transcript divider can show the before → after amounts.
*/
#projectCompactedContextTokens(args: {
summary: string;
shortSummary: string | undefined;
tokensBefore: number;
firstKeptEntryId: string;
preserveData: Record<string, unknown> | undefined;
}): number {
const archive = snapcompact.getPreservedArchive(args.preserveData);
const blocks = archive
? snapcompact.historyBlocks(archive, { maxFrameDataBytes: snapcompact.FRAME_DATA_BYTES_BUDGET })
: undefined;
const summaryMessage = createCompactionSummaryMessage(args.summary, args.tokensBefore, new Date().toISOString(), {
shortSummary: args.shortSummary,
blocks,
});
let tokens =
computeNonMessageTokens(this.#host.nonMessageTokenSource(), this.#tokenizer) +
this.#tokenizer.countMessage(summaryMessage);
let inKeptRegion = false;
for (const entry of this.#host.sessionManager.getBranch()) {
if (entry.id === args.firstKeptEntryId) inKeptRegion = true;
if (!inKeptRegion) continue;
if (entry.type === "message") tokens += this.#tokenizer.countMessage(entry.message);
}
return tokens;
}
/**
* Post-maintenance progress check for the context-full / snapcompact tail.
*
@@ -2499,9 +2540,18 @@ export class SessionMaintenance {
result.shortSummary,
result.firstKeptEntryId,
result.tokensBefore,
result.details,
false,
result.preserveData,
{
details: result.details,
preserveData: result.preserveData,
method: "snapcompact",
tokensAfter: this.#projectCompactedContextTokens({
summary: result.summary,
shortSummary: result.shortSummary,
tokensBefore: result.tokensBefore,
firstKeptEntryId: result.firstKeptEntryId,
preserveData: result.preserveData,
}),
},
);
const sessionContext = this.#host.buildDisplaySessionContext();
this.#host.agent.replaceMessages(sessionContext.messages);
@@ -2698,6 +2748,7 @@ export class SessionMaintenance {
preserveData: armedSpec.result.preserveData,
fromExtension: false,
codexCompaction: armedSpec.codexCompaction,
method: armedSpec.method,
action,
reason,
willRetry,
@@ -3284,6 +3335,7 @@ export class SessionMaintenance {
preserveData,
fromExtension,
codexCompaction,
method: fromExtension ? undefined : method,
action,
reason,
willRetry,
@@ -3375,6 +3427,7 @@ export class SessionMaintenance {
preserveData: Record<string, unknown> | undefined;
fromExtension: boolean;
codexCompaction: CodexCompactionContext | undefined;
method: CompactionMethod | undefined;
action: "context-full" | "handoff" | "snapcompact" | "remote";
reason: "overflow" | "threshold" | "idle" | "incomplete";
willRetry: boolean;
@@ -3412,6 +3465,7 @@ export class SessionMaintenance {
fromExtension: args.fromExtension,
preserveData: args.preserveData,
codexCompaction: args.codexCompaction,
method: args.method,
advisorResetReason: "auto-compaction",
detachExtensionEmit: detachPostCommit,
});
@@ -21,6 +21,7 @@ import {
import type { StructuredSubagentSchemaMode } from "../task/types";
import { ArtifactManager } from "./artifacts";
import { type BlobPutOptions, type BlobPutResult, BlobStore } from "./blob-store";
import type { CompactionMethod } from "./compaction-methods";
import {
type BashExecutionMessage,
type CustomMessage,
@@ -2240,9 +2241,13 @@ export class SessionManager {
shortSummary: string | undefined,
firstKeptEntryId: string,
tokensBefore: number,
details?: T,
fromExtension?: boolean,
preserveData?: Record<string, unknown>,
options: {
details?: T;
fromExtension?: boolean;
preserveData?: Record<string, unknown>;
method?: CompactionMethod;
tokensAfter?: number;
} = {},
): string {
const entry: CompactionEntry<T> = {
type: "compaction",
@@ -2251,9 +2256,11 @@ export class SessionManager {
shortSummary,
firstKeptEntryId,
tokensBefore,
details,
fromExtension,
preserveData,
tokensAfter: options.tokensAfter,
method: options.method,
details: options.details,
fromExtension: options.fromExtension,
preserveData: options.preserveData,
};
this.#recordEntry(entry);
return entry.id;
@@ -160,9 +160,10 @@ describe("AgentSession snapcompact frame dead-end rescue", () => {
"stale snapcompact archive",
userEntryId,
150_000,
{ readFiles: ["src/a.ts"], modifiedFiles: ["src/b.ts"] },
false,
makeArchivePreserveData(options.frameCount),
{
details: { readFiles: ["src/a.ts"], modifiedFiles: ["src/b.ts"] },
preserveData: makeArchivePreserveData(options.frameCount),
},
);
}
@@ -20,7 +20,7 @@ const SUMMARY = "Earlier the user fixed the login TTL bug.";
function makeComponent(images?: ImageContent[]): CompactionSummaryMessageComponent {
return new CompactionSummaryMessageComponent(
createCompactionSummaryMessage(SUMMARY, 84000, new Date().toISOString(), undefined, undefined, images),
createCompactionSummaryMessage(SUMMARY, 84000, new Date().toISOString(), { images }),
);
}
@@ -36,6 +36,29 @@ describe("CompactionSummaryMessageComponent", () => {
expect(rule).not.toContain(SUMMARY);
});
it("names the compaction method and the before → after amounts on the divider", () => {
const component = new CompactionSummaryMessageComponent(
createCompactionSummaryMessage(SUMMARY, 256_000, new Date().toISOString(), {
method: "remote",
tokensAfter: 20_000,
}),
);
const rule = Bun.stripANSI(component.render(80)[1]);
expect(rule).toContain("remote-compacted");
expect(rule).toContain("256K→20K");
expect(rule).toContain("ctrl+o");
});
it("labels a handoff-method compaction as handed-off", () => {
const component = new CompactionSummaryMessageComponent(
createCompactionSummaryMessage(SUMMARY, 84_000, new Date().toISOString(), { method: "handoff" }),
);
const rule = Bun.stripANSI(component.render(80)[1]);
expect(rule).toContain("handed-off");
// No tokensAfter recorded → no amount badge.
expect(rule).not.toContain("→");
});
it("expanded: reveals the summary (and snapcompact frame count) below the divider", () => {
const component = makeComponent([{ type: "image", data: "ZmFrZQ==", mimeType: "image/png" }]);
component.setExpanded(true);
@@ -39,7 +39,7 @@ describe("handoff summary divider", () => {
expect(component).toBeInstanceOf(HandoffSummaryMessageComponent);
const collapsed = Bun.stripANSI(component!.render(80).join("\n"));
expect(collapsed).toContain("handoff");
expect(collapsed).toContain("handed-off");
expect(collapsed).toContain("ctrl+o");
expect(collapsed).not.toContain("[handoff]");
expect(collapsed).not.toContain("Continue the resize fix");
@@ -82,15 +82,9 @@ describe("large session memory guards", () => {
textTail: "frame",
},
};
const firstCompactionId = session.appendCompaction(
firstSummary,
undefined,
firstKeptEntryId,
1000,
undefined,
undefined,
firstPreserve,
);
const firstCompactionId = session.appendCompaction(firstSummary, undefined, firstKeptEntryId, 1000, {
preserveData: firstPreserve,
});
const rewindId = session.appendMessage({ role: "user", content: "between compactions", timestamp: 3 });
session.appendCompaction(secondSummary, undefined, rewindId, 2000);
await session.flush();
@@ -188,15 +182,9 @@ describe("large session memory guards", () => {
const branchACompactionSummary = `branch-a-${"x".repeat(1024)}`;
const branchAPreserve = { openaiRemoteCompaction: { provider: "anthropic", replacementHistory: [] } };
session.appendCompaction(
branchACompactionSummary,
undefined,
rootId,
1000,
undefined,
undefined,
branchAPreserve,
);
session.appendCompaction(branchACompactionSummary, undefined, rootId, 1000, {
preserveData: branchAPreserve,
});
const branchACompactionId = session.getLeafId();
if (!branchACompactionId) throw new Error("Expected branch A compaction id");
+7 -4
View File
@@ -219,10 +219,13 @@ describe("AgentSession shake", () => {
seedHeavyToolResult("X".repeat(20_000));
const firstKeptEntryId = sessionManager.getBranch()[0]?.id;
if (!firstKeptEntryId) throw new Error("Expected seeded branch");
sessionManager.appendCompaction("remote summary", undefined, firstKeptEntryId, 10_000, {}, false, {
openaiRemoteCompaction: {
provider: "openai",
replacementHistory: [],
sessionManager.appendCompaction("remote summary", undefined, firstKeptEntryId, 10_000, {
details: {},
preserveData: {
openaiRemoteCompaction: {
provider: "openai",
replacementHistory: [],
},
},
});
sessionManager.appendMessage({