diff --git a/src/adapters/kiro.ts b/src/adapters/kiro.ts index 200b1edb77..35b33a922e 100644 --- a/src/adapters/kiro.ts +++ b/src/adapters/kiro.ts @@ -1,5 +1,5 @@ import { decodeEventStream } from "../lib/eventstream-decoder"; -import { estimateTokens } from "../lib/token-estimate"; +import { estimateTokens, estimateTokensFromCharacterCounts } from "../lib/token-estimate"; import { debugProviderDiagnostic } from "../lib/debug"; import { resolveKiroApiRegion, resolveKiroRequestProfile } from "../oauth/kiro"; import { KIRO_MODEL_CONTEXT_WINDOWS, normalizeKiroModelId } from "../providers/kiro-models"; @@ -209,8 +209,9 @@ function estimateKiroWireTokens(text: string, modelId: string): number { if (!text) return 0; const cjk = kiroCjkCount(text); if (cjk === 0) return Math.ceil(estimateKiroTokens(text, modelId) * KIRO_LATIN_WIRE_EXPANSION); - const latinTokens = estimateKiroTokens("x".repeat(text.length - cjk), modelId); - const cjkTokens = estimateKiroTokens("\uac00".repeat(cjk), modelId); + const prefixedModelId = `kiro/${modelId}`; + const latinTokens = estimateTokensFromCharacterCounts(text.length - cjk, 0, prefixedModelId); + const cjkTokens = estimateTokensFromCharacterCounts(0, cjk, prefixedModelId); return Math.ceil(latinTokens * KIRO_LATIN_WIRE_EXPANSION + cjkTokens); } diff --git a/src/lib/token-estimate.ts b/src/lib/token-estimate.ts index bf507e1cf2..0a9e4a7440 100644 --- a/src/lib/token-estimate.ts +++ b/src/lib/token-estimate.ts @@ -142,12 +142,24 @@ export function estimateTokens(text: string, modelId?: string, contextWindow?: n if (!text) return 0; const len = text.length; if (len === 0) return 0; - const latinRatio = charsPerToken(modelId); const cjk = countCjk(text); + return estimateTokensFromCharacterCounts(len - cjk, cjk, modelId, contextWindow); +} + +/** Estimate tokens from already-counted script buckets without materializing replacement text. */ +export function estimateTokensFromCharacterCounts( + latin: number, + cjk: number, + modelId?: string, + contextWindow?: number, +): number { + const len = latin + cjk; + if (len === 0) return 0; + const latinRatio = charsPerToken(modelId); // 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); + : Math.ceil(latin / latinRatio + cjk / CJK_CHARS_PER_TOKEN); return capEstimateAtContextWindow(Math.max(1, estimate), contextWindow); } diff --git a/tests/lib/token-estimate.test.ts b/tests/lib/token-estimate.test.ts index a7ba989da4..fa058fb0b9 100644 --- a/tests/lib/token-estimate.test.ts +++ b/tests/lib/token-estimate.test.ts @@ -1,5 +1,10 @@ import { describe, expect, test } from "bun:test"; -import { capEstimateAtContextWindow, charsPerToken, estimateTokens } from "../../src/lib/token-estimate"; +import { + capEstimateAtContextWindow, + charsPerToken, + estimateTokens, + estimateTokensFromCharacterCounts, +} from "../../src/lib/token-estimate"; describe("script-segmented ratio", () => { const korean = "한국어 텍스트는 토큰 밀도가 높아서 영어 기준 추정이 과소계산됩니다 ".repeat(10); @@ -10,6 +15,13 @@ describe("script-segmented ratio", () => { return Math.ceil((s.length - cjk) / latin + cjk / 1.5); }; + test("pre-counted script buckets preserve estimates without replacement strings", () => { + expect(estimateTokensFromCharacterCounts(10, 3, "kiro/kiro-auto")) + .toBe(estimateTokens("xxxxxxxxxx한한한", "kiro/kiro-auto")); + expect(estimateTokensFromCharacterCounts(250_000_000, 1, "kiro/kiro-auto")) + .toBe(Math.ceil(250_000_000 / 2.8 + 1 / 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));