From b195750f34a6942a2ee205b05f9b9b6b289c9f7c Mon Sep 17 00:00:00 2001
From: digua <i@digua.me>
Date: Thu, 14 May 2026 21:27:39 +0800
Subject: [PATCH] refactor: extract tokenizer to @openchatlab/node-runtime

Move countTokens/countMessagesTokens from Electron to shared package.
Electron side becomes a thin re-export wrapper.
---
 electron/main/ai/tokenizer/index.ts       | 48 ++---------------------
 packages/node-runtime/package.json        |  3 +-
 packages/node-runtime/src/ai/index.ts     |  3 ++
 packages/node-runtime/src/ai/tokenizer.ts | 47 ++++++++++++++++++++++
 packages/node-runtime/src/index.ts        |  1 +
 pnpm-lock.yaml                            |  3 ++
 6 files changed, 60 insertions(+), 45 deletions(-)
 create mode 100644 packages/node-runtime/src/ai/tokenizer.ts

diff --git a/electron/main/ai/tokenizer/index.ts b/electron/main/ai/tokenizer/index.ts
index 2f436647..337d99ab 100644
--- a/electron/main/ai/tokenizer/index.ts
+++ b/electron/main/ai/tokenizer/index.ts
@@ -1,47 +1,7 @@
 /**
- * Token 计数模块
- * 使用 js-tiktoken 的 cl100k_base 编码进行近似 token 计数。
- * 该编码是 GPT-4/Claude 系列的近似值，对国内模型有一定误差，
- * 因此阈值计算时预留了余量。
+ * Token 计数模块 — 薄包装
+ *
+ * 实际实现已提取到 @openchatlab/node-runtime。
  */
 
-import { encodingForModel } from 'js-tiktoken'
-
-let encoder: ReturnType<typeof encodingForModel> | null = null
-
-function getEncoder() {
-  if (!encoder) {
-    encoder = encodingForModel('gpt-4o')
-  }
-  return encoder
-}
-
-/**
- * 计算单段文本的 token 数
- */
-export function countTokens(text: string): number {
-  if (!text) return 0
-  return getEncoder().encode(text).length
-}
-
-/**
- * 计算消息列表的总 token 数（含 systemPrompt）
- * 每条消息额外计 4 tokens 的格式开销（role + 分隔符）
- */
-export function countMessagesTokens(messages: Array<{ role: string; content: string }>, systemPrompt?: string): number {
-  const enc = getEncoder()
-  let total = 0
-
-  if (systemPrompt) {
-    total += enc.encode(systemPrompt).length + 4
-  }
-
-  for (const msg of messages) {
-    total += enc.encode(msg.content).length + 4
-  }
-
-  // 回复引导 token
-  total += 3
-
-  return total
-}
+export { countTokens, countMessagesTokens } from '@openchatlab/node-runtime'
diff --git a/packages/node-runtime/package.json b/packages/node-runtime/package.json
index 648ab24a..a21c50c1 100644
--- a/packages/node-runtime/package.json
+++ b/packages/node-runtime/package.json
@@ -7,6 +7,7 @@
   "types": "./src/index.ts",
   "dependencies": {
     "@openchatlab/core": "workspace:*",
-    "gray-matter": "^4.0.3"
+    "gray-matter": "^4.0.3",
+    "js-tiktoken": "^1.0.21"
   }
 }
diff --git a/packages/node-runtime/src/ai/index.ts b/packages/node-runtime/src/ai/index.ts
index 53460096..2d7bde59 100644
--- a/packages/node-runtime/src/ai/index.ts
+++ b/packages/node-runtime/src/ai/index.ts
@@ -16,3 +16,6 @@ export type {
   TokenUsageData,
   ConversationManagerLogger,
 } from './conversations'
+
+// Tokenizer
+export { countTokens, countMessagesTokens } from './tokenizer'
diff --git a/packages/node-runtime/src/ai/tokenizer.ts b/packages/node-runtime/src/ai/tokenizer.ts
new file mode 100644
index 00000000..2f436647
--- /dev/null
+++ b/packages/node-runtime/src/ai/tokenizer.ts
@@ -0,0 +1,47 @@
+/**
+ * Token 计数模块
+ * 使用 js-tiktoken 的 cl100k_base 编码进行近似 token 计数。
+ * 该编码是 GPT-4/Claude 系列的近似值，对国内模型有一定误差，
+ * 因此阈值计算时预留了余量。
+ */
+
+import { encodingForModel } from 'js-tiktoken'
+
+let encoder: ReturnType<typeof encodingForModel> | null = null
+
+function getEncoder() {
+  if (!encoder) {
+    encoder = encodingForModel('gpt-4o')
+  }
+  return encoder
+}
+
+/**
+ * 计算单段文本的 token 数
+ */
+export function countTokens(text: string): number {
+  if (!text) return 0
+  return getEncoder().encode(text).length
+}
+
+/**
+ * 计算消息列表的总 token 数（含 systemPrompt）
+ * 每条消息额外计 4 tokens 的格式开销（role + 分隔符）
+ */
+export function countMessagesTokens(messages: Array<{ role: string; content: string }>, systemPrompt?: string): number {
+  const enc = getEncoder()
+  let total = 0
+
+  if (systemPrompt) {
+    total += enc.encode(systemPrompt).length + 4
+  }
+
+  for (const msg of messages) {
+    total += enc.encode(msg.content).length + 4
+  }
+
+  // 回复引导 token
+  total += 3
+
+  return total
+}
diff --git a/packages/node-runtime/src/index.ts b/packages/node-runtime/src/index.ts
index 11c4c46e..ec0d9d68 100644
--- a/packages/node-runtime/src/index.ts
+++ b/packages/node-runtime/src/index.ts
@@ -34,6 +34,7 @@ export {
 export type { AssistantConfig, AssistantSummary, SkillDef, SkillSummary } from './ai'
 export { parseAssistantFile, serializeAssistant, parseSkillFile, extractSkillId } from './ai'
 export { AIConversationManager } from './ai'
+export { countTokens, countMessagesTokens } from './ai'
 export type {
   AIConversation,
   AIMessage,
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index a4bf498d..e4feec67 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -161,6 +161,9 @@ importers:
       gray-matter:
         specifier: ^4.0.3
         version: 4.0.3
+      js-tiktoken:
+        specifier: ^1.0.21
+        version: 1.0.21
 
   packages/parser:
     dependencies: