@vellumai/assistant 0.6.4 → 0.6.5

package/src/context/estimator-calibration.ts

@@ -0,0 +1,136 @@
+/**
+ * Self-calibrating correction ratios for the token estimator.
+ *
+ * Every successful provider call returns ground-truth `usage.inputTokens`.
+ * By comparing that to the pre-send estimate, we can maintain a per-model
+ * EWMA correction ratio that multiplies future estimates — catching
+ * miscalibration proactively instead of waiting for provider overflow
+ * errors to reverse-engineer a token count from the error message.
+ *
+ * State is process-local; correction resets on restart. That is acceptable
+ * because the ratio converges quickly (EWMA alpha = 0.2, ~5 samples).
+ */
+
+interface CalibrationState {
+  /** EWMA of (actual / estimated) — multiply estimates by this to correct. */
+  ratio: number;
+  /** Total samples recorded, for observability. */
+  sampleCount: number;
+}
+
+const CALIBRATIONS: Map<string, CalibrationState> = new Map();
+
+/** Fast-adapting EWMA — converges to a steady state in ~5 consistent samples. */
+const EWMA_ALPHA = 0.2;
+
+/**
+ * Below this magnitude both numbers are noisy — a 500-token prompt with a
+ * 600-token usage is within normal overhead fluctuation and should not
+ * move the correction ratio.
+ */
+const MIN_SAMPLE_MAGNITUDE = 500;
+
+/**
+ * Outlier guard — discard samples where the ratio is more than 3x off in
+ * either direction. A ratio that extreme is almost always a bug (wrong
+ * estimate site, wrong provider, double-counting) rather than a genuine
+ * estimation error the calibrator should learn from.
+ */
+const MIN_ACCEPTABLE_RATIO = 1 / 3;
+const MAX_ACCEPTABLE_RATIO = 3;
+
+function key(provider: string, model: string): string {
+  return `${provider}::${model}`;
+}
+
+/** Apply a single EWMA update at an exact (provider, model) key. */
+function applyEwmaUpdate(provider: string, model: string, ratio: number): void {
+  const k = key(provider, model);
+  const prev = CALIBRATIONS.get(k);
+  const next: CalibrationState = prev
+    ? {
+        ratio: prev.ratio + EWMA_ALPHA * (ratio - prev.ratio),
+        sampleCount: prev.sampleCount + 1,
+      }
+    : { ratio, sampleCount: 1 };
+  CALIBRATIONS.set(k, next);
+}
+
+/**
+ * Fold a new (estimated, actual) observation into the EWMA ratio for this
+ * (provider, model). No-op when either number is too small to be reliable,
+ * or when the ratio is an outlier.
+ *
+ * When `model` is non-empty, the same observation is also folded into the
+ * per-provider aggregate key `(provider, "")` so callers that cannot resolve
+ * a specific model (early-init paths, provider-only estimate sites) still
+ * pick up a reasonable correction via {@link getCorrection}.
+ */
+export function recordEstimate(
+  provider: string,
+  model: string,
+  estimated: number,
+  actual: number,
+): void {
+  if (estimated < MIN_SAMPLE_MAGNITUDE || actual < MIN_SAMPLE_MAGNITUDE) return;
+  const ratio = actual / estimated;
+  if (ratio < MIN_ACCEPTABLE_RATIO || ratio > MAX_ACCEPTABLE_RATIO) return;
+
+  applyEwmaUpdate(provider, model, ratio);
+
+  // Also fold into the per-provider aggregate so callers without a modelId
+  // fall back to a meaningful rolling correction instead of the 1.0 default.
+  // Skip when the caller already passed an empty model (avoids double-counting).
+  if (model.length > 0) {
+    applyEwmaUpdate(provider, "", ratio);
+  }
+}
+
+/**
+ * Correction factor to multiply a raw estimate by. Defaults to 1.0 for any
+ * unseen (provider, model) tuple, so first-call behavior is unchanged.
+ *
+ * When the model-specific key has no recorded samples, falls back to the
+ * per-provider aggregate `(provider, "")`. This covers model-alias
+ * mismatches where the configured model string (e.g. `claude-opus-4-6`)
+ * differs from the model the provider echoes back in its response
+ * (e.g. `claude-opus-4-6-20250514`).
+ */
+export function getCorrection(provider: string, model: string): number {
+  const specific = CALIBRATIONS.get(key(provider, model));
+  if (specific) return specific.ratio;
+  if (model.length > 0) {
+    return CALIBRATIONS.get(key(provider, ""))?.ratio ?? 1.0;
+  }
+  return 1.0;
+}
+
+/** Test helper — clears all calibration state. */
+export function resetCalibrations(): void {
+  CALIBRATIONS.clear();
+}
+
+/** Observability — list current calibrations for logging/debugging. */
+export function getCalibrationSnapshot(): Array<{
+  provider: string;
+  model: string;
+  ratio: number;
+  samples: number;
+}> {
+  const out: Array<{
+    provider: string;
+    model: string;
+    ratio: number;
+    samples: number;
+  }> = [];
+  for (const [k, state] of CALIBRATIONS) {
+    const [provider, model] = k.split("::", 2) as [string, string];
+    out.push({
+      provider,
+      model,
+      ratio: state.ratio,
+      samples: state.sampleCount,
+    });
+  }
+  return out;
+}

package/src/context/microcompact.ts

@@ -0,0 +1,443 @@
+import type {
+  ContentBlock,
+  Message,
+  ToolResultContent,
+} from "../providers/types.js";
+import { estimateTextTokens } from "./token-estimator.js";
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+
+/** Default number of recent user-assistant exchanges to keep untouched. */
+const DEFAULT_PROTECT_RECENT_TURNS = 4;
+
+/** Default minimum reclaimable tokens required before we commit a pass. */
+const DEFAULT_MIN_GAIN_TOKENS = 2_000;
+
+/**
+ * Default list of tool names whose results are never wholesale-cleared.
+ * Matches opencode's `PRUNE_PROTECTED_TOOLS` and Claude Code's equivalent —
+ * protects subagent outputs and curated skill results, which are expensive
+ * to regenerate and usually load-bearing for the rest of the conversation.
+ *
+ * Note: ax-tree stripping still applies to protected tool results. Only the
+ * full-body replacement is skipped.
+ */
+const DEFAULT_PROTECTED_TOOLS: readonly string[] = [
+  "Task",
+  "subagent",
+  "skill",
+];
+
+/** Replacement body for cleared tool-result content. */
+const CLEARED_TOOL_RESULT_TEXT = "[Old tool result content cleared]";
+
+/** Replacement text for stubbed image blocks. */
+const CLEARED_IMAGE_TEXT = "[image omitted]";
+
+/** Replacement text for stubbed file blocks. */
+const CLEARED_FILE_TEXT = "[file omitted]";
+
+/**
+ * Regex that matches `<ax-tree>...</ax-tree>` blocks (non-greedy).
+ * Kept in sync with `AX_TREE_PATTERN` in `assistant/src/agent/loop.ts`; this
+ * module subsumes that function's responsibility for stale tool results.
+ */
+const AX_TREE_PATTERN = /<ax-tree>[\s\S]*?<\/ax-tree>/g;
+
+/** Placeholder inserted in place of a stripped ax-tree block. */
+const AX_TREE_PLACEHOLDER = "<ax_tree_omitted />";
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+export interface MicrocompactOptions {
+  /** Preserve the last N user-assistant exchanges verbatim. Default 4. */
+  protectRecentTurns?: number;
+  /** Tool names whose results are never microcompacted (e.g. "Task", sub-agents). */
+  protectedTools?: string[];
+  /** Minimum reclaimable tokens required to bother — skip no-op passes. Default 2000. */
+  minGainTokens?: number;
+}
+
+export interface MicrocompactResult {
+  messages: Message[];
+  reclaimedTokens: number;
+  clearedToolResults: number;
+  /** Count of image + file blocks that were replaced with text stubs. */
+  clearedMedia: number;
+}
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+/**
+ * Deterministically compact stale content in the message history.
+ *
+ * Walks `messages` from newest to oldest. The most recent
+ * `protectRecentTurns` user-assistant exchanges are left untouched. In the
+ * older region:
+ *
+ *  - Each `tool_result` whose owning tool is NOT in `protectedTools` has its
+ *    `content` replaced with a short placeholder.
+ *  - Every `image` / `file` block is replaced with a text stub.
+ *  - `<ax-tree>...</ax-tree>` blocks inside any tool_result (including
+ *    protected ones) are collapsed to a placeholder, subsuming the old
+ *    `compactAxTreeHistory` in `assistant/src/agent/loop.ts`.
+ *
+ * The `tool_use` / `tool_result` block structure is preserved — we only
+ * mutate block bodies, never the block types or their pairing — so provider
+ * serialization remains valid after compaction.
+ *
+ * If the estimated savings (`reclaimedTokens`) is below `minGainTokens`, the
+ * original `messages` reference is returned unchanged. The pass is
+ * idempotent: re-invoking on a previously compacted history produces zero
+ * incremental reclaim.
+ */
+export function microcompact(
+  messages: Message[],
+  options?: MicrocompactOptions,
+): MicrocompactResult {
+  const protectRecentTurns =
+    options?.protectRecentTurns ?? DEFAULT_PROTECT_RECENT_TURNS;
+  const minGainTokens = options?.minGainTokens ?? DEFAULT_MIN_GAIN_TOKENS;
+  const protectedToolSet = new Set(
+    options?.protectedTools ?? DEFAULT_PROTECTED_TOOLS,
+  );
+
+  if (messages.length === 0) {
+    return {
+      messages,
+      reclaimedTokens: 0,
+      clearedToolResults: 0,
+      clearedMedia: 0,
+    };
+  }
+
+  // Build a lookup of tool_use_id -> tool name by scanning every assistant
+  // message. Tool_use_ids are globally unique within a conversation so a
+  // single map is valid for the whole history.
+  const toolNameById = buildToolNameLookup(messages);
+
+  // Determine the index up to which messages are "protected" (left alone).
+  // Anything at index >= firstProtectedIdx is in the protected tail; the
+  // older region (indices < firstProtectedIdx) is where we clear.
+  const firstProtectedIdx = computeProtectedBoundary(
+    messages,
+    protectRecentTurns,
+  );
+
+  if (firstProtectedIdx <= 0) {
+    // Every message is protected — nothing to clear.
+    return {
+      messages,
+      reclaimedTokens: 0,
+      clearedToolResults: 0,
+      clearedMedia: 0,
+    };
+  }
+
+  let reclaimedTokens = 0;
+  let clearedToolResults = 0;
+  let clearedMedia = 0;
+  let anyChange = false;
+
+  const nextMessages: Message[] = messages.map((msg, idx) => {
+    if (idx >= firstProtectedIdx) return msg;
+
+    let changed = false;
+    const nextContent: ContentBlock[] = msg.content.map((block) => {
+      // guard:allow-tool-result-only — compaction here operates on locally-
+      // executed `tool_result` bodies (string `.content`, possible
+      // `.contentBlocks` media). `web_search_tool_result` has an opaque
+      // encrypted content shape and is never microcompacted; it's treated as
+      // a tool-response only in `isToolResultOnlyUserMessage` above.
+      if (block.type === "tool_result") {
+        const tr = block as ToolResultContent;
+        const toolName = toolNameById.get(tr.tool_use_id);
+        const isProtected = toolName != null && protectedToolSet.has(toolName);
+
+        const { replacement, tokensSaved, didChange, cleared } =
+          compactToolResult(tr, isProtected);
+        if (didChange) {
+          reclaimedTokens += tokensSaved;
+          if (cleared) clearedToolResults += 1;
+          changed = true;
+          return replacement;
+        }
+        return block;
+      }
+
+      if (block.type === "image") {
+        const saved = estimateImageReclaim(block);
+        // Only count a reclaim if stubbing actually shrinks the block.
+        // The stub is always valid, but avoid double-counting no-op cases.
+        reclaimedTokens += saved;
+        clearedMedia += 1;
+        changed = true;
+        return { type: "text", text: CLEARED_IMAGE_TEXT };
+      }
+
+      if (block.type === "file") {
+        const saved = estimateFileReclaim(block);
+        reclaimedTokens += saved;
+        clearedMedia += 1;
+        changed = true;
+        return { type: "text", text: CLEARED_FILE_TEXT };
+      }
+
+      return block;
+    });
+
+    if (!changed) return msg;
+    anyChange = true;
+    return { ...msg, content: nextContent };
+  });
+
+  if (!anyChange || reclaimedTokens < minGainTokens) {
+    return {
+      messages,
+      reclaimedTokens: 0,
+      clearedToolResults: 0,
+      clearedMedia: 0,
+    };
+  }
+
+  return {
+    messages: nextMessages,
+    reclaimedTokens,
+    clearedToolResults,
+    clearedMedia,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Internals
+// ---------------------------------------------------------------------------
+
+function buildToolNameLookup(messages: Message[]): Map<string, string> {
+  const map = new Map<string, string>();
+  for (const msg of messages) {
+    if (msg.role !== "assistant") continue;
+    for (const block of msg.content) {
+      if (block.type === "tool_use") {
+        map.set(block.id, block.name);
+      }
+    }
+  }
+  return map;
+}
+
+/**
+ * Return the index of the first message that is inside the "protected tail".
+ *
+ * A user-assistant exchange begins at a user message that carries real user
+ * content (i.e. not a tool_result-only follow-up, which is how the provider
+ * encodes the tool response turn). We walk newest-to-oldest and, once we've
+ * seen `protectRecentTurns` such user-turn starts, return the index of the
+ * oldest such start. Everything at or after that index is protected.
+ *
+ * If the history has fewer than `protectRecentTurns` user turns, protect the
+ * entire history by returning 0.
+ */
+function computeProtectedBoundary(
+  messages: Message[],
+  protectRecentTurns: number,
+): number {
+  if (protectRecentTurns <= 0) return messages.length;
+
+  let seen = 0;
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const msg = messages[i];
+    if (msg.role !== "user") continue;
+    if (isToolResultOnlyUserMessage(msg)) continue;
+    seen += 1;
+    if (seen >= protectRecentTurns) {
+      return i;
+    }
+  }
+  // Fewer user turns than the protection budget — protect everything.
+  return 0;
+}
+
+/**
+ * A user message that contains only tool-response blocks (or system-injected
+ * metadata) represents the response side of a tool call rather than a fresh
+ * user turn.
+ *
+ * Qualifying blocks:
+ *  - `tool_result` — locally-executed tool responses.
+ *  - `web_search_tool_result` — server-side web-search responses. Same
+ *    semantic as `tool_result` (paired with a prior `server_tool_use`), just
+ *    a distinct discriminant. Missing this variant would cause these
+ *    messages to masquerade as real user turns and eat `protectRecentTurns`
+ *    budget on tool churn.
+ *  - `text` blocks whose body is wholly inside `<system_notice>...</system_notice>`.
+ *    Those are system-injected reminders/progress checks, not user-authored.
+ */
+function isToolResultOnlyUserMessage(message: Message): boolean {
+  if (message.content.length === 0) return false;
+  return message.content.every(isToolResponseOrSystemNoticeBlock);
+}
+
+function isToolResponseOrSystemNoticeBlock(block: ContentBlock): boolean {
+  if (block.type === "tool_result" || block.type === "web_search_tool_result") {
+    return true;
+  }
+  if (block.type === "text") {
+    const text = block.text;
+    return (
+      text.startsWith("<system_notice>") && text.endsWith("</system_notice>")
+    );
+  }
+  return false;
+}
+
+/**
+ * Compact a single tool_result block in the stripped region.
+ *
+ *  - If `isProtected` is true, we keep the body but strip `<ax-tree>` blocks
+ *    and drop media (image/file) entries from `contentBlocks`, preserving
+ *    any text entries so meaningful tool output isn't silently removed.
+ *  - Otherwise we replace the body with a short placeholder and drop
+ *    `contentBlocks` entirely.
+ *
+ * `tokensSaved` is the delta between the original and the replacement as
+ * measured by `estimateTextTokens`, floored at 0. `cleared` is true only when
+ * we fully replaced the body (i.e. counted against `clearedToolResults`).
+ */
+function compactToolResult(
+  block: ToolResultContent,
+  isProtected: boolean,
+): {
+  replacement: ToolResultContent;
+  tokensSaved: number;
+  didChange: boolean;
+  cleared: boolean;
+} {
+  const originalContent = block.content;
+  const originalContentTokens = estimateTextTokens(originalContent);
+
+  // Rich contentBlocks may include text (meaningful tool output) alongside
+  // media (images/files). For UNPROTECTED results we drop the entire array
+  // (the body is being wholesale-replaced). For PROTECTED results we keep
+  // text entries but strip media — images aren't the expensive part of a
+  // subagent result and are rarely load-bearing once the turn is stale, but
+  // text entries can carry real content we must not silently erase.
+  const originalContentBlocks = block.contentBlocks;
+  const hadContentBlocks =
+    originalContentBlocks != null && originalContentBlocks.length > 0;
+
+  let preservedContentBlocks: ContentBlock[] | undefined;
+  let droppedBlocksTokens = 0;
+  if (hadContentBlocks) {
+    if (isProtected) {
+      const kept: ContentBlock[] = [];
+      for (const cb of originalContentBlocks) {
+        if (cb.type === "text") {
+          kept.push(cb);
+        } else {
+          droppedBlocksTokens += estimateContentBlockTokenCost(cb);
+        }
+      }
+      preservedContentBlocks = kept;
+    } else {
+      for (const cb of originalContentBlocks) {
+        droppedBlocksTokens += estimateContentBlockTokenCost(cb);
+      }
+      preservedContentBlocks = undefined;
+    }
+  }
+
+  let newContent: string;
+  let cleared: boolean;
+  if (isProtected) {
+    // Strip ax-tree blocks but keep the rest of the text body.
+    newContent = originalContent.replace(AX_TREE_PATTERN, AX_TREE_PLACEHOLDER);
+    cleared = false;
+  } else {
+    newContent = CLEARED_TOOL_RESULT_TEXT;
+    cleared = true;
+  }
+
+  const newContentTokens = estimateTextTokens(newContent);
+  const bodyTokensSaved = originalContentTokens - newContentTokens;
+
+  const bodyChanged = newContent !== originalContent;
+  // A protected result with only text contentBlocks would be a no-op — nothing
+  // meaningful would be dropped. Only count contentBlocks as a "change" when
+  // we actually drop at least one entry.
+  const contentBlocksChanged = hadContentBlocks && droppedBlocksTokens > 0;
+  const didChange = bodyChanged || contentBlocksChanged;
+
+  if (!didChange) {
+    return {
+      replacement: block,
+      tokensSaved: 0,
+      didChange: false,
+      cleared: false,
+    };
+  }
+
+  const tokensSaved = Math.max(0, bodyTokensSaved + droppedBlocksTokens);
+
+  const replacement: ToolResultContent = {
+    type: "tool_result",
+    tool_use_id: block.tool_use_id,
+    content: newContent,
+  };
+  if (block.is_error != null) {
+    replacement.is_error = block.is_error;
+  }
+  if (preservedContentBlocks != null && preservedContentBlocks.length > 0) {
+    replacement.contentBlocks = preservedContentBlocks;
+  }
+
+  return {
+    replacement,
+    tokensSaved,
+    didChange,
+    cleared: cleared && bodyChanged,
+  };
+}
+
+function estimateImageReclaim(
+  block: Extract<ContentBlock, { type: "image" }>,
+): number {
+  // Base64 payloads are the expensive part of image blocks for our naive
+  // char/4 heuristic. The Anthropic-aware estimator in `token-estimator.ts`
+  // charges per-pixel — we deliberately use the simpler heuristic here so
+  // the reclaim number is a conservative lower bound on actual savings.
+  const payloadTokens = estimateTextTokens(block.source.data);
+  const stubTokens = estimateTextTokens(CLEARED_IMAGE_TEXT);
+  return Math.max(0, payloadTokens - stubTokens);
+}
+
+function estimateFileReclaim(
+  block: Extract<ContentBlock, { type: "file" }>,
+): number {
+  const payloadTokens =
+    estimateTextTokens(block.source.data) +
+    estimateTextTokens(block.extracted_text ?? "");
+  const stubTokens = estimateTextTokens(CLEARED_FILE_TEXT);
+  return Math.max(0, payloadTokens - stubTokens);
+}
+
+function estimateContentBlockTokenCost(block: ContentBlock): number {
+  switch (block.type) {
+    case "text":
+      return estimateTextTokens(block.text);
+    case "image":
+      return estimateTextTokens(block.source.data);
+    case "file":
+      return (
+        estimateTextTokens(block.source.data) +
+        estimateTextTokens(block.extracted_text ?? "")
+      );
+    default:
+      return 0;
+  }
+}

package/src/context/prompts/compact.md

@@ -0,0 +1,12 @@
+You compress long assistant conversations into durable working memory.
+Focus on actionable state, not prose.
+Preserve concrete facts: goals, constraints, decisions, pending questions, file paths, commands, errors, and TODOs.
+Remove repetition and stale details that were superseded.
+Thread anchors: when a "Retained Thread References" section is present, each listed reply cites its parent via `→ Mxxxxxx`. If that parent appears in the Transcript, preserve its text verbatim (reactions may be aggregated as "N users reacted"). Omit when the section is absent.
+Return concise markdown using these section headers exactly:
+## Goals
+## Constraints
+## Decisions
+## Open Conversations
+## Key Artifacts
+## Recent Progress

package/src/context/token-estimator.ts

@@ -1,10 +1,27 @@
 import type {
   ContentBlock,
   Message,
+  Provider,
   ToolDefinition,
 } from "../providers/types.js";
+import { getCorrection } from "./estimator-calibration.js";
 import { parseImageDimensions } from "./image-dimensions.js";
 
+/**
+ * Canonical provider key used for calibration lookups and updates. Wrapper
+ * providers (e.g. OpenRouter routing `anthropic/*` traffic to the Messages
+ * API) set `tokenEstimationProvider` to the upstream provider name so the
+ * calibration key matches the one used when the provider actually produces
+ * the response. Falls back to `name` when the wrapper hint is unset.
+ *
+ * Every caller that records a sample or applies a correction must use this
+ * helper — otherwise wrapper-provider data is scattered across mismatched
+ * keys and the calibration becomes a no-op.
+ */
+export function getCalibrationProviderKey(provider: Provider): string {
+  return provider.tokenEstimationProvider ?? provider.name;
+}
+
 const CHARS_PER_TOKEN = 4;
 const MESSAGE_OVERHEAD_TOKENS = 4;
 const TEXT_BLOCK_OVERHEAD_TOKENS = 2;
@@ -103,7 +120,9 @@ function estimateAnthropicImageTokens(width: number, height: number): number {
     scaledHeight = Math.round(scaledHeight * mpScale);
   }
 
-  return Math.ceil(scaledWidth * scaledHeight * ANTHROPIC_IMAGE_TOKENS_PER_PIXEL);
+  return Math.ceil(
+    scaledWidth * scaledHeight * ANTHROPIC_IMAGE_TOKENS_PER_PIXEL,
+  );
 }
 
 function estimateImageTokens(
@@ -139,13 +158,28 @@ export function estimateContentBlockTokens(
         estimateTextTokens(stableJson(block.input))
       );
     case "tool_result": {
+      // Mirror the Anthropic serializer in providers/anthropic/client.ts
+      // (toAnthropicBlockSafe): block.content is always sent as the first
+      // text part, and contentBlocks are appended — but only `image` and
+      // `text` sub-blocks survive, and `image` is filtered out when
+      // is_error is true. Counting every contentBlocks entry regardless
+      // of type overestimates the wire size and can trigger spurious
+      // compaction on conversations that carry e.g. thinking sub-blocks.
+      // OpenAI and Gemini forward error-result images normally, so the
+      // is_error image drop is Anthropic-specific.
+      const anthropicDropsErrorImage =
+        options?.providerName === "anthropic" && block.is_error === true;
       let tokens =
         TOOL_BLOCK_OVERHEAD_TOKENS +
         estimateTextTokens(block.tool_use_id) +
         estimateTextTokens(block.content);
       if (block.contentBlocks) {
         for (const cb of block.contentBlocks) {
-          tokens += estimateContentBlockTokens(cb, options);
+          if (cb.type === "text") {
+            tokens += estimateContentBlockTokens(cb, options);
+          } else if (cb.type === "image" && !anthropicDropsErrorImage) {
+            tokens += estimateContentBlockTokens(cb, options);
+          }
         }
       }
       return tokens;
@@ -224,7 +258,13 @@ export function estimateToolsTokens(tools: ToolDefinition[]): number {
   return total;
 }
 
-export function estimatePromptTokens(
+/**
+ * Raw (uncorrected) prompt-token estimate — exposed so the calibrator
+ * can record (raw, actual) pairs. Applying calibration to the estimate
+ * it uses for training would create a feedback loop that eventually
+ * drives the correction ratio back to 1.0 regardless of true bias.
+ */
+export function estimatePromptTokensRaw(
   messages: Message[],
   systemPrompt?: string,
   options?: TokenEstimatorOptions,
@@ -236,6 +276,24 @@ export function estimatePromptTokens(
   return systemTokens + toolTokens + estimateMessagesTokens(messages, options);
 }
 
+export function estimatePromptTokens(
+  messages: Message[],
+  systemPrompt?: string,
+  options?: TokenEstimatorOptions,
+): number {
+  const raw = estimatePromptTokensRaw(messages, systemPrompt, options);
+
+  // Apply the self-calibration correction. Default is 1.0 for any
+  // (provider, model) pair we haven't recorded a sample for, so first-call
+  // behavior is unchanged. As usage data accumulates, the correction ratio
+  // pulls estimates toward the provider's ground-truth token count. Lookup
+  // uses the per-provider aggregate key — `getCorrection` falls back to
+  // `(provider, "")` when a model-specific sample is not available.
+  const providerName = options?.providerName ?? "";
+  const correction = getCorrection(providerName, "");
+  return correction === 1.0 ? raw : Math.ceil(raw * correction);
+}
+
 function stableJson(value: unknown): string {
   try {
     return JSON.stringify(value);