From 69cd900a1f4d3f90af33167abeed7fab13b473b6 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 16:01:50 +1100 Subject: [PATCH 1/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20stop=20counting=20pro?= =?UTF-8?q?vider=20replay=20metadata=20in=20the=20context=20estimate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The budget walker counted providerOptions on assistant parts as input text. Anthropic thinking signatures are large base64 strings that Anthropic does not bill, so a 1M Opus chat hit its 20% handoff target at about 10% real usage. Skip providerOptions on every SDK message and message part. Tool-output content items keep it because some adapters JSON-stringify output.value. --- src/common/utils/compaction/contextBudget.ts | 90 +++++++------------ .../services/contextBudgetCounting.test.ts | 68 ++++++++------ 2 files changed, 71 insertions(+), 87 deletions(-) diff --git a/src/common/utils/compaction/contextBudget.ts b/src/common/utils/compaction/contextBudget.ts index 2b68196d881..d8fb5fc2eaf 100644 --- a/src/common/utils/compaction/contextBudget.ts +++ b/src/common/utils/compaction/contextBudget.ts @@ -156,6 +156,25 @@ export function estimateToolResultSize(result: unknown): { return measureBudgetContent(result); } +type BudgetContentKind = + | "json" + | "messages" + | "message" + | "parts" + | "part" + | "assistant-parts" + | "assistant-part" + | "output" + | "output-parts" + | "output-part"; + +const BUDGET_ITEM_KIND: Partial> = { + messages: "message", + parts: "part", + "assistant-parts": "assistant-part", + "output-parts": "output-part", +}; + function measureBudgetContent( result: unknown, textParts?: string[], @@ -167,23 +186,9 @@ function measureBudgetContent( let toolResultChars = 0; let imageParts = 0; const ancestors = new Set(); - const stack: Array<{ - value: unknown; - leave?: boolean; - kind?: - | "json" - | "messages" - | "message" - | "parts" - | "part" - | "assistant-parts" - | "assistant-part" - | "reasoning-options" - | "openai-reasoning-options" - | "tool-result-options" - | "mux-tool-result-options" - | "output"; - }> = [{ value: result, kind }]; + const stack: Array<{ value: unknown; leave?: boolean; kind?: BudgetContentKind }> = [ + { value: result, kind }, + ]; while (stack.length > 0) { // eslint-disable-next-line @typescript-eslint/no-non-null-assertion -- legacy site from before the ban; remove when touching this line const entry = stack.pop()!; @@ -220,18 +225,8 @@ function measureBudgetContent( stack.push({ value, leave: true }); toolResultChars += 2; if (Array.isArray(value)) { - for (const child of value) - stack.push({ - value: child, - kind: - entry.kind === "messages" - ? "message" - : entry.kind === "assistant-parts" - ? "assistant-part" - : entry.kind === "parts" - ? "part" - : "json", - }); + const itemKind = (entry.kind && BUDGET_ITEM_KIND[entry.kind]) ?? "json"; + for (const child of value) stack.push({ value: child, kind: itemKind }); toolResultChars += value.length; continue; } @@ -241,7 +236,8 @@ function measureBudgetContent( // Tool JSON can impersonate SDK part shapes. Only direct model-message/fresh // attachment parts get SDK media semantics; canonical tool wrappers are also // safe because the shared attachment sanitizer removes their data recursively. - const isPart = entry.kind === "part" || entry.kind === "assistant-part"; + const messagePart = entry.kind === "part" || entry.kind === "assistant-part"; + const isPart = messagePart || entry.kind === "output-part"; const image = isPart && record.type === "image" && "image" in record; const inlineText = typeof record.data === "object" && @@ -262,22 +258,10 @@ function measureBudgetContent( (urlMedia && key === "url") ) continue; - // Ciphertext length does not measure reasoning tokens. Provider usage tracks the underlying reasoning. - // Exclude only the SDK replay field; identical keys in user/tool JSON must still count. - if ( - entry.kind === "openai-reasoning-options" && - key === "reasoningEncryptedContent" && - typeof child === "string" - ) - continue; - // The dedupe projection's restoration stash (toolCatalog.ts) never reaches - // the wire: adapters serialize only their own providerOptions namespace. - if ( - entry.kind === "mux-tool-result-options" && - key === "rawToolSearchOutput" && - typeof child === "string" - ) - continue; + // providerOptions on a message or message part holds replay metadata, not input text: + // reasoning signatures and ciphertext, item IDs, cache hints, and Xum's own stash. + // Tool-output items keep it: some adapters send output.value as JSON text. + if (key === "providerOptions" && (messagePart || entry.kind === "message")) continue; toolResultChars += JSON.stringify(key).length + 2; textParts?.push(key); let childKind: typeof entry.kind = "json"; @@ -286,20 +270,8 @@ function measureBudgetContent( else if (record.role === "user" || record.role === "tool") childKind = "parts"; } else if (isPart && record.type === "tool-result" && key === "output") { childKind = "output"; - } else if (isPart && record.type === "tool-result" && key === "providerOptions") { - childKind = "tool-result-options"; - } else if (entry.kind === "tool-result-options" && key === "mux") { - childKind = "mux-tool-result-options"; } else if (entry.kind === "output" && record.type === "content" && key === "value") { - childKind = "parts"; - } else if ( - entry.kind === "assistant-part" && - record.type === "reasoning" && - key === "providerOptions" - ) { - childKind = "reasoning-options"; - } else if (entry.kind === "reasoning-options" && key === "openai") { - childKind = "openai-reasoning-options"; + childKind = "output-parts"; } stack.push({ value: child, kind: childKind }); } diff --git a/src/node/services/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index 484dedfbebd..9fc4042972c 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -194,41 +194,62 @@ describe("real-encoding budget guards", () => { ).toBeUndefined(); }); - test("encrypted OpenAI reasoning does not inflate either count or mutate the request", async () => { - const payload = (encrypted: string) => ({ + test.each([ + { + name: "OpenAI encrypted reasoning", + type: "reasoning", + options: (v: string) => ({ openai: { reasoningEncryptedContent: v } }), + }, + { + name: "an Anthropic thinking signature", + type: "reasoning", + options: (v: string) => ({ anthropic: { signature: v } }), + }, + { + name: "Anthropic redacted thinking", + type: "reasoning", + options: (v: string) => ({ anthropic: { redactedData: v } }), + }, + { + name: "a Google thought signature", + type: "text", + options: (v: string) => ({ google: { thoughtSignature: v } }), + }, + { + name: "OpenRouter reasoning details", + type: "tool-call", + options: (v: string) => ({ + openrouter: { reasoning_details: [{ type: "reasoning.encrypted", data: v }] }, + }), + }, + ])("$name does not inflate either count or mutate the request", async ({ type, options }) => { + const payload = (secret: string) => ({ messages: [ { role: "assistant", content: [ { - type: "reasoning", - text: "Check the result.", - providerOptions: { openai: { reasoningEncryptedContent: encrypted } }, + ...(type === "tool-call" + ? { type, toolCallId: "call", toolName: "read", input: {} } + : { type, text: "Check the result." }), + providerOptions: options(secret), }, ], }, ], }); - const encrypted = "a0b1c2d3e4f5".repeat(4000); - const request = payload(encrypted); + const secret = "a0b1c2d3e4f5".repeat(4000); + const request = payload(secret); expect(estimateAssembledRequestTokens(request)).toBe( estimateAssembledRequestTokens(payload("short")) ); expect( await checkAssembledRequestBudgetForModel(request, { model, modelContextLimit: 10000 }) ).toBeUndefined(); - expect(request.messages[0].content[0].providerOptions.openai.reasoningEncryptedContent).toBe( - encrypted - ); + expect(request.messages[0].content[0].providerOptions).toEqual(options(secret)); }); - test.each([ - "text", - "other-openai-field", - "other-provider", - "nested-field", - "non-string", - ] as const)("reasoning still counts %s beside encrypted OpenAI metadata", async (kind) => { + test("reasoning text still counts beside replay metadata", async () => { const large = "a0b1c2d3e4f5".repeat(4000); const request = { messages: [ @@ -237,17 +258,8 @@ describe("real-encoding budget guards", () => { content: [ { type: "reasoning", - text: kind === "text" ? large : "Check the result.", - providerOptions: { - openai: { - reasoningEncryptedContent: kind === "non-string" ? { text: large } : large, - ...(kind === "other-openai-field" ? { extra: large } : {}), - ...(kind === "nested-field" - ? { extra: { reasoningEncryptedContent: large } } - : {}), - }, - ...(kind === "other-provider" ? { xai: { reasoningEncryptedContent: large } } : {}), - }, + text: large, + providerOptions: { openai: { reasoningEncryptedContent: large } }, }, ], }, From 23c4a7873da0e920bfca24926d4949984e25c2fd Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 16:35:11 +1100 Subject: [PATCH 2/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20keep=20counting=20req?= =?UTF-8?q?uest=20options=20on=20user=20and=20tool=20content?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/common/utils/compaction/contextBudget.ts | 15 ++++++++----- .../services/contextBudgetCounting.test.ts | 21 +++++++++++++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/src/common/utils/compaction/contextBudget.ts b/src/common/utils/compaction/contextBudget.ts index d8fb5fc2eaf..0f38c8b3351 100644 --- a/src/common/utils/compaction/contextBudget.ts +++ b/src/common/utils/compaction/contextBudget.ts @@ -166,7 +166,8 @@ type BudgetContentKind = | "assistant-part" | "output" | "output-parts" - | "output-part"; + | "output-part" + | "request-options"; const BUDGET_ITEM_KIND: Partial> = { messages: "message", @@ -258,10 +259,12 @@ function measureBudgetContent( (urlMedia && key === "url") ) continue; - // providerOptions on a message or message part holds replay metadata, not input text: - // reasoning signatures and ciphertext, item IDs, cache hints, and Xum's own stash. - // Tool-output items keep it: some adapters send output.value as JSON text. - if (key === "providerOptions" && (messagePart || entry.kind === "message")) continue; + // On model output, providerOptions replays provider metadata (reasoning signatures and + // ciphertext, item IDs, cache hints), which providers do not bill as input text. + if (key === "providerOptions" && entry.kind === "assistant-part") continue; + // Options on user and tool content can reach the model as text (an Anthropic + // document's context and title), but adapters never send Xum's own namespace. + if (entry.kind === "request-options" && key === "mux") continue; toolResultChars += JSON.stringify(key).length + 2; textParts?.push(key); let childKind: typeof entry.kind = "json"; @@ -272,6 +275,8 @@ function measureBudgetContent( childKind = "output"; } else if (entry.kind === "output" && record.type === "content" && key === "value") { childKind = "output-parts"; + } else if (key === "providerOptions" && (messagePart || entry.kind === "message")) { + childKind = "request-options"; } stack.push({ value: child, kind: childKind }); } diff --git a/src/node/services/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index 9fc4042972c..e84dc9de5b1 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -274,6 +274,27 @@ describe("real-encoding budget guards", () => { ).toBe("context_budget_exceeded"); }); + test("request options on a user part still count", async () => { + // Anthropic sends a document part's context option to the model as text. + const attachment = { + type: "file", + mediaType: "application/pdf", + data: "JVBERi0xLjQ=", + providerOptions: { anthropic: { context: "a0b1c2d3e4f5".repeat(4000) } }, + }; + const request = { messages: [{ role: "user", content: [attachment] }] }; + expect( + (await checkAssembledRequestBudgetForModel(request, { model, modelContextLimit: 10000 })) + ?.type + ).toBe("context_budget_exceeded"); + expect( + await estimateFreshRequestTokensForModel( + { userText: "Summarize it.", attachments: [attachment], modelContextLimit: 10000 }, + { model } + ) + ).toBeGreaterThan(getContextBudgetHardCeiling(10000)); + }); + test.each(["tool-json", "tool-content", "tool-input", "user-text"] as const)( "reasoning metadata lookalikes remain counted in %s", async (kind) => { From ce3ccc3e0a4a58c3ea93cdc0dc03878801636812 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 17:10:03 +1100 Subject: [PATCH 3/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20count=20the=20turn-st?= =?UTF-8?q?art=20budget=20from=20the=20previous=20reply's=20provider=20usa?= =?UTF-8?q?ge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The turn-start check and step zero now count E as the newest reply's provider input plus the messages after it, when that reply is a completed one from the same model, route and agent. Growth of the sent system text since that reply is added. Otherwise they keep the full estimate. Stage prompts are decided on the same E, so a stage published at turn start is never refused at step zero. The loopback fixture now reports provider usage that tracks the request and runs the real tokenizer. --- .../services/contextBudgetCounting.test.ts | 137 +++++++++++++++ src/node/services/contextBudgetCounting.ts | 162 ++++++++++++++++-- .../strategies/tokenBudget.ts | 12 +- .../streamManager.contextBudget.test.ts | 86 +++++++++- .../services/turnEngine/ProviderRequest.ts | 4 + .../services/turnEngine/RefusalFallback.ts | 1 + .../services/turnEngine/StepPreparation.ts | 15 +- src/node/services/turnEngine/types.ts | 9 +- src/node/services/turnRequestBuilder.test.ts | 31 ++++ src/node/services/turnRequestBuilder.ts | 101 ++++++++--- ...enBudgetStages.fixture.integration.test.ts | 95 ++++++---- 11 files changed, 562 insertions(+), 91 deletions(-) diff --git a/src/node/services/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index e84dc9de5b1..dec0407bdac 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -8,10 +8,13 @@ import { type AssembledRequestBudgetInput, } from "@/common/utils/compaction/contextBudget"; import type { JSONSchema7 } from "@ai-sdk/provider"; +import { createMuxMessage, type MuxMetadata } from "@/common/types/message"; +import { CLAUDE_ENCODING_BUDGET_FACTOR } from "@/constants/tokenizerCorrection"; import calibration from "./__fixtures__/contextBudgetClaudeTokenizer.json"; import { checkAssembledRequestBudgetForModel, createContextBudgetAnchor, + createTurnBudgetAnchor, estimateAnchoredRequestTokensForModel, estimateAssembledRequestTokensForModel, estimateFreshRequestTokensForModel, @@ -838,6 +841,132 @@ describe("anchored request estimate (#4858)", () => { }); }); +describe("turn-start anchor (#7662)", () => { + const history = [{ role: "user", content: "long history ".repeat(4000) }]; + const reply = { role: "assistant", content: "answer ".repeat(2000) }; + const prompt = { role: "user", content: "follow-up question ".repeat(20) }; + const payload = { messages: [...history, reply, prompt] }; + const options = { model, modelContextLimit: 200_000 }; + const system = "Answer briefly and cite the files you read."; + const systemTokens = async (forModel: string) => + (await tokenizerModule.getTokenizerForModel(forModel)).countTokens(system); + const request = { model, agentId: "exec", routeProvider: "openai", system }; + const anchorFor = async ( + metadata?: Partial, + overrides?: Partial + ) => { + const reply = createMuxMessage("a1", "assistant", "answer", { + model, + agentId: "exec", + routeProvider: "openai", + systemMessageTokens: await systemTokens(overrides?.model ?? model), + contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, + ...metadata, + }); + return createTurnBudgetAnchor([createMuxMessage("u1", "user", "hi"), reply], { + ...request, + ...overrides, + }); + }; + + test("costs the previous reply's input plus the messages from its last step on", async () => { + const anchored = await estimateAnchoredRequestTokensForModel( + payload, + options, + await anchorFor() + ); + const delta = (await estimateAssembledRequestTokensForModel( + { messages: [reply, prompt] }, + options + ))!.estimate; + expect(anchored).toEqual({ + estimate: 1000 + delta, + hardCeiling: getContextBudgetHardCeiling(200_000), + }); + expect(anchored!.estimate).toBeLessThan( + (await estimateAssembledRequestTokensForModel(payload, options))!.estimate + ); + }); + + test("counts cache reads and writes once, adds system growth, and leaves out the reply's reasoning", async () => { + const contextUsage = { + inputTokens: 1000, + outputTokens: 80, + cachedInputTokens: 600, + reasoningTokens: 50, + totalTokens: 1080, + }; + // Rows keep cache writes in provider metadata: 300 uncached + 600 read + 100 written. + const contextProviderMetadata = { anthropic: { cacheCreationInputTokens: 100 } }; + const sent = await systemTokens(model); + expect((await anchorFor({ contextUsage, contextProviderMetadata }))?.providerTokens).toBe(1000); + const grown = { contextUsage, systemMessageTokens: sent - 30 }; + expect((await anchorFor(grown))?.providerTokens).toBe(1030); + const shrunk = { contextUsage, systemMessageTokens: sent + 30 }; + expect((await anchorFor(shrunk))?.providerTokens).toBe(1000); + // Claude-encoded growth is scaled like every other claude count (#5219). Tests default to + // the approximate tokenizer, which has no claude encoding. + const original = process.env.XUM_FORCE_REAL_TOKENIZER; + process.env.XUM_FORCE_REAL_TOKENIZER = "1"; + try { + const claude = "anthropic:claude-opus-5-5"; + const claudeGrown = { + model: claude, + contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, + systemMessageTokens: (await systemTokens(claude)) - 30, + }; + expect((await anchorFor(claudeGrown, { model: claude }))?.providerTokens).toBe( + 1000 + Math.ceil(30 * CLAUDE_ENCODING_BUDGET_FACTOR) + ); + } finally { + if (original == null) delete process.env.XUM_FORCE_REAL_TOKENIZER; + else process.env.XUM_FORCE_REAL_TOKENIZER = original; + } + }); + + test.each<[string, Partial, Partial]>([ + ["a partial reply", { partial: true }, {}], + ["a synthetic reply", { synthetic: true }, {}], + ["a failed reply", { error: "boom" }, {}], + ["a reply with an error type", { errorType: "unknown" }, {}], + ["a compaction summary", { compacted: "user" }, {}], + ["a context boundary", { compactionBoundary: true }, {}], + ["a reply without usage", { contextUsage: undefined }, {}], + [ + "a reply without input", + { contextUsage: { inputTokens: 0, outputTokens: 5, totalTokens: 5 } }, + {}, + ], + ["a reply without its system count", { systemMessageTokens: undefined }, {}], + ["a reply without its model", { model: undefined }, {}], + ["another model", {}, { model: "openai:gpt-4.1" }], + ["another route", {}, { routeProvider: "openrouter" }], + ["another agent", {}, { agentId: "plan" }], + ])("%s yields no turn anchor", async (_name, metadata, overrides) => { + expect(await anchorFor(metadata, overrides)).toBeUndefined(); + }); + + test("a reply marked as not compacted still anchors", async () => { + expect((await anchorFor({ compacted: false }))?.providerTokens).toBe(1000); + }); + + test("a row without a route or agent still anchors; history without a reply does not", async () => { + expect(await anchorFor({ agentId: undefined, routeProvider: undefined })).toBeDefined(); + expect( + await createTurnBudgetAnchor([createMuxMessage("u1", "user", "hi")], request) + ).toBeUndefined(); + }); + + test.each([ + ["no assistant message", { messages: [...history, prompt] }, options], + ["another model", payload, { ...options, model: "openai:gpt-4.1" }], + ])("%s falls back to the full estimate", async (_name, input, opts) => { + expect(await estimateAnchoredRequestTokensForModel(input, opts, await anchorFor())).toEqual( + await estimateAssembledRequestTokensForModel(input, opts) + ); + }); +}); + describe("repeated native tool references (#5413)", () => { // Anthropic expands every tool_reference into the full definition, repeats included. const claude = "anthropic:claude-opus-5-5"; @@ -912,4 +1041,12 @@ describe("repeated native tool references (#5413)", () => { await estimateAssembledRequestTokensForModel(payload, options) ); }); + + test("a reference after the turn anchor's reply falls back to the full estimate", async () => { + const turnAnchor = { model: claude, providerTokens: 1000 }; + const payload = { tools, messages: once }; + expect(await estimateAnchoredRequestTokensForModel(payload, options, turnAnchor)).toEqual( + await estimateAssembledRequestTokensForModel(payload, options) + ); + }); }); diff --git a/src/node/services/contextBudgetCounting.ts b/src/node/services/contextBudgetCounting.ts index 5dca3f21379..b12dc99fb42 100644 --- a/src/node/services/contextBudgetCounting.ts +++ b/src/node/services/contextBudgetCounting.ts @@ -14,6 +14,8 @@ import { BUDGET_TOKEN_CHUNK_SLACK, REQUEST_FRAMING_TOKENS, } from "@/common/constants/contextBudget"; +import type { MuxMessage } from "@/common/types/message"; +import { normalizeToCanonical } from "@/common/utils/ai/models"; import { createDisplayUsage } from "@/common/utils/tokens/displayUsage"; import { normalizeUsage, type AiSdkUsageLike } from "@/common/utils/tokens/usageHelpers"; import { countToolReferences } from "@/common/utils/tools/toolCatalog"; @@ -162,6 +164,28 @@ export interface ContextBudgetAnchor extends ContextBudgetAnchorRequest { providerTokens: number; } +/** + * A turn-start request's anchor (#7662): the provider-counted input of the previous turn's last + * step. That step's request is this request up to its last assistant message (the step's own + * output), so no request object is kept: the delta is found in the messages at count time. + */ +export interface TurnBudgetAnchor extends BudgetModel { + providerTokens: number; +} + +/** A step's usage normalized for display, or undefined without a positive integer input total. */ +function reportedUsage( + usage: AiSdkUsageLike, + model: string, + providerMetadata: Record | undefined +) { + const normalized = normalizeUsage(usage); + // Cache or reasoning details without an input total do not count the whole request. + const input = normalized.inputTokens; + if (input == null || !Number.isSafeInteger(input) || input <= 0) return undefined; + return createDisplayUsage(normalized, model, providerMetadata); +} + /** * Pairs a step's request with the usage the provider reported for it. Reasoning output is * added because the next request replays it and the local estimator skips encrypted @@ -179,11 +203,7 @@ export function createContextBudgetAnchor( | undefined ): ContextBudgetAnchor | undefined { if (request == null || step?.usage == null) return undefined; - const normalized = normalizeUsage(step.usage); - // Cache or reasoning details without an input total do not count the whole request. - const input = normalized.inputTokens; - if (input == null || !Number.isSafeInteger(input) || input <= 0) return undefined; - const usage = createDisplayUsage(normalized, request.model, step.providerMetadata); + const usage = reportedUsage(step.usage, request.model, step.providerMetadata); if (usage == null) return undefined; // Display recovery also reads provider-metadata-only reasoning counts. Zero cannot tell // "no reasoning" from "not reported", so a step that emitted reasoning needs a count. @@ -196,6 +216,89 @@ export function createContextBudgetAnchor( : undefined; } +/** Both set and different: a row without the field does not count as a switch. */ +function differs(before: string | undefined, now: string | undefined): boolean { + return before != null && now != null && before !== now; +} + +type ReplyMetadata = NonNullable; + +/** A completed reply: not partial, synthetic, failed, a compaction summary or a boundary. */ +function isCompletedReply(metadata: ReplyMetadata): boolean { + // A compaction summary carries an estimated usage, and starts a new segment. + const compacted = metadata.compacted != null && metadata.compacted !== false; + return !( + metadata.partial === true || + metadata.synthetic === true || + metadata.error != null || + metadata.errorType != null || + compacted || + metadata.compactionBoundary === true || + metadata.contextBoundaryKind != null + ); +} + +/** The model, route and agent switches that re-freeze the system prompt. */ +function sameRequestRoute( + metadata: ReplyMetadata & { model: string }, + request: BudgetModel & { routeProvider?: string; agentId?: string } +): boolean { + // Rows record the canonical model (Coder ids raw), so compare both sides canonicalized. + return ( + normalizeToCanonical(metadata.model) === normalizeToCanonical(request.model) && + !differs(metadata.routeProvider, request.routeProvider) && + !differs(metadata.agentId, request.agentId) + ); +} + +/** The newest reply's metadata when its saved usage can anchor this request, else undefined. */ +function anchoringReply( + rows: readonly MuxMessage[], + request: BudgetModel & { routeProvider?: string; agentId?: string } +) { + const metadata = rows.findLast((row) => row.role === "assistant")?.metadata; + const { contextUsage, model, systemMessageTokens } = metadata ?? {}; + if (contextUsage == null || model == null || systemMessageTokens == null) return undefined; + const reply = { ...metadata, contextUsage, model, systemMessageTokens }; + return isCompletedReply(reply) && sameRequestRoute(reply, request) ? reply : undefined; +} + +/** + * The anchor for a turn-start request built from `rows`, its provider request rows (#7662). + * The newest assistant row must be a completed reply from the same model, route and agent (the + * switches that re-freeze the system prompt), with the provider's input total for its last step + * and its sent system token count. Undefined (so the caller uses the full estimate) otherwise. + * + * Reasoning output is not added, unlike the in-turn anchor: the request carries that step's + * output, reasoning text included, after the anchored prefix, so the estimator counts it. + * Growth of `system`, the sent system text, since that request (a live system prompt changes + * between turns when the provider takes no system rows) is added, because the provider count + * holds the old prompt. New tail rows go after the newest reply, where the delta counts them. + */ +export async function createTurnBudgetAnchor( + rows: readonly MuxMessage[], + request: BudgetModel & { routeProvider?: string; agentId?: string; system: string } +): Promise { + const reply = anchoringReply(rows, request); + if (reply == null) return undefined; + const usage = reportedUsage(reply.contextUsage, request.model, reply.contextProviderMetadata); + if (usage == null) return undefined; + // The row counted its system text, tail rows included, with this tokenizer (measureSentSystem), + // so its rows only make the growth smaller. Claude-encoded growth is scaled like every other + // claude count (#5219). + const tokenizer = await getTokenizerForModel(request.model, request.metadataModel); + const growth = (await tokenizer.countTokens(request.system)) - reply.systemMessageTokens; + const factor = tokenizer.encoding === "claude" ? CLAUDE_ENCODING_BUDGET_FACTOR : 1; + const providerTokens = + usage.input.tokens + + usage.cached.tokens + + usage.cacheCreate.tokens + + Math.ceil(Math.max(0, growth) * factor); + return Number.isSafeInteger(providerTokens) + ? { model: request.model, metadataModel: request.metadataModel, providerTokens } + : undefined; +} + export function isExactAppend( payload: AssembledRequestBudgetInput, options: BudgetModel & { activeTools?: readonly string[] }, @@ -218,12 +321,41 @@ export function isExactAppend( } /** - * Anchored estimate (#4858). Contract: the anchor is used ONLY when this request is an exact - * append of the step whose usage it carries: same model, the same system prompt and tool set - * objects, the same advertised tool names, a message list that starts with that step's sent + * The messages the anchor's provider count does not hold, or undefined when the anchor does + * not apply. A turn anchor's delta starts at the last assistant message: the previous turn's + * last step output, which Xum's history rebuild sends as its own message (stepStartPartIndices). + */ +function anchoredDelta( + payload: AssembledRequestBudgetInput, + options: BudgetModel & { activeTools?: readonly string[] }, + anchor: ContextBudgetAnchor | TurnBudgetAnchor +): readonly unknown[] | undefined { + if ("messages" in anchor) { + return isExactAppend(payload, options, anchor) + ? payload.messages.slice(anchor.messages.length) + : undefined; + } + if (anchor.model !== options.model || anchor.metadataModel !== options.metadataModel) { + return undefined; + } + const start = payload.messages.findLastIndex( + (message) => (message as { role?: unknown } | null)?.role === "assistant" + ); + if (start < 0) return undefined; + const delta = payload.messages.slice(start); + // A tool_reference expands into a definition the delta never charges (#5413), as above. + return countToolReferences(delta).size === 0 ? delta : undefined; +} + +/** + * Anchored estimate (#4858). Contract: an in-turn anchor is used ONLY when this request is an + * exact append of the step whose usage it carries: same model, the same system prompt and tool + * set objects, the same advertised tool names, a message list that starts with that step's sent * messages (element identity, so any rewrite of the prefix fails), and a positive integer - * provider input total plus a reasoning count whenever the step emitted reasoning. The estimate is then the provider's count plus the existing estimator applied to the - * appended messages only. ANY doubt means the full estimate; edge cases get a new full-estimate + * provider input total plus a reasoning count whenever the step emitted reasoning. A turn + * anchor (createTurnBudgetAnchor) needs the same model and a last assistant message. The + * estimate is then the provider's count plus the existing estimator applied to the appended + * messages only. ANY doubt means the full estimate; edge cases get a new full-estimate * condition here, never new mechanism. */ export async function estimateAnchoredRequestTokensForModel( @@ -232,15 +364,15 @@ export async function estimateAnchoredRequestTokensForModel( modelContextLimit: number | null | undefined; activeTools?: readonly string[]; }, - anchor: ContextBudgetAnchor | undefined + anchor: ContextBudgetAnchor | TurnBudgetAnchor | undefined ): Promise<{ estimate: number; hardCeiling: number } | undefined> { - if (anchor == null || !isExactAppend(payload, options, anchor)) { + const delta = anchor == null ? undefined : anchoredDelta(payload, options, anchor); + if (anchor == null || delta == null) { return estimateAssembledRequestTokensForModel(payload, options); } const limit = options.modelContextLimit; if (limit == null || !Number.isFinite(limit) || limit <= 0) return undefined; const hardCeiling = getContextBudgetHardCeiling(limit); - const delta = payload.messages.slice(anchor.messages.length); // Same early-exit lower bound as the full estimate: above the ceiling, not exact. const deltaEstimate = await countBudgetInput( prepareAssembledRequestTokenCount({ messages: delta }), @@ -261,7 +393,7 @@ export async function measureAssembledRequestBudgetForModel( modelContextLimit: number | null | undefined; activeTools?: readonly string[]; }, - anchor?: ContextBudgetAnchor + anchor?: ContextBudgetAnchor | TurnBudgetAnchor ): Promise< { estimate: number; hardCeiling: number; exceeded?: ContextBudgetExceeded } | undefined > { @@ -283,7 +415,7 @@ export async function checkAssembledRequestBudgetForModel( modelContextLimit: number | null | undefined; activeTools?: readonly string[]; }, - anchor?: ContextBudgetAnchor + anchor?: ContextBudgetAnchor | TurnBudgetAnchor ): Promise { return (await measureAssembledRequestBudgetForModel(payload, options, anchor))?.exceeded; } diff --git a/src/node/services/contextManagement/strategies/tokenBudget.ts b/src/node/services/contextManagement/strategies/tokenBudget.ts index be6bed3ee42..52645c0d9d3 100644 --- a/src/node/services/contextManagement/strategies/tokenBudget.ts +++ b/src/node/services/contextManagement/strategies/tokenBudget.ts @@ -509,11 +509,13 @@ export class TokenBudgetStrategy { /** * Turn-start stage decision (#5286). A stage prompt is decided and published by the send that - * delivers it, in TurnRequestBuilder.start(), from E, the full estimate the turn-start check - * counted on the built request. It opens when it is unclaimed in this window (warning rows in - * the turn's own messages) and E reaches its point. The builder then counts the request with - * the row (E′) and publishes the row only if `fits(E′)`. Due and fit never read saved provider - * usage or in-memory claims, so a restart, an edit or a resume decides the same way. + * delivers it, in TurnRequestBuilder.start(), from E, the count the turn-start check enforced + * on the built request: the newest reply's saved provider usage plus the messages after it + * (createTurnBudgetAnchor), else the full estimate. It opens when it is unclaimed in this + * window (warning rows in the turn's own messages) and E reaches its point. The builder then + * counts the request with the row (E′) and publishes the row only if `fits(E′)`. Due and fit + * read only the turn's history rows, never in-memory claims, so a restart, an edit or a resume + * decides the same way. */ async prepareTurnStartStage( budget: Parameters[0], diff --git a/src/node/services/streamManager.contextBudget.test.ts b/src/node/services/streamManager.contextBudget.test.ts index 9917aaec577..2b1ebf53f76 100644 --- a/src/node/services/streamManager.contextBudget.test.ts +++ b/src/node/services/streamManager.contextBudget.test.ts @@ -14,7 +14,10 @@ import { z } from "zod"; import * as fs from "node:fs/promises"; import * as path from "node:path"; import { createMuxMessage } from "@/common/types/message"; -import { evaluateStepBudget } from "@/common/utils/compaction/contextBudget"; +import { + evaluateStepBudget, + getContextBudgetHardCeiling, +} from "@/common/utils/compaction/contextBudget"; import { LocalRuntime } from "@/node/runtime/LocalRuntime"; import { StreamManager, type SettledStepBudget } from "./streamManager"; import { createTestHistoryService } from "./testHistoryService"; @@ -224,6 +227,87 @@ describe("settled context hard ceiling", () => { } ); + // #7662: step zero has no settled step, so it counts against the turn anchor, as the + // turn-start check did; a full estimate here could refuse what turn start accepted. + test.each([false, true])( + "step zero counts against the turn anchor (anchored=%p)", + async (anchored) => { + const h = await createTestHistoryService(); + const workspaceId = "turn-anchor-budget"; + const messageId = "turn-anchor-assistant"; + let providerCalls = 0; + const model = new MockLanguageModelV3({ + doStream: () => { + providerCalls++; + return Promise.resolve({ + stream: simulateReadableStream({ + chunks: [ + { type: "stream-start", warnings: [] }, + { type: "text-start", id: "answer" }, + { type: "text-delta", id: "answer", delta: "Done" }, + { type: "text-end", id: "answer" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 100, noCache: 100, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 10, text: 10, reasoning: 0 }, + }, + }, + ] satisfies LanguageModelV3StreamPart[], + }), + }); + }, + }); + const manager = new StreamManager(h.historyService); + const runtimeDir = await fs.mkdtemp(path.join(tmpdir(), "turn-anchor-budget-")); + try { + expect( + ( + await h.historyService.appendManyToHistory(workspaceId, [ + createMuxMessage("user", "user", "Hello"), + createMuxMessage(messageId, "assistant", ""), + ]) + ).success + ).toBe(true); + const started = await manager.startStream({ + workspaceId, + messageId, + historySequence: 1, + model, + modelString: "openai:gpt-4o", + messages: [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi" }, + { role: "user", content: "Again" }, + ], + system: "Answer briefly", + runtime: new LocalRuntime(h.tempDir), + providedRuntimeTempDir: runtimeDir, + contextBudgetLimit: 10000, + // The previous reply's provider count alone is at the ceiling. + ...(anchored + ? { + contextBudgetTurnAnchor: { + model: "openai:gpt-4o", + metadataModel: "openai:gpt-4o", + providerTokens: getContextBudgetHardCeiling(10000), + }, + } + : {}), + }); + if (!started.success) throw new Error("Expected stream construction"); + const completion = await started.data.completion; + expect(completion.status).toBe(anchored ? "failed" : "completed"); + expect(providerCalls).toBe(anchored ? 0 : 1); + } finally { + await manager.stopStream(workspaceId); + await fs.rm(runtimeDir, { recursive: true, force: true }); + await h.cleanup(); + } + } + ); + test.each(["fits", "overflow", "stale-activation"] as const)( "native tool search sends every tool each step and budgets only loaded schemas (%s)", async (mode) => { diff --git a/src/node/services/turnEngine/ProviderRequest.ts b/src/node/services/turnEngine/ProviderRequest.ts index 5bc5be564f8..86ac0993cf8 100644 --- a/src/node/services/turnEngine/ProviderRequest.ts +++ b/src/node/services/turnEngine/ProviderRequest.ts @@ -33,6 +33,7 @@ import { getModelStats, getModelStatsResolved } from "@/common/utils/tokens/mode import type { ToolSearchStreamState } from "@/common/utils/tools/toolCatalog"; import type { ToolPolicy } from "@/common/utils/tools/toolPolicy"; +import type { TurnBudgetAnchor } from "@/node/services/contextBudgetCounting"; import { stripMessageCacheControl, type ContinuousPrefixSwap, @@ -76,6 +77,7 @@ export interface StreamRequestOptions { onStepMessages?: (messages: ModelMessage[]) => void; onStepSettled?: OnStepSettled; contextBudgetLimit?: number; + contextBudgetTurnAnchor?: TurnBudgetAnchor; toolSearchState?: ToolSearchStreamState; thinkingOverrideState?: ActiveTurnThinkingOverride; rebuildProviderOptionsForThinkingLevel?: RebuildProviderOptionsForThinkingLevel; @@ -315,6 +317,7 @@ export class ProviderRequest { onStepMessages, onStepSettled, contextBudgetLimit, + contextBudgetTurnAnchor, toolSearchState, onToolExecutionStart, thinkingOverrideState, @@ -370,6 +373,7 @@ export class ProviderRequest { onStepMessages, onStepSettled, contextBudgetLimit, + contextBudgetTurnAnchor, toolPolicy, toolSearchState, thinkingOverrideState, diff --git a/src/node/services/turnEngine/RefusalFallback.ts b/src/node/services/turnEngine/RefusalFallback.ts index 15c534074ca..5f1c68ff6a4 100644 --- a/src/node/services/turnEngine/RefusalFallback.ts +++ b/src/node/services/turnEngine/RefusalFallback.ts @@ -531,6 +531,7 @@ export class RefusalFallback { onStepMessages: streamInfo.request.onStepMessages, onStepSettled: streamInfo.request.onStepSettled, contextBudgetLimit: prepared.contextBudgetLimit, + contextBudgetTurnAnchor: prepared.contextBudgetTurnAnchor, // Same state object: aiService's fallback prepare() rebuilt it in place // against the fallback toolset, so prepareStep keeps reading live state. toolSearchState: streamInfo.request.toolSearchState, diff --git a/src/node/services/turnEngine/StepPreparation.ts b/src/node/services/turnEngine/StepPreparation.ts index e138ad8d3a5..818303e3252 100644 --- a/src/node/services/turnEngine/StepPreparation.ts +++ b/src/node/services/turnEngine/StepPreparation.ts @@ -36,6 +36,7 @@ import { checkAssembledRequestBudgetForModel, createContextBudgetAnchor, type ContextBudgetAnchorRequest, + type TurnBudgetAnchor, } from "@/node/services/contextBudgetCounting"; import { ContextBudgetBlockedError } from "@/node/services/contextBudgetError"; import type { HistoryService } from "@/node/services/historyService"; @@ -434,17 +435,22 @@ export function buildContextBudgetRequest( }; } -/** Hard-stop the step when its assembled request exceeds the safe context budget. */ +/** + * Hard-stop the step when its assembled request exceeds the safe context budget. Step zero has + * no settled step, so it counts against the turn anchor, as the turn-start check did (#7662). + */ async function enforceContextBudgetPreflight( request: Pick, modelContextLimit: number, budgetRequest: ContextBudgetAnchorRequest, previousRequest: ContextBudgetAnchorRequest | undefined, lastStep: Parameters[1], - stepNumber: number + stepNumber: number, + turnAnchor: TurnBudgetAnchor | undefined ): Promise { const budgetOptions = { ...budgetRequest, modelContextLimit }; - const anchor = createContextBudgetAnchor(previousRequest, lastStep); + const anchor = + stepNumber === 0 ? turnAnchor : createContextBudgetAnchor(previousRequest, lastStep); const exceeded = await checkAssembledRequestBudgetForModel(budgetRequest, budgetOptions, anchor); // Step zero can follow executed tools on a fallback. This late hard stop // preserves settled results; it must not reset/replay the activated catalog. @@ -715,7 +721,8 @@ export class StepPreparation { budgetRequest, previousRequest, steps.at(-1), - stepNumber + stepNumber, + request.contextBudgetTurnAnchor ); } return buildPrepareStepResult({ diff --git a/src/node/services/turnEngine/types.ts b/src/node/services/turnEngine/types.ts index 3b060d1517f..f37c4c6996f 100644 --- a/src/node/services/turnEngine/types.ts +++ b/src/node/services/turnEngine/types.ts @@ -13,7 +13,10 @@ import type { QueuedInputStopCause, StreamStopCause } from "@/common/types/strea import type { ToolSearchStreamState } from "@/common/utils/tools/toolCatalog"; import type { ToolPolicy } from "@/common/utils/tools/toolPolicy"; import type { AutoThinkingEscalationState } from "@/node/services/autoThinkingEscalation"; -import type { ContextBudgetAnchorRequest } from "@/node/services/contextBudgetCounting"; +import type { + ContextBudgetAnchorRequest, + TurnBudgetAnchor, +} from "@/node/services/contextBudgetCounting"; import type { ContinuousPrefixSwap } from "@/node/services/continuousCompactionJournal"; import type { ProvidersConfigMap } from "@/common/orpc/types"; import type { SendMessageError, StreamErrorType } from "@/common/types/errors"; @@ -131,6 +134,8 @@ export type OnStepSettled = (step: SettledStepBudget) => Promise void; onStepSettled?: OnStepSettled; contextBudgetLimit?: number; + /** Step 0's budget anchor: the previous turn's provider count (createTurnBudgetAnchor). */ + contextBudgetTurnAnchor?: TurnBudgetAnchor; toolPolicy?: ToolPolicy; /** * Tool-search deferral state. Owned and mutated by diff --git a/src/node/services/turnRequestBuilder.test.ts b/src/node/services/turnRequestBuilder.test.ts index 97e87b006ff..8a42eb3a891 100644 --- a/src/node/services/turnRequestBuilder.test.ts +++ b/src/node/services/turnRequestBuilder.test.ts @@ -446,6 +446,37 @@ describe("TurnRequestBuilder assembled preflight", () => { const off = await assembleBudgetCheckedPromptPayload(request, { enabled: false }); expect(off.contextBudgetEstimate).toBeUndefined(); }); + + // #7662: after a completed reply, E is that reply's provider input plus the messages after it. + it("anchors the enforced estimate to the newest reply's provider input", async () => { + const base = options("openai:large-context-model"); + const reply = createMuxMessage("reply", "assistant", "done", { + model: base.modelString, + agentId: "exec", + systemMessageTokens: 0, + contextUsage: { inputTokens: 50_000, outputTokens: 10, totalTokens: 50_010 }, + }); + const request = { + ...base, + history: [createMuxMessage("earlier", "user", "earlier request"), reply, ...base.history], + }; + const payload = await assembleBudgetCheckedPromptPayload(request, { enabled: true }); + const full = (await estimateAssembledRequestTokensForModel(payload, { + model: request.modelString, + modelContextLimit: payload.contextBudgetLimit, + }))!.estimate; + expect(full).toBeLessThan(50_000); + expect(payload.contextBudgetTurnAnchor!.providerTokens).toBeGreaterThan(50_000); + expect(payload.contextBudgetEstimate).toBeGreaterThan( + payload.contextBudgetTurnAnchor!.providerTokens + ); + // A compaction request's trim can drop rows the provider count holds. + const compaction = await assembleBudgetCheckedPromptPayload( + { ...request, compactionInputBudgetTokens: 90_000 }, + { enabled: true } + ); + expect(compaction.contextBudgetTurnAnchor).toBeUndefined(); + }); }); describe("TurnRequestBuilder MCP overrides deadline", () => { diff --git a/src/node/services/turnRequestBuilder.ts b/src/node/services/turnRequestBuilder.ts index 220461c36cc..e486f7742fc 100644 --- a/src/node/services/turnRequestBuilder.ts +++ b/src/node/services/turnRequestBuilder.ts @@ -7,7 +7,11 @@ import type { QueuedInputStopCause } from "@/common/types/streamStopCause"; import { execBuffered } from "@/node/utils/runtime/helpers"; import { shellQuote } from "@/common/utils/shell"; import type { OnStepSettled } from "./streamManager"; -import { measureAssembledRequestBudgetForModel } from "./contextBudgetCounting"; +import { + createTurnBudgetAnchor, + measureAssembledRequestBudgetForModel, + type TurnBudgetAnchor, +} from "./contextBudgetCounting"; import { ContextBudgetExceededError } from "./contextBudgetError"; import type { ContextBudgetExceeded } from "@/common/utils/compaction/contextBudget"; import { getEffectiveContextLimit } from "@/common/utils/compaction/contextLimit"; @@ -588,7 +592,7 @@ interface FoldedStreamState { messages: TMessages; systemRecords: TTail; systemMessageTokens: number; - /** The full estimate of `messages`, the stage decision's E (#5286). */ + /** The count of `messages` the budget check enforced, the stage decision's E (#5286). */ contextBudgetEstimate?: number; } @@ -661,7 +665,7 @@ async function foldPendingThinkingOverrides( return state; } -/** Called once per primary start() with E, the full estimate of the built request (#5286). */ +/** Called once per primary start() with E, the enforced count of the built request (#5286). */ export type OnTurnStartStage = (budget: { estimate: number; limit: number; @@ -755,10 +759,50 @@ async function countAndPublishStage[0], + payload: PromptPayload, + metadataModel: string | undefined +): Promise { + if (options.compactionInputBudgetTokens != null) return undefined; + return createTurnBudgetAnchor(payload.providerRequestMessages, { + model: options.modelString, + metadataModel, + routeProvider: options.routeProvider, + agentId: options.effectiveAgentId, + system: payload.sentSystem.message, + }); +} + +/** The model's effective context limit, or undefined (logged) when it is unknown. */ +function contextBudgetLimitFor( + options: Parameters[0], + providerOptions: MuxProviderOptions | undefined +): number | undefined { + const limit = getEffectiveContextLimit( + options.modelString, + isAnthropic1MEffectivelyEnabled(options.modelString, providerOptions, options.providersConfig), + options.providersConfig, + { openaiWireFormat: providerOptions?.openai?.wireFormat } + ); + if (limit == null) { + log.warn("Context budget preflight unavailable: model context limit is unknown", { + workspaceId: options.workspaceId, + model: options.modelString, + }); + } + return limit ?? undefined; +} + /** * Shared assembly path for primary, fallback, and thinking-rebuild provider attempts. - * `contextBudgetEstimate` is the full estimate the check enforced (undefined when the budget is - * off or the limit is unknown). + * `contextBudgetEstimate` is the estimate the check enforced (undefined when the budget is off + * or the limit is unknown): the previous turn's provider count plus the messages after it when + * `contextBudgetTurnAnchor` applies, else the full estimate. */ export async function assembleBudgetCheckedPromptPayload( options: Parameters[0], @@ -771,43 +815,40 @@ export async function assembleBudgetCheckedPromptPayload( Awaited> & { contextBudgetLimit?: number; contextBudgetEstimate?: number; + contextBudgetTurnAnchor?: TurnBudgetAnchor; } > { const payload = await assemblePromptPayload(options); let contextBudgetLimit: number | undefined; let contextBudgetEstimate: number | undefined; + let contextBudgetTurnAnchor: TurnBudgetAnchor | undefined; // Check after provider transforms and system/schema assembly: history-only // estimates cannot prevent oversized requests from reaching the provider. if (budget.enabled) { - contextBudgetLimit = - getEffectiveContextLimit( - options.modelString, - isAnthropic1MEffectivelyEnabled( - options.modelString, - budget.providerOptions, - options.providersConfig - ), - options.providersConfig, - { openaiWireFormat: budget.providerOptions?.openai?.wireFormat } - ) ?? undefined; - if (contextBudgetLimit == null) { - log.warn("Context budget preflight unavailable: model context limit is unknown", { - workspaceId: options.workspaceId, - model: options.modelString, - }); + contextBudgetLimit = contextBudgetLimitFor(options, budget.providerOptions); + const metadataModel = resolveModelForMetadata( + options.modelString, + options.providersConfig ?? null + ); + if (contextBudgetLimit != null) { + contextBudgetTurnAnchor = await turnBudgetAnchorFor(options, payload, metadataModel); } - const measured = await measureAssembledRequestBudgetForModel(payload, { - model: options.modelString, - metadataModel: resolveModelForMetadata(options.modelString, options.providersConfig ?? null), - modelContextLimit: contextBudgetLimit, - activeTools: budget.activeTools, - }); + const measured = await measureAssembledRequestBudgetForModel( + payload, + { + model: options.modelString, + metadataModel, + modelContextLimit: contextBudgetLimit, + activeTools: budget.activeTools, + }, + contextBudgetTurnAnchor + ); contextBudgetEstimate = measured?.estimate; if (measured?.exceeded) { throw new ContextBudgetExceededError(measured.exceeded); } } - return { ...payload, contextBudgetLimit, contextBudgetEstimate }; + return { ...payload, contextBudgetLimit, contextBudgetEstimate, contextBudgetTurnAnchor }; } /** @@ -3617,6 +3658,7 @@ export class TurnRequestBuilder { engineSystem: attemptPayload.system, systemRecords: rowSystemRecords(attemptPayload, options.partialContinuationMessage), contextBudgetLimit: attemptPayload.contextBudgetLimit, + contextBudgetTurnAnchor: attemptPayload.contextBudgetTurnAnchor, contextBudgetEstimate: attemptPayload.contextBudgetEstimate, ...compactionInputBudget, systemMessageTokens: sent.tokens, @@ -3646,6 +3688,7 @@ export class TurnRequestBuilder { systemRecords: rowSystemRecords(payload, options.partialContinuationMessage), messages: payload.messages, contextBudgetLimit: payload.contextBudgetLimit, + contextBudgetTurnAnchor: payload.contextBudgetTurnAnchor, contextBudgetEstimate: payload.contextBudgetEstimate, }; }, @@ -4135,6 +4178,7 @@ export class TurnRequestBuilder { system: nextRequest.engineSystem, tools: nextRequest.engineTools, contextBudgetLimit: nextRequest.contextBudgetLimit, + contextBudgetTurnAnchor: nextRequest.contextBudgetTurnAnchor, providerOptions: nextRequest.providerOptions, headers: nextHeaders, callSettingsOverrides: nextRequest.resolvedOverrides.standard, @@ -4173,6 +4217,7 @@ export class TurnRequestBuilder { abortSignal: combinedAbortSignal, tools: toolsForStream, contextBudgetLimit: primaryRequest.contextBudgetLimit, + contextBudgetTurnAnchor: primaryRequest.contextBudgetTurnAnchor, initialMetadata: { ...(requestHistorySequence >= 0 ? { requestHistorySequence } : {}), systemMessageTokens, diff --git a/tests/ipc/tokenBudgetStages.fixture.integration.test.ts b/tests/ipc/tokenBudgetStages.fixture.integration.test.ts index 0b38ad7f3f2..4ab68eb454c 100644 --- a/tests/ipc/tokenBudgetStages.fixture.integration.test.ts +++ b/tests/ipc/tokenBudgetStages.fixture.integration.test.ts @@ -7,22 +7,23 @@ * last user text carries a Token Budget warning with a memory write to the session checkpoint, * and answers everything else (tool results, "Continue", side requests) with short text. * - * The fixture reports provider usage through a per-test function of the request size. That is the - * gap under test: stage prompts must be decided from the built request's estimate (E), not from - * saved provider usage, or they are published and then refused at the ceiling (#5286). + * The fixture reports provider usage through a per-test function of the request size. Stage + * prompts are decided on E, the count the turn-start check enforces: the newest reply's provider + * usage plus the messages after it (#7662), else the full estimate. Step zero counts the same + * way, so a published prompt is never refused at the ceiling (#5286). * * Numbers (limit 100,000): hard ceiling C = 91,808, final point F = 81,616, slider 70 → handoff - * point H = 70,000. Measured with these fixtures: E ≈ 43.8k for the first request (system prompt - * and tool schemas), and each "read 4 files" turn adds ≈ 15.7k, so the third turn starts at - * E ≈ 75k, inside the handoff band. + * point H = 70,000. Measured with these fixtures and the real tokenizer: the system prompt and + * tool schemas count ≈ 43.8k once padded (below), and each file read adds ≈ 3.9k, so after two + * "read 4 files" turns the third turn starts at E ≈ 73k, inside the handoff band. * * The system prompt also lists global skills and instructions found under HOME, so the first E - * differs between machines (≈ 37.1k with an empty HOME, as in CI). After the first turn each - * scenario pads the workspace scratchpad by the difference, so every later E matches the - * measured bands. The pad and the later growth go in the scratchpad, not AGENTS.md: inside a - * context segment an AGENTS.md change is sent as a row with the whole new block while the - * system prompt keeps the old one (prompt-cache mechanism 10), so it would grow E by more - * than its own size. + * differs between machines (≈ 25k with an empty HOME, as in CI). After the first turn each + * scenario pads the workspace scratchpad up to CALIBRATED_FIRST_ESTIMATE, so every later E stays + * within about 1k of the measured bands. The pad and the later growth go in the scratchpad, not + * AGENTS.md: inside a context segment an AGENTS.md change is sent as a row with the whole new + * block while the system prompt keeps the old one (prompt-cache mechanism 10), so it would grow + * E by more than its own size. */ import assert from "node:assert/strict"; @@ -58,7 +59,7 @@ const TEST_TIMEOUT_MS = 180_000; const WARNING_MARKER = "Context budget ~"; const CHECKPOINT_PATH = "/memories/session/checkpoint.md"; const NOTE_FILES = 40; -/** E of the first request on the machine where the bands in this file were measured. */ +/** First-request E that each scenario is padded to; the bands in this file were measured from it. */ const CALIBRATED_FIRST_ESTIMATE = 43_800; interface Block { @@ -377,26 +378,41 @@ async function startScenario(options: { type Scenario = Awaited>; -/** Saved usage stays far below every stage point; only the request estimate crosses them. */ -const LOW_USAGE: UsageFor = () => 1_000; +/** Provider usage that tracks the full estimate, so the anchored E keeps the measured bands. */ +const TRACKING_USAGE: UsageFor = (chars) => Math.round(chars / 2.4); + +/** Three turns; only the last turn start crosses H (E ≈ 76.5k), and the handoff is published. */ +async function sendUntilHandoff(scenario: Scenario) { + await scenario.send("Turn 0: read 4 files."); + // One file keeps turn 1's estimate below H, so its settlement opens no stage. + await scenario.send("Turn 1: read 1 files."); + await scenario.send(`Turn 2: read 4 files. Pasted deploy log:\n${filler("Paste", 14_000)}`); +} describeIntegration("Token Budget stages (loopback fixture)", () => { + // Rows record systemMessageTokens with the real tokenizer, as in production: the turn anchor + // compares the sent system text against that count (#7662). + const forceRealTokenizer = process.env.XUM_FORCE_REAL_TOKENIZER; beforeAll(async () => { + process.env.XUM_FORCE_REAL_TOKENIZER = "1"; await preloadTestModules(); }); + afterAll(() => { + if (forceRealTokenizer == null) delete process.env.XUM_FORCE_REAL_TOKENIZER; + else process.env.XUM_FORCE_REAL_TOKENIZER = forceRealTokenizer; + }); - describe("estimate-only crossing (R2a)", () => { + describe("turn-start crossing (R2a)", () => { let scenario: Scenario; beforeAll(async () => { - scenario = await startScenario({ slider: 70, usageFor: LOW_USAGE }); - for (const turn of [0, 1, 2]) await scenario.send(`Turn ${turn}: read 4 files.`); + scenario = await startScenario({ slider: 70, usageFor: TRACKING_USAGE }); + await sendUntilHandoff(scenario); }, TEST_TIMEOUT_MS); afterAll(async () => { await scenario?.close(); }); test("publishes the handoff from the request estimate and the checkpoint write runs", async () => { - // Settlement schedules in provider terms (1,000 tokens), so only turn start can see E ≥ H. const stage = scenario.stageRequests(); expect(stage).toHaveLength(1); expect(stage[0].lastUserText).toContain("Turn 2: read 4 files."); @@ -457,9 +473,10 @@ describeIntegration("Token Budget stages (loopback fixture)", () => { }); test( - "realistic usage gap (R2b): no warning row is followed by a rollover without a reply", + "realistic usage gap (R2b): a warning decided on the provider count is answered, not refused", async () => { - // E/usage ≈ 1.5, as when provider usage trails the full request estimate. + // E/usage ≈ 1.5, as when provider usage trails the full request estimate. E follows the + // provider count (#7662), and step zero counts the same way. const scenario = await startScenario({ slider: 70, usageFor: (chars) => Math.round(chars / 3.6), @@ -467,11 +484,17 @@ describeIntegration("Token Budget stages (loopback fixture)", () => { try { await scenario.send("Turn 0: read 4 files."); await scenario.send("Turn 1: read 4 files."); - // A large paste takes E past the ceiling while saved usage plus the paste reads only ~74k. + // Saved usage plus the paste reads ≈76k, inside the handoff band. await scenario.send(`Turn 2: read 0 files. Pasted deploy log:\n${filler("Paste", 25_000)}`); const history = await scenario.history(); - // The ceiling was reached (non-vacuous): the paste rolled the window over. - expect(history.some(isReset)).toBe(true); + expect(history.filter(isWarning).map((row) => row.metadata?.muxMetadata?.handoff)).toEqual([ + true, + ]); + // Non-vacuous: the full estimate (≈2.4 request chars per token) of the sent stage request + // is past the hard ceiling, so a full count at step zero would have refused it. + const [stage] = scenario.stageRequests(); + expect(stage.chars / 2.4).toBeGreaterThan(91_808); + expect(history.some(isReset)).toBe(false); const unanswered = history.flatMap((row, index) => { if (!isWarning(row)) return []; const next = history.slice(index + 1).find((r) => r.role === "assistant"); @@ -489,17 +512,14 @@ describeIntegration("Token Budget stages (loopback fixture)", () => { "instruction growth: no warning when the row and its reserve no longer fit, and the turn runs", async () => { // Usage tracks E here, so an old usage-based decision would also open the handoff stage. - const scenario = await startScenario({ - slider: 70, - usageFor: (chars) => Math.round(chars / 2.4), - }); + const scenario = await startScenario({ slider: 70, usageFor: TRACKING_USAGE }); try { await scenario.send("Turn 0: read 4 files."); // One file keeps turn 1's estimate below H, so its settlement opens no stage (#5223). await scenario.send("Turn 1: read 1 files."); - // Grow the instructions so E ≈ 90.8k: below C (91,808), but E′ + 2,048 ≥ C. The assertions bracket + // Grow the instructions so E ≈ 90.6k: below C (91,808), but E′ + 2,048 ≥ C. The assertions bracket // the band: less growth publishes a warning, more is refused and rolls the window over. - await scenario.writeInstructions(filler("Instruction", 26_270)); + await scenario.writeInstructions(filler("Instruction", 29_400)); await scenario.send("Turn 2: read 0 files."); const history = await scenario.history(); expect(history.filter(isWarning)).toEqual([]); @@ -519,13 +539,13 @@ describeIntegration("Token Budget stages (loopback fixture)", () => { test( "after a restart the next send decides from its own estimate and does not repeat the handoff", async () => { - const scenario = await startScenario({ slider: 70, usageFor: LOW_USAGE }); + const scenario = await startScenario({ slider: 70, usageFor: TRACKING_USAGE }); try { - for (const turn of [0, 1, 2]) await scenario.send(`Turn ${turn}: read 4 files.`); + await sendUntilHandoff(scenario); expect((await scenario.history()).filter(isWarning)).toHaveLength(1); await scenario.restart(); // E ≈ 84k after restart: inside the final band [81,616, C − 4,048), handoff already claimed. - await scenario.writeInstructions(filler("Instruction", 7_500)); + await scenario.writeInstructions(filler("Instruction", 5_500)); await scenario.send("Turn 3: read 0 files."); const warnings = (await scenario.history()).filter(isWarning); expect( @@ -549,13 +569,13 @@ describeIntegration("Token Budget stages (loopback fixture)", () => { test( "a refusal fallback with a stage due is sent without one, and the next turn start offers it once", async () => { - // The primary's window (200k, H = 140k) has no stage due at E ≈ 75k. The fallback's window + // The primary's window (200k, H = 140k) has no stage due at E ≈ 73k. The fallback's window // (100k, H = 70k) does. The fallback request is sent as built, with no stage decision: the // warning cannot sit before the turn's placeholder row, and a handover would need a new // mechanism (#5286). The stage comes at the next turn start that reaches it, one turn late. const scenario = await startScenario({ slider: 70, - usageFor: LOW_USAGE, + usageFor: TRACKING_USAGE, refusalFallback: { primaryLimit: 200_000 }, }); try { @@ -593,10 +613,11 @@ describeIntegration("Token Budget stages (loopback fixture)", () => { test( "a rollover whose first request is past the handoff point publishes the warning in the new window", async () => { - const scenario = await startScenario({ slider: 70, usageFor: LOW_USAGE }); + const scenario = await startScenario({ slider: 70, usageFor: TRACKING_USAGE }); try { await scenario.send("Turn 0: read 4 files."); - await scenario.send("Turn 1: read 4 files."); + // Two files keep turn 1's estimate below H, so the old window publishes no stage. + await scenario.send("Turn 1: read 2 files."); // The paste takes the old window past C, so the send rolls it over. The fresh window's // own first request (system prompt, lead-in, paste) is past H (70k) and still has room // for the row and its reserve below C (#7215). From 115c8c7522e533e8fbf9a1adf8bb247ff268fec9 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 18:12:44 +1100 Subject: [PATCH 4/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20anchor=20the=20turn?= =?UTF-8?q?=20start=20only=20on=20the=20same=20resolved=20model,=20with=20?= =?UTF-8?q?replayed=20reasoning?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The turn anchor now adds the last step's reasoning output, like the in-turn anchor. Claude Opus 5.5 replays that step's thinking at its full count, and the estimator sees at most the summary text. Redacted thinking and encrypted reasoning have no text at all. A last step that emitted reasoning without a reported count does not anchor. A row from another resolved model (a providers.jsonc mapping change), or from a build that did not record it, does not anchor. --- .../services/contextBudgetCounting.test.ts | 72 ++++++++--- src/node/services/contextBudgetCounting.ts | 121 ++++++++++-------- src/node/services/turnRequestBuilder.test.ts | 1 + 3 files changed, 128 insertions(+), 66 deletions(-) diff --git a/src/node/services/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index dec0407bdac..0a84c5f1288 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -8,7 +8,7 @@ import { type AssembledRequestBudgetInput, } from "@/common/utils/compaction/contextBudget"; import type { JSONSchema7 } from "@ai-sdk/provider"; -import { createMuxMessage, type MuxMetadata } from "@/common/types/message"; +import { createMuxMessage, type MuxMessage, type MuxMetadata } from "@/common/types/message"; import { CLAUDE_ENCODING_BUDGET_FACTOR } from "@/constants/tokenizerCorrection"; import calibration from "./__fixtures__/contextBudgetClaudeTokenizer.json"; import { @@ -850,19 +850,32 @@ describe("turn-start anchor (#7662)", () => { const system = "Answer briefly and cite the files you read."; const systemTokens = async (forModel: string) => (await tokenizerModule.getTokenizerForModel(forModel)).countTokens(system); - const request = { model, agentId: "exec", routeProvider: "openai", system }; + const request: { + model: string; + metadataModel?: string; + agentId: string; + routeProvider: string; + system: string; + } = { model, agentId: "exec", routeProvider: "openai", system }; const anchorFor = async ( metadata?: Partial, - overrides?: Partial + overrides?: Partial, + parts?: MuxMessage["parts"] ) => { - const reply = createMuxMessage("a1", "assistant", "answer", { - model, - agentId: "exec", - routeProvider: "openai", - systemMessageTokens: await systemTokens(overrides?.model ?? model), - contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, - ...metadata, - }); + const reply = createMuxMessage( + "a1", + "assistant", + parts == null ? "answer" : "", + { + model, + agentId: "exec", + routeProvider: "openai", + systemMessageTokens: await systemTokens(overrides?.model ?? model), + contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, + ...metadata, + }, + parts + ); return createTurnBudgetAnchor([createMuxMessage("u1", "user", "hi"), reply], { ...request, ...overrides, @@ -888,7 +901,7 @@ describe("turn-start anchor (#7662)", () => { ); }); - test("counts cache reads and writes once, adds system growth, and leaves out the reply's reasoning", async () => { + test("counts cache reads and writes once, and adds the reply's reasoning and system growth", async () => { const contextUsage = { inputTokens: 1000, outputTokens: 80, @@ -899,11 +912,11 @@ describe("turn-start anchor (#7662)", () => { // Rows keep cache writes in provider metadata: 300 uncached + 600 read + 100 written. const contextProviderMetadata = { anthropic: { cacheCreationInputTokens: 100 } }; const sent = await systemTokens(model); - expect((await anchorFor({ contextUsage, contextProviderMetadata }))?.providerTokens).toBe(1000); + expect((await anchorFor({ contextUsage, contextProviderMetadata }))?.providerTokens).toBe(1050); const grown = { contextUsage, systemMessageTokens: sent - 30 }; - expect((await anchorFor(grown))?.providerTokens).toBe(1030); + expect((await anchorFor(grown))?.providerTokens).toBe(1080); const shrunk = { contextUsage, systemMessageTokens: sent + 30 }; - expect((await anchorFor(shrunk))?.providerTokens).toBe(1000); + expect((await anchorFor(shrunk))?.providerTokens).toBe(1050); // Claude-encoded growth is scaled like every other claude count (#5219). Tests default to // the approximate tokenizer, which has no claude encoding. const original = process.env.XUM_FORCE_REAL_TOKENIZER; @@ -912,10 +925,12 @@ describe("turn-start anchor (#7662)", () => { const claude = "anthropic:claude-opus-5-5"; const claudeGrown = { model: claude, + metadataModel: claude, contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, systemMessageTokens: (await systemTokens(claude)) - 30, }; - expect((await anchorFor(claudeGrown, { model: claude }))?.providerTokens).toBe( + const claudeRequest = { model: claude, metadataModel: claude }; + expect((await anchorFor(claudeGrown, claudeRequest))?.providerTokens).toBe( 1000 + Math.ceil(30 * CLAUDE_ENCODING_BUDGET_FACTOR) ); } finally { @@ -940,12 +955,37 @@ describe("turn-start anchor (#7662)", () => { ["a reply without its system count", { systemMessageTokens: undefined }, {}], ["a reply without its model", { model: undefined }, {}], ["another model", {}, { model: "openai:gpt-4.1" }], + // A providers.jsonc mapping changes the resolved model and keeps the model. + ["another resolved model", { metadataModel: "openai:gpt-4.1" }, { metadataModel: model }], + ["a reply without its resolved model", {}, { metadataModel: model }], ["another route", {}, { routeProvider: "openrouter" }], ["another agent", {}, { agentId: "plan" }], ])("%s yields no turn anchor", async (_name, metadata, overrides) => { expect(await anchorFor(metadata, overrides)).toBeUndefined(); }); + // The next request replays the last step's thinking at its full count; the estimator sees at + // most its summary text (Anthropic redacted thinking and OpenAI encrypted reasoning have none). + const thinking = { type: "reasoning" as const, text: "" }; + const text = (value: string) => ({ type: "text" as const, text: value }); + test.each<[string, MuxMessage["parts"], Partial, number | undefined]>([ + ["last-step reasoning without a reported count", [thinking], {}, undefined], + [ + "last-step reasoning counted only in provider metadata", + [thinking], + { contextProviderMetadata: { openai: { reasoningTokens: 40 } } }, + 1040, + ], + [ + "reasoning only before the last step", + [thinking, text("step one"), text("answer")], + { stepStartPartIndices: [0, 2] }, + 1000, + ], + ])("a reply with %s", async (_name, parts, metadata, providerTokens) => { + expect((await anchorFor(metadata, {}, parts))?.providerTokens).toBe(providerTokens); + }); + test("a reply marked as not compacted still anchors", async () => { expect((await anchorFor({ compacted: false }))?.providerTokens).toBe(1000); }); diff --git a/src/node/services/contextBudgetCounting.ts b/src/node/services/contextBudgetCounting.ts index b12dc99fb42..77248818894 100644 --- a/src/node/services/contextBudgetCounting.ts +++ b/src/node/services/contextBudgetCounting.ts @@ -16,8 +16,12 @@ import { } from "@/common/constants/contextBudget"; import type { MuxMessage } from "@/common/types/message"; import { normalizeToCanonical } from "@/common/utils/ai/models"; -import { createDisplayUsage } from "@/common/utils/tokens/displayUsage"; -import { normalizeUsage, type AiSdkUsageLike } from "@/common/utils/tokens/usageHelpers"; +import { isTrustedStepStartList } from "@/common/utils/messages/stepStartPartIndices"; +import { + getPromptTokens, + readUsageTokens, + type AiSdkUsageLike, +} from "@/common/utils/tokens/usageHelpers"; import { countToolReferences } from "@/common/utils/tools/toolCatalog"; import { getTokenizerForModel } from "@/node/utils/main/tokenizer"; import { @@ -173,24 +177,33 @@ export interface TurnBudgetAnchor extends BudgetModel { providerTokens: number; } -/** A step's usage normalized for display, or undefined without a positive integer input total. */ -function reportedUsage( +/** + * The provider's count of a step's request (cache reads and writes included) plus the step's + * reasoning output: the next request replays that reasoning, and the local estimator skips its + * encrypted form (signatures, redacted data, encrypted content). The visible reasoning text the + * next request also carries is counted again by its estimate, a small overcount. Undefined (so + * the caller uses the full estimate) without a positive integer input total, or when the step + * emitted reasoning without a reported count. + */ +function replayedStepTokens( usage: AiSdkUsageLike, - model: string, - providerMetadata: Record | undefined -) { - const normalized = normalizeUsage(usage); + providerMetadata: Record | undefined, + emittedReasoning: boolean +): number | undefined { // Cache or reasoning details without an input total do not count the whole request. - const input = normalized.inputTokens; + const input = usage.inputTokens; if (input == null || !Number.isSafeInteger(input) || input <= 0) return undefined; - return createDisplayUsage(normalized, model, providerMetadata); + // Provider-metadata-only reasoning counts are read too. Zero cannot tell "no reasoning" from + // "not reported", so a step that emitted reasoning needs a count. + const tokens = readUsageTokens(usage, providerMetadata); + if (emittedReasoning && tokens.reasoning === 0) return undefined; + const total = getPromptTokens(tokens) + tokens.reasoning; + return Number.isSafeInteger(total) ? total : undefined; } /** - * Pairs a step's request with the usage the provider reported for it. Reasoning output is - * added because the next request replays it and the local estimator skips encrypted - * reasoning. Undefined (so the caller uses the full estimate) when the provider reported no - * positive integer input total, or when the step emitted reasoning without a reported count. + * Pairs a step's request with its provider count (replayedStepTokens). Undefined (so the caller + * uses the full estimate) when the step has none. */ export function createContextBudgetAnchor( request: ContextBudgetAnchorRequest | undefined, @@ -203,17 +216,9 @@ export function createContextBudgetAnchor( | undefined ): ContextBudgetAnchor | undefined { if (request == null || step?.usage == null) return undefined; - const usage = reportedUsage(step.usage, request.model, step.providerMetadata); - if (usage == null) return undefined; - // Display recovery also reads provider-metadata-only reasoning counts. Zero cannot tell - // "no reasoning" from "not reported", so a step that emitted reasoning needs a count. - const reasoning = usage.reasoning.tokens; - if (reasoning === 0 && (step.reasoning?.length ?? 0) > 0) return undefined; - const providerTokens = - usage.input.tokens + usage.cached.tokens + usage.cacheCreate.tokens + reasoning; - return Number.isSafeInteger(providerTokens) && providerTokens > 0 - ? { ...request, providerTokens } - : undefined; + const emittedReasoning = (step.reasoning?.length ?? 0) > 0; + const providerTokens = replayedStepTokens(step.usage, step.providerMetadata, emittedReasoning); + return providerTokens == null ? undefined : { ...request, providerTokens }; } /** Both set and different: a row without the field does not count as a switch. */ @@ -238,62 +243,78 @@ function isCompletedReply(metadata: ReplyMetadata): boolean { ); } -/** The model, route and agent switches that re-freeze the system prompt. */ -function sameRequestRoute( - metadata: ReplyMetadata & { model: string }, - request: BudgetModel & { routeProvider?: string; agentId?: string } -): boolean { - // Rows record the canonical model (Coder ids raw), so compare both sides canonicalized. +type AnchorRequest = BudgetModel & { routeProvider?: string; agentId?: string }; + +/** + * The model the provider counted with, and the route and agent switches that re-freeze the + * system prompt. + */ +function sameRequestRoute(metadata: ReplyMetadata & { model: string }, request: AnchorRequest) { + // Rows record the canonical model (Coder ids raw), so compare both sides canonicalized. A + // providers.jsonc mapping keeps the model and changes the resolved model, which picks the + // tokenizer and the request transforms, so a row from another one (or from a build that did + // not record it) does not anchor. return ( normalizeToCanonical(metadata.model) === normalizeToCanonical(request.model) && + metadata.metadataModel === request.metadataModel && !differs(metadata.routeProvider, request.routeProvider) && !differs(metadata.agentId, request.agentId) ); } /** The newest reply's metadata when its saved usage can anchor this request, else undefined. */ -function anchoringReply( - rows: readonly MuxMessage[], - request: BudgetModel & { routeProvider?: string; agentId?: string } -) { - const metadata = rows.findLast((row) => row.role === "assistant")?.metadata; +function anchoringReply(metadata: ReplyMetadata | undefined, request: AnchorRequest) { const { contextUsage, model, systemMessageTokens } = metadata ?? {}; if (contextUsage == null || model == null || systemMessageTokens == null) return undefined; const reply = { ...metadata, contextUsage, model, systemMessageTokens }; return isCompletedReply(reply) && sameRequestRoute(reply, request) ? reply : undefined; } +/** Whether the reply's last step (from its last trusted step start) emitted reasoning. */ +function lastStepEmittedReasoning(reply: Pick): boolean { + const starts: unknown = reply.metadata?.stepStartPartIndices; + const first = isTrustedStepStartList(starts) + ? (starts.findLast((index) => index < reply.parts.length) ?? 0) + : 0; + return reply.parts.slice(first).some((part) => part.type === "reasoning"); +} + /** * The anchor for a turn-start request built from `rows`, its provider request rows (#7662). - * The newest assistant row must be a completed reply from the same model, route and agent (the - * switches that re-freeze the system prompt), with the provider's input total for its last step - * and its sent system token count. Undefined (so the caller uses the full estimate) otherwise. + * The newest assistant row must be a completed reply from the same model, resolved model, route + * and agent, with the provider's input total for its last step and its sent system token count. + * Undefined (so the caller uses the full estimate) otherwise. + * + * Like the in-turn anchor, the last step's reasoning output is added (replayedStepTokens): the + * request replays it after the anchored prefix, and a model that keeps earlier thinking (Claude + * Opus 5.5 does) counts all of it, not the summary text the estimator sees. For a model that + * drops it, this overcounts by that step's reasoning. * - * Reasoning output is not added, unlike the in-turn anchor: the request carries that step's - * output, reasoning text included, after the anchored prefix, so the estimator counts it. * Growth of `system`, the sent system text, since that request (a live system prompt changes * between turns when the provider takes no system rows) is added, because the provider count * holds the old prompt. New tail rows go after the newest reply, where the delta counts them. */ export async function createTurnBudgetAnchor( rows: readonly MuxMessage[], - request: BudgetModel & { routeProvider?: string; agentId?: string; system: string } + request: AnchorRequest & { system: string } ): Promise { - const reply = anchoringReply(rows, request); + const row = rows.findLast((candidate) => candidate.role === "assistant"); + if (row == null) return undefined; + const reply = anchoringReply(row.metadata, request); if (reply == null) return undefined; - const usage = reportedUsage(reply.contextUsage, request.model, reply.contextProviderMetadata); - if (usage == null) return undefined; + const replayed = replayedStepTokens( + reply.contextUsage, + reply.contextProviderMetadata, + lastStepEmittedReasoning(row) + ); + if (replayed == null) return undefined; // The row counted its system text, tail rows included, with this tokenizer (measureSentSystem), // so its rows only make the growth smaller. Claude-encoded growth is scaled like every other // claude count (#5219). const tokenizer = await getTokenizerForModel(request.model, request.metadataModel); const growth = (await tokenizer.countTokens(request.system)) - reply.systemMessageTokens; const factor = tokenizer.encoding === "claude" ? CLAUDE_ENCODING_BUDGET_FACTOR : 1; - const providerTokens = - usage.input.tokens + - usage.cached.tokens + - usage.cacheCreate.tokens + - Math.ceil(Math.max(0, growth) * factor); + const providerTokens = replayed + Math.ceil(Math.max(0, growth) * factor); return Number.isSafeInteger(providerTokens) ? { model: request.model, metadataModel: request.metadataModel, providerTokens } : undefined; diff --git a/src/node/services/turnRequestBuilder.test.ts b/src/node/services/turnRequestBuilder.test.ts index 8a42eb3a891..a955a922b10 100644 --- a/src/node/services/turnRequestBuilder.test.ts +++ b/src/node/services/turnRequestBuilder.test.ts @@ -452,6 +452,7 @@ describe("TurnRequestBuilder assembled preflight", () => { const base = options("openai:large-context-model"); const reply = createMuxMessage("reply", "assistant", "done", { model: base.modelString, + metadataModel: base.modelString, agentId: "exec", systemMessageTokens: 0, contextUsage: { inputTokens: 50_000, outputTokens: 10, totalTokens: 50_010 }, From 5ccad0f70bb81e639121964c45f445091bd8b7cb Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 18:30:36 +1100 Subject: [PATCH 5/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20reject=20negative=20u?= =?UTF-8?q?sage=20counts=20and=20compare=20system=20text=20with=20its=20ro?= =?UTF-8?q?ws=20in=20the=20turn=20anchor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - replayedStepTokens returns no anchor when a persisted input, cache or reasoning count is negative or not a safe integer. - The turn anchor measures system growth against the sent system message plus its system rows, the same text the reply's systemMessageTokens counted. --- .../services/contextBudgetCounting.test.ts | 24 +++++++++++ src/node/services/contextBudgetCounting.ts | 21 ++++++--- src/node/services/turnRequestBuilder.test.ts | 43 +++++++++++++++++++ src/node/services/turnRequestBuilder.ts | 3 +- 4 files changed, 83 insertions(+), 8 deletions(-) diff --git a/src/node/services/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index 0a84c5f1288..1541e69c29e 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -952,6 +952,30 @@ describe("turn-start anchor (#7662)", () => { { contextUsage: { inputTokens: 0, outputTokens: 5, totalTokens: 5 } }, {}, ], + [ + "a negative reasoning count", + { + contextUsage: { + inputTokens: 1000, + outputTokens: 5, + totalTokens: 1005, + reasoningTokens: -900, + }, + }, + {}, + ], + [ + "a negative cache-read count", + { + contextUsage: { + inputTokens: 1000, + outputTokens: 5, + totalTokens: 1005, + cachedInputTokens: -900, + }, + }, + {}, + ], ["a reply without its system count", { systemMessageTokens: undefined }, {}], ["a reply without its model", { model: undefined }, {}], ["another model", {}, { model: "openai:gpt-4.1" }], diff --git a/src/node/services/contextBudgetCounting.ts b/src/node/services/contextBudgetCounting.ts index 77248818894..dec1a8ef78d 100644 --- a/src/node/services/contextBudgetCounting.ts +++ b/src/node/services/contextBudgetCounting.ts @@ -177,6 +177,11 @@ export interface TurnBudgetAnchor extends BudgetModel { providerTokens: number; } +/** Persisted usage is unchecked JSON: a negative count would shrink the anchor. */ +function isTokenCount(count: unknown): count is number { + return typeof count === "number" && Number.isSafeInteger(count) && count >= 0; +} + /** * The provider's count of a step's request (cache reads and writes included) plus the step's * reasoning output: the next request replays that reasoning, and the local estimator skips its @@ -192,10 +197,12 @@ function replayedStepTokens( ): number | undefined { // Cache or reasoning details without an input total do not count the whole request. const input = usage.inputTokens; - if (input == null || !Number.isSafeInteger(input) || input <= 0) return undefined; + if (!isTokenCount(input) || input === 0) return undefined; // Provider-metadata-only reasoning counts are read too. Zero cannot tell "no reasoning" from // "not reported", so a step that emitted reasoning needs a count. const tokens = readUsageTokens(usage, providerMetadata); + const counts = [tokens.input, tokens.cached, tokens.cacheCreate, tokens.reasoning]; + if (!counts.every(isTokenCount)) return undefined; if (emittedReasoning && tokens.reasoning === 0) return undefined; const total = getPromptTokens(tokens) + tokens.reasoning; return Number.isSafeInteger(total) ? total : undefined; @@ -290,9 +297,10 @@ function lastStepEmittedReasoning(reply: Pick) * Opus 5.5 does) counts all of it, not the summary text the estimator sees. For a model that * drops it, this overcounts by that step's reasoning. * - * Growth of `system`, the sent system text, since that request (a live system prompt changes - * between turns when the provider takes no system rows) is added, because the provider count - * holds the old prompt. New tail rows go after the newest reply, where the delta counts them. + * Growth of `system`, the sent system text with its system rows, since that request (a live + * system prompt changes between turns when the provider takes no system rows) is added, because + * the provider count holds the old prompt. A row this request appends after the newest reply is + * also in the delta, a small overcount. */ export async function createTurnBudgetAnchor( rows: readonly MuxMessage[], @@ -308,9 +316,8 @@ export async function createTurnBudgetAnchor( lastStepEmittedReasoning(row) ); if (replayed == null) return undefined; - // The row counted its system text, tail rows included, with this tokenizer (measureSentSystem), - // so its rows only make the growth smaller. Claude-encoded growth is scaled like every other - // claude count (#5219). + // The row counted the same text (message and rows) with this tokenizer (measureSentSystem). + // Claude-encoded growth is scaled like every other claude count (#5219). const tokenizer = await getTokenizerForModel(request.model, request.metadataModel); const growth = (await tokenizer.countTokens(request.system)) - reply.systemMessageTokens; const factor = tokenizer.encoding === "claude" ? CLAUDE_ENCODING_BUDGET_FACTOR : 1; diff --git a/src/node/services/turnRequestBuilder.test.ts b/src/node/services/turnRequestBuilder.test.ts index a955a922b10..53ebd9dd7d4 100644 --- a/src/node/services/turnRequestBuilder.test.ts +++ b/src/node/services/turnRequestBuilder.test.ts @@ -4,6 +4,7 @@ import { z } from "zod"; import { getContextBudgetHardCeiling } from "@/common/utils/compaction/contextBudget"; import { ContextBudgetExceededError } from "./contextBudgetError"; import { estimateAssembledRequestTokensForModel } from "./contextBudgetCounting"; +import { getTokenizerForModel } from "@/node/utils/main/tokenizer"; import { afterEach, describe, expect, it, mock, spyOn } from "bun:test"; import { CONTEXT_BOUNDARY_KINDS } from "@/common/constants/contextBoundary"; import { MULTI_PROJECT_CONFIG_KEY } from "@/common/constants/multiProject"; @@ -478,6 +479,48 @@ describe("TurnRequestBuilder assembled preflight", () => { ); expect(compaction.contextBudgetTurnAnchor).toBeUndefined(); }); + + // #7662: the reply's systemMessageTokens counted the system message and its rows, so a system + // message that grows by less than the retained rows still adds its growth. + it("adds system growth when the earlier request also sent system rows", async () => { + const opus = { + ...options(), + modelString: "anthropic:claude-opus-5-5", + providerForMessages: "anthropic", + routeProvider: "anthropic", + }; + const stable = "stable instructions"; + const tail = `\n\n[Warning: ${"retained row ".repeat(300)}]`; + const withSystem = (system: string, history: MuxMessage[]) => + assembleBudgetCheckedPromptPayload( + { ...opus, history, systemMessage: system + tail, volatileSystemSuffixLength: tail.length }, + { enabled: true } + ); + const u1 = createMuxMessage("u1", "user", "first"); + const a1 = createMuxMessage("a1", "assistant", "answer a1", { + volatileSystemTail: { text: "\n\n[Warning: A]", placement: "system" }, + }); + const u2 = createMuxMessage("u2", "user", "second"); + const earlier = await withSystem(stable, [u1, a1, u2]); + expect(earlier.sentSystem.rowsText).not.toBe(""); + const tokenizer = await getTokenizerForModel(opus.modelString); + const reply = createMuxMessage("a2", "assistant", "answer a2", { + model: opus.modelString, + metadataModel: opus.modelString, + agentId: "exec", + routeProvider: "anthropic", + volatileSystemTail: { text: tail, placement: "row" }, + systemMessageTokens: await tokenizer.countTokens( + earlier.sentSystem.message + earlier.sentSystem.rowsText + ), + contextUsage: { inputTokens: 50_000, outputTokens: 10, totalTokens: 50_010 }, + }); + const history = [u1, a1, u2, reply, createMuxMessage("u3", "user", "third")]; + const same = await withSystem(stable, history); + expect(same.contextBudgetTurnAnchor!.providerTokens).toBe(50_000); + const grown = await withSystem(`${stable} and more instructions`, history); + expect(grown.contextBudgetTurnAnchor!.providerTokens).toBeGreaterThan(50_000); + }); }); describe("TurnRequestBuilder MCP overrides deadline", () => { diff --git a/src/node/services/turnRequestBuilder.ts b/src/node/services/turnRequestBuilder.ts index e486f7742fc..207cc5cc24d 100644 --- a/src/node/services/turnRequestBuilder.ts +++ b/src/node/services/turnRequestBuilder.ts @@ -774,7 +774,8 @@ async function turnBudgetAnchorFor( metadataModel, routeProvider: options.routeProvider, agentId: options.effectiveAgentId, - system: payload.sentSystem.message, + // The text the reply's systemMessageTokens counted (measureSentSystem), rows included. + system: payload.sentSystem.message + payload.sentSystem.rowsText, }); } From dd159a55d9bf37f08ec183fd2b0bd3f652fd1f3d Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 18:48:46 +1100 Subject: [PATCH 6/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20skip=20the=20turn=20a?= =?UTF-8?q?nchor=20after=20a=20thinking-level=20change=20or=20a=20split=20?= =?UTF-8?q?last=20step?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The turn anchor needs the reply's recorded thinking level to equal the request's effective level. Message preparation keeps or strips earlier Anthropic reasoning by that level. - The turn anchor is skipped when content follows a tool call in the reply's last step. The rebuild then splits the step, and a delta from the last assistant message would miss the tool call and its result. --- .../services/contextBudgetCounting.test.ts | 30 ++++++++++++- src/node/services/contextBudgetCounting.ts | 43 ++++++++++++++----- src/node/services/turnRequestBuilder.test.ts | 2 + src/node/services/turnRequestBuilder.ts | 1 + 4 files changed, 64 insertions(+), 12 deletions(-) diff --git a/src/node/services/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index 1541e69c29e..7b87fc19653 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -9,6 +9,7 @@ import { } from "@/common/utils/compaction/contextBudget"; import type { JSONSchema7 } from "@ai-sdk/provider"; import { createMuxMessage, type MuxMessage, type MuxMetadata } from "@/common/types/message"; +import type { ThinkingLevel } from "@/common/types/thinking"; import { CLAUDE_ENCODING_BUDGET_FACTOR } from "@/constants/tokenizerCorrection"; import calibration from "./__fixtures__/contextBudgetClaudeTokenizer.json"; import { @@ -855,8 +856,9 @@ describe("turn-start anchor (#7662)", () => { metadataModel?: string; agentId: string; routeProvider: string; + thinkingLevel: ThinkingLevel; system: string; - } = { model, agentId: "exec", routeProvider: "openai", system }; + } = { model, agentId: "exec", routeProvider: "openai", thinkingLevel: "high", system }; const anchorFor = async ( metadata?: Partial, overrides?: Partial, @@ -870,6 +872,7 @@ describe("turn-start anchor (#7662)", () => { model, agentId: "exec", routeProvider: "openai", + thinkingLevel: "high", systemMessageTokens: await systemTokens(overrides?.model ?? model), contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, ...metadata, @@ -984,6 +987,9 @@ describe("turn-start anchor (#7662)", () => { ["a reply without its resolved model", {}, { metadataModel: model }], ["another route", {}, { routeProvider: "openrouter" }], ["another agent", {}, { agentId: "plan" }], + // Anthropic history keeps signed reasoning with thinking on and strips it at "off". + ["another thinking level", {}, { thinkingLevel: "off" }], + ["a reply without its thinking level", { thinkingLevel: undefined }, {}], ])("%s yields no turn anchor", async (_name, metadata, overrides) => { expect(await anchorFor(metadata, overrides)).toBeUndefined(); }); @@ -992,6 +998,14 @@ describe("turn-start anchor (#7662)", () => { // most its summary text (Anthropic redacted thinking and OpenAI encrypted reasoning have none). const thinking = { type: "reasoning" as const, text: "" }; const text = (value: string) => ({ type: "text" as const, text: value }); + const toolCall = { + type: "dynamic-tool" as const, + toolCallId: "call-1", + toolName: "web_search", + state: "output-available" as const, + input: { query: "docs" }, + output: "results ".repeat(500), + }; test.each<[string, MuxMessage["parts"], Partial, number | undefined]>([ ["last-step reasoning without a reported count", [thinking], {}, undefined], [ @@ -1006,6 +1020,20 @@ describe("turn-start anchor (#7662)", () => { { stepStartPartIndices: [0, 2] }, 1000, ], + // The rebuild sends the tool call and its result before the trailing text message. + [ + "text after a last-step tool call", + [text("searching"), toolCall, text("answer")], + {}, + undefined, + ], + ["a tool call that ends the last step", [text("checking"), toolCall], {}, 1000], + [ + "a tool call only before the last step", + [toolCall, text("answer")], + { stepStartPartIndices: [0, 1] }, + 1000, + ], ])("a reply with %s", async (_name, parts, metadata, providerTokens) => { expect((await anchorFor(metadata, {}, parts))?.providerTokens).toBe(providerTokens); }); diff --git a/src/node/services/contextBudgetCounting.ts b/src/node/services/contextBudgetCounting.ts index dec1a8ef78d..ea585506f28 100644 --- a/src/node/services/contextBudgetCounting.ts +++ b/src/node/services/contextBudgetCounting.ts @@ -15,6 +15,8 @@ import { REQUEST_FRAMING_TOKENS, } from "@/common/constants/contextBudget"; import type { MuxMessage } from "@/common/types/message"; +import type { ThinkingLevel } from "@/common/types/thinking"; +import { isDynamicToolPart } from "@/common/types/toolParts"; import { normalizeToCanonical } from "@/common/utils/ai/models"; import { isTrustedStepStartList } from "@/common/utils/messages/stepStartPartIndices"; import { @@ -250,11 +252,15 @@ function isCompletedReply(metadata: ReplyMetadata): boolean { ); } -type AnchorRequest = BudgetModel & { routeProvider?: string; agentId?: string }; +type AnchorRequest = BudgetModel & { + routeProvider?: string; + agentId?: string; + thinkingLevel: ThinkingLevel; +}; /** - * The model the provider counted with, and the route and agent switches that re-freeze the - * system prompt. + * The model the provider counted with, the route and agent switches that re-freeze the system + * prompt, and the thinking level that decides which earlier reasoning replays. */ function sameRequestRoute(metadata: ReplyMetadata & { model: string }, request: AnchorRequest) { // Rows record the canonical model (Coder ids raw), so compare both sides canonicalized. A @@ -264,6 +270,9 @@ function sameRequestRoute(metadata: ReplyMetadata & { model: string }, request: return ( normalizeToCanonical(metadata.model) === normalizeToCanonical(request.model) && metadata.metadataModel === request.metadataModel && + // Message preparation keeps or strips earlier Anthropic reasoning by the effective level, + // which the row records, so another level replays history the count does not hold. + metadata.thinkingLevel === request.thinkingLevel && !differs(metadata.routeProvider, request.routeProvider) && !differs(metadata.agentId, request.agentId) ); @@ -277,20 +286,31 @@ function anchoringReply(metadata: ReplyMetadata | undefined, request: AnchorRequ return isCompletedReply(reply) && sameRequestRoute(reply, request) ? reply : undefined; } -/** Whether the reply's last step (from its last trusted step start) emitted reasoning. */ -function lastStepEmittedReasoning(reply: Pick): boolean { +/** The reply's last step: its parts from its last trusted step start. */ +function lastStepParts(reply: Pick): MuxMessage["parts"] { const starts: unknown = reply.metadata?.stepStartPartIndices; const first = isTrustedStepStartList(starts) ? (starts.findLast((index) => index < reply.parts.length) ?? 0) : 0; - return reply.parts.slice(first).some((part) => part.type === "reasoning"); + return reply.parts.slice(first); +} + +/** + * Whether content follows a tool call in the step. The rebuild then splits the step into an + * assistant message, its tool results and a trailing assistant message (a provider-run web + * search does this), so a delta from the last assistant message would miss the first two. + */ +function splitsAfterToolCall(parts: MuxMessage["parts"]): boolean { + const firstTool = parts.findIndex(isDynamicToolPart); + return firstTool >= 0 && parts.slice(firstTool).some((part) => !isDynamicToolPart(part)); } /** * The anchor for a turn-start request built from `rows`, its provider request rows (#7662). - * The newest assistant row must be a completed reply from the same model, resolved model, route - * and agent, with the provider's input total for its last step and its sent system token count. - * Undefined (so the caller uses the full estimate) otherwise. + * The newest assistant row must be a completed reply from the same model, resolved model, route, + * agent and thinking level, with the provider's input total for its last step and its sent + * system token count, and its last step must end in one assistant message. Undefined (so the + * caller uses the full estimate) otherwise. * * Like the in-turn anchor, the last step's reasoning output is added (replayedStepTokens): the * request replays it after the anchored prefix, and a model that keeps earlier thinking (Claude @@ -309,11 +329,12 @@ export async function createTurnBudgetAnchor( const row = rows.findLast((candidate) => candidate.role === "assistant"); if (row == null) return undefined; const reply = anchoringReply(row.metadata, request); - if (reply == null) return undefined; + const lastStep = lastStepParts(row); + if (reply == null || splitsAfterToolCall(lastStep)) return undefined; const replayed = replayedStepTokens( reply.contextUsage, reply.contextProviderMetadata, - lastStepEmittedReasoning(row) + lastStep.some((part) => part.type === "reasoning") ); if (replayed == null) return undefined; // The row counted the same text (message and rows) with this tokenizer (measureSentSystem). diff --git a/src/node/services/turnRequestBuilder.test.ts b/src/node/services/turnRequestBuilder.test.ts index 53ebd9dd7d4..e7bbda924d3 100644 --- a/src/node/services/turnRequestBuilder.test.ts +++ b/src/node/services/turnRequestBuilder.test.ts @@ -455,6 +455,7 @@ describe("TurnRequestBuilder assembled preflight", () => { model: base.modelString, metadataModel: base.modelString, agentId: "exec", + thinkingLevel: base.effectiveThinkingLevel, systemMessageTokens: 0, contextUsage: { inputTokens: 50_000, outputTokens: 10, totalTokens: 50_010 }, }); @@ -509,6 +510,7 @@ describe("TurnRequestBuilder assembled preflight", () => { metadataModel: opus.modelString, agentId: "exec", routeProvider: "anthropic", + thinkingLevel: opus.effectiveThinkingLevel, volatileSystemTail: { text: tail, placement: "row" }, systemMessageTokens: await tokenizer.countTokens( earlier.sentSystem.message + earlier.sentSystem.rowsText diff --git a/src/node/services/turnRequestBuilder.ts b/src/node/services/turnRequestBuilder.ts index 207cc5cc24d..c2625acbb76 100644 --- a/src/node/services/turnRequestBuilder.ts +++ b/src/node/services/turnRequestBuilder.ts @@ -774,6 +774,7 @@ async function turnBudgetAnchorFor( metadataModel, routeProvider: options.routeProvider, agentId: options.effectiveAgentId, + thinkingLevel: options.effectiveThinkingLevel, // The text the reply's systemMessageTokens counted (measureSentSystem), rows included. system: payload.sentSystem.message + payload.sentSystem.rowsText, }); From 2e70fc7a078a4773c23d427279c0c209fd6c20a1 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 22:23:01 +1100 Subject: [PATCH 7/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20use=20the=20turn=20an?= =?UTF-8?q?chor=20of=20the=20thinking=20level=20that=20is=20sent?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A thinking-level change during PREPARING, or one that races stream setup, rebuilds the first request at the new level. Step zero still counted that request against the anchor captured at the old level, so a high-to-off switch could reuse a count with replayed reasoning and refuse a request that fits. - The pre-stream fold now carries the anchor of its rebuild to step zero. - Step zero uses the full estimate when it applies a thinking override itself. --- src/node/services/aiService.test.ts | 32 ++++++++++++++++ .../streamManager.contextBudget.test.ts | 38 +++++++++++++------ .../services/turnEngine/StepPreparation.ts | 14 ++++++- src/node/services/turnRequestBuilder.ts | 9 ++++- 4 files changed, 79 insertions(+), 14 deletions(-) diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index 283d44436c0..0715e0fb05b 100644 --- a/src/node/services/aiService.test.ts +++ b/src/node/services/aiService.test.ts @@ -2638,6 +2638,8 @@ describe("AIService.streamMessage compaction boundary slicing", () => { duringPlaceholderAppend?: () => void; abortController?: AbortController; activeTurnThinkingOverride?: ActiveTurnThinkingOverride; + /** History rows before the turn's user row. */ + earlier?: MuxMessage[]; } = {} ) { const metadata = createLocalWorkspaceMetadata("turn-start-stage", xumHomePath); @@ -2658,6 +2660,9 @@ describe("AIService.streamMessage compaction boundary slicing", () => { return appended; }); } + for (const row of options.earlier ?? []) { + expect((await harness.historyService.appendToHistory(metadata.id, row)).success).toBe(true); + } const user = createMuxMessage("user-1", "user", USER_TEXT); expect((await harness.historyService.appendToHistory(metadata.id, user)).success).toBe(true); const read = await harness.historyService.getHistoryFromLatestBoundary(metadata.id); @@ -2851,6 +2856,33 @@ describe("AIService.streamMessage compaction boundary slicing", () => { expect(turn.measured).toEqual([await countRequest(turn, sent.messages)]); expect(turn.countOptions.length).toBe(3); }); + + // #7662: the turn anchor holds the provider's count at the reply's thinking level. The fold + // rebuilds the request at the pending level, so step zero gets the anchor of that rebuild. + it.each([ + { replyLevel: "off", anchored: false }, + { replyLevel: "high", anchored: true }, + ] as const)( + "a pre-stream thinking fold dispatches the turn anchor of its rebuild (reply at $replyLevel)", + async ({ replyLevel, anchored }) => { + using xumHome = new DisposableTempDir("ai-turn-start-thinking-fold-anchor"); + const reply = createMuxMessage("reply-0", "assistant", "done", { + model: KNOWN_MODELS.SONNET.id, + metadataModel: KNOWN_MODELS.SONNET.id, + thinkingLevel: replyLevel, + systemMessageTokens: 0, + contextUsage: { inputTokens: 500, outputTokens: 10, totalTokens: 510 }, + }); + const override: ActiveTurnThinkingOverride = { pending: "high" }; + const turn = await runStageTurn(xumHome.path, { + stageCallback: false, + activeTurnThinkingOverride: override, + earlier: [createMuxMessage("user-0", "user", "first request"), reply], + }); + expect(turn.sent!.thinkingLevel).toBe("high"); + expect(turn.sent!.contextBudgetTurnAnchor != null).toBe(anchored); + } + ); }); it("emits startup breadcrumbs as runtime-status events before stream start", async () => { diff --git a/src/node/services/streamManager.contextBudget.test.ts b/src/node/services/streamManager.contextBudget.test.ts index 2b1ebf53f76..1d2b9cf54fa 100644 --- a/src/node/services/streamManager.contextBudget.test.ts +++ b/src/node/services/streamManager.contextBudget.test.ts @@ -9,7 +9,7 @@ import { createTestToolConfig } from "./tools/testHelpers"; import { describe, expect, spyOn, test } from "bun:test"; import { MockLanguageModelV3, simulateReadableStream } from "ai/test"; import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; -import { tool } from "ai"; +import { tool, type ModelMessage } from "ai"; import { z } from "zod"; import * as fs from "node:fs/promises"; import * as path from "node:path"; @@ -228,10 +228,13 @@ describe("settled context hard ceiling", () => { ); // #7662: step zero has no settled step, so it counts against the turn anchor, as the - // turn-start check did; a full estimate here could refuse what turn start accepted. - test.each([false, true])( - "step zero counts against the turn anchor (anchored=%p)", - async (anchored) => { + // turn-start check did; a full estimate here could refuse what turn start accepted. A + // thinking change that races stream setup rebuilds step zero at a level the anchor's + // provider count does not hold, so that step uses the full estimate. + test.each(["none", "anchored", "rethought"] as const)( + "step zero counts against the turn anchor (%s)", + async (mode) => { + const anchored = mode === "anchored"; const h = await createTestHistoryService(); const workspaceId = "turn-anchor-budget"; const messageId = "turn-anchor-assistant"; @@ -270,29 +273,40 @@ describe("settled context hard ceiling", () => { ]) ).success ).toBe(true); + const messages: ModelMessage[] = [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi" }, + { role: "user", content: "Again" }, + ]; const started = await manager.startStream({ workspaceId, messageId, historySequence: 1, model, modelString: "openai:gpt-4o", - messages: [ - { role: "user", content: "Hello" }, - { role: "assistant", content: "Hi" }, - { role: "user", content: "Again" }, - ], + messages, system: "Answer briefly", runtime: new LocalRuntime(h.tempDir), providedRuntimeTempDir: runtimeDir, contextBudgetLimit: 10000, // The previous reply's provider count alone is at the ceiling. - ...(anchored - ? { + ...(mode === "none" + ? {} + : { contextBudgetTurnAnchor: { model: "openai:gpt-4o", metadataModel: "openai:gpt-4o", providerTokens: getContextBudgetHardCeiling(10000), }, + }), + ...(mode === "rethought" + ? { + thinkingOverrideState: { pending: "off" as const }, + rebuildProviderOptionsForThinkingLevel: () => ({ + providerOptions: {}, + effectiveLevel: "off" as const, + }), + rebuildFirstStepForThinkingLevel: () => Promise.resolve({ messages }), } : {}), }); diff --git a/src/node/services/turnEngine/StepPreparation.ts b/src/node/services/turnEngine/StepPreparation.ts index 818303e3252..8ff652b04c8 100644 --- a/src/node/services/turnEngine/StepPreparation.ts +++ b/src/node/services/turnEngine/StepPreparation.ts @@ -435,6 +435,18 @@ export function buildContextBudgetRequest( }; } +/** + * The turn anchor counts the request at the level the turn started at. A thinking override the + * step applied (the step-0 rebuild, or options only when it failed) changes which reasoning is + * sent, so the step uses the full estimate then (#7662). + */ +function turnAnchorAfterOverride( + request: Pick, + thinkingOverride: Record | undefined +): TurnBudgetAnchor | undefined { + return thinkingOverride === undefined ? request.contextBudgetTurnAnchor : undefined; +} + /** * Hard-stop the step when its assembled request exceeds the safe context budget. Step zero has * no settled step, so it counts against the turn anchor, as the turn-start check did (#7662). @@ -722,7 +734,7 @@ export class StepPreparation { previousRequest, steps.at(-1), stepNumber, - request.contextBudgetTurnAnchor + turnAnchorAfterOverride(request, thinkingOverride) ); } return buildPrepareStepResult({ diff --git a/src/node/services/turnRequestBuilder.ts b/src/node/services/turnRequestBuilder.ts index c2625acbb76..0f59e62af13 100644 --- a/src/node/services/turnRequestBuilder.ts +++ b/src/node/services/turnRequestBuilder.ts @@ -594,6 +594,8 @@ interface FoldedStreamState { systemMessageTokens: number; /** The count of `messages` the budget check enforced, the stage decision's E (#5286). */ contextBudgetEstimate?: number; + /** The turn anchor that count used; a rebuild at another level carries its own (#7662). */ + contextBudgetTurnAnchor?: TurnBudgetAnchor; } /** @@ -620,6 +622,7 @@ async function foldPendingThinkingOverrides( systemRecords: TTail; systemMessageTokens: number; contextBudgetEstimate?: number; + contextBudgetTurnAnchor?: TurnBudgetAnchor; }>; currentEffectiveLevelRef: { current: ThinkingLevel }; /** Leave pending changes for StreamManager's step 0 (after a published stage, #5286). */ @@ -650,6 +653,7 @@ async function foldPendingThinkingOverrides( systemRecords: rebuilt.systemRecords, systemMessageTokens: rebuilt.systemMessageTokens, contextBudgetEstimate: rebuilt.contextBudgetEstimate, + contextBudgetTurnAnchor: rebuilt.contextBudgetTurnAnchor, }; } catch (error) { if (error instanceof ContextBudgetExceededError) return { exceeded: error.details }; @@ -3639,6 +3643,7 @@ export class TurnRequestBuilder { return { messages: rebuiltPayload.messages, contextBudgetEstimate: rebuiltPayload.contextBudgetEstimate, + contextBudgetTurnAnchor: rebuiltPayload.contextBudgetTurnAnchor, systemMessageTokens: sent.tokens, systemRecords: rowSystemRecords(rebuiltPayload, options.partialContinuationMessage), }; @@ -3812,6 +3817,7 @@ export class TurnRequestBuilder { systemRecords: primaryRequest.systemRecords, systemMessageTokens, contextBudgetEstimate: primaryRequest.contextBudgetEstimate, + contextBudgetTurnAnchor: primaryRequest.contextBudgetTurnAnchor, }); if ("exceeded" in streamState) { runLanguageModelCleanup(modelResult.data.model); @@ -3898,6 +3904,7 @@ export class TurnRequestBuilder { providerOptions: streamProviderOptions, messages: streamFinalMessages, systemRecords: streamSystemRecords, + contextBudgetTurnAnchor: streamTurnAnchor, } = streamState; ({ systemMessageTokens } = streamState); // Debug sinks pair systemMessage with the message list, so when the @@ -4219,7 +4226,7 @@ export class TurnRequestBuilder { abortSignal: combinedAbortSignal, tools: toolsForStream, contextBudgetLimit: primaryRequest.contextBudgetLimit, - contextBudgetTurnAnchor: primaryRequest.contextBudgetTurnAnchor, + contextBudgetTurnAnchor: streamTurnAnchor, initialMetadata: { ...(requestHistorySequence >= 0 ? { requestHistorySequence } : {}), systemMessageTokens, From a93f8e9c0531be7331bd9edd6930b9ac625f004d Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Tue, 27 Oct 2026 22:53:43 +1100 Subject: [PATCH 8/8] =?UTF-8?q?=F0=9F=A4=96=20fix:=20anchor=20only=20on=20?= =?UTF-8?q?a=20reply=20with=20a=20valid=20system=20count=20and=20the=20sam?= =?UTF-8?q?e=20saved=20request=20identity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The turn anchor compared the route and agent only when both sides set them, so a row from a build that did not save them anchored a request on another route or agent. A negative saved system count read as system growth. - One identity rule: the resolved model, thinking level, route and agent a reply saved must each equal the request's, unset included (a direct route saves no route). - The saved system count goes through the same check as the usage counts. --- src/node/services/aiService.test.ts | 1 + .../services/contextBudgetCounting.test.ts | 30 +++++++++----- src/node/services/contextBudgetCounting.ts | 41 ++++++++++--------- 3 files changed, 42 insertions(+), 30 deletions(-) diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index 0715e0fb05b..b818035f4aa 100644 --- a/src/node/services/aiService.test.ts +++ b/src/node/services/aiService.test.ts @@ -2869,6 +2869,7 @@ describe("AIService.streamMessage compaction boundary slicing", () => { const reply = createMuxMessage("reply-0", "assistant", "done", { model: KNOWN_MODELS.SONNET.id, metadataModel: KNOWN_MODELS.SONNET.id, + agentId: "exec", thinkingLevel: replyLevel, systemMessageTokens: 0, contextUsage: { inputTokens: 500, outputTokens: 10, totalTokens: 510 }, diff --git a/src/node/services/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index 7b87fc19653..cbf3f5c31d8 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -849,8 +849,8 @@ describe("turn-start anchor (#7662)", () => { const payload = { messages: [...history, reply, prompt] }; const options = { model, modelContextLimit: 200_000 }; const system = "Answer briefly and cite the files you read."; - const systemTokens = async (forModel: string) => - (await tokenizerModule.getTokenizerForModel(forModel)).countTokens(system); + const systemTokens = async (forModel: string, text = system) => + (await tokenizerModule.getTokenizerForModel(forModel)).countTokens(text); const request: { model: string; metadataModel?: string; @@ -914,12 +914,15 @@ describe("turn-start anchor (#7662)", () => { }; // Rows keep cache writes in provider metadata: 300 uncached + 600 read + 100 written. const contextProviderMetadata = { anthropic: { cacheCreationInputTokens: 100 } }; - const sent = await systemTokens(model); - expect((await anchorFor({ contextUsage, contextProviderMetadata }))?.providerTokens).toBe(1050); + // Long enough that 30 fewer saved tokens is still a valid (non-negative) count. + const longSystem = { system: system.repeat(20) }; + const sent = await systemTokens(model, longSystem.system); + const cached = { contextUsage, contextProviderMetadata, systemMessageTokens: sent }; + expect((await anchorFor(cached, longSystem))?.providerTokens).toBe(1050); const grown = { contextUsage, systemMessageTokens: sent - 30 }; - expect((await anchorFor(grown))?.providerTokens).toBe(1080); + expect((await anchorFor(grown, longSystem))?.providerTokens).toBe(1080); const shrunk = { contextUsage, systemMessageTokens: sent + 30 }; - expect((await anchorFor(shrunk))?.providerTokens).toBe(1050); + expect((await anchorFor(shrunk, longSystem))?.providerTokens).toBe(1050); // Claude-encoded growth is scaled like every other claude count (#5219). Tests default to // the approximate tokenizer, which has no claude encoding. const original = process.env.XUM_FORCE_REAL_TOKENIZER; @@ -930,9 +933,9 @@ describe("turn-start anchor (#7662)", () => { model: claude, metadataModel: claude, contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, - systemMessageTokens: (await systemTokens(claude)) - 30, + systemMessageTokens: (await systemTokens(claude, longSystem.system)) - 30, }; - const claudeRequest = { model: claude, metadataModel: claude }; + const claudeRequest = { model: claude, metadataModel: claude, ...longSystem }; expect((await anchorFor(claudeGrown, claudeRequest))?.providerTokens).toBe( 1000 + Math.ceil(30 * CLAUDE_ENCODING_BUDGET_FACTOR) ); @@ -980,13 +983,18 @@ describe("turn-start anchor (#7662)", () => { {}, ], ["a reply without its system count", { systemMessageTokens: undefined }, {}], + ["a negative system count", { systemMessageTokens: -100_000 }, {}], + ["a fractional system count", { systemMessageTokens: 1.5 }, {}], ["a reply without its model", { model: undefined }, {}], ["another model", {}, { model: "openai:gpt-4.1" }], // A providers.jsonc mapping changes the resolved model and keeps the model. ["another resolved model", { metadataModel: "openai:gpt-4.1" }, { metadataModel: model }], ["a reply without its resolved model", {}, { metadataModel: model }], ["another route", {}, { routeProvider: "openrouter" }], + ["a reply without its route", { routeProvider: undefined }, {}], + ["a direct-route request after a routed reply", {}, { routeProvider: undefined }], ["another agent", {}, { agentId: "plan" }], + ["a reply without its agent", { agentId: undefined }, {}], // Anthropic history keeps signed reasoning with thinking on and strips it at "off". ["another thinking level", {}, { thinkingLevel: "off" }], ["a reply without its thinking level", { thinkingLevel: undefined }, {}], @@ -1042,8 +1050,10 @@ describe("turn-start anchor (#7662)", () => { expect((await anchorFor({ compacted: false }))?.providerTokens).toBe(1000); }); - test("a row without a route or agent still anchors; history without a reply does not", async () => { - expect(await anchorFor({ agentId: undefined, routeProvider: undefined })).toBeDefined(); + test("a direct route on both sides anchors; history without a reply does not", async () => { + expect( + (await anchorFor({ routeProvider: undefined }, { routeProvider: undefined }))?.providerTokens + ).toBe(1000); expect( await createTurnBudgetAnchor([createMuxMessage("u1", "user", "hi")], request) ).toBeUndefined(); diff --git a/src/node/services/contextBudgetCounting.ts b/src/node/services/contextBudgetCounting.ts index ea585506f28..8c89d296b5f 100644 --- a/src/node/services/contextBudgetCounting.ts +++ b/src/node/services/contextBudgetCounting.ts @@ -230,11 +230,6 @@ export function createContextBudgetAnchor( return providerTokens == null ? undefined : { ...request, providerTokens }; } -/** Both set and different: a row without the field does not count as a switch. */ -function differs(before: string | undefined, now: string | undefined): boolean { - return before != null && now != null && before !== now; -} - type ReplyMetadata = NonNullable; /** A completed reply: not partial, synthetic, failed, a compaction summary or a boundary. */ @@ -259,31 +254,37 @@ type AnchorRequest = BudgetModel & { }; /** - * The model the provider counted with, the route and agent switches that re-freeze the system - * prompt, and the thinking level that decides which earlier reasoning replays. + * The saved request fields that decide what a reply's provider count holds (#7662). The + * resolved model picks the tokenizer and the request transforms (a providers.jsonc mapping + * changes it and keeps the model). The thinking level decides which earlier Anthropic reasoning + * message preparation replays. The route and the agent each re-freeze the system prompt and + * tools. A reply anchors only when every field equals the request's, unset included: a direct + * route saves no route, and a row from a build that did not save a field proves no match. */ -function sameRequestRoute(metadata: ReplyMetadata & { model: string }, request: AnchorRequest) { - // Rows record the canonical model (Coder ids raw), so compare both sides canonicalized. A - // providers.jsonc mapping keeps the model and changes the resolved model, which picks the - // tokenizer and the request transforms, so a row from another one (or from a build that did - // not record it) does not anchor. +const ANCHOR_IDENTITY_KEYS = [ + "metadataModel", + "thinkingLevel", + "routeProvider", + "agentId", +] as const satisfies ReadonlyArray; + +function sameRequestIdentity(metadata: ReplyMetadata & { model: string }, request: AnchorRequest) { + // Rows record the canonical model (Coder ids raw), so compare both sides canonicalized. return ( normalizeToCanonical(metadata.model) === normalizeToCanonical(request.model) && - metadata.metadataModel === request.metadataModel && - // Message preparation keeps or strips earlier Anthropic reasoning by the effective level, - // which the row records, so another level replays history the count does not hold. - metadata.thinkingLevel === request.thinkingLevel && - !differs(metadata.routeProvider, request.routeProvider) && - !differs(metadata.agentId, request.agentId) + ANCHOR_IDENTITY_KEYS.every((key) => metadata[key] === request[key]) ); } /** The newest reply's metadata when its saved usage can anchor this request, else undefined. */ function anchoringReply(metadata: ReplyMetadata | undefined, request: AnchorRequest) { const { contextUsage, model, systemMessageTokens } = metadata ?? {}; - if (contextUsage == null || model == null || systemMessageTokens == null) return undefined; + // A negative saved system count would read as system growth and inflate the anchor. + if (contextUsage == null || model == null || !isTokenCount(systemMessageTokens)) { + return undefined; + } const reply = { ...metadata, contextUsage, model, systemMessageTokens }; - return isCompletedReply(reply) && sameRequestRoute(reply, request) ? reply : undefined; + return isCompletedReply(reply) && sameRequestIdentity(reply, request) ? reply : undefined; } /** The reply's last step: its parts from its last trusted step start. */