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

src/adapters/kiro-calibration.ts

@@ -0,0 +1,178 @@
+/**
+ * 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 TERMINAL attempt 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, so a long conversation converges
+ * on the rate it actually exhibits rather than the rate the average conversation exhibits.
+ *
+ * "Terminal" matters: the adapter's bounded completion retry streams a second time for the same
+ * user turn against a rebuilt payload, so the caller commits an observation only once that turn
+ * is genuinely over.
+ *
+ * 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;
+
+/**
+ * One conversation's calibration state.
+ *
+ * `factor` is the learned correction. `rawEstimate` is the uncorrected estimate most recently
+ * applied: the stream records what upstream charged, but by then the value it holds has already
+ * been corrected, and learning from that would make the factor measure its own residual error
+ * instead of the heuristic's — each round would close only the part it had not yet fixed and the
+ * factor would converge short of the true ratio.
+ *
+ * Both live in ONE entry so they cannot desynchronise. Held as two maps, a conversation could be
+ * refreshed in one and evicted from the other, and the pair could retain twice the documented
+ * bound in distinct ids.
+ */
+interface Calibration {
+  factor?: number;
+  rawEstimate?: number;
+}
+
+/** Insertion-ordered, so the first key is the least recently touched entry. */
+const calibrations = new Map<string, Calibration>();
+
+/** Re-insert to move an entry to the most-recently-used end, then evict from the old end. */
+function touch(conversationId: string, update: (current: Calibration) => Calibration): void {
+  const current = calibrations.get(conversationId) ?? {};
+  calibrations.delete(conversationId);
+  calibrations.set(conversationId, update(current));
+  while (calibrations.size > MAX_TRACKED_CONVERSATIONS) {
+    const oldest = calibrations.keys().next();
+    if (oldest.done) break;
+    calibrations.delete(oldest.value);
+  }
+}
+
+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;
+  const entry = calibrations.get(conversationId);
+  // 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 = entry?.rawEstimate ?? estimated;
+  if (!Number.isFinite(baseline) || baseline <= 0) return;
+  const observed = charged / baseline;
+  if (!Number.isFinite(observed) || observed <= 0 || observed > MAX_PLAUSIBLE_OBSERVATION) return;
+
+  // Smooth from 1 (the "no correction" prior) on the FIRST observation too. Jumping straight to
+  // `clamp(observed)` would let a single cache-affected or truncated turn set the factor outright,
+  // which is exactly the swing the smoothing exists to prevent.
+  const previous = entry?.factor ?? 1;
+  const next = clamp(previous + (observed - previous) * SMOOTHING);
+
+  // The observation for this turn is consumed: dropping the raw estimate keeps a later record for
+  // the same conversation from being scored a second time against a stale baseline, which would
+  // re-learn the same correction from a payload that has since grown.
+  touch(conversationId, () => ({ factor: next }));
+}
+
+/**
+ * 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;
+  touch(conversationId, current => ({ ...current, rawEstimate: estimate }));
+  const factor = calibrations.get(conversationId)?.factor;
+  return factor === undefined ? estimate : Math.ceil(estimate * factor);
+}
+
+/**
+ * Carry a conversation's state to the id upstream actually returned.
+ *
+ * The estimate is stored under the id the request was BUILT with, but Kiro can answer with a
+ * different conversation id, and the stream then records under that one. Without this the record
+ * would miss its own raw estimate and silently fall back to the corrected value — the exact
+ * feedback bug the raw estimate exists to avoid.
+ */
+export function rekeyKiroCalibration(fromConversationId: string | undefined, toConversationId: string | undefined): void {
+  if (!fromConversationId || !toConversationId || fromConversationId === toConversationId) return;
+  const entry = calibrations.get(fromConversationId);
+  if (!entry) return;
+  calibrations.delete(fromConversationId);
+  touch(toConversationId, current => ({ ...current, ...entry }));
+}
+
+/** Test seam: drop all learned state. */
+export function resetKiroCalibration(): void {
+  calibrations.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, rekeyKiroCalibration } 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 {
@@ -990,6 +1021,11 @@ async function* parseKiroAttempt(
   // the attempt boundary. Anything the inner parser leaves behind is flushed before the terminal.
   const deferred: AdapterEvent[] = [];
   const retention = createKiroAttemptRetention(budget);
+  // Shared box: the inner parser stages its calibration observation here on the completion path,
+  // and this wrapper decides whether the attempt was terminal enough to commit it. A box rather
+  // than a return field because the completion path has a dozen terminal returns and threading a
+  // field through every one of them is exactly the kind of edit that misses one.
+  const attemptCalibration: { value?: { conversationId: string; estimated: number; charged: number } } = {};
   const attempt = parseKiroAttemptEvents(
     response,
     budget,
@@ -1001,12 +1037,22 @@ async function* parseKiroAttempt(
     conversationId,
     deferred,
     retention,
+    attemptCalibration,
     contextInputEstimate,
     priorEmittedOutput,
   );
   let handedOff = false;
   try {
     const result = yield* attempt;
+    // A staged observation only counts when this attempt is the LAST one for the user turn. An
+    // attempt that asks for the bounded fallback streams again against a rebuilt payload, so
+    // committing here would move the factor twice for one turn and score the second observation
+    // against a payload the first had already inflated.
+    const staged = attemptCalibration.value;
+    attemptCalibration.value = undefined;
+    if (staged && !result.needsFallback) {
+      recordKiroCalibration(staged.conversationId, staged.estimated, staged.charged);
+    }
     for (const event of deferred.splice(0)) {
       try { yield event; } finally { retention.releaseEvent(event); }
     }
@@ -1028,10 +1074,13 @@ async function* parseKiroAttemptEvents(
   conversationId: string | undefined,
   deferred: AdapterEvent[],
   retention: KiroAttemptRetention,
+  attemptCalibration: { value?: { conversationId: string; estimated: number; charged: number } },
   contextInputEstimate?: number,
   priorEmittedOutput = false,
 ): AsyncGenerator<AdapterEvent, KiroAttemptParseResult> {
   const emptyResult = (): KiroAttemptParseResult => ({ assistantText: "", sawReasoning: false });
+  // Every early return below is a failure path that stages nothing; only the completion path
+  // writes `attemptCalibration`, and the wrapper decides whether to commit it.
   if (!response.body) {
     return {
       ...emptyResult(),
@@ -1346,7 +1395,13 @@ async function* parseKiroAttemptEvents(
           if (ev.stopReason !== undefined) stopReason = ev.stopReason;
           break;
         case "message_metadata":
-          if (isValidKiroConversationId(ev.conversationId)) returnedConversationId = ev.conversationId;
+          if (isValidKiroConversationId(ev.conversationId)) {
+            // Kiro can answer under a different conversation id than the request was built with.
+            // Carry the calibration entry across so the record below finds its own raw estimate
+            // instead of silently falling back to the already-corrected value.
+            rekeyKiroCalibration(returnedConversationId, ev.conversationId);
+            returnedConversationId = ev.conversationId;
+          }
           break;
         case "content":
           if (ev.modelId) {
@@ -1475,6 +1530,29 @@ async function* parseKiroAttemptEvents(
         ...(contextWindowState.value ? { upstreamContextWindow: contextWindowState.value } : {}),
       });
     }
+    // Upstream just told us what this payload cost. 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.
+    //
+    // Staged, not recorded. An attempt that sets `needsFallback` is not over: the adapter rebuilds
+    // the payload and streams a second time for the SAME user turn. Learning here would apply the
+    // fresh factor to that rebuild and then learn again from it, so one turn would move the factor
+    // twice and the second observation would score a payload the first had already inflated. Only
+    // the outer parser knows whether an attempt is terminal, so it commits.
+    //
+    // Subtract the output first. `contextUsageTotalFloor` is the absolute context size AFTER the
+    // response (`OcxUsage.contextTotalTokens`, types/request.ts), while `contextInputEstimate`
+    // covers the request payload alone. Dividing one by the other would charge generated tokens to
+    // prompt-tokenization error, so a short prompt answered at length would learn a large factor
+    // and inflate every later request in that conversation — the premature compaction this work
+    // exists to prevent.
+    const chargedTotal = contextUsageTotalFloor();
+    if (chargedTotal !== undefined && contextInputEstimate !== undefined) {
+      const chargedInput = chargedTotal - finalUsage.outputTokens;
+      if (chargedInput > 0 && returnedConversationId) {
+        attemptCalibration.value = { conversationId: returnedConversationId, estimated: contextInputEstimate, charged: chargedInput };
+      }
+    }
     // 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 +2012,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,