feat: enhance token counting with model validation and update tests

Author: cuipengfei

src/routes/generate-content/handler.ts

@@ -300,7 +300,16 @@ export async function handleGeminiCountTokens(c: Context) {
 
   const openAIPayload = translateGeminiCountTokensToOpenAI(geminiPayload, model)
 
-  const tokenCounts = getTokenCount(openAIPayload.messages)
+  // Find the selected model object from state
+  const selectedModel = state.models?.data.find((m) => m.id === model)
+
+  if (!selectedModel) {
+    throw new Error(
+      `Model ${model} not found in available models. Please ensure the model list is loaded.`,
+    )
+  }
+
+  const tokenCounts = await getTokenCount(openAIPayload, selectedModel)
 
   const totalTokens = tokenCounts.input + tokenCounts.output
   const geminiResponse = translateTokenCountToGemini(totalTokens)

tests/generate-content/core-functionality.test.ts

@@ -27,7 +27,21 @@ test("translates request and uses local tokenizer without downstream call", asyn
     },
   }))
   await mock.module("~/lib/tokenizer", () => ({
-    getTokenCount: (_messages: unknown) => ({ input: 2, output: 3 }),
+    getTokenCount: async (_messages: unknown, _model: unknown) =>
+      Promise.resolve({ input: 2, output: 3 }),
+  }))
+  await mock.module("~/lib/state", () => ({
+    state: {
+      models: {
+        data: [
+          {
+            id: "gemini-pro",
+            name: "Gemini Pro",
+            capabilities: { tokenizer: "o200k_base" },
+          },
+        ],
+      },
+    },
   }))
 
   const { server } = (await import("~/server")) as { server: TestServer }

tests/generate-content/route-routing.test.ts

@@ -66,7 +66,21 @@ test("routes to stream endpoint based on URL keyword", async () => {
 
 test("routes to countTokens endpoint based on URL keyword", async () => {
   await mock.module("~/lib/tokenizer", () => ({
-    getTokenCount: (_messages: unknown) => ({ input: 2, output: 3 }),
+    getTokenCount: async (_messages: unknown, _model: unknown) =>
+      Promise.resolve({ input: 2, output: 3 }),
+  }))
+  await mock.module("~/lib/state", () => ({
+    state: {
+      models: {
+        data: [
+          {
+            id: "gemini-pro",
+            name: "Gemini Pro",
+            capabilities: { tokenizer: "o200k_base" },
+          },
+        ],
+      },
+    },
   }))
   await mock.module("~/lib/rate-limit", () => ({
     checkRateLimit: () => {},