From 5aad9c40ec11add6af0df33afe762e885a3489d9 Mon Sep 17 00:00:00 2001 From: Gugu8 Date: Wed, 9 Sep 2026 11:55:12 +1200 Subject: [PATCH 1/2] Improve comments in chat-compression.ts Refactor comments and documentation in chat-compression.ts for clarity. --- cli/src/utils/chat-compression.ts | 421 ++++++++++++++++++++++++++++++ 1 file changed, 421 insertions(+) create mode 100644 cli/src/utils/chat-compression.ts diff --git a/cli/src/utils/chat-compression.ts b/cli/src/utils/chat-compression.ts new file mode 100644 index 0000000000..bd5d59cf2a --- /dev/null +++ b/cli/src/utils/chat-compression.ts @@ -0,0 +1,421 @@ +// With the help of opus 5. I mostly did it tho. Just had opus to fix any bugs and make it easier to read. + +// type +export type MessageRole = 'system' | 'user' | 'assistant' | 'tool' + +// minimal structure +export interface ChatMessage { + role: MessageRole + content: unknown + //set on use by syntehtic data + _compaction?: CompactionMeta + [key: string]: unknown +} + +export interface CompactionMeta { + kind: 'summary' + // hhow many original messages this summary stands in for + replacedCount: number + // token count of it + replacedTokens: number + // goes up everytime a summary is summazried + generation: number + createdAt: number +} + +export interface CompressionOptions { + // ceiling of the token + maxContextTokens: number + /** + * Fraction of `maxContextTokens` at which compression kicks in. + * @default 0.8 + */ + triggerRatio?: number + /** + * Fraction of `maxContextTokens` to target after compression. + * @default 0.5 + */ + targetRatio?: number + /** + * Always keep at least this many trailing messages verbatim. + * @default 6 + */ + minTailMessages?: number + /** + * Never summarize unless at least this many messages would be collapsed + * (avoids burning an LLM call to save nothing). + * @default 4 + */ + minMessagesToSummarize?: number + /** + * Number of leading messages treated as pinned (system prompt, etc.). + * Auto-detected from leading `system` messages if omitted. + */ + pinnedHeadCount?: number + /** Token counter. Defaults to a ~4 chars/token heuristic. */ + countTokens?: (message: ChatMessage) => number + /** Produces the summary text. Required when compression actually runs. */ + summarize?: Summarizer +} + +export type Summarizer = (input: SummarizeInput) => Promise + +export interface SummarizeInput { + /** Messages to be collapsed, in order. */ + messages: ChatMessage[] + /** Text of any prior summary being folded in, if present. */ + previousSummary?: string + /** Soft budget for the produced summary. */ + maxSummaryTokens: number +} + +export type CompressionReason = + | 'under-threshold' + | 'nothing-to-summarize' + | 'compressed' + | 'no-summarizer' + +export interface CompressionResult { + messages: ChatMessage[] + compressed: boolean + reason: CompressionReason + tokensBefore: number + tokensAfter: number + messagesRemoved: number + /** Index in the returned array where the summary lives, if any. */ + summaryIndex?: number +} + +// --------------------------------------------------------------------------- +// Defaults +// --------------------------------------------------------------------------- + +const DEFAULTS = { + triggerRatio: 0.8, + targetRatio: 0.5, + minTailMessages: 6, + minMessagesToSummarize: 4, +} as const + +/** Rough heuristic: ~4 characters per token for English + code. */ +const CHARS_PER_TOKEN = 4 + +/** Overhead per message for role/formatting tokens. */ +const PER_MESSAGE_OVERHEAD_TOKENS = 4 + +// --------------------------------------------------------------------------- +// Token estimation +// --------------------------------------------------------------------------- + +/** + * Flattens arbitrary message content into a string for length estimation. + * Handles strings, content-part arrays, and nested objects. + */ +export function contentToText(content: unknown): string { + if (content == null) return '' + if (typeof content === 'string') return content + if (typeof content === 'number' || typeof content === 'boolean') { + return String(content) + } + if (Array.isArray(content)) { + return content.map(contentToText).join('\n') + } + if (typeof content === 'object') { + const part = content as Record + if (typeof part.text === 'string') return part.text + // tool calls, tool results, images, etc. + try { + return JSON.stringify(part) + } catch { + return '' + } + } + return '' +} + +export function estimateMessageTokens(message: ChatMessage): number { + const text = contentToText(message.content) + return Math.ceil(text.length / CHARS_PER_TOKEN) + PER_MESSAGE_OVERHEAD_TOKENS +} + +export function estimateTotalTokens( + messages: readonly ChatMessage[], + countTokens: (m: ChatMessage) => number = estimateMessageTokens, +): number { + let total = 0 + for (const message of messages) total += countTokens(message) + return total +} + +// --------------------------------------------------------------------------- +// Message classification +// --------------------------------------------------------------------------- + +function isToolResultMessage(message: ChatMessage): boolean { + if (message.role === 'tool') return true + const content = message.content + if (!Array.isArray(content)) return false + return content.some( + (part) => + typeof part === 'object' && + part !== null && + (part as Record).type === 'tool_result', + ) +} + +function isSummaryMessage(message: ChatMessage): boolean { + return message._compaction?.kind === 'summary' +} + +function detectPinnedHeadCount(messages: readonly ChatMessage[]): number { + let i = 0 + while (i < messages.length && messages[i].role === 'system') i++ + return i +} + +// --------------------------------------------------------------------------- +// Boundary selection +// --------------------------------------------------------------------------- + +/** + * Walks `index` backwards until it no longer points at a tool result, so the + * assistant message that issued the tool call stays with its results. + */ +function alignToPairBoundary( + messages: readonly ChatMessage[], + index: number, + lowerBound: number, +): number { + let i = index + while (i > lowerBound && isToolResultMessage(messages[i])) i-- + return i +} + +/** + * Chooses the index where the preserved tail begins. + * + * Grows the tail backwards from the end until it would exceed `tailBudget` + * tokens, honours `minTailMessages`, then snaps to a tool-pair boundary. + */ +export function findTailStart( + messages: readonly ChatMessage[], + opts: { + pinnedHeadCount: number + minTailMessages: number + tailBudgetTokens: number + countTokens: (m: ChatMessage) => number + }, +): number { + const { pinnedHeadCount, minTailMessages, tailBudgetTokens, countTokens } = + opts + + let index = messages.length + let used = 0 + + while (index > pinnedHeadCount) { + const candidate = index - 1 + const cost = countTokens(messages[candidate]) + const kept = messages.length - candidate + const withinBudget = used + cost <= tailBudgetTokens + + if (!withinBudget && kept > minTailMessages) break + + used += cost + index = candidate + } + + return alignToPairBoundary(messages, index, pinnedHeadCount) +} + +// --------------------------------------------------------------------------- +// Summary message construction +// --------------------------------------------------------------------------- + +const SUMMARY_PREAMBLE = + 'The earlier portion of this conversation was compacted to fit the context ' + + 'window. Treat the following as an accurate record of what happened:' + +export function buildSummaryMessage( + summaryText: string, + meta: Omit, +): ChatMessage { + return { + role: 'user', + content: `${SUMMARY_PREAMBLE}\n\n${summaryText.trim()}`, + _compaction: { + kind: 'summary', + createdAt: Date.now(), + ...meta, + }, + } +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +export function shouldCompress( + messages: readonly ChatMessage[], + options: Pick< + CompressionOptions, + 'maxContextTokens' | 'triggerRatio' | 'countTokens' + >, +): boolean { + const { + maxContextTokens, + triggerRatio = DEFAULTS.triggerRatio, + countTokens = estimateMessageTokens, + } = options + return ( + estimateTotalTokens(messages, countTokens) >= + maxContextTokens * triggerRatio + ) +} + +/** + * Compresses a conversation if it has crossed the trigger threshold. + * + * Returns the original array (same reference) when no work was done, so + * callers can cheaply check `result.compressed`. + */ +export async function compressChat( + messages: ChatMessage[], + options: CompressionOptions, +): Promise { + const { + maxContextTokens, + triggerRatio = DEFAULTS.triggerRatio, + targetRatio = DEFAULTS.targetRatio, + minTailMessages = DEFAULTS.minTailMessages, + minMessagesToSummarize = DEFAULTS.minMessagesToSummarize, + countTokens = estimateMessageTokens, + summarize, + } = options + + const pinnedHeadCount = + options.pinnedHeadCount ?? detectPinnedHeadCount(messages) + + const tokensBefore = estimateTotalTokens(messages, countTokens) + + const unchanged = (reason: CompressionReason): CompressionResult => ({ + messages, + compressed: false, + reason, + tokensBefore, + tokensAfter: tokensBefore, + messagesRemoved: 0, + }) + + if (tokensBefore < maxContextTokens * triggerRatio) { + return unchanged('under-threshold') + } + + const headTokens = estimateTotalTokens( + messages.slice(0, pinnedHeadCount), + countTokens, + ) + const targetTokens = maxContextTokens * targetRatio + const maxSummaryTokens = Math.max( + 256, + Math.floor(targetTokens * 0.15), + ) + const tailBudgetTokens = Math.max( + 0, + targetTokens - headTokens - maxSummaryTokens, + ) + + const tailStart = findTailStart(messages, { + pinnedHeadCount, + minTailMessages, + tailBudgetTokens, + countTokens, + }) + + const middle = messages.slice(pinnedHeadCount, tailStart) + if (middle.length < minMessagesToSummarize) { + return unchanged('nothing-to-summarize') + } + if (!summarize) { + return unchanged('no-summarizer') + } + + // Fold any existing summary into the new one instead of nesting them. + const priorSummary = middle.find(isSummaryMessage) + const toSummarize = middle.filter((m) => !isSummaryMessage(m)) + const generation = (priorSummary?._compaction?.generation ?? 0) + 1 + + const summaryText = await summarize({ + messages: toSummarize, + previousSummary: priorSummary + ? contentToText(priorSummary.content) + : undefined, + maxSummaryTokens, + }) + + if (!summaryText.trim()) { + return unchanged('nothing-to-summarize') + } + + const summaryMessage = buildSummaryMessage(summaryText, { + replacedCount: + middle.length + (priorSummary?._compaction?.replacedCount ?? 0), + replacedTokens: + estimateTotalTokens(middle, countTokens) + + (priorSummary?._compaction?.replacedTokens ?? 0), + generation, + }) + + const next = [ + ...messages.slice(0, pinnedHeadCount), + summaryMessage, + ...messages.slice(tailStart), + ] + + return { + messages: next, + compressed: true, + reason: 'compressed', + tokensBefore, + tokensAfter: estimateTotalTokens(next, countTokens), + messagesRemoved: middle.length - 1, + summaryIndex: pinnedHeadCount, + } +} + +/** + * Convenience wrapper: compress repeatedly until under the target, or until + * no further progress is possible. Guards against summarizer no-ops. + */ +export async function compressChatToFit( + messages: ChatMessage[], + options: CompressionOptions, + maxPasses = 3, +): Promise { + let current = messages + let last: CompressionResult | undefined + + for (let pass = 0; pass < maxPasses; pass++) { + const result = await compressChat(current, options) + last = result + if (!result.compressed) break + if (result.tokensAfter >= result.tokensBefore) break + current = result.messages + if ( + result.tokensAfter < + options.maxContextTokens * (options.targetRatio ?? DEFAULTS.targetRatio) + ) { + break + } + } + + return ( + last ?? { + messages, + compressed: false, + reason: 'under-threshold', + tokensBefore: estimateTotalTokens(messages, options.countTokens), + tokensAfter: estimateTotalTokens(messages, options.countTokens), + messagesRemoved: 0, + } + ) +} From dfec4cd2cebb58c630cb605fb3b645aedeeb9729 Mon Sep 17 00:00:00 2001 From: Gugu8 Date: Wed, 9 Sep 2026 12:02:26 +1200 Subject: [PATCH 2/2] Refactor chat-compression.ts for clarity and maintainability Refactor chat compression utility with improved type annotations and documentation. Update compression logic to enhance clarity and maintainability. --- cli/src/utils/chat-compression.ts | 217 ++++++++++++++++++++++++------ 1 file changed, 175 insertions(+), 42 deletions(-) diff --git a/cli/src/utils/chat-compression.ts b/cli/src/utils/chat-compression.ts index bd5d59cf2a..dace233d13 100644 --- a/cli/src/utils/chat-compression.ts +++ b/cli/src/utils/chat-compression.ts @@ -1,31 +1,62 @@ -// With the help of opus 5. I mostly did it tho. Just had opus to fix any bugs and make it easier to read. +// had opus 5 help me. Mostly did it myself for once. + +import { + formatTokens, + freshInputTokens, + totalTokens, +} from '@codebuff/common/util/tokens' // ← verify this specifier + +// types -// type export type MessageRole = 'system' | 'user' | 'assistant' | 'tool' -// minimal structure export interface ChatMessage { role: MessageRole content: unknown - //set on use by syntehtic data + /** Set by us on synthetic summary messages. */ _compaction?: CompactionMeta [key: string]: unknown } export interface CompactionMeta { kind: 'summary' - // hhow many original messages this summary stands in for + /** How many original messages this summary stands in for. */ replacedCount: number - // token count of it + /** Token count of the messages that were replaced. */ replacedTokens: number - // goes up everytime a summary is summazried + /** Incremented each time a summary is re-summarized. */ generation: number createdAt: number } +/** + * A provider usage row, in the same shape `util/tokens` consumes. + * + * `inputTokens` is the provider's `prompt_tokens` and ALREADY INCLUDES + * `cacheReadTokens` — never add them together. See `totalTokens`. + */ +export interface UsageSnapshot { + inputTokens: number + outputTokens: number + cacheReadTokens: number + /** + * How many leading entries of the message array this row accounts for, + * INCLUDING the assistant message the row's `outputTokens` produced. + * + * Anything past this index has never been sent to the provider and must be + * estimated. Get this wrong and the threshold silently drifts. + */ + coversMessageCount: number +} + export interface CompressionOptions { - // ceiling of the token + /** Hard ceiling for the conversation, in tokens. */ maxContextTokens: number + /** + * Most recent provider usage row. When supplied, the threshold check uses + * real numbers for the covered prefix instead of the ~4 chars/token guess. + */ + lastUsage?: UsageSnapshot /** * Fraction of `maxContextTokens` at which compression kicks in. * @default 0.8 @@ -47,12 +78,24 @@ export interface CompressionOptions { * @default 4 */ minMessagesToSummarize?: number + /** + * Skip compression while the provider is prefilling less than this fraction + * of the prompt, i.e. while the cache is doing the work for us. + * + * Compaction rewrites the prefix and invalidates the cached prefill, so a + * conversation at 85% cache hit can get *more* expensive after compressing. + * Requires `lastUsage`; ignored without it. + * + * Set to 0 to always compress on threshold. + * @default 0.15 + */ + minFreshInputRatio?: number /** * Number of leading messages treated as pinned (system prompt, etc.). * Auto-detected from leading `system` messages if omitted. */ pinnedHeadCount?: number - /** Token counter. Defaults to a ~4 chars/token heuristic. */ + /** Fallback estimator for messages the provider has not priced. */ countTokens?: (message: ChatMessage) => number /** Produces the summary text. Required when compression actually runs. */ summarize?: Summarizer @@ -71,9 +114,13 @@ export interface SummarizeInput { export type CompressionReason = | 'under-threshold' + | 'cache-warm' | 'nothing-to-summarize' - | 'compressed' | 'no-summarizer' + | 'compressed' + +/** Where a token figure came from. Worth surfacing — the two differ a lot. */ +export type TokenSource = 'provider' | 'estimated' | 'mixed' export interface CompressionResult { messages: ChatMessage[] @@ -81,6 +128,7 @@ export interface CompressionResult { reason: CompressionReason tokensBefore: number tokensAfter: number + tokenSource: TokenSource messagesRemoved: number /** Index in the returned array where the summary lives, if any. */ summaryIndex?: number @@ -95,6 +143,7 @@ const DEFAULTS = { targetRatio: 0.5, minTailMessages: 6, minMessagesToSummarize: 4, + minFreshInputRatio: 0.15, } as const /** Rough heuristic: ~4 characters per token for English + code. */ @@ -104,7 +153,7 @@ const CHARS_PER_TOKEN = 4 const PER_MESSAGE_OVERHEAD_TOKENS = 4 // --------------------------------------------------------------------------- -// Token estimation +// Token estimation (fallback only) // --------------------------------------------------------------------------- /** @@ -147,6 +196,57 @@ export function estimateTotalTokens( return total } +// --------------------------------------------------------------------------- +// Context measurement +// --------------------------------------------------------------------------- + +export interface ContextMeasurement { + tokens: number + source: TokenSource +} + +/** + * Current conversation size. + * + * With a usage row: `totalTokens(row)` for the covered prefix — the provider's + * own prompt+completion figure, cached prefix counted exactly once — plus an + * estimate for anything appended since. Without one: pure estimate. + */ +export function measureContext( + messages: readonly ChatMessage[], + options: Pick, +): ContextMeasurement { + const { lastUsage, countTokens = estimateMessageTokens } = options + + if (!lastUsage) { + return { tokens: estimateTotalTokens(messages, countTokens), source: 'estimated' } + } + + const covered = Math.min( + Math.max(0, lastUsage.coversMessageCount), + messages.length, + ) + const uncounted = messages.slice(covered) + const tokens = + totalTokens(lastUsage) + estimateTotalTokens(uncounted, countTokens) + + return { + tokens, + source: uncounted.length === 0 ? 'provider' : 'mixed', + } +} + +/** + * Fraction of the last prompt the provider actually had to prefill. + * + * Returns 1 when there is nothing to go on, so a missing usage row never + * suppresses compression. + */ +export function freshInputRatio(usage: UsageSnapshot | undefined): number { + if (!usage || usage.inputTokens <= 0) return 1 + return freshInputTokens(usage) / usage.inputTokens +} + // --------------------------------------------------------------------------- // Message classification // --------------------------------------------------------------------------- @@ -239,9 +339,13 @@ export function buildSummaryMessage( summaryText: string, meta: Omit, ): ChatMessage { + const header = + `${SUMMARY_PREAMBLE}\n` + + `[${meta.replacedCount} messages, ${formatTokens(meta.replacedTokens)} tokens]` + return { role: 'user', - content: `${SUMMARY_PREAMBLE}\n\n${summaryText.trim()}`, + content: `${header}\n\n${summaryText.trim()}`, _compaction: { kind: 'summary', createdAt: Date.now(), @@ -258,18 +362,23 @@ export function shouldCompress( messages: readonly ChatMessage[], options: Pick< CompressionOptions, - 'maxContextTokens' | 'triggerRatio' | 'countTokens' + | 'maxContextTokens' + | 'triggerRatio' + | 'countTokens' + | 'lastUsage' + | 'minFreshInputRatio' >, ): boolean { const { maxContextTokens, triggerRatio = DEFAULTS.triggerRatio, - countTokens = estimateMessageTokens, + minFreshInputRatio = DEFAULTS.minFreshInputRatio, } = options - return ( - estimateTotalTokens(messages, countTokens) >= - maxContextTokens * triggerRatio - ) + + const { tokens } = measureContext(messages, options) + if (tokens < maxContextTokens * triggerRatio) return false + + return freshInputRatio(options.lastUsage) >= minFreshInputRatio } /** @@ -284,10 +393,12 @@ export async function compressChat( ): Promise { const { maxContextTokens, + lastUsage, triggerRatio = DEFAULTS.triggerRatio, targetRatio = DEFAULTS.targetRatio, minTailMessages = DEFAULTS.minTailMessages, minMessagesToSummarize = DEFAULTS.minMessagesToSummarize, + minFreshInputRatio = DEFAULTS.minFreshInputRatio, countTokens = estimateMessageTokens, summarize, } = options @@ -295,30 +406,31 @@ export async function compressChat( const pinnedHeadCount = options.pinnedHeadCount ?? detectPinnedHeadCount(messages) - const tokensBefore = estimateTotalTokens(messages, countTokens) + const before = measureContext(messages, options) const unchanged = (reason: CompressionReason): CompressionResult => ({ messages, compressed: false, reason, - tokensBefore, - tokensAfter: tokensBefore, + tokensBefore: before.tokens, + tokensAfter: before.tokens, + tokenSource: before.source, messagesRemoved: 0, }) - if (tokensBefore < maxContextTokens * triggerRatio) { + if (before.tokens < maxContextTokens * triggerRatio) { return unchanged('under-threshold') } + if (freshInputRatio(lastUsage) < minFreshInputRatio) { + return unchanged('cache-warm') + } const headTokens = estimateTotalTokens( messages.slice(0, pinnedHeadCount), countTokens, ) const targetTokens = maxContextTokens * targetRatio - const maxSummaryTokens = Math.max( - 256, - Math.floor(targetTokens * 0.15), - ) + const maxSummaryTokens = Math.max(256, Math.floor(targetTokens * 0.15)) const tailBudgetTokens = Math.max( 0, targetTokens - headTokens - maxSummaryTokens, @@ -371,12 +483,15 @@ export async function compressChat( ...messages.slice(tailStart), ] + // The usage row described the *old* prefix, which no longer exists. Every + // figure from here on is an estimate until the next response comes back. return { messages: next, compressed: true, reason: 'compressed', - tokensBefore, + tokensBefore: before.tokens, tokensAfter: estimateTotalTokens(next, countTokens), + tokenSource: before.source === 'provider' ? 'mixed' : before.source, messagesRemoved: middle.length - 1, summaryIndex: pinnedHeadCount, } @@ -385,37 +500,55 @@ export async function compressChat( /** * Convenience wrapper: compress repeatedly until under the target, or until * no further progress is possible. Guards against summarizer no-ops. + * + * Only the first pass can use the provider usage row — after that the prefix + * has been rewritten, so subsequent passes drop it and estimate. */ export async function compressChatToFit( messages: ChatMessage[], options: CompressionOptions, maxPasses = 3, ): Promise { + const targetRatio = options.targetRatio ?? DEFAULTS.targetRatio let current = messages + let currentOptions = options let last: CompressionResult | undefined for (let pass = 0; pass < maxPasses; pass++) { - const result = await compressChat(current, options) + const result = await compressChat(current, currentOptions) last = result if (!result.compressed) break if (result.tokensAfter >= result.tokensBefore) break + current = result.messages - if ( - result.tokensAfter < - options.maxContextTokens * (options.targetRatio ?? DEFAULTS.targetRatio) - ) { - break - } + currentOptions = { ...currentOptions, lastUsage: undefined } + + if (result.tokensAfter < options.maxContextTokens * targetRatio) break } + if (last) return last + + const measured = measureContext(messages, options) + return { + messages, + compressed: false, + reason: 'under-threshold', + tokensBefore: measured.tokens, + tokensAfter: measured.tokens, + tokenSource: measured.source, + messagesRemoved: 0, + } +} + +/** One-line rendering for logs and the status bar. */ +export function describeCompression(result: CompressionResult): string { + if (!result.compressed) { + return `no compression (${result.reason}, ${formatTokens(result.tokensBefore)})` + } + const saved = result.tokensBefore - result.tokensAfter return ( - last ?? { - messages, - compressed: false, - reason: 'under-threshold', - tokensBefore: estimateTotalTokens(messages, options.countTokens), - tokensAfter: estimateTotalTokens(messages, options.countTokens), - messagesRemoved: 0, - } + `compacted ${result.messagesRemoved} messages: ` + + `${formatTokens(result.tokensBefore)} → ${formatTokens(result.tokensAfter)} ` + + `(−${formatTokens(saved)})` ) }