🧢 fix: Bound Tokenizer Work Per Input (#15246)

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();
+      }
+    });
   });
 });

packages/api/src/utils/tokenizer.ts

@@ -5,6 +5,20 @@ export type EncodingName = 'o200k_base' | 'claude';
 
 type EncodingData = ConstructorParameters<typeof AiTokenizer>[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<Record<EncodingName, AiTokenizer>> = {};
   private loadingPromises: Partial<Record<EncodingName, Promise<void>>> = {};
@@ -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);
     }
   }