diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 0e8871723..fa173c3a7 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added a per-message estimation cache (`estimateTokens`) keyed by message identity, so settled history is token-counted once and reused until an owner mutates it. Non-assistant roles cache unconditionally; assistants cache only when settled (real `usage` with a terminal, non-`aborted`/`error` `stopReason`) so streaming partials never freeze a mid-stream count. Dual option-split maps keep the default and `excludeEncryptedReasoning` (compaction-floor) estimates from colliding. Prune, shake, and cross-package convert caches invalidate through `invalidateMessageCache` / `registerMessageCacheInvalidator` at their mutation seams ([#5934](https://github.com/can1357/oh-my-pi/issues/5934)). + ## [17.0.2] - 2026-07-17 ### Fixed diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 2b5186010..3781ed431 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -43,6 +43,7 @@ import { V2_RETAINED_MESSAGE_TOKEN_BUDGET, } from "./compaction-v2-streaming"; import type { CompactionEntry, SessionEntry } from "./entries"; +import { isEstimateCacheable, readEstimateCache, writeEstimateCache } from "./message-cache"; import { type ConvertToLlm, createBranchSummaryMessage, createCustomMessage, defaultConvertToLlm } from "./messages"; import { buildOpenAiNativeHistory, @@ -364,6 +365,21 @@ const IMAGE_TOKEN_ESTIMATE = 1200; * content) excludes them to avoid false triggers on thinking-heavy turns. */ export function estimateTokens(message: AgentMessage, options?: { excludeEncryptedReasoning?: boolean }): number { + // Settled historical messages are counted once and reused until an owner + // (prune/shake/strip-images) invalidates them; streaming assistants bypass + // the cache entirely (see message-cache.ts settle-gate invariant). + const cacheable = isEstimateCacheable(message); + const excludeEncryptedReasoning = options?.excludeEncryptedReasoning === true; + if (cacheable) { + const cached = readEstimateCache(message, excludeEncryptedReasoning); + if (cached !== undefined) return cached; + } + const result = computeMessageTokens(message, options); + if (cacheable) writeEstimateCache(message, excludeEncryptedReasoning, result); + return result; +} + +function computeMessageTokens(message: AgentMessage, options?: { excludeEncryptedReasoning?: boolean }): number { const fragments: string[] = []; let extra = 0; if ((message as { role?: string }).role === "bashExecution") { diff --git a/packages/agent/src/compaction/index.ts b/packages/agent/src/compaction/index.ts index 401215724..c0586f18d 100644 --- a/packages/agent/src/compaction/index.ts +++ b/packages/agent/src/compaction/index.ts @@ -6,6 +6,7 @@ export * from "./branch-summarization"; export * from "./compaction"; export * from "./entries"; export * from "./errors"; +export * from "./message-cache"; export * from "./messages"; export * from "./openai"; export * from "./pruning"; diff --git a/packages/agent/src/compaction/message-cache.ts b/packages/agent/src/compaction/message-cache.ts new file mode 100644 index 000000000..a2249f926 --- /dev/null +++ b/packages/agent/src/compaction/message-cache.ts @@ -0,0 +1,92 @@ +/** + * Per-message memoization for the two hot history walks: token estimation + * ({@link estimateTokens}) and LLM conversion (the coding-agent's `convertToLlm`). + * + * Long sessions re-walk a settled `AgentMessage[]` every turn, re-tokenizing and + * re-converting historical objects that only the newest suffix can change. These + * caches key on message *identity* so a settled message is counted/converted once + * and reused until an owner rewrites it. + * + * Correctness rests on two invariants: + * + * 1. **Settle gate.** A streaming assistant is mutated under one identity while + * its `usage`/`stopReason` are provisional (the seed carries zeroed usage and + * a placeholder `stopReason`). Caching it would freeze a mid-stream count, so + * estimation only caches assistants that are settled — real `usage` + * (`totalTokens > 0`) with a terminal `stopReason` that is not `"aborted"` / + * `"error"`. Unsettled assistants never read or insert. Non-assistant roles + * are immutable once appended and cache by identity. + * 2. **Owner invalidation.** `pruneToolOutputs` / `pruneSupersededToolResults`, + * `applyShakeRegion`, and `stripImagesFromMessage` rewrite message content in + * place under a stable identity. Each MUST call {@link invalidateMessageCache} + * on the mutated message before the next convert/estimate pass so both caches + * drop the stale entry. The convert cache lives in another package, so it + * subscribes via {@link registerMessageCacheInvalidator}. + */ +import type { AssistantMessage } from "@oh-my-pi/pi-ai"; +import type { AgentMessage } from "../types"; + +/** External cache invalidators (e.g. the coding-agent `convertToLlm` memo). */ +const externalInvalidators = new Set<(message: AgentMessage) => void>(); + +/** + * Register a cache tied to message identity so owner mutations in this package + * (prune/shake) can invalidate it across the package boundary. Returns an + * unregister function. The coding-agent `convertToLlm` memo registers here. + */ +export function registerMessageCacheInvalidator(invalidate: (message: AgentMessage) => void): () => void { + externalInvalidators.add(invalidate); + return () => { + externalInvalidators.delete(invalidate); + }; +} + +// Dual option-split estimate caches: the compaction floor passes +// `excludeEncryptedReasoning` (dropping opaque provider reasoning), so a message +// has two distinct estimates that must not collide in one map. +// +// These are WeakMaps, not symbol-tagged properties, deliberately: callers spread +// messages to derive throwaway variants for counting — `estimateBranchSummaryTokens` +// does `estimateTokens({ ...message, content: truncated })`. A symbol-keyed cache +// value rides along an object spread, so the truncated clone would inherit (and +// return) the full-content estimate. Keying strictly on identity keeps the cache +// off spread copies, which get their own fresh count. +const estimateCacheDefault = new WeakMap(); +const estimateCacheFloored = new WeakMap(); + +/** + * True when this message's estimate is safe to cache by identity. Non-assistants + * are immutable once appended; assistants are cached only once settled (see the + * settle-gate invariant above). + */ +export function isEstimateCacheable(message: AgentMessage): boolean { + if (message.role !== "assistant") return true; + const assistant = message as AssistantMessage; + return ( + assistant.stopReason !== "aborted" && + assistant.stopReason !== "error" && + assistant.usage != null && + assistant.usage.totalTokens > 0 + ); +} + +/** Read a cached estimate for the given option split, or `undefined` on miss. */ +export function readEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean): number | undefined { + return (excludeEncryptedReasoning ? estimateCacheFloored : estimateCacheDefault).get(message); +} + +/** Store an estimate for the given option split. */ +export function writeEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean, value: number): void { + (excludeEncryptedReasoning ? estimateCacheFloored : estimateCacheDefault).set(message, value); +} + +/** + * Drop every cached derivation of `message` after an in-place rewrite. Owners of + * mutation (prune, shake, strip-images) call this at the mutation seam so the + * next convert/estimate pass recomputes from the new content. + */ +export function invalidateMessageCache(message: AgentMessage): void { + estimateCacheDefault.delete(message); + estimateCacheFloored.delete(message); + for (const invalidate of externalInvalidators) invalidate(message); +} diff --git a/packages/agent/src/compaction/pruning.ts b/packages/agent/src/compaction/pruning.ts index 1ba9a7ac0..83d5c2dd5 100644 --- a/packages/agent/src/compaction/pruning.ts +++ b/packages/agent/src/compaction/pruning.ts @@ -6,6 +6,7 @@ import type { ToolResultMessage } from "@oh-my-pi/pi-ai"; import type { AgentMessage, AgentToolCall } from "../types"; import { estimateTokens } from "./compaction"; import type { SessionEntry, SessionMessageEntry } from "./entries"; +import { invalidateMessageCache } from "./message-cache"; import { collectToolCallsById, isProtectedToolResult, @@ -295,6 +296,7 @@ export function pruneSupersededToolResults(entries: SessionEntry[], config: Supe for (const candidate of toPrune) { candidate.message.content = [{ type: "text", text: candidate.notice }]; candidate.message.prunedAt = prunedAt; + invalidateMessageCache(candidate.message as AgentMessage); tokensSaved += estimatePrunedSavings(candidate.tokens, candidate.notice); } return { prunedCount: toPrune.length, tokensSaved }; @@ -398,6 +400,7 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = : createPrunedNotice(candidate.tokens); message.content = [{ type: "text", text: notice }]; message.prunedAt = prunedAt; + invalidateMessageCache(message as AgentMessage); prunedCount++; } diff --git a/packages/agent/src/compaction/shake.ts b/packages/agent/src/compaction/shake.ts index 7e0f2c5ed..db95486de 100644 --- a/packages/agent/src/compaction/shake.ts +++ b/packages/agent/src/compaction/shake.ts @@ -15,6 +15,7 @@ import { countTokens } from "../tokenizer"; import type { AgentMessage } from "../types"; import { estimateTokens } from "./compaction"; import type { CustomMessageEntry, SessionEntry, SessionMessageEntry } from "./entries"; +import { invalidateMessageCache } from "./message-cache"; import { collectToolCallsById, isProtectedToolResult, @@ -406,12 +407,18 @@ export function applyShakeRegion(region: ShakeRegion, replacement: string): void const message = region.entry.message as ToolResultMessage; message.content = [{ type: "text", text: replacement }]; message.prunedAt = Date.now(); + invalidateMessageCache(message as AgentMessage); return; } const slot = getBlockTextSlot(region.entry, region.blockIndex); if (!slot) return; const text = slot.read(); slot.write(text.slice(0, region.start) + replacement + text.slice(region.end)); + // Message entries keep a stable `entry.message` identity across context + // rebuilds, so an in-place block rewrite must drop its cached estimate/convert. + // Custom-message entries are re-materialized into a fresh AgentMessage on every + // buildSessionContext, so they carry no stable cached identity to invalidate. + if (region.entry.type === "message") invalidateMessageCache(region.entry.message); } /** diff --git a/packages/agent/test/message-cache.test.ts b/packages/agent/test/message-cache.test.ts new file mode 100644 index 000000000..ac3db9997 --- /dev/null +++ b/packages/agent/test/message-cache.test.ts @@ -0,0 +1,159 @@ +import { describe, expect, test } from "bun:test"; +import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import type { SessionMessageEntry } from "@oh-my-pi/pi-agent-core/compaction"; +import { + applyShakeRegion, + collectShakeRegions, + DEFAULT_PRUNE_CONFIG, + estimateTokens, + invalidateMessageCache, + isEstimateCacheable, + pruneToolOutputs, +} from "@oh-my-pi/pi-agent-core/compaction"; +import type { AssistantMessage, ToolResultMessage, Usage } from "@oh-my-pi/pi-ai"; + +let idCounter = 0; +function nextId(): string { + return `mc-${idCounter++}`; +} + +function messageEntry(message: AgentMessage): SessionMessageEntry { + return { type: "message", id: nextId(), parentId: null, timestamp: new Date().toISOString(), message }; +} + +function usage(totalTokens: number): Usage { + return { + input: totalTokens, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function settledAssistant(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "anthropic-messages", + provider: "anthropic", + model: "bench", + usage: usage(120), + stopReason: "stop", + timestamp: 1, + }; +} + +function toolResult(text: string, extra?: Partial): ToolResultMessage { + return { + role: "toolResult", + toolCallId: `call-${idCounter++}`, + toolName: "read", + content: [{ type: "text", text }], + isError: false, + timestamp: Date.now(), + ...extra, + }; +} + +describe("estimate cache settle gate", () => { + test("caches settled assistants (terminal stopReason + real usage)", () => { + expect(isEstimateCacheable(settledAssistant("done"))).toBe(true); + }); + + test("bypasses a streaming assistant (zero usage seed)", () => { + const streaming: AssistantMessage = { ...settledAssistant("partial"), usage: usage(0), stopReason: "stop" }; + expect(isEstimateCacheable(streaming)).toBe(false); + }); + + test("bypasses aborted and error assistants even with usage", () => { + expect(isEstimateCacheable({ ...settledAssistant("x"), stopReason: "aborted" })).toBe(false); + expect(isEstimateCacheable({ ...settledAssistant("x"), stopReason: "error" })).toBe(false); + }); + + test("caches non-assistant roles unconditionally", () => { + expect(isEstimateCacheable(toolResult("out") as AgentMessage)).toBe(true); + expect(isEstimateCacheable({ role: "user", content: "hi", timestamp: 1 } as AgentMessage)).toBe(true); + }); + + test("a streaming assistant re-estimates as its content grows", () => { + const streaming: AssistantMessage = { + ...settledAssistant("first chunk"), + usage: usage(0), + stopReason: "stop", + }; + const before = estimateTokens(streaming as AgentMessage); + streaming.content = [{ type: "text", text: "first chunk plus a much longer continuation of streamed text" }]; + const after = estimateTokens(streaming as AgentMessage); + // Unsettled assistants never read the cache, so the grown content is recounted. + expect(after).toBeGreaterThan(before); + }); +}); + +describe("estimate cache option split", () => { + test("default and floored estimates do not collide in one map", () => { + const blob = "blob ".repeat(4000); + const msg: AssistantMessage = { + ...settledAssistant("thinking heavy"), + content: [ + { type: "text", text: "answer" }, + { type: "thinking", thinking: "reasoning", thinkingSignature: blob }, + ], + }; + // Prime the default map first, then the floored one; the floored estimate + // (which drops the encrypted-reasoning blob) must not read the default entry. + const withBlob = estimateTokens(msg as AgentMessage); + const floored = estimateTokens(msg as AgentMessage, { excludeEncryptedReasoning: true }); + expect(withBlob).toBeGreaterThan(floored + 500); + // Cached reads return the same split values. + expect(estimateTokens(msg as AgentMessage)).toBe(withBlob); + expect(estimateTokens(msg as AgentMessage, { excludeEncryptedReasoning: true })).toBe(floored); + }); +}); + +describe("estimate cache invalidation seams", () => { + test("pruneToolOutputs drops the cached estimate of a pruned result", () => { + const big = toolResult("x".repeat(20_000)); + const entries = [messageEntry(big as AgentMessage)]; + const before = estimateTokens(big as AgentMessage); + expect(before).toBeGreaterThan(1000); + + const result = pruneToolOutputs(entries, { ...DEFAULT_PRUNE_CONFIG, protectTokens: 0, minimumSavings: 0 }); + expect(result.prunedCount).toBe(1); + + // After the in-place prune the estimate must reflect the short placeholder, + // not the stale full-content count. + const after = estimateTokens(big as AgentMessage); + expect(after).toBeLessThan(before); + }); + + test("applyShakeRegion drops the cached estimate of a shaken result", () => { + const big = toolResult(`\`\`\`ts\n${"const value = compute(a, b, c, d, e);\n".repeat(400)}\`\`\``); + const entry = messageEntry(big as AgentMessage); + const before = estimateTokens(big as AgentMessage); + + const regions = collectShakeRegions([entry], { + protectTokens: 0, + minSavings: 0, + protectedTools: [], + fenceMinTokens: 0, + }); + expect(regions.length).toBeGreaterThan(0); + applyShakeRegion(regions[0], "[shaken]"); + + const after = estimateTokens(big as AgentMessage); + expect(after).toBeLessThan(before); + }); + + test("explicit invalidateMessageCache forces a recount", () => { + const result = toolResult("original content here"); + const before = estimateTokens(result as AgentMessage); + // Mutate content directly (simulating an owner rewrite) then invalidate. + result.content = [{ type: "text", text: "a much longer replacement body that should count higher than before" }]; + // Without invalidation the stale cached value would still be returned. + expect(estimateTokens(result as AgentMessage)).toBe(before); + invalidateMessageCache(result as AgentMessage); + expect(estimateTokens(result as AgentMessage)).toBeGreaterThan(before); + }); +}); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ad20be261..cf23e12f8 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Long sessions no longer re-run `convertToLlm` over settled history every turn. Conversion is memoized per message identity (plus the assistant `interruptedNext` neighbor flag): an exact re-convert of the same array reuses the outer `Message[]`, append-only growth reuses the converted prefix via slice-on-growth, and the prune/shake/strip-images/prewalk-scrub rewrite seams invalidate the affected message before the next pass. On the `llm-assembly` bench (N=5000) steady/append convert and repeat estimate are all >10x faster with robust MAD-noise well under 20% ([#5934](https://github.com/can1357/oh-my-pi/issues/5934)). + ## [17.0.3] - 2026-07-17 ### Changed diff --git a/packages/coding-agent/bench/llm-assembly.bench.ts b/packages/coding-agent/bench/llm-assembly.bench.ts new file mode 100644 index 000000000..195ae0ab1 --- /dev/null +++ b/packages/coding-agent/bench/llm-assembly.bench.ts @@ -0,0 +1,242 @@ +/** + * Benchmark: LLM-assembly recompute over settled history (perf/long-session-convert-estimate-memo). + * + * Before each model call and during compaction accounting, the agent walks the + * full live `AgentMessage[]` history through: + * 1. `convertToLlm(messages)` — role-specific conversion into provider `Message[]`. + * 2. `estimateTokens(message)` — cl100k-style token counting for prune/shake/floors. + * + * In a long session those historical objects are settled, yet before the memo + * both paths recompute from scratch on every pass. This bench measures cold + * (fresh identities → cache miss) against steady state (warm cache): + * + * - convert first: cold conversion of a never-before-seen history. + * - convert steady: re-convert of the same (warmed) array + append-only growth + * that reuses the settled prefix. + * - estimate first: cold token count of a never-before-seen history. + * - estimate second: repeat count of the identical warmed history. + * + * Acceptance (issue #5934): on N=5000 the first/steady convert and first/second + * estimate speedups are >=10x, and the absolute noise gate uses robust MAD + * noise <=20% of the median (not raw stddev/median). + * + * Run: `bun run packages/coding-agent/bench/llm-assembly.bench.ts` + * Env: `LLM_ASSEMBLY_N` overrides the history length (default 5000); + * `PI_TOKENIZER_ACCURATE=1` uses the native cl100k tokenizer. + */ +import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { estimateTokens } from "@oh-my-pi/pi-agent-core/compaction"; +import type { AssistantMessage, ToolResultMessage, Usage } from "@oh-my-pi/pi-ai"; +import { convertToLlm } from "../src/session/messages"; + +const N = Number(Bun.env.LLM_ASSEMBLY_N ?? 5000); +const WARMUP = 5; +const SAMPLES = 25; + +function settledUsage(total: number): Usage { + return { + input: total, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: total, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function codeBlob(seed: number): string { + return `\`\`\`typescript\nexport function f${seed}(a: number, b: number): number {\n\treturn a + b + ${seed};\n}\n\`\`\``; +} + +/** Build a settled, mixed history: user / assistant (settled usage + tool call) / tool-result triples. + * Every call mints fresh object identities so it reads as a cold (uncached) workload. */ +function buildHistory(count: number): AgentMessage[] { + const messages: AgentMessage[] = []; + for (let i = 0; i < count; i++) { + const ts = 1_700_000_000_000 + i * 1000; + const kind = i % 3; + if (kind === 0) { + messages.push({ + role: "user", + content: `User turn ${i}: please look at this.\n\n${codeBlob(i)}`, + timestamp: ts, + } as AgentMessage); + } else if (kind === 1) { + const assistant: AssistantMessage = { + role: "assistant", + content: [ + { type: "text", text: `Assistant turn ${i}. ${codeBlob(i)}` }, + { type: "toolCall", id: `call-${i}`, name: "read", arguments: { path: `src/f${i}.ts` } }, + ], + api: "anthropic-messages", + provider: "anthropic", + model: "bench", + usage: settledUsage(200 + (i % 50)), + stopReason: "toolUse", + timestamp: ts, + }; + messages.push(assistant as AgentMessage); + } else { + const toolResult: ToolResultMessage = { + role: "toolResult", + toolCallId: `call-${i - 1}`, + toolName: "read", + content: [{ type: "text", text: `Tool result ${i}.\n${codeBlob(i)}\n${codeBlob(i + 1)}` }], + isError: false, + timestamp: ts, + }; + messages.push(toolResult as AgentMessage); + } + } + return messages; +} + +interface Stats { + median: number; + madNoise: number; +} + +/** Median and robust MAD-based noise (median absolute deviation, normalized). */ +function stats(samples: number[]): Stats { + const sorted = [...samples].sort((a, b) => a - b); + const median = sorted[sorted.length >> 1]; + const deviations = sorted.map(x => Math.abs(x - median)).sort((a, b) => a - b); + const mad = deviations[deviations.length >> 1]; + // 1.4826 scales MAD to a stddev-equivalent for a normal distribution. + const madNoise = median === 0 ? 0 : (1.4826 * mad) / median; + return { median, madNoise }; +} + +/** + * Time `run(workload)` across samples. `makeWorkload` builds inputs OUTSIDE the + * timing window, so a cold phase can hand each sample fresh (uncached) identities + * without allocation noise polluting the measurement; a warm phase hands back one + * shared, already-primed workload. + * + * `batch` runs that many independent workloads inside one timed window and + * reports per-op time. A sub-millisecond cold op sits near the timer/scheduler + * floor where jitter dominates MAD-noise; batching lifts the measured window well + * above that floor while keeping each op a genuine cache miss. + */ +function sample(makeWorkload: () => T, run: (workload: T) => void, batch = 1): Stats { + for (let i = 0; i < WARMUP; i++) run(makeWorkload()); + const samples: number[] = []; + for (let i = 0; i < SAMPLES; i++) { + const workloads: T[] = []; + for (let b = 0; b < batch; b++) workloads.push(makeWorkload()); + // Collect workload-allocation garbage BEFORE timing so a GC pause can't land + // inside the window and inflate MAD-noise. + Bun.gc(true); + const t0 = Bun.nanoseconds(); + for (let b = 0; b < batch; b++) run(workloads[b]); + samples.push((Bun.nanoseconds() - t0) / 1e6 / batch); + } + return stats(samples); +} + +function estimateAll(messages: AgentMessage[]): number { + let total = 0; + for (const m of messages) total += estimateTokens(m); + return total; +} + +console.log(`\nBenchmark: llm-assembly (N=${N}, warmup=${WARMUP}, samples=${SAMPLES})\n`); + +// ─── convertToLlm ───────────────────────────────────────────────────────────── +// Cold: a fresh-identity history per sample → every message is a cache miss. +const convertFirst = sample( + () => buildHistory(N), + history => { + convertToLlm(history); + }, + 16, +); +// Steady: re-convert the same (warmed) array. transformContext re-converts the +// same live array multiple times per turn (prompt assembly, prune/shake +// accounting, context breakdown); the exact-repeat shortcut hands back the same +// outer array. Priming twice warms both the per-message memo and the shortcut. +const warmConvert = buildHistory(N); +convertToLlm(warmConvert); +convertToLlm(warmConvert); +const convertSteady = sample( + () => warmConvert, + history => { + convertToLlm(history); + }, +); + +// Append-growth: push one settled turn onto the same array identity each sample, +// then reconvert. Slice-on-growth reuses the unchanged prefix output and +// reconverts only the boundary message plus the new suffix, so the per-turn cost +// is O(suffix), not O(history). +const growConvert = buildHistory(N); +convertToLlm(growConvert); +let growSeed = N; +const convertGrow = sample( + () => { + growConvert.push({ + role: "user", + content: `User turn ${growSeed}: one more.\n\n${codeBlob(growSeed)}`, + timestamp: 1_700_000_000_000 + growSeed * 1000, + } as AgentMessage); + growSeed++; + return growConvert; + }, + history => { + convertToLlm(history); + }, +); + +// ─── estimateTokens ─────────────────────────────────────────────────────────── +// Cold: fresh-identity history per sample → every estimate is a cache miss. +const estimateFirst = sample( + () => buildHistory(N), + history => { + estimateAll(history); + }, +); +// Warm: one history, primed once, re-counted every sample from the cache. +const warmEstimate = buildHistory(N); +estimateAll(warmEstimate); +const estimateSecond = sample( + () => warmEstimate, + history => { + estimateAll(history); + }, +); + +function report(label: string, s: Stats): void { + console.log( + ` ${label.padEnd(18)} median ${s.median.toFixed(4).padStart(10)} ms MAD-noise ${(s.madNoise * 100).toFixed(1).padStart(5)}%`, + ); +} + +report("convert first", convertFirst); +report("convert steady", convertSteady); +report("convert grow", convertGrow); +report("estimate first", estimateFirst); +report("estimate second", estimateSecond); + +const convertSteadySpeedup = convertFirst.median / convertSteady.median; +const convertGrowSpeedup = convertFirst.median / convertGrow.median; +const estimateSpeedup = estimateFirst.median / estimateSecond.median; +console.log(`\n convert speedup (first / steady): ${convertSteadySpeedup.toFixed(2)}x`); +console.log(` convert speedup (first / grow): ${convertGrowSpeedup.toFixed(2)}x`); +console.log(` estimate speedup (first / second): ${estimateSpeedup.toFixed(2)}x`); + +const noiseGate = 0.2; +const worstNoise = Math.max( + convertFirst.madNoise, + convertSteady.madNoise, + convertGrow.madNoise, + estimateFirst.madNoise, + estimateSecond.madNoise, +); +console.log(` worst MAD-noise: ${(worstNoise * 100).toFixed(1)}% (gate ${(noiseGate * 100).toFixed(0)}%)\n`); + +console.log(`METRIC convert_steady_speedup=${convertSteadySpeedup.toFixed(3)}`); +console.log(`METRIC convert_grow_speedup=${convertGrowSpeedup.toFixed(3)}`); +console.log(`METRIC estimate_speedup=${estimateSpeedup.toFixed(3)}`); +console.log(`METRIC worst_mad_noise=${worstNoise.toFixed(4)}`); + +process.exit(0); diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index aeb4a2e49..0529b7d25 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -63,6 +63,7 @@ import { estimateTokens, generateBranchSummary, generateHandoffFromContext, + invalidateMessageCache, prepareCompaction, renderHandoffPrompt, resolveBudgetReserveTokens, @@ -2524,7 +2525,13 @@ export class AgentSession { const isPlanNudge = (m: AgentMessage): boolean => m.role === "custom" && m.customType === PREWALK_PLAN_MESSAGE_TYPE; for (let i = liveMessages.length - 1; i >= 0; i--) { - if (isPlanNudge(liveMessages[i])) liveMessages.splice(i, 1); + if (isPlanNudge(liveMessages[i])) { + // Interior removal on the live array: drop the scrubbed message from + // the convert/estimate caches so the next convert can't reuse a prefix + // that still carries its fragment (the array shrinks in place). + invalidateMessageCache(liveMessages[i]); + liveMessages.splice(i, 1); + } } const stateMessages = this.agent.state.messages; const filtered = stateMessages.filter(m => !isPlanNudge(m)); diff --git a/packages/coding-agent/src/session/messages.test.ts b/packages/coding-agent/src/session/messages.test.ts index cfc01d73b..5d6680047 100644 --- a/packages/coding-agent/src/session/messages.test.ts +++ b/packages/coding-agent/src/session/messages.test.ts @@ -8,6 +8,7 @@ import { replaceLlmImagesWithText, SKILL_PROMPT_MESSAGE_TYPE, type SkillPromptDetails, + stripImagesFromMessage, } from "./messages"; function customMessage(customType: string, attribution: "agent" | "user"): CustomMessage { @@ -125,6 +126,96 @@ describe("convertToLlm", () => { }); }); +function settledAssistant(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage: { + input: 100, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 120, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 1, + }; +} + +function userMessage(text: string, timestamp: number): AgentMessage { + return { role: "user", content: text, attribution: "user", timestamp } as AgentMessage; +} + +describe("convertToLlm caching", () => { + it("reuses the outer array on an exact repeat of the same history", () => { + const messages: AgentMessage[] = [userMessage("hello", 1), settledAssistant("hi")]; + const first = convertToLlm(messages); + const second = convertToLlm(messages); + expect(second).toBe(first); + }); + + it("reuses the unchanged prefix output on append-only growth", () => { + const messages: AgentMessage[] = [userMessage("one", 1), settledAssistant("reply one")]; + const first = convertToLlm(messages); + messages.push(userMessage("two", 2)); + const grown = convertToLlm(messages); + // New outer array (no held-result aliasing), but the converted prefix is + // byte-identical and the appended turn is present. + expect(grown).not.toBe(first); + expect(grown.length).toBe(first.length + 1); + expect(grown.slice(0, first.length)).toEqual(first); + expect(grown[grown.length - 1]?.role).toBe("user"); + }); + + it("recomputes the boundary assistant when a following interrupted-thinking marker appears on growth", () => { + const messages: AgentMessage[] = [ + abortedAssistant([ + { type: "text", text: "partial answer" }, + { type: "thinking", thinking: "interrupted reasoning" }, + ]), + ]; + const before = convertToLlm(messages); + const beforeAssistant = before.find(entry => entry.role === "assistant"); + expect(Array.isArray(beforeAssistant?.content) && beforeAssistant.content.map(b => b.type)).toEqual([ + "text", + "thinking", + ]); + + // Append the continuity marker on the same array: the assistant is now the + // boundary message and its LLM view must drop the trailing thinking run. + messages.push(interruptedThinkingContinuity()); + const after = convertToLlm(messages); + const afterAssistant = after.find(entry => entry.role === "assistant"); + expect(Array.isArray(afterAssistant?.content) && afterAssistant.content.map(b => b.type)).toEqual(["text"]); + }); + + it("recomputes a message after strip-images invalidates its cache", () => { + const withImage: AgentMessage = { + role: "user", + content: [ + { type: "text", text: "look" }, + { type: "image", data: "aaaa", mimeType: "image/png" }, + ], + attribution: "user", + timestamp: 1, + }; + const messages: AgentMessage[] = [withImage]; + const before = convertToLlm(messages); + const beforeUser = before.find(entry => entry.role === "user"); + expect(Array.isArray(beforeUser?.content) && beforeUser.content.some(b => b.type === "image")).toBe(true); + + // Mutate in place through the owner seam, which must invalidate the cache. + stripImagesFromMessage(withImage); + const after = convertToLlm(messages); + const afterUser = after.find(entry => entry.role === "user"); + expect(Array.isArray(afterUser?.content) && afterUser.content.some(b => b.type === "image")).toBe(false); + }); +}); + describe("replaceLlmImagesWithText", () => { it("replaces image blocks in user, developer, and tool-result messages with the placeholder", () => { const converted = convertToLlm([ diff --git a/packages/coding-agent/src/session/messages.ts b/packages/coding-agent/src/session/messages.ts index 6fd7b24ed..97a607ddf 100644 --- a/packages/coding-agent/src/session/messages.ts +++ b/packages/coding-agent/src/session/messages.ts @@ -5,6 +5,10 @@ * and provides a transformer to convert them to LLM-compatible messages. */ import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { + invalidateMessageCache, + registerMessageCacheInvalidator, +} from "@oh-my-pi/pi-agent-core/compaction/message-cache"; import { type BranchSummaryMessage, type CompactionSummaryMessage, @@ -457,6 +461,14 @@ function stripImagesFromArrayContent(content: (TextContent | ImageContent)[]): S * pure local mutation and intentionally does neither. */ export function stripImagesFromMessage(message: AgentMessage): number { + const removed = stripImagesFromMessageContent(message); + // The mutated message keeps its identity across context rebuilds, so drop its + // cached estimate/convert before the next pass counts/converts the new shape. + if (removed > 0) invalidateMessageCache(message); + return removed; +} + +function stripImagesFromMessageContent(message: AgentMessage): number { switch (message.role) { case "user": case "developer": @@ -750,6 +762,184 @@ function convertImageBearingCustomMessage(message: CustomMessage | HookMessage): return converted; } +/** + * Per-message conversion result, keyed by message identity. `interruptedNext` + * records the neighbor state the fragment was built against so an assistant + * whose following {@link INTERRUPTED_THINKING_MESSAGE_TYPE} marker appears or + * disappears is recomputed (its LLM view strips the trailing thinking run only + * while that marker follows). + * + * WeakMap (not a symbol tag) is deliberate: `wrapSteeringForModel` and + * `deobfuscateAgentMessages` spread messages into fresh variants with different + * content; a symbol-keyed fragment would ride that spread and mis-convert the + * copy. Identity keying keeps the cache off spread copies. + */ +interface ConvertMemoEntry { + interruptedNext: boolean; + fragment: Message[]; +} +const convertCache = new WeakMap(); + +// Array-level shortcuts over the per-message memo. The live agent mutates one +// `AgentMessage[]` identity across a turn: appending new messages and swapping +// the streaming tail (`context.messages[len-1] = partial → trailing`). Between +// owner invalidations (prune/shake/strip bump `convertGeneration`) and for a +// given array identity, only the last index is ever swapped and the array only +// grows — interior prefix messages are immutable. That invariant lets two +// shortcuts skip the O(N) re-walk: +// - exact-repeat: same array, same length, same generation, same tail identity +// → hand back the same outer array. +// - slice-on-growth: same array, same generation, length grew → copy the +// unchanged prefix output and reconvert only the neighbor-sensitive boundary +// message plus the appended suffix. +// The tail-identity guard on exact-repeat catches the streaming snapshot swap +// (partial → trailing is a fresh identity), so a settled tail is never served +// from a stale mid-stream fragment. +let convertGeneration = 0; +let lastConvertInput: AgentMessage[] | undefined; +let lastConvertLength = 0; +let lastConvertOutput: Message[] | undefined; +let lastConvertGeneration = -1; +let lastConvertTail: AgentMessage | undefined; +// Output-message count contributed by messages[0 .. lastConvertLength-1), i.e. +// every message except the last. The last message is neighbor-sensitive (its LLM +// view drops the trailing thinking run only while an interrupted-thinking marker +// follows), so growth reconverts it rather than reusing its old fragment. +let lastConvertPrefixOutputLen = 0; + +registerMessageCacheInvalidator(message => { + convertCache.delete(message); + convertGeneration++; +}); + +/** Convert one message to its LLM fragment. `interruptedNext` is true only for an + * assistant turn immediately followed by its interrupted-thinking marker. */ +function convertOne(m: AgentMessage, interruptedNext: boolean): Message[] { + switch (m.role) { + case "bashExecution": + if (m.excludeFromContext) { + return []; + } + return [ + { + role: "user", + content: [{ type: "text", text: bashExecutionToText(m) }], + attribution: "user", + timestamp: m.timestamp, + }, + ]; + case "pythonExecution": + if (m.excludeFromContext) { + return []; + } + return [ + { + role: "user", + content: [{ type: "text", text: pythonExecutionToText(m) }], + attribution: "user", + timestamp: m.timestamp, + }, + ]; + case "fileMention": { + // One `fileMention` can mix `@notes.md` (text) and `@screenshot.png` (image) + // in the same turn (`generateFileMentionMessages` packs every `@…` into a + // single message). Splitting by image presence keeps text-only mentions on + // the higher-priority `developer` slot while routing image attachments + // through `user`, the only Responses content slot that legitimately accepts + // `input_image` (Codex chatgpt.com /codex/responses rejects everything else + // with `Invalid value: 'input_image'`, #3443). + const wrap = (file: FileMentionMessage["files"][number]): string => { + const inner = file.content ? `\n${file.content}\n` : "\n"; + return `${inner}`; + }; + const textFiles = m.files.filter(file => !file.image); + const imageFiles = m.files.filter(file => file.image); + const out: Message[] = []; + if (textFiles.length > 0) { + out.push({ + role: "developer", + content: [{ type: "text" as const, text: textFiles.map(wrap).join("\n") }], + attribution: "user", + timestamp: m.timestamp, + }); + } + if (imageFiles.length > 0) { + const content: (TextContent | ImageContent)[] = [ + { type: "text" as const, text: imageFiles.map(wrap).join("\n") }, + ]; + for (const file of imageFiles) { + if (file.image) content.push(file.image); + } + out.push({ + role: "user", + content, + attribution: "user", + timestamp: m.timestamp, + }); + } + return out; + } + case "custom": { + if (!isCustomMessageContent(m.content)) return []; + if (isUserInvokedSkillPrompt(m)) { + return [ + { + role: "user", + content: customMessageContentToLlmContent(m.content), + attribution: "user", + timestamp: m.timestamp, + }, + ]; + } + const split = convertImageBearingCustomMessage(m); + if (split) return split; + const converted = convertMessageToLlm(m); + return converted ? [converted] : []; + } + case "hookMessage": { + if (!isCustomMessageContent(m.content)) return []; + const split = convertImageBearingCustomMessage(m); + if (split) return split; + const converted = convertMessageToLlm(m); + return converted ? [converted] : []; + } + case "assistant": { + // A user-interrupted turn keeps its trailing thinking run on the + // persisted/displayed message so reload and Ctrl+L rebuilds still + // show it. That run is incomplete/unsigned and gets rejected on + // resend, so strip it here — LLM path only — when the hidden + // interrupted-thinking continuity message follows. + const source = interruptedNext ? stripDemotedThinkingForLlm(m) : m; + const converted = convertMessageToLlm(source); + return converted ? [converted] : []; + } + case "branchSummary": + case "compactionSummary": + case "user": + case "developer": + case "toolResult": { + // Core roles share one transformer with agent-core — + // duplicating them here is how snapcompact frames once + // silently fell off the provider request. + const converted = convertMessageToLlm(m); + return converted ? [converted] : []; + } + default: + m satisfies never; + return []; + } +} + +/** Cached per-message conversion. Reuses the stored fragment while identity and + * `interruptedNext` neighbor state hold; recomputes on a neighbor flip. */ +function convertOneCached(m: AgentMessage, interruptedNext: boolean): Message[] { + const cached = convertCache.get(m); + if (cached !== undefined && cached.interruptedNext === interruptedNext) return cached.fragment; + const fragment = convertOne(m, interruptedNext); + convertCache.set(m, { interruptedNext, fragment }); + return fragment; +} + /** * Transform AgentMessages (including custom types) to LLM-compatible Messages. * @@ -757,121 +947,69 @@ function convertImageBearingCustomMessage(message: CustomMessage | HookMessage): * - Agent's transormToLlm option (for prompt calls and queued messages) * - Compaction's generateSummary (for summarization) * - Custom extensions and tools + * + * Settled history converts once and is reused per message identity: an + * append-only turn on the same array re-pays only the new suffix, and an + * unchanged re-convert of the same array hands back the same outer `Message[]`. + * Owner mutations (prune/shake/strip-images) invalidate the affected message + * through the shared registry before the next pass. */ export function convertToLlm(messages: AgentMessage[]): Message[] { - return messages.flatMap((m, index): Message[] => { - switch (m.role) { - case "bashExecution": - if (m.excludeFromContext) { - return []; - } - return [ - { - role: "user", - content: [{ type: "text", text: bashExecutionToText(m) }], - attribution: "user", - timestamp: m.timestamp, - }, - ]; - case "pythonExecution": - if (m.excludeFromContext) { - return []; - } - return [ - { - role: "user", - content: [{ type: "text", text: pythonExecutionToText(m) }], - attribution: "user", - timestamp: m.timestamp, - }, - ]; - case "fileMention": { - // One `fileMention` can mix `@notes.md` (text) and `@screenshot.png` (image) - // in the same turn (`generateFileMentionMessages` packs every `@…` into a - // single message). Splitting by image presence keeps text-only mentions on - // the higher-priority `developer` slot while routing image attachments - // through `user`, the only Responses content slot that legitimately accepts - // `input_image` (Codex chatgpt.com /codex/responses rejects everything else - // with `Invalid value: 'input_image'`, #3443). - const wrap = (file: FileMentionMessage["files"][number]): string => { - const inner = file.content ? `\n${file.content}\n` : "\n"; - return `${inner}`; - }; - const textFiles = m.files.filter(file => !file.image); - const imageFiles = m.files.filter(file => file.image); - const out: Message[] = []; - if (textFiles.length > 0) { - out.push({ - role: "developer", - content: [{ type: "text" as const, text: textFiles.map(wrap).join("\n") }], - attribution: "user", - timestamp: m.timestamp, - }); - } - if (imageFiles.length > 0) { - const content: (TextContent | ImageContent)[] = [ - { type: "text" as const, text: imageFiles.map(wrap).join("\n") }, - ]; - for (const file of imageFiles) { - if (file.image) content.push(file.image); - } - out.push({ - role: "user", - content, - attribution: "user", - timestamp: m.timestamp, - }); - } - return out; - } - case "custom": { - if (!isCustomMessageContent(m.content)) return []; - if (isUserInvokedSkillPrompt(m)) { - return [ - { - role: "user", - content: customMessageContentToLlmContent(m.content), - attribution: "user", - timestamp: m.timestamp, - }, - ]; - } - const split = convertImageBearingCustomMessage(m); - if (split) return split; - const converted = convertMessageToLlm(m); - return converted ? [converted] : []; - } - case "hookMessage": { - if (!isCustomMessageContent(m.content)) return []; - const split = convertImageBearingCustomMessage(m); - if (split) return split; - const converted = convertMessageToLlm(m); - return converted ? [converted] : []; - } - case "assistant": { - // A user-interrupted turn keeps its trailing thinking run on the - // persisted/displayed message so reload and Ctrl+L rebuilds still - // show it. That run is incomplete/unsigned and gets rejected on - // resend, so strip it here — LLM path only — when the hidden - // interrupted-thinking continuity message follows. - const source = followedByInterruptedThinking(messages, index) ? stripDemotedThinkingForLlm(m) : m; - const converted = convertMessageToLlm(source); - return converted ? [converted] : []; - } - case "branchSummary": - case "compactionSummary": - case "user": - case "developer": - case "toolResult": { - // Core roles share one transformer with agent-core — - // duplicating them here is how snapcompact frames once - // silently fell off the provider request. - const converted = convertMessageToLlm(m); - return converted ? [converted] : []; - } - default: - m satisfies never; - return []; - } - }); + const len = messages.length; + const sameArray = messages === lastConvertInput && lastConvertGeneration === convertGeneration; + const tail = len > 0 ? messages[len - 1] : undefined; + + // Exact-repeat: same array, same length, same trailing identity → reuse the + // outer array. The tail-identity check rejects the streaming snapshot swap + // (partial → settled trailing keeps array identity/length but mints a fresh + // tail), so a settled tail never reads a stale mid-stream fragment. + if (sameArray && lastConvertOutput !== undefined && len === lastConvertLength && tail === lastConvertTail) { + return lastConvertOutput; + } + + // Slice-on-growth: same array grew by append. Every interior message is + // immutable under one array identity, so copy the unchanged prefix output + // (messages[0 .. lastLen-1)) and reconvert only the old boundary message + // (neighbor-sensitive: a following interrupted-thinking marker may now exist) + // plus the appended suffix. The boundary-identity check (old tail still sits + // at its old index) rejects an in-place interior splice-replace that grew the + // array while swapping earlier identities, forcing a full rebuild. + let out: Message[]; + let start: number; + if ( + sameArray && + lastConvertOutput !== undefined && + len > lastConvertLength && + lastConvertLength > 0 && + messages[lastConvertLength - 1] === lastConvertTail && + lastConvertPrefixOutputLen <= lastConvertOutput.length + ) { + out = lastConvertOutput.slice(0, lastConvertPrefixOutputLen); + start = lastConvertLength - 1; + } else { + out = []; + start = 0; + } + + // Output length contributed by messages[0 .. len-1), captured when the loop + // reaches the final index so the next growth can reuse this prefix. + let prefixOutputLen = 0; + for (let i = start; i < len; i++) { + if (i === len - 1) prefixOutputLen = out.length; + const m = messages[i]; + const interruptedNext = m.role === "assistant" && followedByInterruptedThinking(messages, i); + const fragment = convertOneCached(m, interruptedNext); + for (const msg of fragment) out.push(msg); + } + if (len === 0) prefixOutputLen = 0; + + // Record for the next call's shortcuts. `out` is a fresh array (slice or new), + // so a prior caller holding the previous `lastConvertOutput` never sees it grow. + lastConvertInput = messages; + lastConvertLength = len; + lastConvertOutput = out; + lastConvertGeneration = convertGeneration; + lastConvertTail = tail; + lastConvertPrefixOutputLen = prefixOutputLen; + return out; }