fix(ai): ignore stale usage after compaction

closes #6464

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

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;
+		if (message.role === "assistant") {
 			const assistant = message as AssistantMessage;
-		if (assistant.stopReason === "aborted" || assistant.stopReason === "error") continue;
-		if (calculateContextTokens(assistant.usage) > 0) return { usage: assistant.usage, index: i };
+			// 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 };
 			}
-	return undefined;
+		}
+		latestPrefixTimestamp = Math.max(latestPrefixTimestamp, message.timestamp);
+	}
+
+	return usageInfo;
 }
 
 function estimateMessages(messages: readonly Message[]): ContextUsageEstimate {

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,
+		});
+	});
+});