feat(api): exact offline token counting for count_tokens fallback via tiktoken (#4087)

Integrated into release/v3.8.28

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 { estimateTokens } from "@/shared/utils/costEstimator";
+import { countTextTokens } 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";
@@ -101,31 +101,31 @@ export async function POST(request) {
 
 function buildEstimatedCountResponse(body) {
   const messages = Array.isArray(body?.messages) ? body.messages : [];
-  let totalChars = 0;
+  let inputTokens = 0;
 
   for (const msg of messages) {
     if (typeof msg?.content === "string") {
-      totalChars += msg.content.length;
+      inputTokens += countTextTokens(msg.content);
       continue;
     }
 
     if (Array.isArray(msg?.content)) {
       for (const part of msg.content) {
         if (part?.type === "text" && typeof part.text === "string") {
-          totalChars += part.text.length;
+          inputTokens += countTextTokens(part.text);
         }
       }
     }
   }
 
   if (typeof body?.system === "string") {
-    totalChars += body.system.length;
+    inputTokens += countTextTokens(body.system);
   }
 
   return new Response(
     JSON.stringify({
-      input_tokens: totalChars > 0 ? Math.ceil(totalChars / 4) : estimateTokens(""),
-      source: "estimated",
+      input_tokens: inputTokens,
+      source: "local",
     }),
     {
       headers: { "Content-Type": "application/json", ...CORS_HEADERS },

src/shared/utils/tiktokenCounter.ts

@@ -0,0 +1,21 @@
+import { getEncoding, type Tiktoken } from "js-tiktoken";
+
+let encoder: Tiktoken | null = null;
+
+function getEncoder(): Tiktoken {
+  if (!encoder) encoder = getEncoding("cl100k_base");
+  return encoder;
+}
+
+/**
+ * Exact token count for a string using cl100k_base (offline, no upstream call).
+ * Defensive: never throws in a counting path — falls back to a char heuristic.
+ */
+export function countTextTokens(text: string): number {
+  if (!text || typeof text !== "string") return 0;
+  try {
+    return getEncoder().encode(text).length;
+  } catch {
+    return Math.ceil(text.length / 4);
+  }
+}

tests/unit/messages-count-tokens-route.test.ts

@@ -97,8 +97,20 @@ test("messages/count_tokens falls back to estimate when model is missing", async
 
   assert.equal(response.status, 200);
   const body = (await response.json()) as any;
-  assert.equal(body.input_tokens, 3);
-  assert.equal(body.source, "estimated");
+  assert.equal(body.input_tokens, 4); // tiktoken: "abcd"=1 + "12345678"=3
+  assert.equal(body.source, "local");
+});
+
+test("count_tokens fallback uses exact tiktoken count with source=local", async () => {
+  const req = new Request("http://localhost/v1/messages/count_tokens", {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ messages: [{ role: "user", content: "hello world" }] }),
+  });
+  const res = await POST(req);
+  const json = (await res.json()) as any;
+  assert.equal(json.source, "local");
+  assert.equal(json.input_tokens, 2); // exact cl100k_base count, not Math.ceil(11/4)=3
 });
 
 test("messages/count_tokens falls back to estimate when real upstream count fails", async () => {
@@ -121,8 +133,8 @@ test("messages/count_tokens falls back to estimate when real upstream count fail
 
     assert.equal(response.status, 200);
     const body = (await response.json()) as any;
-    assert.equal(body.input_tokens, 1);
-    assert.equal(body.source, "estimated");
+    assert.equal(body.input_tokens, 1); // tiktoken: "abcd"=1
+    assert.equal(body.source, "local");
   } finally {
     globalThis.fetch = originalFetch;
   }

tests/unit/tiktoken-counter.test.ts

@@ -0,0 +1,19 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { countTextTokens } 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("countTextTokens handles empty and non-string safely", () => {
+  assert.equal(countTextTokens(""), 0);
+  assert.equal(countTextTokens(undefined as unknown as string), 0);
+});
+
+test("countTextTokens is additive-ish and monotonic for longer text", () => {
+  const short = countTextTokens("the quick brown fox");
+  const long = countTextTokens("the quick brown fox jumps over the lazy dog");
+  assert.ok(long > short);
+  assert.ok(short > 0);
+});