diff --git a/lib/token-utils.ts b/lib/token-utils.ts index bba07fa1..7168b7f6 100644 --- a/lib/token-utils.ts +++ b/lib/token-utils.ts @@ -2,8 +2,16 @@ 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 +const getAnthropicTokenizer = (_anthropicTokenizer.getTokenizer ?? + (_anthropicTokenizer as any).default?.getTokenizer) as typeof _anthropicTokenizer.getTokenizer +// The package's countTokens() builds and frees a new WASM tokenizer on every call (~70 ms each), which +// blocks the host for seconds when many tool outputs are counted. Build it once and reuse it: the +// counts are identical (same encoder, same NFKC normalisation and special-token handling). +let anthropicTokenizer: ReturnType | undefined +function anthropicCountTokens(text: string): number { + anthropicTokenizer ??= getAnthropicTokenizer() + return anthropicTokenizer.encode(text.normalize("NFKC"), "all").length +} import { getLastUserMessage } from "./messages/query" export function getCurrentTokenUsage(state: SessionState, messages: WithParts[]): number { diff --git a/tests/token-counting.test.ts b/tests/token-counting.test.ts index 03e9439c..964c6faa 100644 --- a/tests/token-counting.test.ts +++ b/tests/token-counting.test.ts @@ -4,6 +4,7 @@ import type { WithParts } from "../lib/state" import { COMPACTED_TOOL_OUTPUT_PLACEHOLDER, countAllMessageTokens, + countTokens, countToolTokens, estimateTokensBatch, extractCompletedToolOutput, @@ -172,3 +173,25 @@ 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]) }) + +test("countTokens matches @anthropic-ai/tokenizer and reuses one tokenizer", async () => { + const tokenizer = await import("@anthropic-ai/tokenizer") + const reference = (tokenizer.countTokens ?? + (tokenizer as any).default?.countTokens) as typeof tokenizer.countTokens + 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(10_000), + "<|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 milliseconds. + const start = performance.now() + for (let index = 0; index < 200; index++) countTokens(`tool output ${index} `.repeat(50)) + assert.ok( + performance.now() - start < 2_000, + "countTokens should not rebuild the tokenizer on every call", + ) +})