From 6ce14101e79101f4f9a10fd8d4486aba266d816f Mon Sep 17 00:00:00 2001 From: Tiago Date: Mon, 19 Oct 2026 18:44:57 -0300 Subject: [PATCH] perf: reuse one tokenizer and bound the cost of counting huge inputs countTokens() from @anthropic-ai/tokenizer builds and frees a WASM tokenizer on every call (~70 ms each) and its encode() runs synchronously on the host's event loop with super-linear cost on huge or degenerate inputs. On long sessions this blocked OpenCode's server for seconds per transform, up to the point of being killed by the client health check mid-compression (upstream #638, #659; see also #639 and #660). - build the tokenizer once and share it (~70 ms -> ~0.2 ms per call) - memoize exact counts (LRU, bounded by entries and total chars) because the same summaries, system prompt and tool outputs are re-counted on every transform and after tool-cache evictions - above DCP_TOKEN_COUNT_MAX_EXACT_CHARS (default 16k characters), extrapolate from head/middle/tail samples instead of encoding everything; sequences with a long single-character run are sampled as well, since their encode cost grows super-linearly - count each item in estimateTokensBatch instead of joining all texts into one string - DCP_TOKEN_DEBUG=1 logs each count for measurement; resetTokenCountingForTests() clears the shared state for tests scripts/bench-token-count.ts measures wall time and missed timer ticks against the package's own countTokens; sampled estimates stay within ~0.6% of exact counts on the bench corpus. --- lib/token-utils.ts | 123 +++++++++++++++++++++++++++- scripts/bench-token-accuracy.ts | 56 +++++++++++++ scripts/bench-token-count.ts | 81 ++++++++++++++++++ scripts/bench-token-pathological.ts | 19 +++++ tests/token-counting.test.ts | 95 +++++++++++++++++++++ 5 files changed, 370 insertions(+), 4 deletions(-) create mode 100644 scripts/bench-token-accuracy.ts create mode 100644 scripts/bench-token-count.ts create mode 100644 scripts/bench-token-pathological.ts diff --git a/lib/token-utils.ts b/lib/token-utils.ts index bba07fa1..36d47030 100644 --- a/lib/token-utils.ts +++ b/lib/token-utils.ts @@ -2,10 +2,95 @@ import { SessionState, WithParts } from "./state" import { AssistantMessage, UserMessage } from "@opencode-ai/sdk/v2" import { Logger } from "./logger" import * as _anthropicTokenizer from "@anthropic-ai/tokenizer" -const anthropicCountTokens = (_anthropicTokenizer.countTokens ?? - (_anthropicTokenizer as any).default?.countTokens) as typeof _anthropicTokenizer.countTokens import { getLastUserMessage } from "./messages/query" +const anthropicGetTokenizer = (_anthropicTokenizer.getTokenizer ?? + (_anthropicTokenizer as any).default?.getTokenizer) as typeof _anthropicTokenizer.getTokenizer + +// The package's countTokens() builds and frees a WASM tokenizer on every call (~70 ms each), +// while an encode() on a reused tokenizer costs a fraction of a millisecond (upstream #638/#639). +// Build it once and share it. Never free() the shared instance. +let anthropicTokenizer: ReturnType | undefined + +function tokenizeExact(text: string): number { + anthropicTokenizer ??= anthropicGetTokenizer() + return anthropicTokenizer.encode(text.normalize("NFKC"), "all").length +} + +// encode() runs synchronously on the host's event loop and its cost is super-linear on long +// runs without token boundaries: a single huge input could block OpenCode for seconds +// (upstream #659). Count exactly up to this many characters and extrapolate from samples above +// it, so one call can never block the host. Override with DCP_TOKEN_COUNT_MAX_EXACT_CHARS +// (0 samples every text; a huge value forces exact counting). +const DEFAULT_MAX_EXACT_CHARS = 16 * 1024 +const SAMPLE_CHARS = 2048 + +function maxExactChars(): number { + const raw = process.env.DCP_TOKEN_COUNT_MAX_EXACT_CHARS + if (raw === undefined) return DEFAULT_MAX_EXACT_CHARS + const parsed = Number.parseInt(raw, 10) + if (!Number.isFinite(parsed) || parsed < 0) return DEFAULT_MAX_EXACT_CHARS + return Math.min(parsed, Number.MAX_SAFE_INTEGER) +} + +// A long run of one repeated character has no token boundaries for the encoder, and its cost +// grows super-linearly: 10k characters can already block the host for hundreds of milliseconds +// (upstream #659). Treat such inputs as huge and sample them too. Up to DEGENERATE_RUN_SAFE_CHARS +// a repeated run is cheap enough to count exactly (measured ~20 ms at 4k characters), and the +// run scan is bounded by the exact-count limit with cheap charCodeAt comparisons. +const MAX_CHAR_RUN = 512 +const DEGENERATE_RUN_SAFE_CHARS = 4 * 1024 + +function hasLongCharRun(text: string): boolean { + let run = 1 + for (let i = 1; i < text.length; i++) { + if (text.charCodeAt(i) === text.charCodeAt(i - 1)) { + run++ + if (run >= MAX_CHAR_RUN) return true + } else { + run = 1 + } + } + return false +} + +// The same summaries, system prompt and tool outputs are re-counted on every chat-messages +// transform (and again after tool-cache evictions), so memoize exact counts. Keys are capped by +// size and total bytes to keep the cache bounded; sampled counts are cheap and are not cached. +const CACHE_MAX_ENTRIES = 1024 +const CACHE_MAX_CHARS = 4 * 1024 * 1024 +const exactCountCache = new Map() +let exactCountCacheChars = 0 + +function trimExactCountCache(): void { + while (exactCountCache.size > CACHE_MAX_ENTRIES || exactCountCacheChars > CACHE_MAX_CHARS) { + const oldest = exactCountCache.keys().next().value + if (oldest === undefined) return + exactCountCacheChars -= oldest.length + exactCountCache.delete(oldest) + } +} + +/** Clears the shared tokenizer and the count cache. Test helper only. */ +export function resetTokenCountingForTests(): void { + anthropicTokenizer = undefined + exactCountCache.clear() + exactCountCacheChars = 0 +} + +function estimateTokensFromSamples(text: string): number { + // Disjoint windows: head, middle and tail, each at most SAMPLE_CHARS and at least one third + // of the input, so the sampled density is representative and the sampled area is known. + const window = Math.min(SAMPLE_CHARS, Math.floor(text.length / 3)) + if (window < 16) return tokenizeExact(text) + const middleStart = (text.length - window) >> 1 + const head = tokenizeExact(text.slice(0, window)) + const middle = tokenizeExact(text.slice(middleStart, middleStart + window)) + const tail = tokenizeExact(text.slice(-window)) + const perChar = (head + middle + tail) / (3 * window) + return Math.max(1, Math.round(perChar * text.length)) +} + export function getCurrentTokenUsage(state: SessionState, messages: WithParts[]): number { for (let i = messages.length - 1; i >= 0; i--) { const msg = messages[i] @@ -68,16 +153,46 @@ export function getCurrentParams( export function countTokens(text: string): number { if (!text) return 0 + + const cached = exactCountCache.get(text) + if (cached !== undefined) { + exactCountCache.delete(text) + exactCountCache.set(text, cached) + return cached + } + + const limit = maxExactChars() + const exact = + text.length <= limit && (text.length <= DEGENERATE_RUN_SAFE_CHARS || !hasLongCharRun(text)) + + let tokens: number try { - return anthropicCountTokens(text) + tokens = exact ? tokenizeExact(text) : estimateTokensFromSamples(text) } catch { + anthropicTokenizer = undefined return Math.round(text.length / 4) } + + if (exact) { + exactCountCache.set(text, tokens) + exactCountCacheChars += text.length + trimExactCountCache() + } + + if (process.env.DCP_TOKEN_DEBUG === "1") { + process.stderr.write( + `[dcp-token] chars=${text.length} tokens=${tokens} path=${exact ? "exact" : "sampled"}\n`, + ) + } + + return tokens } export function estimateTokensBatch(texts: string[]): number { if (texts.length === 0) return 0 - return countTokens(texts.join(" ")) + let total = 0 + for (const text of texts) total += countTokens(text) + return total } export const COMPACTED_TOOL_OUTPUT_PLACEHOLDER = "[Old tool result content cleared]" diff --git a/scripts/bench-token-accuracy.ts b/scripts/bench-token-accuracy.ts new file mode 100644 index 00000000..164a29dd --- /dev/null +++ b/scripts/bench-token-accuracy.ts @@ -0,0 +1,56 @@ +// Accuracy of the sampled estimator against exact BPE counting, for texts above the +// exact-count limit. Run: npx tsx scripts/bench-token-accuracy.ts +import * as t from "@anthropic-ai/tokenizer" +import { countTokens } from "../lib/token-utils" + +const getTokenizer = ((t as any).getTokenizer ?? + (t as any).default?.getTokenizer) as typeof t.getTokenizer +const tokenizer = getTokenizer() +const exactCount = (text: string): number => tokenizer.encode(text.normalize("NFKC"), "all").length + +const prose = + "The billing cycle console groups invoices per account and shows the cycle window for each subscription. " +const json = + '{"path":"apps/bud/app/models/invoice.rb","line":42,"text":"scope :open, -> { where(paid_at: nil) },"} ' +const code = + "def cycle_bounds(account)\n [account.cycle_started_at, account.cycle_ends_at]\nend\n\n" + +function build(kind: string, targetChars: number): string { + const unit = + kind === "prose" + ? prose + : kind === "json" + ? json + : kind === "code" + ? code + : prose + json + code + const parts: string[] = [] + let chars = 0 + while (chars < targetChars) { + parts.push(unit) + chars += unit.length + } + return parts.join("") +} + +function mixedHalves(targetChars: number): string { + const half = build("prose", targetChars / 2) + const b64 = "QUJDREVGR0hJSktMTU5PUFFSU1RVVldYWVo=".repeat(Math.ceil(targetChars / 2 / 36)) + return half + b64 +} + +const sizes = [32 * 1024, 128 * 1024, 512 * 1024, 1024 * 1024, 8 * 1024 * 1024] +console.log("kind\tsizeKB\texact\tsampled\terror%") +for (const [kind, text] of [ + ...sizes.map((s) => ["prose", build("prose", s)] as const), + ...sizes.map((s) => ["json", build("json", s)] as const), + ...sizes.map((s) => ["code", build("code", s)] as const), + ...sizes.map((s) => ["mixed-halves", mixedHalves(s)] as const), +]) { + const exact = exactCount(text) + const sampled = countTokens(text) + const error = ((sampled - exact) / exact) * 100 + console.log( + `${kind}\t${(text.length / 1024).toFixed(0)}\t${exact}\t${sampled}\t${error.toFixed(2)}%`, + ) +} diff --git a/scripts/bench-token-count.ts b/scripts/bench-token-count.ts new file mode 100644 index 00000000..89284283 --- /dev/null +++ b/scripts/bench-token-count.ts @@ -0,0 +1,81 @@ +// Bench for countTokens: legacy (the package's own countTokens, i.e. build+encode+free per call) +// versus the patched implementation. Measures wall time and missed 10 ms timer ticks (blocking). +// npx tsx scripts/bench-token-count.ts [sizeKB...] +import { performance } from "node:perf_hooks" +import * as anthropicTokenizer from "@anthropic-ai/tokenizer" +import { countTokens, estimateTokensBatch } from "../lib/token-utils" + +const legacyCountTokens = ((anthropicTokenizer as any).countTokens ?? + (anthropicTokenizer as any).default?.countTokens) as (text: string) => number + +// Realistic mix of prose, code and JSON — what long sessions actually contain. +function synthSessionText(targetBytes: number): string { + const chunk = [ + "user: refactor the billing cycle console so invoices are grouped per account", + "assistant: I will start by reading the console controller and the invoice serializer.", + '{"path":"apps/bud/app/controllers/billing/invoices_controller.rb","line":128,"text":"render json: InvoiceSerializer.new(invoice).as_json"}', + "tool: bash -lc 'pnpm exec vitest run tests/serializers/invoice.test.ts' — 24 passed", + "- [x] read controller\n- [ ] patch serializer\n- [ ] run tests", + "assistant: The serializer now emits `cycle_started_at` and `cycle_ends_at`; tests pass.", + "```ruby\ndef cycle_bounds(account)\n [account.cycle_started_at, account.cycle_ends_at]\nend\n```", + ].join("\n") + const parts: string[] = [] + let bytes = 0 + while (bytes < targetBytes) { + parts.push(chunk) + bytes += chunk.length + 1 + } + return parts.join("\n") +} + +function measure(fn: () => number): { ms: number; tokens: number; missedTicks: number } { + let ticks = 0 + const handle = setInterval(() => ticks++, 10) + const start = performance.now() + const tokens = fn() + const ms = performance.now() - start + clearInterval(handle) + return { ms, tokens, missedTicks: Math.max(0, Math.floor(ms / 10) - ticks) } +} + +const sizesKb = process.argv + .slice(2) + .map(Number) + .filter((n) => Number.isFinite(n) && n > 0) +const kb = sizesKb.length > 0 ? sizesKb : [100, 1024, 8192] + +console.log(`node ${process.version} | sizes(KB): ${kb.join(", ")}`) +console.log( + "sizeKB\tchars\tlegacy ms\tlegacy ticks\tpatched ms\tpatched ticks\tpatched 2nd ms\ttokens", +) +for (const size of kb) { + const text = synthSessionText(size * 1024) + const legacy = measure(() => legacyCountTokens(text)) + const patched = measure(() => countTokens(text)) + const patched2 = measure(() => countTokens(text)) + console.log( + [ + size, + text.length, + legacy.ms.toFixed(1), + legacy.missedTicks, + patched.ms.toFixed(1), + patched.missedTicks, + patched2.ms.toFixed(1), + patched.tokens, + ].join("\t"), + ) +} + +// Batch path: master joined every text into one string before counting. +const batch = Array.from({ length: 200 }, (_, i) => + synthSessionText(2000).replace("user:", `tool ${i}:`), +) +const batchChars = batch.reduce((n, s) => n + s.length, 0) +const legacyBatch = measure(() => legacyCountTokens(batch.join(" "))) +const patchedBatch = measure(() => estimateTokensBatch(batch)) +console.log( + `\nestimateTokensBatch: 200 items (~${(batchChars / 1024).toFixed(0)} KB)\n` + + ` legacy join+count: ${legacyBatch.ms.toFixed(1)} ms (${legacyBatch.missedTicks} ticks)\n` + + ` per-item count: ${patchedBatch.ms.toFixed(1)} ms (${patchedBatch.missedTicks} ticks) -> ${patchedBatch.tokens} tokens`, +) diff --git a/scripts/bench-token-pathological.ts b/scripts/bench-token-pathological.ts new file mode 100644 index 00000000..de37ebf7 --- /dev/null +++ b/scripts/bench-token-pathological.ts @@ -0,0 +1,19 @@ +// Warm-path bench for degenerate runs (upstream #659). Repeated chars in one process with a +// warmup call first, so the numbers separate real encode cost from JS/WASM startup. +// timeout 120 npx tsx scripts/bench-token-pathological.ts [sizes...] +import { performance } from "node:perf_hooks" +import { countTokens } from "../lib/token-utils" + +countTokens("warmup ".repeat(64)) +countTokens("x".repeat(2048)) + +const sizes = process.argv.slice(2).map(Number) +const list = sizes.length > 0 ? sizes : [4096, 10000, 16384, 32768, 100000, 500000] + +console.log("chars\tms\ttokens") +for (const size of list) { + const text = "x".repeat(size) + const start = performance.now() + const tokens = countTokens(text) + console.log(`${size}\t${(performance.now() - start).toFixed(1)}\t${tokens}`) +} diff --git a/tests/token-counting.test.ts b/tests/token-counting.test.ts index 03e9439c..b9abc2a9 100644 --- a/tests/token-counting.test.ts +++ b/tests/token-counting.test.ts @@ -4,10 +4,12 @@ import type { WithParts } from "../lib/state" import { COMPACTED_TOOL_OUTPUT_PLACEHOLDER, countAllMessageTokens, + countTokens, countToolTokens, estimateTokensBatch, extractCompletedToolOutput, extractToolContent, + resetTokenCountingForTests, } from "../lib/token-utils" function buildToolMessage(part: Record): WithParts { @@ -172,3 +174,96 @@ test("counting uses the compacted tool placeholder for completed outputs", () => assert.equal(extractCompletedToolOutput(part), COMPACTED_TOOL_OUTPUT_PLACEHOLDER) assertCounted(part, [JSON.stringify(input), COMPACTED_TOOL_OUTPUT_PLACEHOLDER]) }) + +// --- shared tokenizer, bounded counting and caching (upstream #638/#659) --- + +async function referenceTokenizer() { + const tokenizer = await import("@anthropic-ai/tokenizer") + return ((tokenizer as any).countTokens ?? (tokenizer as any).default?.countTokens) as ( + text: string, + ) => number +} + +function measure(fn: () => number): { ms: number; tokens: number } { + const start = performance.now() + const tokens = fn() + return { ms: performance.now() - start, tokens } +} + +function repeatTo(chars: number, unit: string): string { + return unit.repeat(Math.ceil(chars / unit.length)).slice(0, chars) +} + +test("countTokens matches the reference tokenizer and reuses one tokenizer", async () => { + const reference = await referenceTokenizer() + const samples = [ + "hello world", + "Crème brûlée — fi ligature ① and emoji 🚀", + JSON.stringify({ path: "src/index.ts", edits: [{ line: 1, text: "const x = 1" }] }), + "x".repeat(1000), + "<|endoftext|> special tokens are counted, not rejected", + ] + for (const sample of samples) assert.equal(countTokens(sample), reference(sample)) + + // Building a tokenizer per call took ~70 ms; 200 calls with a shared one take a few ms. + resetTokenCountingForTests() + const unit = "tool output with words, punctuation and identifiers. " + const start = performance.now() + for (let index = 0; index < 200; index++) { + assert.ok(countTokens(`${index} ${unit}`) > 0) + } + const elapsed = performance.now() - start + assert.ok(elapsed < 6_000, `200 counts took ${elapsed.toFixed(0)} ms; is the tokenizer shared?`) +}) + +test("countTokens bounds the cost of huge and degenerate inputs", () => { + const realistic = repeatTo(4 * 1024 * 1024, "counts must stay bounded on huge inputs. ") + const realisticRun = measure(() => countTokens(realistic)) + assert.ok( + realisticRun.ms < 1_000, + `counting 4 MB took ${realisticRun.ms.toFixed(0)} ms; exact counting is unbounded`, + ) + + const degenerate = "x".repeat(2 * 1024 * 1024) + const degenerateRun = measure(() => countTokens(degenerate)) + assert.ok( + degenerateRun.ms < 1_000, + `counting a 2 MB single-character run took ${degenerateRun.ms.toFixed(0)} ms`, + ) +}) + +test("countTokens caches exact counts", () => { + resetTokenCountingForTests() + const text = repeatTo(7000, "cached content with words, punctuation and identifiers. ") + const first = measure(() => countTokens(text)) + assert.ok(first.tokens > 0) + + const start = performance.now() + for (let index = 0; index < 500; index++) assert.equal(countTokens(text), first.tokens) + const elapsed = performance.now() - start + assert.ok(elapsed < 1_000, `500 cached counts took ${elapsed.toFixed(0)} ms`) +}) + +test("sampled estimates stay close to exact counts above the limit", async () => { + const reference = await referenceTokenizer() + process.env.DCP_TOKEN_COUNT_MAX_EXACT_CHARS = "0" + try { + const units = [ + "memory context pruning keeps summaries technical and dense. ", + '{"tool":"read","input":{"filePath":"/tmp/app.log","limit":2000}} ', + "def compress(block)\n summary_tokens = count(block.summary)\nend\n", + ] + for (const unit of units) { + const text = repeatTo(64 * 1024, unit) + const exact = reference(text) + const estimated = countTokens(text) + const error = Math.abs(estimated - exact) / exact + assert.ok( + error < 0.1, + `sampled estimate off by ${(error * 100).toFixed(1)}% for "${unit.slice(0, 24)}"`, + ) + } + } finally { + delete process.env.DCP_TOKEN_COUNT_MAX_EXACT_CHARS + } +})