diff --git a/src/adapters/kiro-calibration.ts b/src/adapters/kiro-calibration.ts new file mode 100644 index 00000000000..4bfc07ba52b --- /dev/null +++ b/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(); + +/** 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(); +} diff --git a/src/adapters/kiro.ts b/src/adapters/kiro.ts index 35a0ff93b32..3da59974278 100644 --- a/src/adapters/kiro.ts +++ b/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, modelId: string): number { const conversationState = (payload as { conversationState?: { @@ -216,7 +245,9 @@ function estimateKiroPayloadInputTokens(payload: Record, 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 { 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, diff --git a/src/lib/token-estimate.ts b/src/lib/token-estimate.ts index 337ae6a5dd5..bf507e1cf27 100644 --- a/src/lib/token-estimate.ts +++ b/src/lib/token-estimate.ts @@ -9,50 +9,112 @@ * ~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. + * + * Applied ONLY to `kiro/`-prefixed ids. The measurement is Kiro's charge for Kiro's payload + * shape, and the id families below (`claude`, `deepseek`, `qwen`, ...) are also routed by + * Cursor, Anthropic direct and Antigravity, whose consumers read this same helper to size + * admission ceilings, `count_tokens` answers, and overflow-vs-429 classification. Widening a + * Kiro-derived constant to those callers would retune three unrelated subsystems from evidence + * that says nothing about them. + */ +const KIRO_CHARS_PER_TOKEN = 2.8; + +/** + * Ratio for the model families Kiro shares with other providers, when NOT routed through Kiro. + * + * These are code-heavy agent models, so 3.5 remains right for them; it is the value every + * non-Kiro consumer of this helper was calibrated against. */ -const KIRO_CHARS_PER_TOKEN = 3.5; +const AGENT_MODEL_CHARS_PER_TOKEN = 3.5; -const KIRO_MODEL_PREFIXES = ["kiro", "claude", "deepseek", "minimax", "glm", "qwen"]; +const AGENT_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(); - if (KIRO_MODEL_PREFIXES.some(p => id.startsWith(p))) return KIRO_CHARS_PER_TOKEN; + // `estimateKiroTokens` always prefixes `kiro/`, so the Kiro-measured ratio reaches Kiro traffic + // and only Kiro traffic. + if (id.startsWith("kiro")) return KIRO_CHARS_PER_TOKEN; + if (AGENT_MODEL_PREFIXES.some(p => id.startsWith(p))) return AGENT_MODEL_CHARS_PER_TOKEN; return DEFAULT_CHARS_PER_TOKEN; } /** - * 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 +142,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); } diff --git a/src/server/responses/input-admission.ts b/src/server/responses/input-admission.ts index 2bd01a67804..ef8b81265e3 100644 --- a/src/server/responses/input-admission.ts +++ b/src/server/responses/input-admission.ts @@ -19,16 +19,23 @@ import type { OcxContentPart, OcxParsedRequest, OcxProviderConfig } from "../../ /** * Multiplier applied to the ceiling before refusing. * - * 2.5, not something tighter, because `estimateTokens` can overshoot by 1.6x on its own. - * `cjkRatio` samples every `stride`-th character, so a payload of fixed-width records whose - * length aligns with the stride samples as 100% CJK while being ~1.6% CJK, firing the - * 2.5-chars/token clamp instead of the 4.0 default. Measured on Bun 1.3.14: 126,046 chars, - * true CJK ratio 0.0161, sampled ratio 1.0, estimate inflated 1.6x. Since 4.0 / 2.5 = 1.6 - * is that branch maximum divergence, a threshold at or under 1.6 would convert the - * estimator error bar into false 413s. + * 2.5, not something tighter, because the estimate is a heuristic and the cost of being wrong + * is asymmetric: a false refusal fails a turn the provider would have answered, while an + * over-admission merely pays for one round trip the provider then rejects itself. * - * 2.5 sits above it with room for the ~10% model-family ratio spread, and still refuses the - * #1412 shape (10x) four times over. + * The original margin was sized against a sampling artifact — `cjkRatio` read every stride-th + * character, so a payload of fixed-width records could sample as 100% CJK while being 1.6% CJK + * and inflate the estimate by 4.0/2.5 = 1.6x. That sampler is gone: CJK characters are now + * counted exactly, so that particular 1.6x divergence cannot occur and the headroom it bought is + * no longer spent on it. + * + * The margin is still 2.5 because the estimator it guards got LARGER, not smaller. Counting the + * two scripts separately raises a pure-Latin estimate by 1.25x and a Korean one by up to 1.67x + * against the previous model, which consumes real headroom: measured against the old estimator's + * scale, 2.5 now behaves like roughly 2.0x for Latin and 1.5x for Korean-dominant input. That is + * the intended direction — the estimates are closer to what providers actually charge, so the + * same multiplier is a tighter and more honest bound — and it still refuses the #1412 shape + * (10x compounding) several times over. */ export const ADMISSION_TOLERANCE = 2.5; diff --git a/tests/kiro-calibration.test.ts b/tests/kiro-calibration.test.ts new file mode 100644 index 00000000000..7a9c1162267 --- /dev/null +++ b/tests/kiro-calibration.test.ts @@ -0,0 +1,174 @@ +import { beforeEach, describe, expect, test } from "bun:test"; +import { calibrateKiroEstimate, recordKiroCalibration, rekeyKiroCalibration, resetKiroCalibration } from "../src/adapters/kiro-calibration"; + +beforeEach(() => resetKiroCalibration()); + +describe("kiro per-conversation calibration", () => { + test("an unseen conversation is returned untouched", () => { + expect(calibrateKiroEstimate("conv-new", 1000)).toBe(1000); + expect(calibrateKiroEstimate(undefined, 1000)).toBe(1000); + }); + + test("an under-estimate is corrected upward on the next turn", () => { + // Upstream charged 1500 for a payload we called 1000: we are reading 2/3 of the truth. + recordKiroCalibration("conv-a", 1000, 1500); + const corrected = calibrateKiroEstimate("conv-a", 1000); + expect(corrected).toBeGreaterThan(1000); + expect(corrected).toBeLessThanOrEqual(1500); + }); + + // A first observation must be smoothed like any other: jumping straight to the observed ratio + // lets one cache-affected or truncated turn set the factor outright. + test("the first observation is smoothed, not adopted outright", () => { + recordKiroCalibration("conv-first", 1000, 2900); + const applied = calibrateKiroEstimate("conv-first", 1000); + expect(applied).toBeLessThan(2000); + }); + + // Kiro may answer under a different conversation id than the request was built with. Without + // carrying the entry across, the record would miss its raw estimate and fall back to the + // already-corrected value — the feedback bug the raw estimate exists to prevent. + test("state follows a conversation id replaced by the provider", () => { + calibrateKiroEstimate("conv-built", 1000); + rekeyKiroCalibration("conv-built", "conv-returned"); + recordKiroCalibration("conv-returned", 4242, 1500); + // Learned against the raw 1000, not the bogus 4242 the caller passed. + const applied = calibrateKiroEstimate("conv-returned", 1000); + expect(applied).toBeGreaterThan(1000); + expect(calibrateKiroEstimate("conv-built", 1000)).toBe(1000); + }); + + test("repeated consistent observations converge toward the observed ratio", () => { + for (let i = 0; i < 12; i++) recordKiroCalibration("conv-b", 1000, 1500); + // Smoothing means it approaches 1.5 rather than jumping there. + expect(calibrateKiroEstimate("conv-b", 1000)).toBeGreaterThan(1400); + expect(calibrateKiroEstimate("conv-b", 1000)).toBeLessThanOrEqual(1500); + }); + + test("a single outlier cannot swing the estimate to the outlier's ratio", () => { + for (let i = 0; i < 10; i++) recordKiroCalibration("conv-c", 1000, 1000); + recordKiroCalibration("conv-c", 1000, 2900); + // One reading moves it partway, not all the way. + expect(calibrateKiroEstimate("conv-c", 1000)).toBeLessThan(1800); + }); + + test("the factor stays inside its clamp under absurd input", () => { + for (let i = 0; i < 50; i++) recordKiroCalibration("conv-hi", 1, 5); + expect(calibrateKiroEstimate("conv-hi", 1000)).toBeLessThanOrEqual(3000); + for (let i = 0; i < 50; i++) recordKiroCalibration("conv-lo", 1000, 1); + // Shrinking is the dangerous direction, so the lower clamp is tighter. + expect(calibrateKiroEstimate("conv-lo", 1000)).toBeGreaterThanOrEqual(700); + }); + + test("implausible and malformed observations are ignored, not clamped in", () => { + for (const [est, charged] of [[1000, 0], [0, 1000], [1000, Number.NaN], [Number.POSITIVE_INFINITY, 1000], [1000, 100_000]] as const) { + recordKiroCalibration("conv-junk", est, charged); + } + expect(calibrateKiroEstimate("conv-junk", 1000)).toBe(1000); + }); + + test("conversations are isolated from each other", () => { + recordKiroCalibration("conv-x", 1000, 2000); + expect(calibrateKiroEstimate("conv-y", 1000)).toBe(1000); + expect(calibrateKiroEstimate("conv-x", 1000)).toBeGreaterThan(1000); + }); + + test("tracking is bounded: old conversations are evicted rather than accumulating", () => { + for (let i = 0; i < 600; i++) recordKiroCalibration("conv-" + i, 1000, 1500); + // The earliest conversation has been evicted and behaves as unseen again. + expect(calibrateKiroEstimate("conv-0", 1000)).toBe(1000); + // The most recent is still remembered. + expect(calibrateKiroEstimate("conv-599", 1000)).toBeGreaterThan(1000); + }); + + // Estimating and recording must share ONE entry. Held separately, a conversation could be + // refreshed in one map and evicted from the other, and the pair could retain twice the bound. + test("estimating refreshes recency, so an active conversation is not evicted by new ones", () => { + recordKiroCalibration("conv-active", 1000, 1500); + const learned = calibrateKiroEstimate("conv-active", 1000); + expect(learned).toBeGreaterThan(1000); + // Push PAST the 256-entry cap so eviction actually runs; touching the active conversation + // each round must keep it alive. A filler count under the cap would pass without eviction + // ever being exercised. + for (let i = 0; i < 400; i++) { + recordKiroCalibration("filler-" + i, 1000, 1500); + calibrateKiroEstimate("conv-active", 1000); + } + expect(calibrateKiroEstimate("conv-active", 1000)).toBe(learned); + // ...while an untouched conversation from the same era is gone. + expect(calibrateKiroEstimate("filler-0", 1000)).toBe(1000); + }); + + test("a conversation that was already accurate is left essentially alone", () => { + for (let i = 0; i < 10; i++) recordKiroCalibration("conv-ok", 1000, 1000); + expect(calibrateKiroEstimate("conv-ok", 1000)).toBe(1000); + }); + + // The correction must be learned against the RAW heuristic output. If it were learned against + // the already-corrected estimate it would only ever measure its own leftover error, so each + // round would close part of the remaining gap and the factor would stall short of the truth. + test("a mis-estimating conversation converges on the real charge, not part of the way to it", () => { + const conv = "conv-converge"; + const trueRatio = 1.35; + let lastRatio = 0; + for (let turn = 1; turn <= 6; turn++) { + const raw = 10_000 * turn; + const applied = calibrateKiroEstimate(conv, raw); + const charged = Math.round(raw * trueRatio); + lastRatio = applied / charged; + recordKiroCalibration(conv, applied, charged); + } + expect(lastRatio).toBeGreaterThan(0.95); + expect(lastRatio).toBeLessThan(1.03); + }); + + test("the raw-estimate map is bounded even when turns are never charged", () => { + for (let i = 0; i < 600; i++) calibrateKiroEstimate("uncharged-" + i, 1000); + recordKiroCalibration("uncharged-0", 1000, 1500); + // Evicted, so it learns from the caller's value and stays inside the clamp. + expect(calibrateKiroEstimate("uncharged-0", 1000)).toBeLessThanOrEqual(1500); + }); + + // A turn can report twice — the bounded completion retry rebuilds and re-streams one turn. + // The raw estimate is CONSUMED by the first report, so the second is scored against the + // caller's own value instead of a baseline that no longer describes the payload in flight. + test("a second report for one turn is scored against the caller's value, not a consumed baseline", () => { + calibrateKiroEstimate("conv-dup", 1000); + recordKiroCalibration("conv-dup", 1000, 1300); + // Second metadata frame for the SAME turn: no rebuild, so no new estimate is recorded. The + // baseline was consumed by the first report, so 12_000/10_000 reads as its true 1.2x rather + // than a 12x absurdity that the plausibility cap would silently discard. + recordKiroCalibration("conv-dup", 10_000, 12_000); + const afterBoth = calibrateKiroEstimate("conv-dup", 1000); + + // Compare against a conversation that saw only the first report. + calibrateKiroEstimate("conv-once", 1000); + recordKiroCalibration("conv-once", 1000, 1300); + const afterOne = calibrateKiroEstimate("conv-once", 1000); + + // The second observation was counted, and it converged rather than jumping. + expect(afterBoth).toBeGreaterThan(afterOne); + expect(afterBoth).toBeLessThan(1300); + }); + + // The upstream checkpoint is the context size AFTER the response, while the estimate covers the + // request alone. Learning from the raw checkpoint would charge generated tokens to + // prompt-tokenization error: a short prompt answered at length would look like a huge + // under-estimate and inflate every later request in that conversation. + test("output tokens must be removed before learning, or a long answer poisons the factor", () => { + const requestEstimate = 1000; + const trulyAccurate = 1000; // our estimate was exactly right for the prompt + const longAnswer = 2000; // ...but the model then wrote a long response + const checkpointAfterResponse = trulyAccurate + longAnswer; + + // Correct: subtract the output, observe a 1.0 ratio, leave the estimate alone. + calibrateKiroEstimate("conv-good", requestEstimate); + recordKiroCalibration("conv-good", requestEstimate, checkpointAfterResponse - longAnswer); + expect(calibrateKiroEstimate("conv-good", requestEstimate)).toBe(requestEstimate); + + // Wrong: feeding the post-response checkpoint straight in learns a 3x error that never existed. + calibrateKiroEstimate("conv-bad", requestEstimate); + recordKiroCalibration("conv-bad", requestEstimate, checkpointAfterResponse); + expect(calibrateKiroEstimate("conv-bad", requestEstimate)).toBeGreaterThan(requestEstimate * 1.5); + }); +}); diff --git a/tests/kiro-stream.test.ts b/tests/kiro-stream.test.ts index 30400b4647e..e0df737cb9b 100644 --- a/tests/kiro-stream.test.ts +++ b/tests/kiro-stream.test.ts @@ -14,6 +14,7 @@ import { } from "../src/adapters/kiro-constants"; import { parseKiroEvent } from "../src/adapters/kiro-events"; import { resetKiroThrottleStateForTests } from "../src/adapters/kiro-retry"; +import { resetKiroCalibration } from "../src/adapters/kiro-calibration"; import { buildResponseJSON } from "../src/bridge"; import { encodeMessage } from "../src/lib/eventstream-decoder"; import { estimateTokens } from "../src/lib/token-estimate"; @@ -452,11 +453,11 @@ describe("kiro adapter — parseStream", () => { const done = events.at(-1); expect(done?.type).toBe("done"); const usage = done?.type === "done" ? done.usage : undefined; - expect(usage?.outputTokens).toBe(estimateTokens(firstText, "claude-sonnet-4.5") + estimateTokens(finalText, "claude-sonnet-4.5")); + expect(usage?.outputTokens).toBe(estimateTokens(firstText, "kiro/claude-sonnet-4.5") + estimateTokens(finalText, "kiro/claude-sonnet-4.5")); expect(usage?.contextTotalTokens).toBeGreaterThan( initialContextEstimate + Math.max( - estimateTokens(firstText, "claude-sonnet-4.5"), - estimateTokens(finalText, "claude-sonnet-4.5"), + estimateTokens(firstText, "kiro/claude-sonnet-4.5"), + estimateTokens(finalText, "kiro/claude-sonnet-4.5"), ), ); }); @@ -489,13 +490,13 @@ describe("kiro adapter — parseStream", () => { expect(largeProgress).toBeGreaterThan(smallProgress); // And the extra pressure must be on the order of the extra progress, not a rounding artefact. expect(largeProgress - smallProgress).toBeGreaterThan( - estimateTokens("p".repeat(20000), "claude-sonnet-4.5"), + estimateTokens("p".repeat(20000), "kiro/claude-sonnet-4.5"), ); }); test("bounded fallback preserves definite growth after an upstream context checkpoint", async () => { const finalText = "f".repeat(3500); - const finalOutputTokens = estimateTokens(finalText, "claude-sonnet-4.5"); + const finalOutputTokens = estimateTokens(finalText, "kiro/claude-sonnet-4.5"); globalThis.fetch = (async () => new Response(streamOf(eventFrame({ content: finalText })))) as typeof fetch; const adapter = createKiroAdapter(provider); await adapter.buildRequest(parsedWith([{ role: "user", content: "do it" }], [bashTool])); @@ -1457,8 +1458,9 @@ describe("kiro adapter — parseStream", () => { const adapter = createKiroAdapter(provider); await adapter.buildRequest(parsedWith([{ role: "user", content: "x".repeat(700) }])); const done = await doneUsage(adapter, eventFrame({ content: "y".repeat(350) })); - expect(done.inputTokens).toBe(200); - expect(done.outputTokens).toBe(100); + // 700 chars at the 2.8 kiro Latin ratio (the turn adds a trailing newline separator). + expect(done.inputTokens).toBe(251); + expect(done.outputTokens).toBe(126); expect(done.estimated).toBe(true); }); @@ -1501,7 +1503,7 @@ describe("kiro adapter — parseStream", () => { ); expect(done).toEqual({ inputTokens: 15, - contextTotalTokens: 204, + contextTotalTokens: 298, cachedInputTokens: 3, cacheReadInputTokens: 3, cacheCreationInputTokens: 2, @@ -1577,8 +1579,8 @@ describe("kiro adapter — parseStream", () => { eventFrame({ contextUsagePercentage: 25 }), ); - expect(done.inputTokens).toBe(200); - expect(done.outputTokens).toBe(100); + expect(done.inputTokens).toBe(251); + expect(done.outputTokens).toBe(126); expect(done.totalTokens).toBeUndefined(); expect(done.estimated).toBe(true); expect(done.contextTotalTokens).toBe(50_000); @@ -1601,10 +1603,10 @@ describe("kiro adapter — parseStream", () => { eventFrame({ contextUsagePercentage: 25 }), ); - expect(done.inputTokens).toBe(200); - expect(done.outputTokens).toBe(100); + expect(done.inputTokens).toBe(251); + expect(done.outputTokens).toBe(126); expect(done.totalTokens).toBeUndefined(); - expect(done.contextTotalTokens).toBe(300); + expect(done.contextTotalTokens).toBe(420); }); test("Kiro auto uses the concrete response model to decode context percentage", async () => { @@ -1619,14 +1621,131 @@ describe("kiro adapter — parseStream", () => { expect(done.contextTotalTokens).toBe(50_000); }); + // End-to-end proof that calibration survives the adapter boundary: one full stream reports a + // context percentage, and the NEXT request for the same conversation is estimated higher for + // identical input. Unit tests cover the arithmetic; this covers the wiring. + // + // The payload is sized so the reported percentage lands within a believable multiple of our + // estimate. A wildly larger ratio is rejected as implausible by design — a 30x reading is a + // compaction or a window mismatch, not a mis-tokenized prompt — so a naive fixture proves + // nothing about the wiring. + test("a reported context percentage calibrates the next request in the same conversation", async () => { + resetKiroCalibration(); + // 28k chars estimates ~11.2k tokens once framing is counted; 10% of the model window is 20k, + // a believable 1.8x under-read rather than a compaction-sized outlier that the plausibility + // guard would rightly discard. + const messages = [{ role: "user", content: "x".repeat(28_000) }]; + // A conversation is only the SAME conversation when the continuation id is carried: a fresh + // request mints a random one (stableConversationId), which is precisely why calibration is + // keyed on it rather than on the adapter instance. + const conversationId = "11111111-2222-3333-4444-555555555555"; + const sameConversation = () => ({ + ...parsedWith(messages), + _providerContinuation: { kiro: { conversationId } }, + }); + + const first = createKiroAdapter(provider); + await first.buildRequest(sameConversation()); + const done = await doneUsage(first, eventFrame({ content: "ok" }), eventFrame({ contextUsagePercentage: 10 })); + expect(done.contextTotalTokens ?? 0).toBeGreaterThan(0); + + const second = createKiroAdapter(provider); + await second.buildRequest(sameConversation()); + const after = await doneUsage(second, eventFrame({ content: "ok" })); + + resetKiroCalibration(); + const uncalibrated = createKiroAdapter(provider); + await uncalibrated.buildRequest(sameConversation()); + const baseline = await doneUsage(uncalibrated, eventFrame({ content: "ok" })); + + expect(after.contextTotalTokens ?? 0).toBeGreaterThan(baseline.contextTotalTokens ?? 0); + }); + + // The negative half of the same contract: a turn that does not complete must teach nothing. A + // stream that fails mid-flight proves nothing about how densely its payload tokenized, and a + // turn that asks for the bounded completion retry is not over — it streams again against a + // rebuilt payload, so learning from the first attempt would move the factor twice for one turn. + test("an unfinished turn records no calibration", async () => { + const messages = [{ role: "user", content: "x".repeat(28_000) }]; + const conversationId = "99999999-8888-7777-6666-555555555555"; + const sameConversation = () => ({ + ...parsedWith(messages), + _providerContinuation: { kiro: { conversationId } }, + }); + + resetKiroCalibration(); + const failing = createKiroAdapter(provider); + await failing.buildRequest(sameConversation()); + // Report a context percentage, then end the stream with an upstream error rather than a + // completion. The percentage is real; the turn is not finished. + const events = await collectAdapterEvents(failing.parseStream(new Response(streamOf( + eventFrame({ contextUsagePercentage: 10 }), + eventFrame({ message: "boom" }, "invalidStateEvent"), + )))); + // Whatever the exact terminal shape, the turn did NOT deliver a completed answer. + expect(events.some(event => event.type === "done")).toBe(false); + + // Compare the ESTIMATE the next request carries, not a second completed turn — completing one + // would itself record and mask what we are asserting. + const after = createKiroAdapter(provider); + await after.buildRequest(sameConversation()); + const afterEstimate = (await doneUsage(after, eventFrame({ content: "ok" }))).contextTotalTokens; + + resetKiroCalibration(); + const baselineAdapter = createKiroAdapter(provider); + await baselineAdapter.buildRequest(sameConversation()); + const baseline = (await doneUsage(baselineAdapter, eventFrame({ content: "ok" }))).contextTotalTokens; + + expect(afterEstimate).toBe(baseline); + }); + + // The bounded completion retry is the case the terminal-only rule exists for: attempt 1 reports + // a real percentage but asks for a fallback, so the SAME user turn streams again against a + // rebuilt payload. Learning from attempt 1 would move the factor twice for one turn. + test("an attempt that falls back to the completion retry does not calibrate", async () => { + const messages = [{ role: "user", content: "x".repeat(28_000) }]; + const conversationId = "abcdabcd-1234-5678-9abc-def012345678"; + const sameConversation = (tools?: unknown[]) => ({ + ...parsedWith(messages, tools), + _providerContinuation: { kiro: { conversationId } }, + }); + + resetKiroCalibration(); + const originalFetch = globalThis.fetch; + globalThis.fetch = (async () => new Response(streamOf(eventFrame({ content: "Final from fallback." })))) as typeof fetch; + try { + const falling = createKiroAdapter(provider); + await falling.buildRequest(sameConversation([bashTool])); + // Text plus a context percentage, but no private completion call: this attempt needs the + // bounded fallback, so its observation must be discarded rather than learned. + await collectAdapterEvents(falling.parseStream(new Response(streamOf( + eventFrame({ content: "I am checking." }), + eventFrame({ contextUsagePercentage: 10 }), + )))); + } finally { + globalThis.fetch = originalFetch; + } + + const after = createKiroAdapter(provider); + await after.buildRequest(sameConversation()); + const afterEstimate = (await doneUsage(after, eventFrame({ content: "ok" }))).contextTotalTokens; + + resetKiroCalibration(); + const baselineAdapter = createKiroAdapter(provider); + await baselineAdapter.buildRequest(sameConversation()); + const baseline = (await doneUsage(baselineAdapter, eventFrame({ content: "ok" }))).contextTotalTokens; + + expect(afterEstimate).toBe(baseline); + }); + test("Kiro GPT routes use the Kiro token ratio without context percentage", async () => { const adapter = createKiroAdapter(provider); await adapter.buildRequest(parsedWith([{ role: "user", content: "x".repeat(3500) }], undefined, "gpt-5.6-sol")); const done = await doneUsage(adapter, eventFrame({ content: "y".repeat(3500) })); - expect(done.inputTokens).toBe(1000); - expect(done.outputTokens).toBe(1000); - expect(done.contextTotalTokens).toBe(2000); + expect(done.inputTokens).toBe(1250); + expect(done.outputTokens).toBe(1250); + expect(done.contextTotalTokens).toBe(2663); }); test("fresh payload includes history while usage counts only the current turn", async () => { @@ -1651,7 +1770,9 @@ describe("kiro adapter — parseStream", () => { const longUsage = await doneUsage(longAdapter, eventFrame({ content: "ok" })); expect(longBody.length).toBeGreaterThan(shortBody.length + 10_000); expect(longUsage.inputTokens).toBe(shortUsage.inputTokens); - expect(longUsage.inputTokens).toBe(estimateTokens(latest, "claude-sonnet-4.5")); + // The adapter estimates through the kiro-prefixed id, which carries the Kiro-measured ratio; + // the bare id is a different provider's route and keeps the shared agent ratio. + expect(longUsage.inputTokens).toBe(estimateTokens(latest, "kiro/claude-sonnet-4.5")); expect(longUsage.contextTotalTokens).toBeGreaterThan(shortUsage.contextTotalTokens ?? 0); }); @@ -1696,10 +1817,16 @@ describe("kiro adapter — parseStream", () => { const request = await adapter.buildRequest(parsedWith(messages)); const usage = await doneUsage(adapter, eventFrame({ content: "ok" })); - expect(usage.inputTokens).toBe(estimateTokens(latest, "claude-sonnet-4.5")); + expect(usage.inputTokens).toBe(estimateTokens(latest, "kiro/claude-sonnet-4.5")); expect(request.usageLog?.estimated).toBe(true); expect(request.usageLog?.inputTokens).toBeGreaterThan(usage.inputTokens + 4000); - expect(usage.contextTotalTokens).toBe((request.usageLog?.inputTokens ?? 0) + usage.outputTokens); + // Both estimates cover the whole conversation, but they are built for different consumers and + // are no longer expected to be equal. The context checkpoint is derived from the SERIALIZED + // payload, so it also carries JSON-escaping expansion and per-entry framing — real charged + // cost that a flat text join cannot see. The log estimate stays a plain text measure. + expect(usage.contextTotalTokens).toBeGreaterThanOrEqual( + (request.usageLog?.inputTokens ?? 0) + usage.outputTokens, + ); }); test("resumed payload preserves the complete locally expanded history", async () => { @@ -1721,7 +1848,7 @@ describe("kiro adapter — parseStream", () => { expect(resumedBody.length).toBe(freshBody.length); expect(cs.history).toHaveLength(4); expect(cs.currentMessage.userInputMessage.content).toBe(latest); - expect(resumedUsage.inputTokens).toBe(estimateTokens(latest, "claude-sonnet-4.5")); + expect(resumedUsage.inputTokens).toBe(estimateTokens(latest, "kiro/claude-sonnet-4.5")); }); test("tool-result follow-up counts new tool output without re-counting prior assistant tool args", async () => { diff --git a/tests/token-estimate.test.ts b/tests/token-estimate.test.ts index f74ab275692..a911a4ac2b7 100644 --- a/tests/token-estimate.test.ts +++ b/tests/token-estimate.test.ts @@ -1,21 +1,70 @@ import { describe, expect, test } from "bun:test"; import { capEstimateAtContextWindow, charsPerToken, estimateTokens } from "../src/lib/token-estimate"; -describe("CJK-aware ratio (devlog 260712 B3)", () => { +describe("script-segmented ratio", () => { const korean = "한국어 텍스트는 토큰 밀도가 높아서 영어 기준 추정이 과소계산됩니다 ".repeat(10); const english = "English text estimates fine at the default four chars per token ratio ".repeat(10); - - test("CJK-heavy text clamps DOWN to 2.5 chars/token (more tokens than English ratio)", () => { - expect(estimateTokens(korean, "gpt-5.6-sol")).toBe(Math.ceil(korean.length / 2.5)); + const cjkOf = (s: string) => [...s].filter(c => /[\uAC00-\uD7A3\u1100-\u11FF\u3130-\u318F\u4E00-\u9FFF\u3400-\u4DBF\u3040-\u30FF]/.test(c)).length; + const expected = (s: string, latin: number) => { + const cjk = cjkOf(s); + return Math.ceil((s.length - cjk) / latin + cjk / 1.5); + }; + + test("CJK characters are counted at their own denser ratio, not the model ratio", () => { + expect(estimateTokens(korean, "gpt-5.6-sol")).toBe(expected(korean, 4)); + expect(estimateTokens(korean, "kiro/claude-opus-5")).toBe(expected(korean, 2.8)); + expect(estimateTokens(korean, "claude-sonnet-4-6")).toBe(expected(korean, 3.5)); }); - test("Claude-shaped models clamp min(3.5, 2.5) -> 2.5 for CJK; keep 3.5 for English (audit R2#7)", () => { - expect(estimateTokens(korean, "claude-sonnet-4-6")).toBe(Math.ceil(korean.length / 2.5)); + test("pure-Latin text uses the model ratio alone", () => { + expect(estimateTokens(english, "gpt-5.6-sol")).toBe(Math.ceil(english.length / 4)); + expect(estimateTokens(english, "kiro/claude-opus-5")).toBe(Math.ceil(english.length / 2.8)); expect(estimateTokens(english, "claude-sonnet-4-6")).toBe(Math.ceil(english.length / 3.5)); }); - test("English text keeps the model ratio (no CJK trigger)", () => { - expect(estimateTokens(english, "gpt-5.6-sol")).toBe(Math.ceil(english.length / 4)); + // The Kiro-measured ratio must not reach the same model families routed by other providers: + // Cursor, Anthropic direct and Antigravity all read this helper for admission ceilings, + // count_tokens answers and overflow classification. + test("the Kiro ratio applies only to kiro-prefixed ids", () => { + expect(charsPerToken("kiro/claude-opus-5")).toBe(2.8); + for (const id of ["claude-4.6-opus-high", "claude-sonnet-4-6", "deepseek-3.2", "qwen3.8-27b", "glm-5", "minimax-m2.5"]) { + expect(charsPerToken(id)).toBe(3.5); + } + }); + + test("Korean costs about twice a Latin character", () => { + const ko = "한".repeat(300); + const en = "a".repeat(300); + const ratio = estimateTokens(ko, "kiro/claude-opus-5") / estimateTokens(en, "kiro/claude-opus-5"); + expect(ratio).toBeGreaterThan(1.5); + expect(ratio).toBeLessThan(2.5); + }); + + // The previous model switched divisor at a 30% sampled CJK share, so a blob at 29% and one at + // 31% differed by ~40% in estimate. Real agent traffic sits inside that band, which is exactly + // where the cliff never fired and the undercount was worst. + test("no cliff: the estimate is continuous across the old 30% threshold", () => { + const at = (share: number) => { + const total = 2000; + const cjk = Math.round(total * share); + return estimateTokens("한".repeat(cjk) + "a".repeat(total - cjk), "kiro/claude-opus-5"); + }; + const below = at(0.29); + const above = at(0.31); + // A 2-point change in composition must move the estimate by only a few percent. + expect(Math.abs(above - below) / below).toBeLessThan(0.05); + // ...and it must still be monotonically increasing in CJK share. + expect(above).toBeGreaterThan(below); + }); + + // Regression for the sampling bug documented in server/responses/input-admission.ts: the old + // stride sampler could read a 1.6%-CJK payload as 100% CJK. An exact count cannot. + test("a mostly-Latin blob with periodic CJK is not counted as CJK-heavy", () => { + const record = "id=0001,name=widget,qty=12,note=".padEnd(63, "x") + "한"; + const blob = record.repeat(400); + const cjk = cjkOf(blob); + expect(cjk / blob.length).toBeLessThan(0.02); + expect(estimateTokens(blob, "kiro/claude-opus-5")).toBe(expected(blob, 2.8)); }); }); @@ -24,8 +73,11 @@ describe("token-estimate sidecar", () => { expect(estimateTokens("", "claude-opus-4.8")).toBe(0); }); - test("kiro text models use the 3.5 ratio", () => { - for (const m of ["kiro-auto", "claude-opus-4.8", "claude-opus-4.5", "deepseek-3.2", "minimax-m2.5", "minimax-m2.1", "glm-5", "qwen3-coder-next"]) { + test("kiro-routed models use the 2.8 Latin ratio; the same families elsewhere keep 3.5", () => { + for (const m of ["kiro-auto", "kiro/claude-opus-4.8", "kiro/deepseek-3.2", "kiro/glm-5"]) { + expect(charsPerToken(m)).toBe(2.8); + } + for (const m of ["claude-opus-4.8", "claude-opus-4.5", "deepseek-3.2", "minimax-m2.5", "minimax-m2.1", "glm-5", "qwen3-coder-next"]) { expect(charsPerToken(m)).toBe(3.5); } }); @@ -40,11 +92,11 @@ describe("token-estimate sidecar", () => { expect(estimateTokens("ab", "claude-opus-4.8")).toBe(1); }); - test("estimate scales with length (ceil(len/3.5))", () => { - // 35 chars / 3.5 = 10 tokens - expect(estimateTokens("x".repeat(35), "claude-opus-4.8")).toBe(10); - // 36 chars / 3.5 = 10.28 -> ceil 11 - expect(estimateTokens("x".repeat(36), "claude-opus-4.8")).toBe(11); + test("estimate scales with length (ceil(len/2.8) on the kiro path)", () => { + // 28 chars / 2.8 = 10 tokens + expect(estimateTokens("x".repeat(28), "kiro/claude-opus-4.8")).toBe(10); + // 29 chars / 2.8 = 10.36 -> ceil 11 + expect(estimateTokens("x".repeat(29), "kiro/claude-opus-4.8")).toBe(11); }); test("lower ratio (kiro) yields more tokens than generic for same text (fail-safe over-count)", () => {