diff --git a/src/shared/utils/tiktokenCounter.ts b/src/shared/utils/tiktokenCounter.ts index ec6af38913..5f4354f045 100644 --- a/src/shared/utils/tiktokenCounter.ts +++ b/src/shared/utils/tiktokenCounter.ts @@ -21,6 +21,29 @@ export function tokenizerContextFromBody(body: unknown): TokenizerContext { const encoders = new Map(); +/** + * Above this many characters the exact tokenizer is skipped in favor of the + * char-heuristic (chars/4). js-tiktoken's pure-JS encoder is near-quadratic on + * large inputs — a 10 MB base64 image payload can block the event loop for + * tens of seconds (OmniRoute worker wedge incident). Token counting is used for + * compression stats/estimates only, so a heuristic on oversized inputs is + * acceptable and keeps the loop responsive. + */ +const MAX_EXACT_TOKEN_COUNT_CHARS = 50_000; + +/** + * Base64 data URIs (e.g. OpenAI-style `image_url.url`) must not be tokenized: + * they are image payloads, not text. Matching a data URI of any `image/*` + * media type and stripping it keeps the count accurate (the raw bytes of an + * image are not meaningful "text" tokens) while avoiding the quadratic encode + * cost on large attachments. + */ +const BASE64_DATA_URI_RE = /data:image\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/gi; + +function stripBase64DataUris(text: string): string { + return text.replace(BASE64_DATA_URI_RE, ""); +} + function normalize(value: unknown): string { return typeof value === "string" ? value.trim().toLowerCase() : ""; } @@ -60,12 +83,19 @@ function getEncoder(encoding: TokenizerEncoding): Tiktoken { * Existing callers retain cl100k_base; Codex callers may pass provider/model context * to use o200k_base. * Defensive: never throws in a counting path — falls back to a char heuristic. + * Oversized inputs (over 50k chars) and base64 image data URIs are never + * tokenized: the encoder is near-quadratic on large strings and would block the + * event loop (worker wedge regression). */ export function countTextTokens(text: string, context?: TokenizerContext): number { if (!text || typeof text !== "string") return 0; + const stripped = stripBase64DataUris(text); + if (stripped.length > MAX_EXACT_TOKEN_COUNT_CHARS) { + return Math.ceil(stripped.length / 4); + } try { - return getEncoder(resolveTokenizerEncoding(context)).encode(text).length; + return getEncoder(resolveTokenizerEncoding(context)).encode(stripped).length; } catch { - return Math.ceil(text.length / 4); + return Math.ceil(stripped.length / 4); } } diff --git a/tests/unit/tiktoken-counter.test.ts b/tests/unit/tiktoken-counter.test.ts index 110618d1b0..d3263dfd31 100644 --- a/tests/unit/tiktoken-counter.test.ts +++ b/tests/unit/tiktoken-counter.test.ts @@ -40,3 +40,34 @@ test("countTextTokens is additive-ish and monotonic for longer text", () => { assert.ok(long > short); assert.ok(short > 0); }); + +test("countTextTokens fast-paths strings over 50k chars without tokenizing (worker wedge regression)", () => { + const big = "user: please review the attached patch\ntext: ".repeat(40_000); + const start = performance.now(); + const tokens = countTextTokens(big); + const elapsed = performance.now() - start; + assert.equal(tokens, Math.ceil(big.length / 4)); + assert.ok(elapsed < 1000, `fast path took ${elapsed.toFixed(0)}ms`); +}); + +test("countTextTokens strips base64 data URIs before tokenizing (images not counted as text)", () => { + const png = + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="; + const b64 = png.repeat(60); + const withImage = countTextTokens( + `{"image_url":{"url":"data:image/png;base64,${b64}"}}`, + { provider: "codex" } + ); + const stripped = countTextTokens('{"image_url":{"url":""}}', { provider: "codex" }); + assert.equal(withImage, stripped); +}); + +test("countTextTokens does not tokenize huge base64 image payloads (wedge repro)", () => { + const b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="; + const body = `{"image_url":{"url":"data:image/png;base64,${b64.repeat(14_000)}"}}`; + const start = performance.now(); + const tokens = countTextTokens(body); + const elapsed = performance.now() - start; + assert.ok(tokens < 1000, `base64 payload inflates token count to ${tokens}`); + assert.ok(elapsed < 1000, `took ${elapsed.toFixed(0)}ms`); +});