From fb6d2a0d776e4588f4b5f8bb082d5d312743f435 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 22 Apr 2026 16:00:12 -0400 Subject: [PATCH] fix(runtime): tokenize context compaction budget --- docs/proposals/C3-real-tokenizer-followup.md | 10 +++--- package.json | 1 + pnpm-lock.yaml | 10 ++++++ src/main/agent-runtime/context-compactor.ts | 21 ++++-------- .../context-compactor-tool-results.test.ts | 32 +++++++++++++++++++ 5 files changed, 55 insertions(+), 19 deletions(-) diff --git a/docs/proposals/C3-real-tokenizer-followup.md b/docs/proposals/C3-real-tokenizer-followup.md index b96a97e7..b7c07c4f 100644 --- a/docs/proposals/C3-real-tokenizer-followup.md +++ b/docs/proposals/C3-real-tokenizer-followup.md @@ -1,7 +1,7 @@ # Proposal C3-followup: 用真实 tokenizer 替换字符估算 > 来源：Roundtable 2026-04-20 surgical compactor 修复；C3 首轮决定先用保守常量 (2 chars/token) 落地 -> 状态：proposed（中间态已落地，真正解决方案待依赖就绪） +> 状态：implemented（2026-04-22：`js-tiktoken` / `cl100k_base` 已接入 `context-compactor.ts`） > 提出：Karpathy (T2) > Owner: TBD > 估计工时：1 天（换依赖）+ 半天（灰度验证） @@ -9,7 +9,7 @@ ## 问题陈述 -`src/main/agent-runtime/context-compactor.ts` 的 `estimateTokens` 现在用 `Math.ceil(chars / 2)` 做 token 估算，作为 120k 阈值的守门人。 +`src/main/agent-runtime/context-compactor.ts` 曾用 `Math.ceil(chars / 2)` 做 token 估算，作为 120k 阈值的守门人。2026-04-22 已改成 `js-tiktoken` 的 `cl100k_base` 本地编码器，并保留每条 chat message 的固定结构开销。 问题： 1. **中文 BPE 实际约 1.5-2 tokens/char**：2 chars/token 的常量对 CJK 大致可用，但 mixed 中英 / emoji / code block 的实际比例会漂移。 @@ -29,9 +29,9 @@ pi-ai 0.67.68 没 export tokenizer。Compactor 作为 `transformContext` hook ## 建议路径 -1. 先等 pi-ai 是否在后续 patch 里暴露 `countTokens(messages, model)`——如果暴露，零成本切过去。 -2. 若 6 周内上游没动作，引入 `tiktoken` + 模型 → encoding 映射表（Anthropic 用 `cl100k_base` 近似；Gemini / GPT 各自对应）。 -3. 保留 `estimateTokens(text: string): number` 的对外签名不变，只换内部实现——B3 配对修复和 hook 调用都无需动。 +1. 已接入 `js-tiktoken` + `cl100k_base`，保持 compactor 同步纯函数。 +2. pi-ai 后续若暴露 provider-native tokenizer，可在 `context-compactor.ts` 内替换编码器实现。 +3. hook 调用保持不变，`compactMessagesIfNeeded(messages)` 继续返回原有 `CompactionOutcome`。 ## 不做的事 diff --git a/package.json b/package.json index 19fc023e..c4239814 100644 --- a/package.json +++ b/package.json @@ -43,6 +43,7 @@ "electron-log": "5.4.3", "electron-updater": "6.8.3", "framer-motion": "12.36.0", + "js-tiktoken": "1.0.21", "react": "19.2.4", "react-dom": "19.2.4", "react-markdown": "^10.1.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 67fa7ab4..02a4aed1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -32,6 +32,9 @@ importers: framer-motion: specifier: 12.36.0 version: 12.36.0(react-dom@19.2.4(react@19.2.4))(react@19.2.4) + js-tiktoken: + specifier: 1.0.21 + version: 1.0.21 react: specifier: 19.2.4 version: 19.2.4 @@ -2944,6 +2947,9 @@ packages: resolution: {integrity: sha512-ekilCSN1jwRvIbgeg/57YFh8qQDNbwDb9xT/qu2DAHbFFZUicIl4ygVaAvzveMhMVr3LnpSKTNnwt8PoOfmKhQ==} hasBin: true + js-tiktoken@1.0.21: + resolution: {integrity: sha512-biOj/6M5qdgx5TKjDnFT1ymSpM5tbd3ylwDtrQvFQSu0Z7bBYko2dF+W/aUkXUPuk6IVpRxk/3Q2sHOzGlS36g==} + js-tokens@4.0.0: resolution: {integrity: sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==} @@ -7998,6 +8004,10 @@ snapshots: jiti@2.6.1: {} + js-tiktoken@1.0.21: + dependencies: + base64-js: 1.5.1 + js-tokens@4.0.0: {} js-yaml@4.1.1: diff --git a/src/main/agent-runtime/context-compactor.ts b/src/main/agent-runtime/context-compactor.ts index a7d04137..3162d14e 100644 --- a/src/main/agent-runtime/context-compactor.ts +++ b/src/main/agent-runtime/context-compactor.ts @@ -8,8 +8,8 @@ * throttle when the context approaches their window. * * This compactor keeps the hook pure and provider-free: no LLM - * round-trip, no mutation, no IO. It estimates token count using a - * 4-chars-per-token heuristic, keeps recent traffic verbatim, and + * round-trip, no mutation, no IO. It counts transcript text with a + * local tiktoken-compatible encoder, keeps recent traffic verbatim, and * collapses older messages into a single "context summary" user turn * that preserves the assistant's narrative plus tool-call headlines. * @@ -19,6 +19,7 @@ */ import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import { getEncoding } from "js-tiktoken"; /** Token-budget threshold that triggers compaction. */ const COMPACTION_TRIGGER_TOKENS = 120_000; @@ -28,17 +29,8 @@ const COMPACTION_TRIGGER_TOKENS = 120_000; * losing mid-stream state (tool-call pairs, user intent). */ const RECENT_MESSAGE_TAIL = 20; -/** - * Conservative tokens-per-char estimator. The former 4 chars/token value - * under-counts CJK traffic by roughly 2x — Chinese BPE tokenizers land - * closer to 1.5-2 tokens per character, so a 120k-token gate would - * silently trip only after the real context already blew past the - * window. 2 chars/token is a deliberate over-estimate for latin text in - * exchange for not under-triggering on Chinese agent runs. The proper - * fix is a real tokenizer — pi-ai 0.67.68 does not export one yet, so - * see docs/proposals/C3-real-tokenizer-followup.md for the upgrade path. - */ -const CHARS_PER_TOKEN = 2; +const CONTEXT_ENCODING = getEncoding("cl100k_base"); +const CHAT_MESSAGE_OVERHEAD_TOKENS = 4; /** Per-assistant-message summary cap to keep the compacted turn readable. */ const MAX_SUMMARY_CHARS_PER_MESSAGE = 160; const MAX_TOOL_RESULT_LINES = 12; @@ -150,7 +142,8 @@ function estimateTokens(messages: AgentMessage[]): number { function estimateMessageTokens(message: AgentMessage): number { const raw = extractTextForEstimate(message as MessageLike); - return Math.ceil(raw.length / CHARS_PER_TOKEN); + if (raw.length === 0) return CHAT_MESSAGE_OVERHEAD_TOKENS; + return CONTEXT_ENCODING.encode(raw).length + CHAT_MESSAGE_OVERHEAD_TOKENS; } function extractTextForEstimate(message: MessageLike): string { diff --git a/tests/unit/context-compactor-tool-results.test.ts b/tests/unit/context-compactor-tool-results.test.ts index 9c9e5b22..e49706f8 100644 --- a/tests/unit/context-compactor-tool-results.test.ts +++ b/tests/unit/context-compactor-tool-results.test.ts @@ -30,6 +30,38 @@ function makeLongTranscript(olderFacts: AgentMessage[]): AgentMessage[] { } describe("context compactor tool result facts", () => { + it("uses tokenizer counts so long CJK transcripts compact before provider overflow", () => { + const messages: AgentMessage[] = [ + { + role: "user", + content: [{ type: "text", text: "Start the Chinese market launch review." }], + timestamp: 1, + } as AgentMessage, + { + role: "assistant", + content: [{ type: "text", text: "I will keep the plan concise." }], + timestamp: 2, + } as AgentMessage, + { + role: "user", + content: [{ type: "text", text: "请完成市场复盘，保留所有客户反馈与行动项。 ".repeat(8_000) }], + timestamp: 3, + } as AgentMessage, + ]; + for (let i = 0; i < 42; i += 1) { + messages.push({ + role: i % 2 === 0 ? "user" : "assistant", + content: [{ type: "text", text: `recent ${i}` }], + timestamp: 100 + i, + } as AgentMessage); + } + + const outcome = compactMessagesIfNeeded(messages); + + expect(outcome.changed).toBe(true); + expect(outcome.telemetry?.estimatedTokens).toBeGreaterThanOrEqual(120_000); + }); + it("preserves task, approval, message, status, decision, and error facts from tool results", () => { const messages = makeLongTranscript([ {