diff --git a/src/common/utils/compaction/contextBudget.ts b/src/common/utils/compaction/contextBudget.ts index 2b68196d881..0f38c8b3351 100644 --- a/src/common/utils/compaction/contextBudget.ts +++ b/src/common/utils/compaction/contextBudget.ts @@ -156,6 +156,26 @@ export function estimateToolResultSize(result: unknown): { return measureBudgetContent(result); } +type BudgetContentKind = + | "json" + | "messages" + | "message" + | "parts" + | "part" + | "assistant-parts" + | "assistant-part" + | "output" + | "output-parts" + | "output-part" + | "request-options"; + +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 +187,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 +226,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 +237,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 +259,12 @@ 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; + // 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"; @@ -286,20 +273,10 @@ 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"; + } else if (key === "providerOptions" && (messagePart || entry.kind === "message")) { + childKind = "request-options"; } stack.push({ value: child, kind: childKind }); } diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index 283d44436c0..b818035f4aa 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,34 @@ 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, + agentId: "exec", + 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/contextBudgetCounting.test.ts b/src/node/services/contextBudgetCounting.test.ts index 484dedfbebd..cbf3f5c31d8 100644 --- a/src/node/services/contextBudgetCounting.test.ts +++ b/src/node/services/contextBudgetCounting.test.ts @@ -8,10 +8,14 @@ import { type AssembledRequestBudgetInput, } 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 { checkAssembledRequestBudgetForModel, createContextBudgetAnchor, + createTurnBudgetAnchor, estimateAnchoredRequestTokensForModel, estimateAssembledRequestTokensForModel, estimateFreshRequestTokensForModel, @@ -194,41 +198,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 +262,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 } }, }, ], }, @@ -262,6 +278,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) => { @@ -805,6 +842,233 @@ 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, text = system) => + (await tokenizerModule.getTokenizerForModel(forModel)).countTokens(text); + const request: { + model: string; + metadataModel?: string; + agentId: string; + routeProvider: string; + thinkingLevel: ThinkingLevel; + system: string; + } = { model, agentId: "exec", routeProvider: "openai", thinkingLevel: "high", system }; + const anchorFor = async ( + metadata?: Partial, + overrides?: Partial, + parts?: MuxMessage["parts"] + ) => { + const reply = createMuxMessage( + "a1", + "assistant", + parts == null ? "answer" : "", + { + model, + agentId: "exec", + routeProvider: "openai", + thinkingLevel: "high", + systemMessageTokens: await systemTokens(overrides?.model ?? model), + contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, + ...metadata, + }, + parts + ); + 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, and adds the reply's reasoning and system growth", 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 } }; + // 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, longSystem))?.providerTokens).toBe(1080); + const shrunk = { contextUsage, systemMessageTokens: sent + 30 }; + 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; + process.env.XUM_FORCE_REAL_TOKENIZER = "1"; + try { + const claude = "anthropic:claude-opus-5-5"; + const claudeGrown = { + model: claude, + metadataModel: claude, + contextUsage: { inputTokens: 1000, outputTokens: 10, totalTokens: 1010 }, + systemMessageTokens: (await systemTokens(claude, longSystem.system)) - 30, + }; + const claudeRequest = { model: claude, metadataModel: claude, ...longSystem }; + expect((await anchorFor(claudeGrown, claudeRequest))?.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 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 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 }, {}], + ])("%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 }); + 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], + [ + "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, + ], + // 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); + }); + + test("a reply marked as not compacted still anchors", async () => { + expect((await anchorFor({ compacted: false }))?.providerTokens).toBe(1000); + }); + + 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(); + }); + + 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"; @@ -879,4 +1143,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..8c89d296b5f 100644 --- a/src/node/services/contextBudgetCounting.ts +++ b/src/node/services/contextBudgetCounting.ts @@ -14,8 +14,16 @@ import { BUDGET_TOKEN_CHUNK_SLACK, REQUEST_FRAMING_TOKENS, } from "@/common/constants/contextBudget"; -import { createDisplayUsage } from "@/common/utils/tokens/displayUsage"; -import { normalizeUsage, type AiSdkUsageLike } from "@/common/utils/tokens/usageHelpers"; +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 { + getPromptTokens, + readUsageTokens, + type AiSdkUsageLike, +} from "@/common/utils/tokens/usageHelpers"; import { countToolReferences } from "@/common/utils/tools/toolCatalog"; import { getTokenizerForModel } from "@/node/utils/main/tokenizer"; import { @@ -163,10 +171,48 @@ export interface ContextBudgetAnchor extends ContextBudgetAnchorRequest { } /** - * 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. + * 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; +} + +/** 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 + * 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, + providerMetadata: Record | undefined, + emittedReasoning: boolean +): number | undefined { + // Cache or reasoning details without an input total do not count the whole request. + const input = usage.inputTokens; + 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; +} + +/** + * 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, @@ -179,20 +225,127 @@ 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); - 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 } + const emittedReasoning = (step.reasoning?.length ?? 0) > 0; + const providerTokens = replayedStepTokens(step.usage, step.providerMetadata, emittedReasoning); + return providerTokens == null ? undefined : { ...request, providerTokens }; +} + +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 + ); +} + +type AnchorRequest = BudgetModel & { + routeProvider?: string; + agentId?: string; + thinkingLevel: ThinkingLevel; +}; + +/** + * 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. + */ +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) && + 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 ?? {}; + // 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) && sameRequestIdentity(reply, request) ? reply : undefined; +} + +/** 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); +} + +/** + * 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, + * 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 + * 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 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[], + request: AnchorRequest & { system: string } +): Promise { + const row = rows.findLast((candidate) => candidate.role === "assistant"); + if (row == null) return undefined; + const reply = anchoringReply(row.metadata, request); + const lastStep = lastStepParts(row); + if (reply == null || splitsAfterToolCall(lastStep)) return undefined; + const replayed = replayedStepTokens( + reply.contextUsage, + reply.contextProviderMetadata, + lastStep.some((part) => part.type === "reasoning") + ); + if (replayed == null) return undefined; + // 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; + const providerTokens = replayed + Math.ceil(Math.max(0, growth) * factor); + return Number.isSafeInteger(providerTokens) + ? { model: request.model, metadataModel: request.metadataModel, providerTokens } : undefined; } @@ -218,12 +371,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 +414,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 +443,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 +465,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..1d2b9cf54fa 100644 --- a/src/node/services/streamManager.contextBudget.test.ts +++ b/src/node/services/streamManager.contextBudget.test.ts @@ -9,12 +9,15 @@ 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"; 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,101 @@ 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. 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"; + 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 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, + system: "Answer briefly", + runtime: new LocalRuntime(h.tempDir), + providedRuntimeTempDir: runtimeDir, + contextBudgetLimit: 10000, + // The previous reply's provider count alone is at the ceiling. + ...(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 }), + } + : {}), + }); + 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..8ff652b04c8 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,34 @@ export function buildContextBudgetRequest( }; } -/** Hard-stop the step when its assembled request exceeds the safe context budget. */ +/** + * 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). + */ 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 +733,8 @@ export class StepPreparation { budgetRequest, previousRequest, steps.at(-1), - stepNumber + stepNumber, + turnAnchorAfterOverride(request, thinkingOverride) ); } 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..e7bbda924d3 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"; @@ -446,6 +447,82 @@ 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, + metadataModel: base.modelString, + agentId: "exec", + thinkingLevel: base.effectiveThinkingLevel, + 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(); + }); + + // #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", + thinkingLevel: opus.effectiveThinkingLevel, + 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 220461c36cc..0f59e62af13 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,8 +592,10 @@ 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; + /** The turn anchor that count used; a rebuild at another level carries its own (#7662). */ + contextBudgetTurnAnchor?: TurnBudgetAnchor; } /** @@ -616,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). */ @@ -646,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 }; @@ -661,7 +669,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 +763,52 @@ 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, + thinkingLevel: options.effectiveThinkingLevel, + // The text the reply's systemMessageTokens counted (measureSentSystem), rows included. + system: payload.sentSystem.message + payload.sentSystem.rowsText, + }); +} + +/** 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 +821,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 }; } /** @@ -3596,6 +3643,7 @@ export class TurnRequestBuilder { return { messages: rebuiltPayload.messages, contextBudgetEstimate: rebuiltPayload.contextBudgetEstimate, + contextBudgetTurnAnchor: rebuiltPayload.contextBudgetTurnAnchor, systemMessageTokens: sent.tokens, systemRecords: rowSystemRecords(rebuiltPayload, options.partialContinuationMessage), }; @@ -3617,6 +3665,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 +3695,7 @@ export class TurnRequestBuilder { systemRecords: rowSystemRecords(payload, options.partialContinuationMessage), messages: payload.messages, contextBudgetLimit: payload.contextBudgetLimit, + contextBudgetTurnAnchor: payload.contextBudgetTurnAnchor, contextBudgetEstimate: payload.contextBudgetEstimate, }; }, @@ -3767,6 +3817,7 @@ export class TurnRequestBuilder { systemRecords: primaryRequest.systemRecords, systemMessageTokens, contextBudgetEstimate: primaryRequest.contextBudgetEstimate, + contextBudgetTurnAnchor: primaryRequest.contextBudgetTurnAnchor, }); if ("exceeded" in streamState) { runLanguageModelCleanup(modelResult.data.model); @@ -3853,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 @@ -4135,6 +4187,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 +4226,7 @@ export class TurnRequestBuilder { abortSignal: combinedAbortSignal, tools: toolsForStream, contextBudgetLimit: primaryRequest.contextBudgetLimit, + contextBudgetTurnAnchor: streamTurnAnchor, 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).