diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 92efd102..6d006c29 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed post-compaction output-token budgeting to ignore stale assistant usage from before the compaction boundary ([#6464](https://github.com/earendil-works/pi/issues/6464)). + ## [0.80.5] - 2026-07-09 ## [0.80.4] - 2026-07-09 diff --git a/packages/ai/src/utils/estimate.ts b/packages/ai/src/utils/estimate.ts index 711bf783..4fe5010f 100644 --- a/packages/ai/src/utils/estimate.ts +++ b/packages/ai/src/utils/estimate.ts @@ -3,11 +3,11 @@ import type { AssistantMessage, Context, ImageContent, Message, TextContent, Usa export interface ContextUsageEstimate { /** Estimated total context tokens. */ tokens: number; - /** Tokens reported by the most recent assistant usage block. */ + /** Tokens reported by the most recent applicable assistant usage block. */ usageTokens: number; - /** Estimated tokens after the most recent assistant usage block. */ + /** Estimated tokens after the most recent applicable assistant usage block. */ trailingTokens: number; - /** Index of the message that provided usage, or null when none exists. */ + /** Index of the applicable message that provided usage, or null when none exists. */ lastUsageIndex: number | null; } @@ -61,14 +61,29 @@ export function estimateMessageTokens(message: Message): number { } function getLastAssistantUsageInfo(messages: readonly Message[]): { usage: Usage; index: number } | undefined { - for (let i = messages.length - 1; i >= 0; i--) { + let latestPrefixTimestamp = Number.NEGATIVE_INFINITY; + let usageInfo: { usage: Usage; index: number } | undefined; + + for (let i = 0; i < messages.length; i++) { const message = messages[i]; - if (message.role !== "assistant") continue; - const assistant = message as AssistantMessage; - if (assistant.stopReason === "aborted" || assistant.stopReason === "error") continue; - if (calculateContextTokens(assistant.usage) > 0) return { usage: assistant.usage, index: i }; + if (message.role === "assistant") { + const assistant = message as AssistantMessage; + // A newer prefix message was inserted after this response (for example, a + // compaction summary), so its usage cannot describe the current prefix. + const usageAppliesToPrefix = assistant.timestamp >= latestPrefixTimestamp; + if ( + usageAppliesToPrefix && + assistant.stopReason !== "aborted" && + assistant.stopReason !== "error" && + calculateContextTokens(assistant.usage) > 0 + ) { + usageInfo = { usage: assistant.usage, index: i }; + } + } + latestPrefixTimestamp = Math.max(latestPrefixTimestamp, message.timestamp); } - return undefined; + + return usageInfo; } function estimateMessages(messages: readonly Message[]): ContextUsageEstimate { diff --git a/packages/ai/test/context-estimate.test.ts b/packages/ai/test/context-estimate.test.ts new file mode 100644 index 00000000..85047309 --- /dev/null +++ b/packages/ai/test/context-estimate.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it } from "vitest"; +import { buildBaseOptions } from "../src/api/simple-options.ts"; +import type { AssistantMessage, Context, Model, Usage } from "../src/types.ts"; +import { estimateContextTokens } from "../src/utils/estimate.ts"; + +function createUsage(totalTokens: number): Usage { + return { + input: totalTokens, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function createAssistant(timestamp: number, totalTokens: number): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text: "kept" }], + api: "openai-responses", + provider: "openai", + model: "test-model", + usage: createUsage(totalTokens), + stopReason: "stop", + timestamp, + }; +} + +const model: Model<"openai-responses"> = { + id: "test-model", + name: "Test Model", + api: "openai-responses", + provider: "openai", + baseUrl: "https://api.openai.com/v1", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 10_000, + maxTokens: 8_000, +}; + +describe("context token estimation", () => { + it("ignores stale assistant usage after a newer message is inserted before it", () => { + const context: Context = { + systemPrompt: "system", + messages: [ + { role: "user", content: "summary", timestamp: 200 }, + createAssistant(100, 9_500), + { role: "user", content: "x".repeat(4_000), timestamp: 300 }, + ], + }; + + expect(estimateContextTokens(context)).toEqual({ + tokens: 1_005, + usageTokens: 0, + trailingTokens: 1_005, + lastUsageIndex: null, + }); + expect(buildBaseOptions(model, context).maxTokens).toBe(4_899); + }); + + it("uses assistant usage again after a response to the inserted context", () => { + const context: Context = { + messages: [ + { role: "user", content: "summary", timestamp: 200 }, + createAssistant(100, 9_500), + { role: "user", content: "new prompt", timestamp: 300 }, + createAssistant(400, 2_000), + { role: "user", content: "tail", timestamp: 500 }, + ], + }; + + expect(estimateContextTokens(context)).toEqual({ + tokens: 2_001, + usageTokens: 2_000, + trailingTokens: 1, + lastUsageIndex: 3, + }); + }); +});