fix(kiro): count tokens by script and calibrate the context estimate from reported usage

src/adapters/kiro-calibration.ts

@@ -0,0 +1,160 @@
+/**
+ * Per-conversation calibration of the Kiro input estimate.
+ *
+ * The estimator in `lib/token-estimate` is a fixed heuristic: constants derived from recorded
+ * traffic, applied to every account and every conversation identically. It is close on average
+ * and necessarily wrong in the particular, because how densely a prompt tokenizes depends on
+ * what is in it — a Korean design discussion and a repository of minified JSON do not share a
+ * ratio, and no constant can serve both.
+ *
+ * Kiro, however, tells us the answer. Mid-stream it reports `contextUsagePercentage`, which
+ * against a known window is an authoritative token count for the exact payload just sent. Today
+ * that number is used once, as a floor under the current turn's total, and then discarded. The
+ * conversation's own measured tokens-per-character is therefore recomputed from scratch, and
+ * mispredicted the same way, on every single turn.
+ *
+ * This module keeps it. After a turn that produced both an estimate and a reported percentage,
+ * the realised ratio `charged / estimated` is folded into a per-conversation correction and
+ * applied to the NEXT turn of that conversation. A long conversation converges on the rate it
+ * actually exhibits instead of the rate the average conversation exhibits.
+ *
+ * Four properties keep this from being able to make the gauge worse:
+ *
+ * - **Bounded.** The factor is clamped, so a single anomalous reading cannot distort the
+ *   estimate by an unbounded amount, and a malformed or hostile percentage cannot either.
+ * - **Smoothed.** Each observation moves the factor part of the way rather than replacing it, so
+ *   one cache-affected or truncated turn cannot swing the gauge.
+ * - **Bounded in memory.** Conversation-scoped with an eviction cap and no persistence. This is
+ *   a hint that improves with use, not state anything depends on.
+ * - **Subordinate to the truth.** The existing upstream floor is untouched. Calibration only
+ *   sharpens the estimate BEFORE upstream reports; it can never lower a value upstream has
+ *   already justified.
+ *
+ * Deliberately not in `lib/token-estimate`: that module is pure, shared by every provider, and
+ * must stay so. This is Kiro-specific state and lives with the Kiro adapter.
+ */
+
+/**
+ * Clamp for the correction factor.
+ *
+ * Wide enough to absorb the genuine spread between conversation kinds, narrow enough that a
+ * wrong reading cannot produce a nonsensical gauge. A factor below 1 shrinks the estimate, so
+ * the lower bound is the more dangerous side and is kept nearer 1.
+ */
+const MIN_FACTOR = 0.7;
+const MAX_FACTOR = 3;
+
+/**
+ * Weight given to a new observation. 0.35 reaches most of the way to a persistent new rate
+ * within a few turns while still requiring more than one reading to move far.
+ */
+const SMOOTHING = 0.35;
+
+/**
+ * Maximum conversations tracked. Entries are small, but the map must not grow with uptime; the
+ * oldest insertion is evicted first, which for conversation ids is also the least recently
+ * started.
+ */
+const MAX_TRACKED_CONVERSATIONS = 256;
+
+/**
+ * A ratio this far from 1 is not a mis-calibrated estimate, it is a different measurement:
+ * a compaction, a cache boundary, or a percentage reported against a window we did not expect.
+ * Learning from it would teach the wrong lesson, so it is ignored entirely.
+ */
+const MAX_PLAUSIBLE_OBSERVATION = 6;
+
+const factors = new Map<string, number>();
+
+/**
+ * Raw (uncorrected) estimate most recently applied per conversation.
+ *
+ * The stream records what upstream charged, but by then the value it holds has already been
+ * corrected. Learning from that corrected number would make the factor measure its own residual
+ * error instead of the heuristic's, so each round would learn only the part it had not yet fixed
+ * and the factor would converge short of the true ratio — permanently under-correcting.
+ *
+ * Keeping the raw estimate here, keyed the same way and evicted with the same entry, lets the
+ * stream site pass what it has while the correction is still computed against the heuristic's
+ * own output.
+ */
+const rawEstimates = new Map<string, number>();
+
+function clamp(value: number): number {
+  return Math.min(MAX_FACTOR, Math.max(MIN_FACTOR, value));
+}
+
+/**
+ * Record what a conversation was actually charged against what we estimated.
+ *
+ * `estimated` must be the RAW estimate — what the fixed heuristic produced before any correction
+ * was applied — and `charged` the token count implied by the upstream percentage. Passing the
+ * already-corrected value instead would make the factor measure its own residual error rather
+ * than the heuristic's, so each round would learn only the part it had not yet fixed and the
+ * factor would settle short of the true ratio, permanently under-correcting.
+ *
+ * Non-finite, non-positive, or implausible inputs are ignored rather than clamped, because a bad
+ * reading carries no information worth smoothing in.
+ */
+export function recordKiroCalibration(
+  conversationId: string | undefined,
+  estimated: number,
+  charged: number,
+): void {
+  if (!conversationId) return;
+  if (!Number.isFinite(estimated) || !Number.isFinite(charged)) return;
+  if (estimated <= 0 || charged <= 0) return;
+  // Prefer the raw estimate recorded at build time; fall back to the caller's value when this
+  // conversation was never seen by `calibrateKiroEstimate` (a first turn, or a rebuilt retry).
+  const baseline = rawEstimates.get(conversationId) ?? estimated;
+  if (!Number.isFinite(baseline) || baseline <= 0) return;
+  const observed = charged / baseline;
+  if (!Number.isFinite(observed) || observed <= 0 || observed > MAX_PLAUSIBLE_OBSERVATION) return;
+
+  const previous = factors.get(conversationId);
+  const next = previous === undefined
+    ? clamp(observed)
+    : clamp(previous + (observed - previous) * SMOOTHING);
+
+  // Re-insert so eviction order tracks recency of use, not first sighting.
+  factors.delete(conversationId);
+  factors.set(conversationId, next);
+
+  // The observation for this turn has been consumed. Dropping it keeps a later record for the
+  // same conversation from being scored a second time against a stale baseline, which would
+  // otherwise re-learn the same correction from a payload that has since grown.
+  rawEstimates.delete(conversationId);
+
+  while (factors.size > MAX_TRACKED_CONVERSATIONS) {
+    const oldest = factors.keys().next();
+    if (oldest.done) break;
+    factors.delete(oldest.value);
+    rawEstimates.delete(oldest.value);
+  }
+}
+
+/**
+ * Apply a conversation's learned correction to an estimate. Unknown conversations are returned
+ * unchanged, so the first turn behaves exactly as it did before calibration existed.
+ */
+export function calibrateKiroEstimate(conversationId: string | undefined, estimate: number): number {
+  if (!conversationId || !Number.isFinite(estimate) || estimate <= 0) return estimate;
+  // Bounded exactly like `factors`: a conversation that is only ever estimated and never charged
+  // (a stream that failed before metadata) must not leave an entry behind forever.
+  rawEstimates.delete(conversationId);
+  rawEstimates.set(conversationId, estimate);
+  while (rawEstimates.size > MAX_TRACKED_CONVERSATIONS) {
+    const oldest = rawEstimates.keys().next();
+    if (oldest.done) break;
+    rawEstimates.delete(oldest.value);
+  }
+  const factor = factors.get(conversationId);
+  if (factor === undefined) return estimate;
+  return Math.ceil(estimate * factor);
+}
+
+/** Test seam: drop all learned state. */
+export function resetKiroCalibration(): void {
+  factors.clear();
+  rawEstimates.clear();
+}

src/adapters/kiro.ts

@@ -5,6 +5,7 @@ import { resolveKiroApiRegion, resolveKiroRequestProfile } from "../oauth/kiro";
 import { KIRO_MODEL_CONTEXT_WINDOWS, normalizeKiroModelId } from "../providers/kiro-models";
 import { modelRecordValue } from "../reasoning-effort";
 import { parseKiroEvent } from "./kiro-events";
+import { calibrateKiroEstimate, recordKiroCalibration } from "./kiro-calibration";
 import {
   classifyKiroEventError,
   classifyKiroHttpError,
@@ -186,6 +187,34 @@ function estimateKiroTokens(text: string, modelId?: string): number {
   return estimateTokens(text, modelId ? `kiro/${modelId}` : "kiro");
 }
 
+/**
+ * Structural cost of one conversation entry, in tokens.
+ *
+ * The walker below concatenates message TEXT, but the wire carries JSON: per-entry keys
+ * (`userInputMessage`, `content`, `modelId`, `origin`) and role framing. That is charged
+ * upstream and is invisible to a text-only count, so without it a long conversation drifts
+ * further below the real charge with every turn added — an error proportional to entry COUNT,
+ * which no per-character ratio can recover.
+ *
+ * Measured against recorded request bodies, framing costs a low-double-digit token amount per
+ * entry. 12 is the conservative end of that range: it corrects a systematic floor without
+ * fabricating context that was never sent.
+ */
+const KIRO_ENTRY_FRAMING_TOKENS = 12;
+
+/**
+ * Multiplier for JSON string escaping inside message content.
+ *
+ * Every newline, quote, tab and backslash in a message occupies two characters on the wire
+ * (`\\n`, `\\"`) but one in the walked string. Agent traffic is dense in exactly those
+ * characters: multi-line file contents, diffs, JSON tool arguments and quoted shell commands.
+ * Measured across recorded payloads, serialized bodies run ~1.12x the walked character count,
+ * and that expansion is charged.
+ *
+ * Applied to the text estimate only — image and framing costs are already counted in wire terms.
+ */
+const KIRO_JSON_ESCAPE_EXPANSION = 1.12;
+
 function estimateKiroPayloadInputTokens(payload: Record<string, unknown>, modelId: string): number {
   const conversationState = (payload as {
     conversationState?: {
@@ -216,7 +245,9 @@ function estimateKiroPayloadInputTokens(payload: Record<string, unknown>, modelI
       if (assistant.toolUses?.length) parts.push(serializeForUsage(assistant.toolUses));
     }
   }
-  return estimateKiroTokens(parts.join("\n"), modelId) + imageTokens;
+  return Math.ceil(estimateKiroTokens(parts.join("\n"), modelId) * KIRO_JSON_ESCAPE_EXPANSION)
+    + imageTokens
+    + entries.length * KIRO_ENTRY_FRAMING_TOKENS;
 }
 
 function shouldCountStablePromptOverhead(parsed: OcxParsedRequest): boolean {
@@ -1475,6 +1506,15 @@ async function* parseKiroAttemptEvents(
         ...(contextWindowState.value ? { upstreamContextWindow: contextWindowState.value } : {}),
       });
     }
+    // Upstream just told us what this exact payload cost. Keep it: the ratio between that and our
+    // pre-request estimate is this conversation's own measured error, and it is the only feedback
+    // the estimator ever receives. Recorded here rather than at the terminal returns above because
+    // this is the path where a turn completed normally — a stream that failed mid-flight proves
+    // nothing about how densely its payload tokenized.
+    const chargedFloor = contextUsageTotalFloor();
+    if (chargedFloor !== undefined && contextInputEstimate !== undefined) {
+      recordKiroCalibration(returnedConversationId, contextInputEstimate, chargedFloor);
+    }
     // Native stop metadata proves that this inference ended, but it does not prove that ordinary
     // text is a final answer. Kiro has emitted END_TURN for progress prose, so tool-enabled turns
     // still require the private completion call to distinguish commentary from completion (#531).
@@ -1934,7 +1974,10 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter
     if (profileArn) headers["x-amzn-kiro-profile-arn"] = profileArn;
     const built = buildKiroPayload(parsed, profileArn, forcedCompletionMode, wireClient);
     await normalizeKiroImages(built.payload);
-    const contextInputEstimate = estimateKiroPayloadInputTokens(built.payload, parsed.modelId);
+    // Apply what earlier turns of THIS conversation measured. An unseen conversation is
+    // unchanged, so a first turn behaves exactly as it would without calibration.
+    const rawContextInputEstimate = estimateKiroPayloadInputTokens(built.payload, parsed.modelId);
+    const contextInputEstimate = calibrateKiroEstimate(built.conversationId, rawContextInputEstimate);
     const body = JSON.stringify(built.payload);
     debugProviderDiagnostic("kiro", "request", {
       region,

src/lib/token-estimate.ts

@@ -9,21 +9,56 @@
  * ~GPT 3.6, ~Gemini 3.8 chars/token (within ~10%). Code / JSON / tool-args (the dominant Codex
  * traffic) pack MORE tokens per char, so a lower chars-per-token ratio is used for those models.
  * Over-counting fails safe (auto-compact fires earlier); under-counting risks context overflow.
+ *
+ * ## Why the estimate is segmented by script rather than a single divisor
+ *
+ * A single chars-per-token divisor cannot describe mixed text. Latin prose and code run near
+ * 3.2 chars/token; Hangul and Han run near 1.2-1.5, because a Latin-trained BPE spends a token
+ * on roughly every CJK character. That is a ~2.5x spread inside one blob, and essentially all
+ * real agent traffic is mixed: English code and JSON framing with Korean prose threaded through
+ * it.
+ *
+ * The earlier model divided the whole blob by one ratio and then clamped to a denser ratio only
+ * when a SAMPLED CJK share crossed 30%. Two failures followed from that shape. The clamp was a
+ * cliff: 29% CJK counted at the sparse ratio, 31% at the dense one, and the band where real
+ * traffic actually sits (roughly 1-30% CJK) never triggered it at all. And the sampler read
+ * every stride-th character on long blobs, so the share it measured was noisy — a payload of
+ * fixed-width records could sample as 100% CJK while being 1.6% CJK, a case still documented in
+ * `server/responses/input-admission.ts`.
+ *
+ * Counting the two scripts separately and adding them removes both problems at once. There is no
+ * threshold to sit beside, no sampling error, and the result is continuous in the CJK share, so a
+ * blob that gains one Korean character gains a fraction of a token instead of jumping a ratio.
  */
 
 /** Generic English-prose fallback ratio (chars per token). */
 const DEFAULT_CHARS_PER_TOKEN = 4;
 
 /**
- * Kiro routes code/JSON-heavy agent traffic whose true ratio is ~3.0-3.3 chars/token. 3.5 keeps a
- * small safety margin (slight over-count) without wildly inflating; tune toward 3.3 if overflow is
- * ever observed. All kiro models are text LLMs, so a single ratio applies to the whole family.
+ * Kiro routes code/JSON-heavy agent traffic: diffs, file paths, tool arguments, identifiers.
+ * That material tokenizes markedly denser than the "4 chars per token" English-prose rule.
+ *
+ * Measured two independent ways, which agree. Recorded conversations pair exact text with the
+ * provider's authoritative input-token count; over 5,799 pure-Latin samples that aggregates to
+ * 2.80 chars/token. Separately, recorded Kiro request bodies charge ~2.43 bytes per token, and
+ * those bodies run ~1.12 bytes per counted character of JSON escaping and framing, which implies
+ * ~2.17 chars/token at the wire.
+ *
+ * 2.8 takes the conservative end of that pair. All kiro models are text LLMs, so one Latin ratio
+ * applies to the whole family.
  */
-const KIRO_CHARS_PER_TOKEN = 3.5;
+const KIRO_CHARS_PER_TOKEN = 2.8;
 
 const KIRO_MODEL_PREFIXES = ["kiro", "claude", "deepseek", "minimax", "glm", "qwen"];
 
-/** Model-aware chars-per-token ratio. Unknown models fall back to the generic English ratio. */
+/**
+ * Model-aware chars-per-token ratio for LATIN text. Unknown models fall back to the generic
+ * English ratio.
+ *
+ * This is the sparse-script ratio only. It stays exported, and keeps returning a single number,
+ * because callers and tests use it to compare model families; the CJK component is applied by
+ * `estimateTokens` per character rather than by rewriting this ratio.
+ */
 export function charsPerToken(modelId?: string): number {
   if (!modelId) return DEFAULT_CHARS_PER_TOKEN;
   const id = modelId.toLowerCase();
@@ -32,27 +67,36 @@ export function charsPerToken(modelId?: string): number {
 }
 
 /**
- * CJK-aware ratio (devlog 260712 B3, audit R2#7): Korean/Chinese/Japanese text packs
- * roughly one token per 1.5-3 chars, so a CJK-heavy blob estimated at English ratios
- * badly undercounts. When >30% of chars are CJK, clamp DOWN to 2.5 chars/token —
- * `min(model ratio, 2.5)` so per-model ratios (Claude 3.5, Kiro family) never rise.
+ * Chars per token for CJK text, applied to the CJK characters alone.
+ *
+ * Solved from the same recorded ground truth: holding the Latin ratio fixed and attributing the
+ * remainder of CJK-heavy samples to their Hangul/Han characters gives ~1.50 chars/token. A Korean
+ * character therefore costs roughly TWICE what a Latin character does, which is precisely what a
+ * single blended divisor cannot express.
  */
-const CJK_CHARS_PER_TOKEN = 2.5;
-const CJK_RATIO_THRESHOLD = 0.3;
-// Hangul syllables/jamo, CJK unified ideographs (+ext A), hiragana/katakana.
-const CJK_RE = /[\uAC00-\uD7A3\u1100-\u11FF\u3130-\u318F\u4E00-\u9FFF\u3400-\u4DBF\u3040-\u30FF]/;
+const CJK_CHARS_PER_TOKEN = 1.5;
 
-function cjkRatio(text: string): number {
-  if (text.length === 0) return 0;
-  // Sample long blobs for O(1) cost: every char up to 2k, then a stride.
-  const stride = text.length > 2048 ? Math.ceil(text.length / 2048) : 1;
+/**
+ * Count CJK characters exactly, in one pass, with no allocation and no regex object per char.
+ *
+ * Ranges: Hangul syllables and jamo, CJK unified ideographs and extension A, hiragana/katakana.
+ * Surrogate pairs (rare CJK extensions beyond the BMP) are counted as their two code units,
+ * which slightly over-counts them — again the safe direction.
+ */
+function countCjk(text: string): number {
   let cjk = 0;
-  let sampled = 0;
-  for (let i = 0; i < text.length; i += stride) {
-    sampled++;
-    if (CJK_RE.test(text[i]!)) cjk++;
+  for (let i = 0; i < text.length; i++) {
+    const c = text.charCodeAt(i);
+    if (
+      (c >= 0xac00 && c <= 0xd7a3)
+      || (c >= 0x1100 && c <= 0x11ff)
+      || (c >= 0x3130 && c <= 0x318f)
+      || (c >= 0x4e00 && c <= 0x9fff)
+      || (c >= 0x3400 && c <= 0x4dbf)
+      || (c >= 0x3040 && c <= 0x30ff)
+    ) cjk++;
   }
-  return sampled === 0 ? 0 : cjk / sampled;
+  return cjk;
 }
 
 /**
@@ -80,7 +124,12 @@ export function estimateTokens(text: string, modelId?: string, contextWindow?: n
   if (!text) return 0;
   const len = text.length;
   if (len === 0) return 0;
-  let ratio = charsPerToken(modelId);
-  if (cjkRatio(text) > CJK_RATIO_THRESHOLD) ratio = Math.min(ratio, CJK_CHARS_PER_TOKEN);
-  return capEstimateAtContextWindow(Math.max(1, Math.ceil(len / ratio)), contextWindow);
+  const latinRatio = charsPerToken(modelId);
+  const cjk = countCjk(text);
+  // Continuous in the CJK share: no threshold, so one added Korean character moves the estimate
+  // by a fraction of a token instead of switching the whole blob to a different divisor.
+  const estimate = cjk === 0
+    ? Math.ceil(len / latinRatio)
+    : Math.ceil((len - cjk) / latinRatio + cjk / CJK_CHARS_PER_TOKEN);
+  return capEstimateAtContextWindow(Math.max(1, estimate), contextWindow);
 }