From 6d8b1cb01356f9145b04354d1b0521ffbfa4a55e Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 26 Aug 2026 07:37:19 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=A2=20fix:=20Bound=20Tokenizer=20Work?= =?UTF-8?q?=20Per=20Input=20(#15246)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/api/src/utils/tokenizer.spec.ts | 32 ++++++++++++++++++++++++ packages/api/src/utils/tokenizer.ts | 26 ++++++++++++++++--- 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/packages/api/src/utils/tokenizer.spec.ts b/packages/api/src/utils/tokenizer.spec.ts index d2a4ccb6a5..1b55b31ca6 100644 --- a/packages/api/src/utils/tokenizer.spec.ts +++ b/packages/api/src/utils/tokenizer.spec.ts @@ -14,6 +14,12 @@ describe('Tokenizer', () => { }); describe('initEncoding', () => { + it('uses the established estimate while the encoding initializes', () => { + const text = 'Cold start token count'; + + expect(Tokenizer.getTokenCount(text, 'o200k_base')).toBe(Math.ceil(text.length / 4)); + }); + it('should load o200k_base encoding', async () => { await Tokenizer.initEncoding('o200k_base'); const count = Tokenizer.getTokenCount('Hello, world!', 'o200k_base'); @@ -53,6 +59,20 @@ describe('Tokenizer', () => { const count = Tokenizer.getTokenCount('Hello, world!', 'claude'); expect(count).toBeGreaterThan(0); }); + + it.each([ + { label: 'an uninterrupted non-whitespace run', text: '_'.repeat(4 * 1024 + 1) }, + { label: 'an oversized input', text: 'word '.repeat(1024) }, + { label: 'multibyte input', text: '界'.repeat(4 * 1024 + 1) }, + ])('uses a conservative estimate without tokenizing $label', ({ text }) => { + const count = jest.spyOn(AiTokenizer.prototype, 'count'); + try { + expect(Tokenizer.getTokenCount(text, 'o200k_base')).toBe(Buffer.byteLength(text, 'utf8')); + expect(count).not.toHaveBeenCalled(); + } finally { + count.mockRestore(); + } + }); }); describe('createExactTokenCounter', () => { @@ -68,5 +88,17 @@ describe('Tokenizer', () => { await Tokenizer.initEncoding('o200k_base'); } }); + + it('uses a conservative estimate for unsafe input without invoking the tokenizer', async () => { + const counter = await Tokenizer.createExactTokenCounter('o200k_base'); + const count = jest.spyOn(AiTokenizer.prototype, 'count'); + try { + const text = '_'.repeat(4 * 1024 + 1); + expect(counter(text)).toBe(Buffer.byteLength(text, 'utf8')); + expect(count).not.toHaveBeenCalled(); + } finally { + count.mockRestore(); + } + }); }); }); diff --git a/packages/api/src/utils/tokenizer.ts b/packages/api/src/utils/tokenizer.ts index 83a3c2691f..9bbd9591ee 100644 --- a/packages/api/src/utils/tokenizer.ts +++ b/packages/api/src/utils/tokenizer.ts @@ -5,6 +5,20 @@ export type EncodingName = 'o200k_base' | 'claude'; type EncodingData = ConstructorParameters[0]; +const MAX_TOKENIZER_INPUT_LENGTH = 4 * 1024; + +function estimateBoundedTokenCount(text: string): number { + return Buffer.byteLength(text, 'utf8'); +} + +function estimateUnavailableTokenCount(text: string): number { + return Math.ceil(text.length / 4); +} + +function requiresTokenEstimate(text: string): boolean { + return text.length > MAX_TOKENIZER_INPUT_LENGTH; +} + class Tokenizer { private tokenizersCache: Partial> = {}; private loadingPromises: Partial>> = {}; @@ -27,7 +41,7 @@ class Tokenizer { return this.loadingPromises[encoding]; } - /** Returns a counter that never substitutes an estimate for tokenizer errors. */ + /** Returns a counter that avoids expensive tokenization for oversized content. */ async createExactTokenCounter(encoding: EncodingName): Promise<(text: string) => number> { await this.initEncoding(encoding); const tokenizer = this.tokenizersCache[encoding]; @@ -35,6 +49,9 @@ class Tokenizer { throw new Error(`Tokenizer encoding failed to initialize: ${encoding}`); } return (text: string): number => { + if (requiresTokenEstimate(text)) { + return estimateBoundedTokenCount(text); + } try { return tokenizer.count(text); } catch (error) { @@ -45,16 +62,19 @@ class Tokenizer { } getTokenCount(text: string, encoding: EncodingName = 'o200k_base'): number { + if (requiresTokenEstimate(text)) { + return estimateBoundedTokenCount(text); + } const tokenizer = this.tokenizersCache[encoding]; if (!tokenizer) { this.initEncoding(encoding); - return Math.ceil(text.length / 4); + return estimateUnavailableTokenCount(text); } try { return tokenizer.count(text); } catch (error) { this.handleCountError(encoding, tokenizer, error); - return Math.ceil(text.length / 4); + return estimateUnavailableTokenCount(text); } }