feat: density-calibrated countTokens (Phase 2 of token calibration)

src/compress-tool.ts

@@ -59,21 +59,25 @@ async function handleCompress(args: CompressArgs, runtime: AcpRuntime, ctx: Exte
   if (ranges.length === 0) return "No ranges provided.";
   const { state, coreMessages } = await runtime.stateFor(ctx);
   const config = runtime.configFor(ctx);
-
-  const beforeTokens = estimateTokens(coreMessages, collectCoveredMessageIds(state));
+  // 显示层对齐（文档 §3.3）：beforeTokens 走密度校准口径，与 kernel 的
+  // countTokens（已乘 density）同口径，模型看到的数字接近真实占用。
+  const modelId = (ctx.model as { id?: string } | undefined)?.id ?? "default";
+  const density = runtime.density.densityFor(modelId);
+  const beforeTokens = Math.round(estimateTokens(coreMessages, collectCoveredMessageIds(state)) * density);
   const summaryMaxChars = args.summaryMaxChars;
   const topLevelTopic = args.topic;
 
   debug.event("compress-in", {
     sid: ctx.sessionManager.getSessionId(),
+    modelId,
+    density,
     ranges: ranges.length,
     spans: ranges.map((r) => ({ span: `${r.startId}..${r.endId}`, summaryLen: r.summary.length, summary: r.summary, topic: r.topic ?? topLevelTopic ?? null })),
     blocksBefore: state.blocks.length,
     activeBefore: state.blocks.filter((b) => b.active).length,
     beforeMsgCount: coreMessages.length,
     beforeTokens,
   });
-
   const applied = runtime.core.applyCompression({
     ranges: ranges.map((r) => ({ startRef: r.startId, endRef: r.endId, summary: r.summary, topic: r.topic ?? topLevelTopic, summaryMaxChars, compressCallId: toolCallId })),
     messages: coreMessages,

src/density.ts

@@ -0,0 +1,117 @@
+/**
+ * 累积锚点密度估计器（Phase 2 of token calibration）。
+ *
+ * 校准 ACP 对"消息字符 → token"的估算密度。kernel 的 T1 pending 用 chars/4
+ * 估算中文会低估 2-4 倍（见 docs/token-calibration-plan.md §1），本模块用
+ * provider 真实 usage（realTotal）与本地估算（estTotal）的**同窗口累积增量**
+ * 算出实时密度系数，注入 countTokens，让 pending/nudge 与实际占用对齐。
+ *
+ * 设计要点（文档 §3.2/§5 定稿）：
+ * - 累积锚点法而非 EMA：Δreal/Δest 同窗口，天然对齐无滞后（评审 D1/D2）
+ * - clamp [0.5, 2.5]：没有自然语言密度能超过 2.5 token/char（评审 D3）
+ * - 最小 Δest=50 门槛：微消息比值抖动（评审 D4）
+ * - ±20% 连续 2 轮确认才采纳：防单轮异常污染锚点（评审 C1）
+ * - 压缩后跳过一轮（postCompressionSkip）：Δest 为负 + provider usage 滞后（评审 D7/F1）
+ * - per-model 存储 + 模型切换重置（评审 D5/D6）
+ */
+import { defaultCountTokens } from "acp-kernel";
+
+export const DENSITY_MIN = 0.5;
+export const DENSITY_MAX = 2.5;
+export const MIN_DELTA_EST = 50;
+export const CONFIRM_RATIO = 0.2; // ±20% 确认带
+export const INITIAL_DENSITY = 1;
+
+interface Estimator {
+  density: number;
+  anchorReal: number | null;
+  anchorEst: number | null;
+  pendingDensity: number | null;
+  confirmCount: number;
+  postCompressionSkip: boolean;
+}
+
+export class DensityEstimator {
+  private models = new Map<string, Estimator>();
+
+  /** 重置指定模型（模型切换/会话开始时调用）。 */
+  resetModel(modelId: string): void {
+    this.models.delete(modelId);
+  }
+
+  /** 返回当前密度系数（未知模型返回初始 1）。 */
+  densityFor(modelId: string): number {
+    return this.models.get(modelId)?.density ?? INITIAL_DENSITY;
+  }
+
+  /**
+   * 每轮 context 事件调用。realTotal 为 provider 真实 usage（可空），
+   * estTotal 为本地估算总 token。postCompression 为压缩刚发生标志。
+   */
+  update(modelId: string, realTotal: number | null, estTotal: number, postCompression = false): void {
+    if (realTotal === null) return; // 无 provider usage，锚点冻结（§5.9）
+    let est = this.models.get(modelId);
+    if (!est) {
+      est = {
+        density: INITIAL_DENSITY,
+        anchorReal: null,
+        anchorEst: null,
+        pendingDensity: null,
+        confirmCount: 0,
+        postCompressionSkip: false,
+      };
+      this.models.set(modelId, est);
+    }
+
+    if (postCompression) {
+      // 压缩后第一轮跳过（D7/F1）：provider usage 滞后，Δest 可能为负
+      est.postCompressionSkip = true;
+      return;
+    }
+    if (est.postCompressionSkip) {
+      est.postCompressionSkip = false;
+      return;
+    }
+    if (est.anchorReal === null || est.anchorEst === null) {
+      // 首轮建立锚点，不产生样本
+      est.anchorReal = realTotal;
+      est.anchorEst = estTotal;
+      return;
+    }
+
+    const dReal = realTotal - est.anchorReal;
+    const dEst = estTotal - est.anchorEst;
+    if (dEst < MIN_DELTA_EST) return; // 增量太小或为负（压缩轮）跳过
+    // 每轮采样后都推进锚点（文档 §3.2）：instant 是相邻轮差值，单轮异常只污染该轮
+    est.anchorReal = realTotal;
+    est.anchorEst = estTotal;
+
+    const instant = clamp(dReal / dEst, DENSITY_MIN, DENSITY_MAX);
+    // C1 加固：连续 2 轮 ±20% 内才采纳，防单轮异常污染锚点
+    if (est.pendingDensity === null) {
+      est.pendingDensity = instant;
+      est.confirmCount = 1;
+    } else if (Math.abs(instant - est.pendingDensity) / est.pendingDensity <= CONFIRM_RATIO) {
+      est.confirmCount += 1;
+    } else {
+      est.pendingDensity = instant;
+      est.confirmCount = 1;
+    }
+    if (est.confirmCount >= 2) {
+      est.density = est.pendingDensity;
+      est.confirmCount = 0;
+      est.pendingDensity = null;
+    }
+  }
+
+  /** 注入器：估算文本 token = defaultCountTokens × density。 */
+  estimateWithDensity(modelId: string, text: string): number {
+    const d = this.densityFor(modelId);
+    if (d === 1) return defaultCountTokens(text); // 未校准时不引入浮点误差
+    return Math.round(defaultCountTokens(text) * d);
+  }
+}
+
+function clamp(v: number, lo: number, hi: number): number {
+  return v < lo ? lo : v > hi ? hi : v;
+}

src/index.ts

@@ -63,6 +63,12 @@ function wireSessionLifecycle(pi: ExtensionAPI, runtime: AcpRuntime): void {
   pi.on("session_start", async (_event, ctx) => {
     runtime.store.invalidate();
     runtime.clearNudgeTracking();
+    // 新会话重置该模型的密度校准（文档 §5.3：模型/窗口切换时重新收敛）
+    const modelId = (ctx.model as { id?: string } | undefined)?.id ?? "default";
+    runtime.density.resetModel(modelId);
+    // Load user config (~/.pi/acp.json + project .pi/acp.json) and apply it so
+    // debug/terminalNudge/autoUpdate/modelContextLimit are runtime-configurable
+    // without env vars or reinstalling.
     const sid = ctx.sessionManager.getSessionId();
     logInfo("session", { event: "start", sid, cwd: ctx.cwd, debug: runtime.adapter.debug ?? null, version: typeof CURRENT_VERSION !== "undefined" ? CURRENT_VERSION : null });
     try {
@@ -104,6 +110,9 @@ function wireContextTransform(pi: ExtensionAPI, runtime: AcpRuntime): void {
     const sid = ctx.sessionManager.getSessionId();
     const release = await runtime.acquireLock(sid);
     try {
+      // 每轮绑定 countTokens 使用的模型（密度校准按 model 隔离）
+      const modelId = (ctx.model as { id?: string } | undefined)?.id ?? "default";
+      runtime.setCountModel(modelId);
       const { state, coreMessages, entries } = await runtime.stateFor(ctx, event.messages);
       const config = runtime.configFor(ctx);
       const coveredIds = collectCoveredMessageIds(state);
@@ -118,6 +127,8 @@ function wireContextTransform(pi: ExtensionAPI, runtime: AcpRuntime): void {
 
       debug.event("context-in", {
         sid,
+        modelId,
+        density: runtime.density.densityFor(modelId),
         eventMsgs: event.messages?.length ?? 0,
         entries: entries.length,
         coreMsgs: coreMessages.length,
@@ -132,6 +143,12 @@ function wireContextTransform(pi: ExtensionAPI, runtime: AcpRuntime): void {
 
       const turn = runtime.core.processTurn({ messages: coreMessages, state, config, tokenCount });
       await runtime.save(turn.state, ctx);
+      // 更新密度校准（Phase 2）：processTurn 后调用，countTokens 用上一轮 density（1 轮延迟可忽略）。
+      // postCompression = 本次新增了 active block（模型刚压缩过）。
+      const postCompression = turn.state.blocks.some(
+        (b) => b.active && !state.blocks.some((o) => o.blockId === b.blockId)
+      );
+      runtime.density.update(modelId, realUsage?.tokens ?? null, estimated, postCompression);
 
       logInfo("turn", {
         sid,
@@ -145,8 +162,9 @@ function wireContextTransform(pi: ExtensionAPI, runtime: AcpRuntime): void {
         blocks: turn.state.blocks.length,
         activeBlocks: turn.state.blocks.filter((b) => b.active).length,
       });
-
       debug.event("processTurn", {
+        modelId,
+        density: runtime.density.densityFor(modelId),
         outMsgs: turn.messages.length,
         summaryMsgs: turn.messages.filter((m) => m.id.startsWith("acp_summary")).length,
         prunedMsgs: coreMessages.length - turn.messages.length + turn.messages.filter((m) => m.id.startsWith("acp_summary")).length,

src/runtime.ts

@@ -1,16 +1,15 @@
 import type { ExtensionContext, SessionEntry, SessionMessageEntry } from "@earendil-works/pi-coding-agent";
 import {
   createCore,
-  defaultCountTokens,
   type CompressionCore,
   type CompressionState,
   type Config,
 } from "acp-kernel";
 import { resolveConfig, type AdapterConfig } from "./config.js";
+import { DensityEstimator } from "./density.js";
 import { entriesToCoreMessages } from "./messages.js";
 import { SessionStateStore } from "./state.js";
 import { logInfo, logWarn } from "./log.js";
-
 // pi exposes `sessionManager.buildContextEntries()`; omp (oh-my-pi) only has
 // `getBranch()`. Both return chronological SessionEntry[]; feature-detect so the
 // adapter runs under either host (omp's runner silently swallows the TypeError).
@@ -39,6 +38,9 @@ export function isPiHost(sm: ExtensionContext["sessionManager"]): boolean {
 export interface AcpRuntime {
   core: CompressionCore;
   store: SessionStateStore;
+  density: DensityEstimator;
+  /** 设置 countTokens 闭包使用的 modelId（每轮 context 事件调用）。 */
+  setCountModel(modelId: string): void;
   adapter: AdapterConfig;
   setAdapter(adapter: AdapterConfig): void;
   /** Record that a nudge was already shown for the turn keyed by last user msg
@@ -55,7 +57,6 @@ export interface AcpRuntime {
   save(state: CompressionState, ctx: ExtensionContext): Promise<void>;
   acquireLock(sid: string): Promise<() => void>;
 }
-
 // omp fires the context event before the current user message is persisted to
 // the session branch (its agent-loop emits message_end only after
 // prepareProviderCall → transformContext), so getBranch() lags one message
@@ -109,9 +110,13 @@ function sameMessage(a: AgentMessage, b: AgentMessage): boolean {
     return a === b;
   }
 }
-
 export function createRuntime(adapter: AdapterConfig): AcpRuntime {
-  const core = createCore({ countTokens: defaultCountTokens });
+  const density = new DensityEstimator();
+  let countModelId = "default";
+  const core = createCore({
+    // 密度校准版 countTokens（Phase 2）：默认回落 defaultCountTokens（density=1）
+    countTokens: (text) => density.estimateWithDensity(countModelId, text),
+  });
   const store = new SessionStateStore();
   const locks = new Map<string, Promise<void>>();
   let adapterRef = adapter;
@@ -165,5 +170,5 @@ export function createRuntime(adapter: AdapterConfig): AcpRuntime {
     await store.save(state, sm.getSessionFile() ?? undefined, sm.getSessionId());
   }
 
-  return { core, store, get adapter() { return adapterRef; }, setAdapter: (a) => { adapterRef = a; }, markNudgeShown: (k) => { nudgeShownTurns.add(k); }, nudgeShownFor: (k) => nudgeShownTurns.has(k), clearNudgeTracking: () => { nudgeShownTurns.clear(); }, liveContextLimit, configFor, stateFor, save, acquireLock };
+  return { core, store, density, setCountModel: (m) => { countModelId = m; }, get adapter() { return adapterRef; }, setAdapter: (a) => { adapterRef = a; }, markNudgeShown: (k) => { nudgeShownTurns.add(k); }, nudgeShownFor: (k) => nudgeShownTurns.has(k), clearNudgeTracking: () => { nudgeShownTurns.clear(); }, liveContextLimit, configFor, stateFor, save, acquireLock };
 }

src/state.ts

@@ -66,6 +66,7 @@
   return {
     blocks: parsed.blocks ?? fresh.blocks,
     messageRefs: parsed.messageRefs ?? fresh.messageRefs,
+    tokenSnapshot: parsed.tokenSnapshot ?? fresh.tokenSnapshot,
     nudge: { ...fresh.nudge, ...(parsed.nudge ?? {}) },
     stats: { ...fresh.stats, ...(parsed.stats ?? {}) },
     nextBlockId: parsed.nextBlockId ?? fresh.nextBlockId,

tests/compress-tool.test.ts

@@ -0,0 +1,109 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { rm } from "node:fs/promises";
+import { createAcpExtension } from "../src/index.js";
+
+// ─── helpers (mirror decompress-tool.test.ts) ──────────────────────────────
+
+function captureApi() {
+  const handlers = new Map<string, ((event: any, ctx: any) => any)[]>();
+  const api = {
+    on(event: string, handler: (e: any, ctx: any) => any) {
+      const list = handlers.get(event) ?? [];
+      list.push(handler);
+      handlers.set(event, list);
+    },
+    tools: [] as any[],
+    commands: new Map<string, any>(),
+    registerTool(tool: any) { this.tools.push(tool); },
+    registerCommand(name: string, options: any) { this.commands.set(name, options); },
+  };
+  return { api, handlers };
+}
+
+function userMsg(id: string, text: string) {
+  return { type: "message", id, parentId: null, timestamp: "", message: { role: "user", content: text, timestamp: Date.now() } };
+}
+
+function fakeCtx(entries: any[], stateFile: string) {
+  let usage: { tokens: number; percent: number } | null = null;
+  return {
+    mode: "rpc",
+    hasUI: false,
+    ui: { notify: () => {}, confirm: async () => true, select: async () => undefined, input: async () => "", setStatus: () => {} },
+    model: { contextWindow: 200_000, id: "test-model" },
+    getContextUsage: () => usage,
+    __setUsage(t: number) { usage = { tokens: t, percent: t / 200_000 }; },
+    sessionManager: {
+      buildContextEntries: () => entries,
+      getSessionId: () => "test-session",
+      getSessionFile: () => stateFile,
+    },
+  };
+}
+
+const ZH = "中".repeat(300);   // 300 CJK tokens
+const ZH2 = "中".repeat(150);  // 150 CJK tokens
+
+function beforeTokensFrom(out: string): number {
+  const m = /▣ ACP \| (\d+) →/.exec(out);
+  assert.ok(m, `no beforeTokens in output: ${out}`);
+  return Number(m![1]!);
+}
+
+async function runContextRound(handlers: Map<string, any[]>, ctx: any) {
+  await handlers.get("context")![0]!({ type: "context", messages: [] }, ctx);
+}
+
+// ─── tests ─────────────────────────────────────────────────────────────────
+
+test("compress beforeTokens at density=1 is uncalibrated estimateTokens", async () => {
+  const { api, handlers } = captureApi();
+  createAcpExtension({ modelContextLimit: 200_000 })(api as any);
+  const stateFile = "/tmp/pai-acp-compress-density-a.session.json";
+  await rm(`${stateFile}.acp.json`, { force: true });
+  const entries = [userMsg("e1", "hello world"), userMsg("e2", ZH)];
+  const ctx = fakeCtx(entries, stateFile);
+  ctx.__setUsage(100_000);
+  await runContextRound(handlers, ctx); // 只锚点，无样本 → density=1
+
+  const compressTool = api.tools.find((t: any) => t.name === "compress")!;
+  const out = await compressTool.execute(
+    "tc1",
+    { content: [{ startId: "m00001", endId: "m00001", summary: "compressed" }] },
+    undefined, undefined, ctx,
+  );
+  const text = typeof out === "string" ? out : out.content?.[0]?.text ?? String(out);
+  assert.equal(beforeTokensFrom(text), 303); // 3 (hello world) + 300 (ZH)
+});
+
+test("compress beforeTokens scales with calibrated density (Phase 2)", async () => {
+  const { api, handlers } = captureApi();
+  createAcpExtension({ modelContextLimit: 200_000 })(api as any);
+  const stateFile = "/tmp/pai-acp-compress-density-b.session.json";
+  await rm(`${stateFile}.acp.json`, { force: true });
+  const entries = [userMsg("e1", "hello world"), userMsg("e2", ZH)];
+  const ctx = fakeCtx(entries, stateFile);
+
+  // 轮1：锚点（real=100k, est=303）
+  ctx.__setUsage(100_000);
+  await runContextRound(handlers, ctx);
+  // 轮2：Δreal=240 / Δest=150 → instant 1.6, pending（追加 e3）
+  entries.push(userMsg("e3", ZH2));
+  ctx.__setUsage(100_240);
+  await runContextRound(handlers, ctx);
+  // 轮3：instant 1.6 → 确认采纳（追加 e4）
+  entries.push(userMsg("e4", ZH2));
+  ctx.__setUsage(100_480);
+  await runContextRound(handlers, ctx);
+
+  const compressTool = api.tools.find((t: any) => t.name === "compress")!;
+  const out = await compressTool.execute(
+    "tc1",
+    { content: [{ startId: "m00001", endId: "m00001", summary: "compressed" }] },
+    undefined, undefined, ctx,
+  );
+  const text = typeof out === "string" ? out : out.content?.[0]?.text ?? String(out);
+  // estTotal = 3+300+150+150 = 603；×1.6 = 964.8 → 965
+  assert.equal(beforeTokensFrom(text), 965);
+});