feat(compression): select model-aware tokenizers (#8009)

* Add model-aware tokenizer selection

* fix: recognize cx Codex model prefix

open-sse/services/compression/hardBudget.ts

@@ -10,7 +10,11 @@
 
 import type { CompressionResult } from "./types.ts";
 import { scoreToken } from "./ultraHeuristic.ts";
-import { countTextTokens } from "../../../src/shared/utils/tiktokenCounter.ts";
+import {
+  countTextTokens,
+  tokenizerContextFromBody,
+  type TokenizerContext,
+} from "../../../src/shared/utils/tiktokenCounter.ts";
 import { createCompressionStats } from "./stats.ts";
 
 interface HardBudgetOptions {
@@ -69,11 +73,11 @@ interface TaggedUnit {
   preserve: boolean;
 }
 
-function tagUnits(units: string[]): TaggedUnit[] {
+function tagUnits(units: string[], tokenizerContext: TokenizerContext): TaggedUnit[] {
   return units.map((u, i) => ({
     i,
     u,
-    tokens: countTextTokens(u),
+    tokens: countTextTokens(u, tokenizerContext),
     score: scoreUnit(u),
     preserve: mustPreserve(u),
   }));
@@ -84,9 +88,7 @@ function dropToTarget(tagged: TaggedUnit[], targetTokens: number): Set<number> {
   let tokCount = tagged.reduce((s, x) => s + x.tokens, 0);
 
   // Sort droppable candidates by score ascending (lowest first = drop first)
-  const candidates = tagged
-    .filter((x) => !x.preserve)
-    .sort((a, b) => a.score - b.score);
+  const candidates = tagged.filter((x) => !x.preserve).sort((a, b) => a.score - b.score);
 
   for (const candidate of candidates) {
     if (tokCount <= targetTokens) break;
@@ -104,14 +106,18 @@ function rebuildText(tagged: TaggedUnit[], dropped: Set<number>): string {
     .join("\n");
 }
 
-function compressText(text: string, targetTokens: number): string {
-  const currentTokens = countTextTokens(text);
+function compressText(
+  text: string,
+  targetTokens: number,
+  tokenizerContext: TokenizerContext
+): string {
+  const currentTokens = countTextTokens(text, tokenizerContext);
   if (currentTokens <= targetTokens) return text;
 
   const units = splitUnits(text);
   if (units.length <= 1) return text;
 
-  const tagged = tagUnits(units);
+  const tagged = tagUnits(units, tokenizerContext);
   const dropped = dropToTarget(tagged, targetTokens);
   if (dropped.size === 0) return text;
 
@@ -137,17 +143,17 @@ export function applyHardBudget(
   const messages = extractMessages(body);
   if (messages.length === 0) return { body, compressed: false, stats: null };
 
+  const tokenizerContext = tokenizerContextFromBody(body);
+
   // Measure total tokens across all messages
   const totalText = messages
     .map((m) => (typeof m.content === "string" ? m.content : JSON.stringify(m.content)))
     .join(" ");
-  const totalTokens = countTextTokens(totalText);
+  const totalTokens = countTextTokens(totalText, tokenizerContext);
 
   // targetTokens wins when both are set
   const effectiveTarget =
-    targetTokens != null
-      ? targetTokens
-      : Math.floor(totalTokens * (targetRatio as number));
+    targetTokens != null ? targetTokens : Math.floor(totalTokens * (targetRatio as number));
 
   if (totalTokens <= effectiveTarget) {
     return { body, compressed: false, stats: null };
@@ -158,16 +164,14 @@ export function applyHardBudget(
   // come back N× over budget).
   const newMessages = messages.map((m) => {
     if (typeof m.content !== "string") return m;
-    const msgTokens = countTextTokens(m.content);
+    const msgTokens = countTextTokens(m.content, tokenizerContext);
     const perMsgTarget =
       totalTokens > 0 ? Math.floor(effectiveTarget * (msgTokens / totalTokens)) : effectiveTarget;
-    const out = compressText(m.content, perMsgTarget);
+    const out = compressText(m.content, perMsgTarget, tokenizerContext);
     return out === m.content ? m : { ...m, content: out };
   });
 
-  const changed = newMessages.some(
-    (m, i) => JSON.stringify(m) !== JSON.stringify(messages[i])
-  );
+  const changed = newMessages.some((m, i) => JSON.stringify(m) !== JSON.stringify(messages[i]));
 
   // Measure the result to detect when preserve-guarded content makes the target
   // unreachable, so callers are not silently left over budget.
@@ -175,7 +179,8 @@ export function applyHardBudget(
   const resultTokens = countTextTokens(
     usedMessages
       .map((m) => (typeof m.content === "string" ? m.content : JSON.stringify(m.content)))
-      .join(" ")
+      .join(" "),
+    tokenizerContext
   );
   const overBudget = resultTokens > effectiveTarget;
 

open-sse/services/compression/stats.ts

@@ -7,6 +7,11 @@ import {
   DEFAULT_RTK_CONFIG,
   DEFAULT_COMPRESSION_LANGUAGE_CONFIG,
 } from "./types.ts";
+import {
+  countTextTokens,
+  isCodexTokenizerContext,
+  tokenizerContextFromBody,
+} from "../../../src/shared/utils/tiktokenCounter.ts";
 import { anthropicImageTokens, ANTHROPIC_IMAGE_BLOCK_OVERHEAD_TOKENS } from "omniglyph";
 
 const CHARS_PER_TOKEN = 4;
@@ -29,9 +34,7 @@ function isAnthropicPngImageBlock(value: unknown): value is AnthropicImageBlock
   const source = block.source as Record<string, unknown> | undefined;
   if (!source || typeof source !== "object") return false;
   return (
-    source.type === "base64" &&
-    source.media_type === "image/png" &&
-    typeof source.data === "string"
+    source.type === "base64" && source.media_type === "image/png" && typeof source.data === "string"
   );
 }
 
@@ -119,17 +122,24 @@ function blankImageBlocksAndSumImageTokens(body: Record<string, unknown>): {
 export function estimateCompressionTokens(text: string | object | null | undefined): number {
   if (!text) return 0;
   if (typeof text === "string") {
-    return Math.ceil(text.length / CHARS_PER_TOKEN);
+    return charTokensOf(text);
   }
   try {
+    const tokenizerContext = tokenizerContextFromBody(text);
+    const useExactTokenizer = isCodexTokenizerContext(tokenizerContext);
     const { clone, imageTokens } = blankImageBlocksAndSumImageTokens(
       text as Record<string, unknown>
     );
     if (imageTokens === 0) {
-      // No recognized image blocks — byte-identical to the legacy behavior.
-      return Math.ceil(JSON.stringify(text).length / CHARS_PER_TOKEN);
+      // Keep the legacy character estimate for generic payloads. Codex payloads use
+      // the model-appropriate tokenizer so their compression stats match hard budgets.
+      return useExactTokenizer
+        ? countTextTokens(JSON.stringify(text), tokenizerContext)
+        : charTokensOf(text);
     }
-    return Math.ceil(JSON.stringify(clone).length / CHARS_PER_TOKEN) + imageTokens;
+    return useExactTokenizer
+      ? countTextTokens(JSON.stringify(clone), tokenizerContext) + imageTokens
+      : charTokensOf(clone) + imageTokens;
   } catch {
     // Non-serializable/unexpected shape → fall back to the legacy char-count,
     // never throw out of an estimator.

src/app/api/v1/messages/count_tokens/route.ts

@@ -1,7 +1,7 @@
 import { CORS_HEADERS } from "@/shared/utils/cors";
 import { v1CountTokensSchema } from "@/shared/validation/schemas";
 import { isValidationFailure, validateBody } from "@/shared/validation/helpers";
-import { countTextTokens } from "@/shared/utils/tiktokenCounter";
+import { countTextTokens, type TokenizerContext } from "@/shared/utils/tiktokenCounter";
 import { getExecutor } from "@omniroute/open-sse/executors/index.ts";
 import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch.ts";
 import { getModelInfo } from "@/sse/services/model";
@@ -40,7 +40,10 @@ export async function POST(request) {
   }
   const body = validation.data;
 
-  const estimated = buildEstimatedCountResponse(body);
+  const tokenizerContext: TokenizerContext = {
+    model: typeof body.model === "string" ? body.model : undefined,
+  };
+  const estimated = buildEstimatedCountResponse(body, tokenizerContext);
   const requestedModel = typeof body.model === "string" ? body.model : "";
   if (!requestedModel) {
     return estimated;
@@ -117,22 +120,24 @@ function safeStringify(value) {
 // content, and `thinking` blocks — counting only `text` (as before) reported
 // near-zero for those messages and silently broke Claude Code's auto-compaction
 // (#2337). Image / redacted_thinking blocks are not text-estimable and count 0.
-function estimateContentBlockTokens(part) {
+function estimateContentBlockTokens(part, tokenizerContext: TokenizerContext) {
   if (!part || typeof part !== "object") return 0;
   let tokens = 0;
   switch (part.type) {
     case "text":
-      if (typeof part.text === "string") tokens += countTextTokens(part.text);
+      if (typeof part.text === "string") tokens += countTextTokens(part.text, tokenizerContext);
       break;
     case "tool_use":
-      if (typeof part.name === "string") tokens += countTextTokens(part.name);
-      if (part.input !== undefined) tokens += countTextTokens(safeStringify(part.input));
+      if (typeof part.name === "string") tokens += countTextTokens(part.name, tokenizerContext);
+      if (part.input !== undefined)
+        tokens += countTextTokens(safeStringify(part.input), tokenizerContext);
       break;
     case "tool_result":
-      tokens += estimateToolResultTokens(part.content);
+      tokens += estimateToolResultTokens(part.content, tokenizerContext);
       break;
     case "thinking":
-      if (typeof part.thinking === "string") tokens += countTextTokens(part.thinking);
+      if (typeof part.thinking === "string")
+        tokens += countTextTokens(part.thinking, tokenizerContext);
       break;
     default:
       break;
@@ -142,13 +147,13 @@ function estimateContentBlockTokens(part) {
 
 // A `tool_result` content can be a plain string or an array of nested blocks
 // (text / image). Count string content and nested text blocks.
-function estimateToolResultTokens(content) {
-  if (typeof content === "string") return countTextTokens(content);
+function estimateToolResultTokens(content, tokenizerContext: TokenizerContext) {
+  if (typeof content === "string") return countTextTokens(content, tokenizerContext);
   if (Array.isArray(content)) {
     let tokens = 0;
     for (const block of content) {
       if (block?.type === "text" && typeof block.text === "string") {
-        tokens += countTextTokens(block.text);
+        tokens += countTextTokens(block.text, tokenizerContext);
       }
     }
     return tokens;
@@ -156,29 +161,29 @@ function estimateToolResultTokens(content) {
   return 0;
 }
 
-function buildEstimatedCountResponse(body) {
+function buildEstimatedCountResponse(body, tokenizerContext: TokenizerContext = {}) {
   const messages = Array.isArray(body?.messages) ? body.messages : [];
   let inputTokens = 0;
 
   for (const msg of messages) {
     if (typeof msg?.content === "string") {
-      inputTokens += countTextTokens(msg.content);
+      inputTokens += countTextTokens(msg.content, tokenizerContext);
       continue;
     }
 
     if (Array.isArray(msg?.content)) {
       for (const part of msg.content) {
-        inputTokens += estimateContentBlockTokens(part);
+        inputTokens += estimateContentBlockTokens(part, tokenizerContext);
       }
     }
   }
 
   if (typeof body?.system === "string") {
-    inputTokens += countTextTokens(body.system);
+    inputTokens += countTextTokens(body.system, tokenizerContext);
   } else if (Array.isArray(body?.system)) {
     for (const block of body.system) {
       if (block?.type === "text" && typeof block.text === "string") {
-        inputTokens += countTextTokens(block.text);
+        inputTokens += countTextTokens(block.text, tokenizerContext);
       }
     }
   }

src/shared/utils/tiktokenCounter.ts

@@ -1,20 +1,61 @@
 import { getEncoding, type Tiktoken } from "js-tiktoken";
 
-let encoder: Tiktoken | null = null;
+export type TokenizerEncoding = "cl100k_base" | "o200k_base";
 
-function getEncoder(): Tiktoken {
-  if (!encoder) encoder = getEncoding("cl100k_base");
-  return encoder;
+export interface TokenizerContext {
+  provider?: string | null;
+  model?: string | null;
+}
+
+export function tokenizerContextFromBody(body: unknown): TokenizerContext {
+  if (!body || typeof body !== "object" || Array.isArray(body)) return {};
+  const record = body as Record<string, unknown>;
+  return {
+    provider: typeof record.provider === "string" ? record.provider : undefined,
+    model: typeof record.model === "string" ? record.model : undefined,
+  };
+}
+
+const encoders = new Map<TokenizerEncoding, Tiktoken>();
+
+function normalize(value: unknown): string {
+  return typeof value === "string" ? value.trim().toLowerCase() : "";
+}
+
+export function isCodexTokenizerContext(context?: TokenizerContext): boolean {
+  const provider = normalize(context?.provider);
+  const model = normalize(context?.model);
+  return (
+    provider === "codex" ||
+    provider === "cx" ||
+    model.startsWith("codex/") ||
+    model.startsWith("cx/") ||
+    model.includes("codex")
+  );
+}
+
+export function resolveTokenizerEncoding(context?: TokenizerContext): TokenizerEncoding {
+  return isCodexTokenizerContext(context) ? "o200k_base" : "cl100k_base";
+}
+
+function getEncoder(encoding: TokenizerEncoding): Tiktoken {
+  const cached = encoders.get(encoding);
+  if (cached) return cached;
+  const created = getEncoding(encoding);
+  encoders.set(encoding, created);
+  return created;
 }
 
 /**
- * Exact token count for a string using cl100k_base (offline, no upstream call).
+ * Exact token count for a string using the selected offline tokenizer.
+ * 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.
  */
-export function countTextTokens(text: string): number {
+export function countTextTokens(text: string, context?: TokenizerContext): number {
   if (!text || typeof text !== "string") return 0;
   try {
-    return getEncoder().encode(text).length;
+    return getEncoder(resolveTokenizerEncoding(context)).encode(text).length;
   } catch {
     return Math.ceil(text.length / 4);
   }

tests/unit/tiktoken-counter.test.ts

@@ -1,11 +1,34 @@
 import { test } from "node:test";
 import assert from "node:assert/strict";
-import { countTextTokens } from "../../src/shared/utils/tiktokenCounter.ts";
+import {
+  countTextTokens,
+  isCodexTokenizerContext,
+  resolveTokenizerEncoding,
+} from "../../src/shared/utils/tiktokenCounter.ts";
 
 test("countTextTokens returns exact tiktoken count for a known string", () => {
   assert.equal(countTextTokens("hello world"), 2); // cl100k_base
 });
 
+test("Codex context selects o200k_base without changing the default", () => {
+  assert.equal(resolveTokenizerEncoding(), "cl100k_base");
+  assert.equal(resolveTokenizerEncoding({ provider: "codex" }), "o200k_base");
+  assert.equal(resolveTokenizerEncoding({ provider: "cx" }), "o200k_base");
+  assert.equal(resolveTokenizerEncoding({ model: "codex/gpt-5.6-sol" }), "o200k_base");
+  assert.equal(resolveTokenizerEncoding({ model: "cx/gpt-5.6-sol" }), "o200k_base");
+  assert.equal(resolveTokenizerEncoding({ provider: "openai", model: "gpt-5.6" }), "cl100k_base");
+  assert.equal(isCodexTokenizerContext({ provider: "codex" }), true);
+  assert.equal(isCodexTokenizerContext({ provider: "openai" }), false);
+});
+
+test("Codex token counting uses the o200k encoder", () => {
+  const text = "antidisestablishmentarianism 中文ภาษาไทย";
+  assert.notEqual(
+    countTextTokens(text, { provider: "codex" }),
+    countTextTokens(text, { provider: "openai" })
+  );
+});
+
 test("countTextTokens handles empty and non-string safely", () => {
   assert.equal(countTextTokens(""), 0);
   assert.equal(countTextTokens(undefined as unknown as string), 0);