From 21d27b40ca5a1d68588637849d8a498b68ca9f26 Mon Sep 17 00:00:00 2001 From: Kartik Date: Fri, 14 Aug 2026 07:27:17 +0530 Subject: [PATCH] fix(usage): cap token estimates at the model context window Port codex-router PR #140: a request the provider answered cannot have exceeded its context window, so the token estimate must never claim it did. estimateTokens gains an optional contextWindow cap. The request log caps the ESTIMATE field (attempt.inputTokenEstimate and the usage-fallback inputTokens) at the routed model's window while the combined inputTokens field keeps its max(reported, estimate) behavior; provider-reported positive counts are never reduced. Kiro auto and unknown adapters/models stay uncapped - a window is never invented. Window sources: KIRO_MODEL_CONTEXT_WINDOWS (src/providers/kiro-models.ts, mirroring adapters/kiro.ts kiroUpstreamContextWindow) and inferCursorContextWindow (src/adapters/cursor/discovery.ts), picked by route adapter in request-log.ts contextWindowForModel. Tests: extended tests/token-estimate.test.ts, added tests/request-log-estimate-cap.test.ts, updated the one request-log test that asserted an above-window estimate. --- src/lib/token-estimate.ts | 21 +++++- src/server/request-log.ts | 40 ++++++++++- tests/request-log-estimate-cap.test.ts | 99 ++++++++++++++++++++++++++ tests/request-log.test.ts | 6 +- tests/token-estimate.test.ts | 41 ++++++++++- 5 files changed, 200 insertions(+), 7 deletions(-) create mode 100644 tests/request-log-estimate-cap.test.ts diff --git a/src/lib/token-estimate.ts b/src/lib/token-estimate.ts index 10b535548..337ae6a5d 100644 --- a/src/lib/token-estimate.ts +++ b/src/lib/token-estimate.ts @@ -55,15 +55,32 @@ function cjkRatio(text: string): number { return sampled === 0 ? 0 : cjk / sampled; } +/** + * Cap an estimated token count at a model's context window (codex-router PR #140). + * + * A request the provider answered cannot have exceeded the window, so the estimate is never + * allowed to claim it did: the estimate only ever substitutes a reported-zero (or missing) + * input count, and a value above the window would be provably false. Positive or missing + * reported counts are never rewritten here — this bounds the heuristic estimate itself. + * Non-positive or non-integer windows (unknown model) leave the estimate untouched. + */ +export function capEstimateAtContextWindow(estimate: number, contextWindow: number | undefined): number { + return typeof contextWindow === "number" && Number.isInteger(contextWindow) && contextWindow > 0 + ? Math.min(estimate, contextWindow) + : estimate; +} + /** * Estimate the token count of a text blob. Pure and deterministic. * Returns 0 for empty/whitespace-free-empty input; otherwise ceil(length / ratio), min 1. + * When `contextWindow` is a positive integer the estimate is capped at it: a request the + * provider answered cannot have exceeded the window, so the estimate must not claim it did. */ -export function estimateTokens(text: string, modelId?: string): number { +export function estimateTokens(text: string, modelId?: string, contextWindow?: number): number { 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 Math.max(1, Math.ceil(len / ratio)); + return capEstimateAtContextWindow(Math.max(1, Math.ceil(len / ratio)), contextWindow); } diff --git a/src/server/request-log.ts b/src/server/request-log.ts index b4b92f93f..5e4592360 100644 --- a/src/server/request-log.ts +++ b/src/server/request-log.ts @@ -38,6 +38,10 @@ import { } from "../usage/debug"; import { matchesLogConversationId } from "./request-log-conversation"; import { enforceAppOwnedMemoryBudget, type RetainedStoreSnapshot } from "../lib/app-owned-memory"; +import { capEstimateAtContextWindow } from "../lib/token-estimate"; +import { inferCursorContextWindow } from "../adapters/cursor/discovery"; +import { KIRO_MODEL_CONTEXT_WINDOWS, normalizeKiroModelId } from "../providers/kiro-models"; +import { modelRecordValue } from "../reasoning-effort"; export interface RequestLogContext { model: string; @@ -839,6 +843,7 @@ export function addFinalRequestLog( logCtx.providerAdapter ?? logCtx.provider, logCtx.usage, logCtx.usageLogInputTokens, + contextWindowForModel(logCtx.providerAdapter ?? logCtx.provider, logCtx.model), ); const attempts = logCtx.attempts?.map(attempt => ({ ...attempt, @@ -960,15 +965,40 @@ interface FinalizedUsageResult { totalTokens?: number; } +/** + * Context window for the routed model, used to cap the token estimate (codex-router PR #140): + * a request the provider answered cannot have exceeded the window, so the estimate must never + * claim it did. The family is picked by the route ADAPTER, not the model id alone, because + * claude-family ids are shared between Kiro and Cursor with different windows. Kiro "auto" is + * a router with no fixed window and is never guessed; unknown adapters/models stay uncapped. + */ +function contextWindowForModel(adapter: string, modelId: string | undefined): number | undefined { + if (!modelId) return undefined; + if (adapter === "kiro" || adapter.startsWith("kiro-")) { + const normalized = normalizeKiroModelId(modelId); + if (normalized === "auto") return undefined; + return modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, modelId) + ?? modelRecordValue(KIRO_MODEL_CONTEXT_WINDOWS, normalized); + } + if (adapter === "cursor" || adapter.startsWith("cursor-")) { + return inferCursorContextWindow(modelId); + } + return undefined; +} + function finalizedUsage( adapter: string, usage: OcxUsage | undefined, inputTokenEstimate: number | undefined, + contextWindow: number | undefined, ): FinalizedUsageResult { + // The ESTIMATE itself is capped at the model's context window (codex-router PR #140). The + // combined value below keeps its max(inputTokens, estimate) behavior — a provider-reported + // positive count is never reduced by this cap, only the estimate that could substitute it. const estimate = typeof inputTokenEstimate === "number" && Number.isFinite(inputTokenEstimate) && inputTokenEstimate >= 0 - ? inputTokenEstimate + ? capEstimateAtContextWindow(inputTokenEstimate, contextWindow) : undefined; const finalUsage = usageForFinalLog(adapter, usage); const usageFallback = !finalUsage && estimate !== undefined @@ -1030,7 +1060,12 @@ export function noteAttemptSend( if (typeof inputTokenEstimate === "number" && Number.isFinite(inputTokenEstimate) && inputTokenEstimate >= 0) { - attempt.inputTokenEstimate = inputTokenEstimate; + // Store the ESTIMATE field already capped at the model's window (codex-router PR #140): + // what gets persisted, and later merged into usage, never claims a count above the window. + attempt.inputTokenEstimate = capEstimateAtContextWindow( + inputTokenEstimate, + contextWindowForModel(attempt.adapter, attempt.model), + ); } if (recovery && !attempt.recoveryKinds.includes(recovery)) { attempt.recoveryKinds.push(recovery); @@ -1047,6 +1082,7 @@ export function finishRequestAttempt( attempt.adapter, usage ?? attempt.usage, attempt.inputTokenEstimate, + contextWindowForModel(attempt.adapter, attempt.model), ); attempt.status = status; attempt.durationMs = Math.max(0, durationMs); diff --git a/tests/request-log-estimate-cap.test.ts b/tests/request-log-estimate-cap.test.ts new file mode 100644 index 000000000..b51f7aeb6 --- /dev/null +++ b/tests/request-log-estimate-cap.test.ts @@ -0,0 +1,99 @@ +import { describe, expect, test } from "bun:test"; +import { + beginRequestAttempt, + finishRequestAttempt, + noteAttemptSend, +} from "../src/server/request-log"; + +/** + * The token-estimate context-window cap (codex-router PR #140) in the request log: + * the ESTIMATE field is bounded by the routed model's context window, while the combined + * inputTokens field keeps its max(reported, estimate) behavior. Provider-reported positive + * counts are never reduced by the cap. + */ +describe("request log token-estimate context-window cap (codex-router PR #140)", () => { + const DEEPSEEK_WINDOW = 128_000; // KIRO_MODEL_CONTEXT_WINDOWS["deepseek-3.2"] + const CURSOR_CLAUDE_WINDOW = 200_000; // inferCursorContextWindow("claude-4.6-opus-high") + + test("the estimate is capped at the kiro model's context window", () => { + const attempt = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(attempt, 200_000); + finishRequestAttempt(attempt, 200, 10); + expect(attempt.inputTokenEstimate).toBe(DEEPSEEK_WINDOW); + expect(attempt.usage).toEqual({ inputTokens: DEEPSEEK_WINDOW, outputTokens: 0, estimated: true }); + }); + + test("below-window estimates pass through unchanged", () => { + const attempt = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(attempt, 100_000); + finishRequestAttempt(attempt, 200, 10); + expect(attempt.inputTokenEstimate).toBe(100_000); + expect(attempt.usage?.inputTokens).toBe(100_000); + }); + + test("a positive provider-reported input count is never reduced by the cap", () => { + const attempt = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(attempt, 200_000); // estimate would exceed the window + finishRequestAttempt(attempt, 200, 10, { inputTokens: 150_000, outputTokens: 500 }); + // combined = max(reported 150k, capped estimate 128k) = 150k: the reported count wins. + expect(attempt.usage?.inputTokens).toBe(150_000); + }); + + test("a missing report substitutes the capped estimate (the ESTIMATE field)", () => { + const attempt = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(attempt, 200_000); + finishRequestAttempt(attempt, 200, 10); + // No usage at all: the fallback inputTokens IS the capped estimate — never above the window. + expect(attempt.usage?.inputTokens).toBe(DEEPSEEK_WINDOW); + expect(attempt.usage?.estimated).toBe(true); + expect(attempt.usageStatus).toBe("estimated"); + }); + + test("the ESTIMATE field vs the combined inputTokens field distinction is preserved", () => { + // Reported zero: combined max(0, capped estimate) equals the capped estimate. + const zero = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(zero, 200_000); + finishRequestAttempt(zero, 200, 10, { inputTokens: 0, outputTokens: 5 }); + expect(zero.usage?.inputTokens).toBe(DEEPSEEK_WINDOW); + + // Reported positive above the capped estimate: combined keeps the reported count. + const positive = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(positive, 200_000); + finishRequestAttempt(positive, 200, 10, { inputTokens: 200_000, outputTokens: 5 }); + expect(positive.usage?.inputTokens).toBe(200_000); + + // Reported positive below the capped estimate: combined takes the capped estimate. + const middle = beginRequestAttempt(1, "kiro", "deepseek-3.2", "kiro"); + noteAttemptSend(middle, 100_000); + finishRequestAttempt(middle, 200, 10, { inputTokens: 50_000, outputTokens: 5 }); + expect(middle.usage?.inputTokens).toBe(100_000); + }); + + test("cursor adapter uses the cursor window inference", () => { + const attempt = beginRequestAttempt(1, "cursor", "claude-4.6-opus-high", "cursor"); + noteAttemptSend(attempt, 500_000); + finishRequestAttempt(attempt, 200, 10); + expect(attempt.inputTokenEstimate).toBe(CURSOR_CLAUDE_WINDOW); + }); + + test("unknown adapters stay uncapped (a window is never invented)", () => { + const attempt = beginRequestAttempt(1, "anthropic", "claude-sonnet-4-6", "anthropic"); + noteAttemptSend(attempt, 500_000); + finishRequestAttempt(attempt, 200, 10); + expect(attempt.inputTokenEstimate).toBe(500_000); + }); + + test("without a window the existing max(reported, estimate) behavior is unchanged", () => { + const attempt = beginRequestAttempt(1, "anthropic", "claude-sonnet-4-6", "anthropic"); + noteAttemptSend(attempt, 500_000); + finishRequestAttempt(attempt, 200, 10, { inputTokens: 100_000, outputTokens: 5 }); + expect(attempt.usage?.inputTokens).toBe(500_000); + }); + + test("kiro auto has no fixed window and is never guessed", () => { + const attempt = beginRequestAttempt(1, "kiro", "kiro-auto", "kiro"); + noteAttemptSend(attempt, 500_000); + finishRequestAttempt(attempt, 200, 10); + expect(attempt.inputTokenEstimate).toBe(500_000); + }); +}); diff --git a/tests/request-log.test.ts b/tests/request-log.test.ts index a8e750794..e1f6eefab 100644 --- a/tests/request-log.test.ts +++ b/tests/request-log.test.ts @@ -1053,10 +1053,12 @@ describe("request log metadata", () => { const text = await response.text(); expect(text).toContain("\"input_tokens\":9"); expect(entries).toHaveLength(1); + // The 240k estimate exceeds claude-sonnet-4.5's 200k window; a request the provider + // answered cannot have exceeded the window, so the estimate is capped (codex-router PR #140). expect(entries[0]).toMatchObject({ usageStatus: "estimated", - totalTokens: 240_004, - usage: { inputTokens: 240_000, outputTokens: 4, estimated: true }, + totalTokens: 200_004, + usage: { inputTokens: 200_000, outputTokens: 4, estimated: true }, }); }); diff --git a/tests/token-estimate.test.ts b/tests/token-estimate.test.ts index fa402f5ea..f74ab2756 100644 --- a/tests/token-estimate.test.ts +++ b/tests/token-estimate.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { charsPerToken, estimateTokens } from "../src/lib/token-estimate"; +import { capEstimateAtContextWindow, charsPerToken, estimateTokens } from "../src/lib/token-estimate"; describe("CJK-aware ratio (devlog 260712 B3)", () => { const korean = "한국어 텍스트는 토큰 밀도가 높아서 영어 기준 추정이 과소계산됩니다 ".repeat(10); @@ -61,3 +61,42 @@ describe("token-estimate sidecar", () => { } }); }); + +describe("context-window cap (codex-router PR #140)", () => { + const GPT56_SOL_WINDOW = 272_000; + const DEEPSEEK_WINDOW = 128_000; + // gpt-5.6-sol uses the generic 4-char ratio: 1.2M chars ~= 300k tokens, far over 272k. + const OVER_WINDOW_TEXT = "x".repeat(1_200_000); + + test("estimate is capped at the model's context window", () => { + expect(estimateTokens(OVER_WINDOW_TEXT, "gpt-5.6-sol", GPT56_SOL_WINDOW)).toBe(GPT56_SOL_WINDOW); + expect(estimateTokens(OVER_WINDOW_TEXT, "gpt-5.6-sol", DEEPSEEK_WINDOW)).toBe(DEEPSEEK_WINDOW); + }); + + test("below-window estimates are unchanged", () => { + const text = "x".repeat(100_000); // ~28.5k tokens at the kiro ratio + const plain = estimateTokens(text, "gpt-5.6-sol"); + expect(plain).toBeLessThan(GPT56_SOL_WINDOW); + expect(estimateTokens(text, "gpt-5.6-sol", GPT56_SOL_WINDOW)).toBe(plain); + }); + + test("only a positive integer window caps (unknown window leaves the estimate untouched)", () => { + const text = "x".repeat(100_000); + const plain = estimateTokens(text, "gpt-5.6-sol"); + for (const window of [undefined, 0, -1, 1.5, Number.NaN]) { + expect(estimateTokens(text, "gpt-5.6-sol", window)).toBe(plain); + } + }); + + test("the min-1 floor survives the cap for tiny inputs", () => { + expect(estimateTokens("a", "claude-opus-4.8", 1)).toBe(1); + expect(estimateTokens("a", "claude-opus-4.8", GPT56_SOL_WINDOW)).toBe(1); + }); + + test("capEstimateAtContextWindow caps only at positive integer windows", () => { + expect(capEstimateAtContextWindow(300_000, GPT56_SOL_WINDOW)).toBe(GPT56_SOL_WINDOW); + expect(capEstimateAtContextWindow(100_000, GPT56_SOL_WINDOW)).toBe(100_000); + expect(capEstimateAtContextWindow(300_000, undefined)).toBe(300_000); + expect(capEstimateAtContextWindow(300_000, 0)).toBe(300_000); + }); +});