fix(coding-agent/session): fixed context usage breakdown to prefer completed in-turn anchors

- Fixed context breakdown to anchor estimates on the latest completed assistant usage message after compaction. - Adjusted pending-context usage selection to prefer an in-turn provider anchor when available at/after cutoff. - Added a contextUsageRevision cache token so status-line context memo invalidates after snapshot clear.
2026-06-17 12:19:49 +02:00
parent 740ced3742
commit 409196bf2e
8 changed files with 260 additions and 66 deletions
@@ -142,14 +142,14 @@ export async function getOAuthApiKey(
 		provider === "alibaba-coding-plan";
 	const apiKey = needsStructuredApiKey
 		? JSON.stringify({
-			token: creds.access,
-			enterpriseUrl: creds.enterpriseUrl,
-			projectId: creds.projectId,
-			refreshToken: creds.refresh,
-			expiresAt: creds.expires,
-			email: creds.email,
-			accountId: creds.accountId,
-		})
+				token: creds.access,
+				enterpriseUrl: creds.enterpriseUrl,
+				projectId: creds.projectId,
+				refreshToken: creds.refresh,
+				expiresAt: creds.expires,
+				email: creds.email,
+				accountId: creds.accountId,
+			})
 		: creds.access;
 	return { newCredentials: creds, apiKey };
 }
@@ -1,7 +1,6 @@
 # Changelog

 ## [Unreleased]
-
 ### Added

 - Added `images.describeForTextModels` option (default `true`) to control automatic image description for attachments sent to models without vision input
@@ -14,9 +13,11 @@

 ### Fixed

+- Fixed context usage breakdown to use a completed assistant usage anchor from the current turn instead of a pending prompt snapshot so totals no longer overcount when a large in-turn tool step returns usage
 - Fixed context token accounting to keep branch-local anchors during branching so sibling-branch messages no longer pollute context estimates
 - Fixed context usage consistency so `/context`, status line, and idle compaction logic now report the same used-token totals
 - Fixed status-line context cache invalidation when assistant reasoning signature data grows so displayed context usage updates accurately
+- Fixed the status-line context% reading inflated during long tool turns and then dropping sharply on the next message even though no compaction ran. While a request was in flight `getContextBreakdown` summed a cl100k estimate of the entire tail on top of the stale turn-start prompt and never re-anchored to completed in-turn steps; it now prefers the real provider prompt-token count of any step that resolves at or after the pending cutoff. The status-line memo also keys on a `contextUsageRevision` that bumps when the in-flight snapshot is set/cleared, so a mid-turn estimate is invalidated on turn end/abort instead of surviving into idle until the next message
 - Fixed image attachment handling for text-only models by saving attachments to `local://` and injecting generated descriptions so they are no longer lost when the target model cannot process images
 - Fixed the ssh tool rejecting valid Windows identity files before invoking OpenSSH by skipping Unix mode-bit key validation on native Windows ([#2850](https://github.com/can1357/oh-my-pi/issues/2850)).

@@ -415,7 +415,7 @@ export class CollabHost {
 		// render exactly the same anchored, provider-real count the host's own
 		// status line shows.
 		const breakdown = this.#ctx.statusLine.getCachedContextBreakdown();
-		const tokens = breakdown.usedTokens;
+		const tokens = breakdown.usedTokens ?? 0;
 		return {
 			isStreaming: session.isStreaming,
 			isAborting: session.isAborting,
@@ -145,6 +145,7 @@ interface ContextUsageMemo {
 	length: number;
 	lastFingerprint: string | undefined;
 	modelContextWindow: number;
+	contextUsageRevision: number;
 	usedTokens: number;
 	contextWindow: number;
 	systemPromptRef: readonly string[] | undefined;
@@ -603,6 +604,10 @@ export class StatusLineComponent implements Component {
 		const modelContextWindow = this.session.model?.contextWindow ?? 0;
 		const length = messages.length;
 		const lastFingerprint = length > 0 ? messageFingerprint(messages[length - 1]!) : undefined;
+		// Bumps when the in-flight pending snapshot is set/cleared. Without it a
+		// value computed mid-turn (estimate of the active tail) would survive after
+		// the turn ends/aborts, since clearing the snapshot touches no message.
+		const contextUsageRevision = this.session.contextUsageRevision ?? 0;

 		const systemPrompt = this.session.systemPrompt;
 		const tools = this.session.agent?.state?.tools;
@@ -615,6 +620,7 @@ export class StatusLineComponent implements Component {
 			cache.length === length &&
 			cache.lastFingerprint === lastFingerprint &&
 			cache.modelContextWindow === modelContextWindow &&
+			cache.contextUsageRevision === contextUsageRevision &&
 			cache.systemPromptRef === systemPrompt &&
 			cache.toolsRef === tools &&
 			cache.skillsRef === skills
@@ -630,6 +636,7 @@ export class StatusLineComponent implements Component {
 			length,
 			lastFingerprint,
 			modelContextWindow,
+			contextUsageRevision,
 			usedTokens,
 			contextWindow,
 			systemPromptRef: systemPrompt,
@@ -1247,6 +1247,11 @@ export class AgentSession {
 				cutoffCount: number;
 		  }
 		| undefined = undefined;
+	// Bumped whenever the pending in-flight snapshot is set/cleared. The
+	// status-line context memo includes this so clearing the snapshot on
+	// turn-end/abort invalidates the cache even though the message list is
+	// unchanged — otherwise a mid-turn estimate would survive into idle.
+	#contextUsageRevision = 0;
 	#obfuscator: SecretObfuscator | undefined;
 	#checkpointState: CheckpointState | undefined = undefined;
 	#pendingRewindReport: string | undefined = undefined;
@@ -5593,15 +5598,15 @@ export class AgentSession {
 				nonMessageTokens +
 					this.messages.reduce((sum, msg) => sum + estimateTokens(msg), 0) +
 					messages.reduce((sum, msg) => sum + estimateTokens(msg), 0);
-			this.#pendingContextSnapshot = {
+			this.#setPendingContextSnapshot({
 				promptTokens,
 				nonMessageTokens,
 				cutoffCount: this.messages.length + messages.length,
-			};
+			});
 			try {
 				await this.#promptAgentWithIdleRetry(messages, agentPromptOptions);
 			} finally {
-				this.#pendingContextSnapshot = undefined;
+				this.#setPendingContextSnapshot(undefined);
 			}
 			if (!options?.skipPostPromptRecoveryWait) {
 				await this.#waitForPostPromptRecovery(generation);
@@ -11276,76 +11281,80 @@ export class AgentSession {

 		const pendingMessages = options?.pendingMessages ?? [];

-		let anchorEntry: SessionMessageEntry | undefined;
-		let isPending = false;
+		const pending = this.#pendingContextSnapshot;

-		if (this.#pendingContextSnapshot) {
-			isPending = true;
-		} else {
-			for (let i = branchEntries.length - 1; i > compactionIndex; i--) {
-				const entry = branchEntries[i];
-				if (entry.type === "message" && entry.message.role === "assistant") {
-					const assistant = entry.message;
-					if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error" && assistant.usage) {
-						anchorEntry = entry;
-						break;
-					}
+		// Always locate the latest real assistant-usage anchor after the last
+		// compaction. Its provider-reported promptTokens is ground truth for
+		// everything up to that point; only the tail after it is estimated.
+		let anchorEntry: SessionMessageEntry | undefined;
+		for (let i = branchEntries.length - 1; i > compactionIndex; i--) {
+			const entry = branchEntries[i];
+			if (entry.type === "message" && entry.message.role === "assistant") {
+				const assistant = entry.message;
+				if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error" && assistant.usage) {
+					anchorEntry = entry;
+					break;
 				}
 			}
 		}

-		if (isPending && this.#pendingContextSnapshot) {
-			const anchor = this.#pendingContextSnapshot;
-			anchored = true;
-
-			const resolvedActiveMessages = this.messages;
-			let tailTokens = 0;
-
-			if (resolvedActiveMessages.length > anchor.cutoffCount) {
-				for (let i = anchor.cutoffCount; i < resolvedActiveMessages.length; i++) {
-					tailTokens += estimateTokens(resolvedActiveMessages[i]);
-				}
+		const resolvedActiveMessages = this.messages;
+		let resolvedAnchorIndex = -1;
+		let anchorAssistant: AssistantMessage | undefined;
+		if (anchorEntry) {
+			const a = anchorEntry.message as AssistantMessage;
+			anchorAssistant = a;
+			resolvedAnchorIndex = resolvedActiveMessages.indexOf(a);
+			if (resolvedAnchorIndex === -1) {
+				resolvedAnchorIndex = resolvedActiveMessages.findIndex(
+					msg => msg.role === "assistant" && msg.timestamp === a.timestamp,
+				);
 			}
+		}

-			usedTokens =
-				anchor.promptTokens +
-				Math.max(0, currentNonMessageTokens - anchor.nonMessageTokens) +
-				tailTokens +
-				pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0);
-		} else if (anchorEntry) {
-			const anchorAssistant = anchorEntry.message as AssistantMessage;
+		// A real anchor supersedes the in-flight estimate only once a step of the
+		// CURRENT turn has produced provider usage — i.e. it resolves at or after
+		// the pending cutoff. While the turn's first response is still pending (or
+		// the newest real anchor predates this turn) the pending snapshot is the
+		// only thing accounting for the just-submitted prompt, so it wins. This
+		// keeps a long tool turn from stacking an estimate of the entire tail on
+		// top of a stale turn-start prompt.
+		const useAnchor =
+			anchorAssistant !== undefined &&
+			resolvedAnchorIndex !== -1 &&
+			(!pending || resolvedAnchorIndex >= pending.cutoffCount);
+
+		if (useAnchor && anchorAssistant) {
 			const promptTokens =
 				anchorAssistant.contextSnapshot?.promptTokens ?? calculatePromptTokens(anchorAssistant.usage);
 			const nonMessageTokens = anchorAssistant.contextSnapshot?.nonMessageTokens ?? computeNonMessageTokens(this);
-			const anchor = { promptTokens, nonMessageTokens };
 			anchored = true;
-
-			const resolvedActiveMessages = this.messages;
-			let resolvedAnchorIndex = resolvedActiveMessages.indexOf(anchorAssistant);
-			if (resolvedAnchorIndex === -1) {
-				resolvedAnchorIndex = resolvedActiveMessages.findIndex(
-					msg => msg.role === "assistant" && msg.timestamp === anchorAssistant.timestamp,
-				);
+			let tailTokens = 0;
+			for (let i = resolvedAnchorIndex + 1; i < resolvedActiveMessages.length; i++) {
+				tailTokens += estimateTokens(resolvedActiveMessages[i]);
 			}
-
-			if (resolvedAnchorIndex !== -1) {
-				let tailTokens = 0;
-				for (let i = resolvedAnchorIndex + 1; i < resolvedActiveMessages.length; i++) {
+			usedTokens =
+				promptTokens +
+				Math.max(0, currentNonMessageTokens - nonMessageTokens) +
+				tailTokens +
+				pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0);
+		} else if (pending) {
+			anchored = true;
+			let tailTokens = 0;
+			if (resolvedActiveMessages.length > pending.cutoffCount) {
+				for (let i = pending.cutoffCount; i < resolvedActiveMessages.length; i++) {
 					tailTokens += estimateTokens(resolvedActiveMessages[i]);
 				}
-				usedTokens =
-					anchor.promptTokens +
-					Math.max(0, currentNonMessageTokens - anchor.nonMessageTokens) +
-					tailTokens +
-					pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0);
-			} else {
-				anchored = false;
 			}
+			usedTokens =
+				pending.promptTokens +
+				Math.max(0, currentNonMessageTokens - pending.nonMessageTokens) +
+				tailTokens +
+				pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0);
 		}

-		if (!anchored && !isPending && branchEntries.length === 0) {
+		if (!anchored && !pending && branchEntries.length === 0) {
 			// Fallback: look for the latest assistant message with usage/snapshot in this.messages (for branchless/fake sessions in tests)
-			const resolvedActiveMessages = this.messages;
 			for (let i = resolvedActiveMessages.length - 1; i >= 0; i--) {
 				const msg = resolvedActiveMessages[i];
 				if (msg.role === "assistant" && msg.stopReason !== "aborted" && msg.stopReason !== "error" && msg.usage) {
@@ -11368,7 +11377,6 @@ export class AgentSession {
 			}
 		}
 		if (!anchored) {
-			const resolvedActiveMessages = this.messages;
 			let messagesTokens = 0;
 			for (const msg of resolvedActiveMessages) {
 				messagesTokens += estimateTokens(msg);
@@ -11403,6 +11411,22 @@ export class AgentSession {
 		};
 	}

+	/**
+	 * Monotonic counter that changes whenever the in-flight pending context
+	 * snapshot is set or cleared. Status-line context memoization keys on this so
+	 * a value computed mid-turn cannot persist after the turn ends/aborts.
+	 */
+	get contextUsageRevision(): number {
+		return this.#contextUsageRevision;
+	}
+
+	#setPendingContextSnapshot(
+		snapshot: { promptTokens: number; nonMessageTokens: number; cutoffCount: number } | undefined,
+	): void {
+		this.#pendingContextSnapshot = snapshot;
+		this.#contextUsageRevision++;
+	}
+
 	#ingestProviderUsageHeaders(response: ProviderResponseMetadata, model?: Model): void {
 		if (model?.provider !== "anthropic") return;
 		this.#modelRegistry.authStorage.ingestUsageHeaders("anthropic", response.headers, {
@@ -409,6 +409,105 @@ describe("Context usage consolidation", () => {
 		await tempDir.remove();
 	});

+	it("prefers a completed in-turn provider anchor over the pending snapshot", async () => {
+		const tempDir = TempDir.createSync("@inturn-anchor-");
+		const { session, sessionManager, agent } = createSession(tempDir);
+
+		const { promise, resolve } = Promise.withResolvers<void>();
+		const promptSpy = vi.spyOn(agent, "prompt").mockImplementation(async () => {
+			await promise;
+		});
+
+		const promptPromise = session.prompt("new question");
+		await Bun.sleep(10);
+
+		// While the request hangs (pending snapshot active, cutoff at the new
+		// prompt), simulate the turn's user message landing plus a completed tool
+		// step carrying a real, large provider prompt count.
+		sessionManager.appendMessage({ role: "user", content: "new question", timestamp: 5000 } as Message);
+		sessionManager.appendMessage({
+			role: "assistant",
+			content: [{ type: "text", text: "step" }],
+			api: mockModel.api,
+			provider: mockModel.provider,
+			model: mockModel.id,
+			stopReason: "toolUse",
+			usage: {
+				input: 9000,
+				output: 20,
+				cacheRead: 0,
+				cacheWrite: 0,
+				totalTokens: 9020,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			},
+			contextSnapshot: { promptTokens: 9000, nonMessageTokens: 10 },
+			timestamp: 6000,
+		} as AssistantMessage);
+		syncSession(session, agent);
+
+		// The in-turn anchor (index >= pending cutoff) must win: usage reflects the
+		// real 9000-token prompt, not the tiny turn-start pending estimate that
+		// would otherwise stack an estimate of the whole tail on top.
+		const breakdown = session.getContextBreakdown();
+		expect(breakdown?.anchored).toBe(true);
+		expect(breakdown?.usedTokens).toBeGreaterThanOrEqual(9000);
+
+		resolve();
+		await promptPromise;
+		promptSpy.mockRestore();
+		await tempDir.remove();
+	});
+
+	it("keeps the pending snapshot (not a pre-cutoff anchor) while the first turn response is pending", async () => {
+		const tempDir = TempDir.createSync("@pending-precutoff-");
+		const { session, sessionManager, agent } = createSession(tempDir);
+
+		// Prior completed turn establishes a real anchor that PREDATES the new
+		// prompt (resolves before the pending cutoff).
+		sessionManager.appendMessage({ role: "user", content: "old", timestamp: 1000 } as Message);
+		sessionManager.appendMessage({
+			role: "assistant",
+			content: [{ type: "text", text: "old response" }],
+			api: mockModel.api,
+			provider: mockModel.provider,
+			model: mockModel.id,
+			stopReason: "stop",
+			usage: {
+				input: 5000,
+				output: 20,
+				cacheRead: 0,
+				cacheWrite: 0,
+				totalTokens: 5020,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			},
+			contextSnapshot: { promptTokens: 5000, nonMessageTokens: 10 },
+			timestamp: 2000,
+		} as AssistantMessage);
+		syncSession(session, agent);
+
+		const { promise, resolve } = Promise.withResolvers<void>();
+		const promptSpy = vi.spyOn(agent, "prompt").mockImplementation(async () => {
+			await promise;
+		});
+
+		// Large prompt in flight, no in-turn step has produced provider usage yet.
+		const bigPrompt = "explain this in detail ".repeat(200);
+		const promptPromise = session.prompt(bigPrompt);
+		await Bun.sleep(10);
+
+		// Pending must win: it accounts for the just-submitted prompt on top of the
+		// prior anchor. The stale pre-cutoff anchor alone (5000) would omit it, so a
+		// regression to "always prefer the latest real anchor" would read ~5000.
+		const breakdown = session.getContextBreakdown();
+		expect(breakdown?.anchored).toBe(true);
+		expect(breakdown?.usedTokens).toBeGreaterThan(5500);
+
+		resolve();
+		await promptPromise;
+		promptSpy.mockRestore();
+		await tempDir.remove();
+	});
+
 	it("guarantees always numeric nullable-vs-speculative contract", async () => {
 		const tempDir = TempDir.createSync("@always-numeric-");
 		const { session, sessionManager, agent } = createSession(tempDir);
@@ -120,4 +120,40 @@ describe("SessionManager usage statistics", () => {
 		const usage = session.getUsageStatistics();
 		expect(usage.premiumRequests).toBe(0);
 	});
+
+	it("accumulates the full billed cost across turns, including cache-read cost", () => {
+		// Contract: the session cost aggregate sums each turn's full `cost.total`
+		// (input+output+cacheRead+cacheWrite), not a cache-excluded "new-work"
+		// subset. Cache-read cost is real billed spend — the cached context is
+		// re-read at the cache-read rate every turn — so it must stay in the
+		// ledger that /usage, ACP usage_update, and hooks consume. Two turns with
+		// nonzero cacheRead make the readings diverge: full total = 18 vs the
+		// excluded subset (input+output+cacheWrite) = 8.
+		const session = SessionManager.inMemory();
+
+		session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
+		for (const timestamp of [2, 3]) {
+			session.appendMessage({
+				role: "assistant",
+				content: [{ type: "text", text: "hi" }],
+				api: "anthropic-messages",
+				provider: "anthropic",
+				model: "claude-sonnet-4",
+				usage: {
+					input: 1,
+					output: 2,
+					cacheRead: 100,
+					cacheWrite: 10,
+					totalTokens: 113,
+					cost: { input: 1, output: 2, cacheRead: 5, cacheWrite: 1, total: 9 },
+				},
+				stopReason: "stop",
+				timestamp,
+			});
+		}
+
+		const usage = session.getUsageStatistics();
+		expect(usage.cacheRead).toBe(200);
+		expect(usage.cost).toBeCloseTo(18, 8);
+	});
 });
@@ -36,12 +36,15 @@ interface Fake {
 	usageCalls: () => number;
 	/** Swap the value the next `getContextUsage()` query returns. */
 	setUsage: (usage: ContextUsage | undefined) => void;
+	/** Bump the in-flight pending revision the next `getCachedContextBreakdown()` reads. */
+	setRevision: (n: number) => void;
 }

 function makeSession(opts: { messages: unknown[]; contextWindow?: number; usage?: ContextUsage | undefined }): Fake {
 	const contextWindow = opts.contextWindow ?? 200_000;
 	let usage: ContextUsage | undefined = "usage" in opts ? opts.usage : { tokens: 1234, contextWindow, percent: 0.6 };
 	let calls = 0;
+	let revision = 0;
 	const session = {
 		messages: opts.messages,
 		systemPrompt: ["You are a helpful assistant."],
@@ -65,6 +68,9 @@ function makeSession(opts: { messages: unknown[]; contextWindow?: number; usage?
 			calls++;
 			return usage;
 		},
+		get contextUsageRevision() {
+			return revision;
+		},
 	} as unknown as AgentSession;
 	return {
 		session,
@@ -72,6 +78,9 @@ function makeSession(opts: { messages: unknown[]; contextWindow?: number; usage?
 		setUsage: next => {
 			usage = next;
 		},
+		setRevision: (n: number) => {
+			revision = n;
+		},
 	};
 }

@@ -155,6 +164,24 @@ describe("StatusLineComponent context breakdown", () => {
 		expect(usageCalls()).toBe(2);
 	});

+	it("re-queries when only the in-flight pending revision changes (no message change)", () => {
+		const fake = makeSession({
+			messages: [userMessage("hi")],
+			usage: { tokens: 190_000, contextWindow: 272_000, percent: 69.9 },
+		});
+		const comp = new StatusLineComponent(fake.session);
+		expect(comp.getCachedContextBreakdown().usedTokens).toBe(190_000);
+
+		// Turn ends/aborts: the message list and last-message fingerprint are
+		// unchanged, but clearing the pending snapshot recalibrates usage to the
+		// real provider anchor. The memo must not keep serving the stale estimate.
+		fake.setUsage({ tokens: 117_000, contextWindow: 272_000, percent: 43.0 });
+		fake.setRevision(1);
+
+		expect(comp.getCachedContextBreakdown().usedTokens).toBe(117_000);
+		expect(fake.usageCalls()).toBe(2);
+	});
+
 	it("propagates a speculative/numeric token count, e.g. right after compaction", () => {
 		const { session } = makeSession({
 			messages: [userMessage("compaction summary")],