Repository navigation
Conversation
convertCbToModelMessages ran on every agent step over the full message history, and almost all of its cost was per-message work that is a pure function of the message: cloneDeep of every content part, the lone-surrogate scrub across every string, and the zod safeParse per converted part (sdk/run.ts measured the clone alone at ~230ms for an 8MB history, ~575ms at 20MB — paid every step, quadratic over a turn). Memoize the per-message product on a WeakMap keyed by message identity: each history message is cloned, scrubbed, converted and schema-checked exactly once per process; a step's remaining cost is the envelope logic (media buffering, same-aux aggregation, cache-control placement) plus O(1) per message. Keyed on identity this is sound because history messages are immutable once appended (fit-tool-results, trim, compaction all allocate replacements — verified at every rewrite site), so a cache entry can only go stale if the message object itself is replaced, which is precisely the invalidation case. Safety within a request: cached converted forms are shared across requests, so both mutation passes were made non-destructive. withCacheControl no longer cloneDeeps — it shallow-copies down to the providerOptions leaf it marks — and the cache-control pass replaces entries instead of writing through aggregated arrays, so repeated conversions of the same history are byte-identical and no mark leaks across requests. Aggregation copies once per request (flagged __aggregatedCopy, non-enumerable) instead of merging into the cached form. The lone-surrogate fast path pre-scans without copying; messages that need a scrub are cloned once and the scrubbed form is cached, so the caller's history is never rewritten. The orphan-tool-call repair moves out of the converter: dropUnansweredToolCalls ran twice per repaired step history (once at run-agent-step, again here — provably a no-op the second time). Rebuild-at-request-time paths that cannot guarantee a repaired history (sdk/impl/llm.ts auxiliary calls) now repair at their own boundary; dropUnansweredToolCalls returns the input array unchanged when there is nothing to repair, so the common case costs one boundary scan. expireMessages likewise returns the input array by reference when nothing expired instead of allocating a full-history copy every step. Equivalence was verified by a randomized differential harness against the pre-change implementation (tool calls, parallel batches, media results, tags, providerOptions, timeToLive, lone/valid unicode, histories to 60 messages x 20 trials x both cache-control modes): every fixture produced byte-identical request messages.
The context estimator (estimateContextTokens and the step loop's input/output estimates beside it) called countTokensMessages over the full history 3-5x per step, plus per-message recounts from the truncation and compaction guards. The walk itself is string-length arithmetic, so a whole-history recount was comparatively cheap but still re-scanned every byte of every tool result every step. Split countTokensMessages into a per-message estimate memoized on a WeakMap keyed by message identity; the list count remains the exact sum of per-message counts, so numbers are unchanged (same integer terms, associative regrouping only). Histories are immutable after append, so identity invalidation (message replacement) is the only staleness path. The one post-count rewrite site: fitToolResults clips its defensive copies in place while binary-searching the budget, and re-probing each intermediate form with the memoized estimator could pin a probe's count onto a message the caller still owns (context-size-guard's original would have reported 71 tokens for a 2.5MB result). fitToolResults now probes with the uncached probeTokensMessages; the already-fits fast path still memoizes the accurate final count. Tests: cached-vs-fresh recount equivalence at several history sizes, per-message sum == list count, replacement invalidation, and the expireMessages no-op identity from the previous commit.
The identity fast path returned messages unchanged when no TTL/expiresAt was set, but INSTRUCTIONS_PROMPT messages carry no TTL — they rely purely on keepLastTags, which is pruned downstream of the filter. Two such messages in history therefore piled up instead of superseding. Track keepLastTags alongside the expiring count; only return by reference when neither is present. Also align the identity test: legacy timeToLives always take the copy path under the new expiresAt machinery (the snitch marker forces one), so assert content equality, not reference identity.
582ec0f to
bf3ac82
Compare
|
The underlying diagnosis is correct and valuable: re-running That said, a few things make this hard to port as-is:
Splitting into (a) the token-counter memoization (self-contained, lower risk) and (b) the conversion cache + repair relocation (needs its own scrutiny and ideally lands after #1587) would make this portable. |
codebuff-team
added
bot:triaged
What
convertCbToModelMessagesran on every agent step over the full history, and almost all of its cost was per-message work that is a pure function of the message:cloneDeepof every content part, the lone-surrogate scrub across every string, and a zodsafeParseper converted part.sdk/run.tsmeasured the clone alone at ~230 ms for an 8 MB history (~575 ms at 20 MB) — paid every step, quadratic over a turn. The per-step token estimate had the same shape: the full history re-walked 3–5× per step.Commit 1 — convert each message exactly once
Memoize the per-message product on a
WeakMapkeyed by message identity: each history message is cloned, scrubbed, converted and schema-checked exactly once per process; a step's remaining cost is the envelope logic plus O(1) per message. Identity keying is sound because history messages are immutable once appended — verified at every rewrite site (fit-tool-resultsclones before clipping,trimMessagescloneDeeps, compaction allocates,dropUnansweredToolCallsspreads).Because cached forms are shared across requests, both mutation passes were made non-destructive:
withCacheControlshallow-copies down to the marked leaf instead ofcloneDeep-ing whole messages, and cache-control placement replaces entries rather than writing through aggregated arrays. The orphan-tool-call repair moved out of the converter (provably a no-op there on the hot path —run-agent-steprepairs and assigns the history before the request is built); auxiliary SDK paths insdk/impl/llm.tsrepair at their own boundary viarepairToolCallHistory.expireMessagesreturns the input array by reference when nothing expired and nothing is pruneable by keepLastTags — this PR is rebased on #1587, so the fast path now tracks keepLastTags too (previously two INSTRUCTIONS_PROMPT messages could be wrongly identity-stable).Equivalence: verified by a randomized differential harness against the pre-change implementation — tool calls, parallel batches with interleaved media, tags/providerOptions/timeToLive, lone/valid unicode, histories to 60 messages × 20 trials × both cache-control modes — byte-identical request JSON for every fixture. (The harness caught and I fixed a real bug in the first draft: deep strings inside tool-call inputs escaping the scrub fast-path.)
Commit 2 — memoize the per-message token estimate
countTokensMessagessplit into per-messagecountTokensMessagememoized on a WeakMap; list totals remain the exact same sum. One semantics clash surfaced and is handled:fitToolResultsclips defensive copies in place during binary search, so the budget loop probes with uncached counting; the already-fits fast path still memoizes the accurate final count.Merge-time change (post-#1587 rebase)
Two small follow-ups this PR adds on top of the main resolution:
expireMessagesfast path trackskeepLastTagstoo, so twoINSTRUCTIONS_PROMPTmessages in history are still superseded rather than deemed identity-stable.expireMessagesidentity test assertion updated to content equality: legacytimeToLivesalways take the copy path under the newexpiresAtmachinery (the snitch marker forces one), so reference-identity assertions no longer hold for those cases; content equality is what the contract actually guarantees.Reviewer caveats (stated, not hidden)
withCacheControl's shallow copy assumes providerMetadata is plain JSON (true today).Tests
message-conversion-cache.test.ts(warmed-cache byte-identity, no cache-control leakage, replacement invalidation, nested-surrogate scrub),expire-messages.test.ts(updated per above),token-counter-cache.test.ts(sum + recount equivalence at sizes 1–100, replacement). common 51/51 messages tests; agent-runtime util 227/227; full agent-runtime suite 995 pass / 1 pre-existing env failure; sdk + agent-runtime suites match base failure sets exactly; typechecks clean (one pre-existing common error unrelated to this change).