Merge PR #2528: emit per-turn token-usage row as standalone block

2026-06-14 02:56:50 +02:00
parent dddfc3320c 35307b60a7
commit eb5cc85eed
10 changed files with 199 additions and 57 deletions
@@ -39,6 +39,10 @@

 - Fixed tool-call spinners animating out of phase across parallel tool calls — each live tool block advanced its glyph from its own per-instance start time, so concurrent spinners showed different frames. Glyphs now derive from a single shared monotonic clock (`sharedSpinnerFrame`), keeping every live block in lockstep.

+### Fixed
+
+- Fixed the per-turn token-usage row (`display.showTokenUsage`) churning and duplicating in scrollback, most visibly with parallel tool calls. The row was rendered inside the assistant block above the turn's tool blocks, so finalizing the block was deferred and late appends recommitted the already-committed tool rows. The assistant block now always finalizes as soon as a tool-call appears, and the usage row is emitted as a standalone finalized block below the turn's tool blocks across all three render paths (live, transcript rebuild, agent-hub).
+
 ## [15.12.5] - 2026-06-13
 ### Changed

@@ -16,6 +16,7 @@
 import * as fs from "node:fs";
 import * as path from "node:path";
 import type { AgentMessage, AgentTool } from "@oh-my-pi/pi-agent-core";
+import type { Usage } from "@oh-my-pi/pi-ai";
 import { Container, Editor, matchesKey, ScrollView, Text, type TUI } from "@oh-my-pi/pi-tui";
 import { formatAge, formatBytes, formatDuration, formatNumber, getProjectDir, logger } from "@oh-my-pi/pi-utils";
 import { COLLAB_PROMPT_MESSAGE_TYPE, type CollabPromptDetails } from "../../collab/protocol";
@@ -57,6 +58,7 @@ import { SkillMessageComponent } from "./skill-message";
 import { formatContextUsage } from "./status-line/context-thresholds";
 import { ToolExecutionComponent } from "./tool-execution";
 import { TranscriptBlock, TranscriptContainer } from "./transcript-container";
+import { createUsageRowBlock } from "./usage-row";
 import { UserMessageComponent } from "./user-message";

 /** Lines per page for PageUp/PageDown */
@@ -213,6 +215,7 @@ export class AgentHubOverlayComponent extends Container {
 	#chatPendingTools = new Map<string, ToolExecutionComponent | ReadToolGroupComponent>();
 	#chatReadArgs = new Map<string, Record<string, unknown>>();
 	#chatReadGroup: ReadToolGroupComponent | null = null;
+	#pendingUsage: Usage | undefined;
 	#chatWaitingPoll: ToolExecutionComponent | null = null;
 	#chatExpandables: Array<{ setExpanded(expanded: boolean): void }> = [];
 	#chatExpanded = false;
@@ -851,6 +854,7 @@ export class AgentHubOverlayComponent extends Container {
 		this.#chatPendingTools.clear();
 		this.#chatReadArgs.clear();
 		this.#chatReadGroup = null;
+		this.#pendingUsage = undefined;
 		this.#chatWaitingPoll = null;
 		this.#chatExpandables = [];
 		this.#chatLog.dispose();
@@ -870,6 +874,13 @@ export class AgentHubOverlayComponent extends Container {
 			this.#appendChatMessage(entries[i].message);
 		}
 		this.#chatBuiltCount = entries.length;
+		// Flush the trailing turn's usage row only once its tools are materialized.
+		// A read (or any tool) whose toolResult lands in a later debounced sync stays
+		// pending in #chatReadArgs / #chatPendingTools; flushing now would emit the
+		// row above it. The sync that drains the maps flushes it below the tools.
+		if (this.#chatReadArgs.size === 0 && this.#chatPendingTools.size === 0) {
+			this.#flushPendingUsage();
+		}
 	}

 	#trackExpandable(component: { setExpanded(expanded: boolean): void }): void {
@@ -899,7 +910,21 @@ export class AgentHubOverlayComponent extends Container {
 		return this.#chatReadGroup;
 	}

+	// The per-turn token-usage row must land below the turn's tool blocks, but
+	// normal `read` calls only materialize their group in #appendToolResult. Defer
+	// the row: stash it on the assistant message and flush once the turn's tools
+	// are placed — before the next non-toolResult message and at the end of each
+	// sync pass — sealing the read run so the row sits under it.
+	#flushPendingUsage(): void {
+		if (!this.#pendingUsage) return;
+		this.#chatReadGroup?.seal();
+		this.#chatReadGroup = null;
+		this.#chatLog.addChild(createUsageRowBlock(this.#pendingUsage));
+		this.#pendingUsage = undefined;
+	}
+
 	#appendChatMessage(message: AgentMessage): void {
+		if (message.role !== "toolResult") this.#flushPendingUsage();
 		switch (message.role) {
 			case "assistant":
 				this.#appendAssistantMessage(message);
@@ -988,7 +1013,6 @@ export class AgentHubOverlayComponent extends Container {
 		const assistantComponent = new AssistantMessageComponent(message, this.#hideThinkingBlock?.() ?? false, () =>
 			this.#requestRender(),
 		);
-		assistantComponent.setUsageInfo(message.usage);
 		this.#chatLog.addChild(assistantComponent);

 		const hasVisibleAssistantContent = message.content.some(
@@ -1067,6 +1091,8 @@ export class AgentHubOverlayComponent extends Container {
 				this.#chatPendingTools.set(content.id, component);
 			}
 		}
+
+		this.#pendingUsage = settings.get("display.showTokenUsage") ? message.usage : undefined;
 	}

 	#appendToolResult(message: Extract<AgentMessage, { role: "toolResult" }>): void {
@@ -1,7 +1,5 @@
-import type { AssistantMessage, ImageContent, Usage } from "@oh-my-pi/pi-ai";
+import type { AssistantMessage, ImageContent } from "@oh-my-pi/pi-ai";
 import { Container, Image, type ImageBudget, ImageProtocol, Markdown, Spacer, TERMINAL, Text } from "@oh-my-pi/pi-tui";
-import { formatNumber } from "@oh-my-pi/pi-utils";
-import { settings } from "../../config/settings";
 import type { AssistantThinkingRenderer } from "../../extensibility/extensions/types";
 import { getMarkdownTheme, theme } from "../../modes/theme/theme";
 import { resolveAbortLabel, shouldRenderAbortReason } from "../../session/messages";
@@ -24,7 +22,6 @@ export class AssistantMessageComponent extends Container {
 	#contentContainer: Container;
 	#lastMessage?: AssistantMessage;
 	#toolImagesByCallId = new Map<string, ImageContent[]>();
-	#usageInfo?: Usage;
 	#convertedKittyImages = new Map<string, ImageContent>();
 	#kittyConversionsInFlight = new Set<string>();
 	#transcriptBlockFinalized: boolean;
@@ -40,11 +37,9 @@ export class AssistantMessageComponent extends Container {
 	/**
 	 * Monotonic content version reported to the transcript container via
 	 * {@link getTranscriptBlockVersion}. Bumped by {@link updateContent} — the
-	 * choke point every mutator funnels through, including the post-finalize
-	 * ones: `setErrorPinned(false)` restoring the inline error at the next
-	 * turn's `agent_start`, late tool-result images, async Kitty conversions,
-	 * and `setUsageInfo`. Without it, the container's committed-scrollback
-	 * bypass would replay this block's pre-mutation bytes forever.
+	 * choke point every mutator funnels through, including post-finalize changes
+	 * such as `setErrorPinned(false)` restoring the inline error at the next
+	 * turn's `agent_start`, late tool-result images, and async Kitty conversions.
 	 */
 	#blockVersion = 0;
 	/** Whether the last updateContent carried an in-flight streaming partial; such
@@ -185,13 +180,6 @@ export class AssistantMessageComponent extends Container {
 		}
 	}

-	setUsageInfo(usage: Usage): void {
-		this.#usageInfo = usage;
-		if (this.#lastMessage) {
-			this.updateContent(this.#lastMessage, { transient: this.#lastUpdateTransient });
-		}
-	}
-
 	#renderToolImages(): void {
 		const imageEntries = Array.from(this.#toolImagesByCallId.entries()).flatMap(([toolCallId, images]) =>
 			images.map((image, index) => ({ image, key: `${toolCallId}:${index}` })),
@@ -256,12 +244,6 @@ export class AssistantMessageComponent extends Container {
 				parts.push(`O:${content.type}`);
 			}
 		}
-		if (settings.get("display.showTokenUsage") && this.#usageInfo) {
-			const u = this.#usageInfo;
-			parts.push(`u:${u.input + u.cacheWrite}:${u.output}:${u.cacheRead}`);
-		} else {
-			parts.push("u:");
-		}
 		return parts.join("|");
 	}

@@ -416,21 +398,6 @@ export class AssistantMessageComponent extends Container {
 		) {
 			this.#appendErrorBlock(message.errorMessage);
 		}
-
-		// Token usage metadata
-		if (settings.get("display.showTokenUsage") && this.#usageInfo) {
-			const usage = this.#usageInfo;
-			const totalInput = usage.input + usage.cacheWrite;
-			const parts: string[] = [];
-			parts.push(`${theme.icon.input} ${formatNumber(totalInput)}`);
-			parts.push(`${theme.icon.output} ${formatNumber(usage.output)}`);
-			if (usage.cacheRead > 0) {
-				parts.push(`cache: ${formatNumber(usage.cacheRead)}`);
-			}
-			this.#contentContainer.addChild(new Spacer(1));
-			this.#contentContainer.addChild(new Text(theme.fg("dim", parts.join("  ")), 1, 0));
-		}
-
 		// Store fast-path state for next call
 		if (shouldCapture) {
 			this.#fastPathItems = captureItems;
@@ -0,0 +1,18 @@
+import type { Usage } from "@oh-my-pi/pi-ai";
+import { Container, Spacer, Text } from "@oh-my-pi/pi-tui";
+import { formatNumber } from "@oh-my-pi/pi-utils";
+import { theme } from "../../modes/theme/theme";
+
+export function createUsageRowBlock(usage: Usage): Container {
+	const totalInput = usage.input + usage.cacheWrite;
+	const parts: string[] = [];
+	parts.push(`${theme.icon.input} ${formatNumber(totalInput)}`);
+	parts.push(`${theme.icon.output} ${formatNumber(usage.output)}`);
+	if (usage.cacheRead > 0) {
+		parts.push(`cache: ${formatNumber(usage.cacheRead)}`);
+	}
+	const block = new Container();
+	block.addChild(new Spacer(1));
+	block.addChild(new Text(theme.fg("dim", parts.join("  ")), 1, 0));
+	return block;
+}
@@ -13,6 +13,7 @@ import {
 import { TodoReminderComponent } from "../../modes/components/todo-reminder";
 import { ToolExecutionComponent } from "../../modes/components/tool-execution";
 import { TtsrNotificationComponent } from "../../modes/components/ttsr-notification";
+import { createUsageRowBlock } from "../../modes/components/usage-row";
 import { getSymbolTheme, theme } from "../../modes/theme/theme";
 import type { InteractiveModeContext, TodoPhase } from "../../modes/types";
 import type { PlanApprovalDetails } from "../../plan-mode/approved-plan";
@@ -454,14 +455,7 @@ export class EventController {
 			// stream (a big write/edit/eval) sits below a still-live block and
 			// can never reach native scrollback: the head of the preview is
 			// neither committed nor on screen and the transcript reads as cut.
-			// Skipped when the per-turn usage row is enabled: that row is only
-			// known at message_end and appends to this block, which would shift
-			// committed tool rows below it every turn (audit recommit →
-			// duplicated preview copies in scrollback).
-			if (
-				this.ctx.streamingMessage.content.some(content => content.type === "toolCall") &&
-				!settings.get("display.showTokenUsage")
-			) {
+			if (this.ctx.streamingMessage.content.some(content => content.type === "toolCall")) {
 				this.ctx.streamingComponent.markTranscriptBlockFinalized();
 			}
 			for (const content of this.ctx.streamingMessage.content) {
@@ -614,8 +608,10 @@ export class EventController {
 				this.#resolveDisplaceablePoll();
 			}
 			this.#lastAssistantComponent = this.ctx.streamingComponent;
-			this.#lastAssistantComponent.setUsageInfo(event.message.usage);
 			this.#lastAssistantComponent.markTranscriptBlockFinalized();
+			if (settings.get("display.showTokenUsage")) {
+				this.ctx.chatContainer.addChild(createUsageRowBlock(event.message.usage));
+			}
 			this.ctx.streamingComponent = undefined;
 			this.ctx.streamingMessage = undefined;
 			// Pin a turn-ending provider error (e.g. Anthropic content-filter block)
@@ -1,5 +1,5 @@
 import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
-import type { AssistantMessage, ImageContent, Message } from "@oh-my-pi/pi-ai";
+import type { AssistantMessage, ImageContent, Message, Usage } from "@oh-my-pi/pi-ai";
 import { type Component, Spacer, Text, TruncatedText } from "@oh-my-pi/pi-tui";
 import { COLLAB_PROMPT_MESSAGE_TYPE, type CollabPromptDetails } from "../../collab/protocol";
 import { settings } from "../../config/settings";
@@ -24,6 +24,7 @@ import {
 import { SkillMessageComponent } from "../../modes/components/skill-message";
 import { ToolExecutionComponent } from "../../modes/components/tool-execution";
 import { TranscriptBlock } from "../../modes/components/transcript-container";
+import { createUsageRowBlock } from "../../modes/components/usage-row";
 import { UserMessageComponent } from "../../modes/components/user-message";
 import { materializeImageReferenceLinksSync } from "../../modes/image-references";
 import { theme } from "../../modes/theme/theme";
@@ -340,6 +341,22 @@ export class UiHelpers {
 		let readGroup: ReadToolGroupComponent | null = null;
 		const readToolCallArgs = new Map<string, Record<string, unknown>>();
 		const readToolCallAssistantComponents = new Map<string, AssistantMessageComponent>();
+		// The per-turn token-usage row (display.showTokenUsage) must land below the
+		// turn's tool blocks. Read tool blocks are only created when their toolResult
+		// message is processed (below), so appending the row in the assistant branch
+		// would place it above a read run. Defer instead: stash the usage on the
+		// assistant message, then flush it once the turn's tools are placed — right
+		// before the next non-toolResult message and at end of rebuild — sealing the
+		// read run so the row sits under it. Mirrors the live path, where the read
+		// group is created during streaming and the row is appended below it.
+		let pendingUsage: Usage | undefined;
+		const flushPendingUsage = () => {
+			if (!pendingUsage) return;
+			readGroup?.seal();
+			readGroup = null;
+			this.ctx.chatContainer.addChild(createUsageRowBlock(pendingUsage));
+			pendingUsage = undefined;
+		};
 		// Rebuild-time mirror of the event controller's displaceable-poll
 		// bookkeeping: a `job` poll that found every watched job still running is
 		// superseded by the next `job` call, so a rebuilt transcript collapses a
@@ -357,14 +374,12 @@ export class UiHelpers {
 			previous.seal();
 		};
 		for (const message of sessionContext.messages) {
+			if (message.role !== "toolResult") flushPendingUsage();
 			// Assistant messages need special handling for tool calls
 			if (message.role === "assistant") {
 				this.ctx.addMessageToChat(message);
 				const lastChild = this.ctx.chatContainer.children[this.ctx.chatContainer.children.length - 1];
 				const assistantComponent = lastChild instanceof AssistantMessageComponent ? lastChild : undefined;
-				if (assistantComponent) {
-					assistantComponent.setUsageInfo(message.usage);
-				}
 				const hasVisibleAssistantContent = message.content.some(
 					content =>
 						(content.type === "text" && content.text.trim().length > 0) ||
@@ -461,6 +476,7 @@ export class UiHelpers {
 						this.ctx.pendingTools.set(content.id, component);
 					}
 				}
+				pendingUsage = this.ctx.settings.get("display.showTokenUsage") ? message.usage : undefined;
 			} else if (message.role === "toolResult") {
 				const pendingReadComponent = this.ctx.pendingTools.get(message.toolCallId);
 				const isReadGroupResult =
@@ -523,6 +539,7 @@ export class UiHelpers {
 				this.ctx.addMessageToChat(message, options);
 			}
 		}
+		flushPendingUsage();

 		// The trailing read run has no following break to close it; seal so the
 		// rebuilt group freezes (even with a never-persisted result) and commits to
@@ -16,8 +16,9 @@
 *       → `updateContent` receives a message with `stopReason: "stop"`;
 *         `errorMessage` is NOT set (TTSR existing behavior unchanged).
 */
-import { describe, expect, it, vi } from "bun:test";
+import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
 import type { AssistantMessage } from "@oh-my-pi/pi-ai";
+import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
 import { EventController } from "@oh-my-pi/pi-coding-agent/modes/controllers/event-controller";
 import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types";
 import type { AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session";
@@ -50,10 +51,9 @@ function createFixture(opts: {
 	retryAttempt?: number;
 }) {
 	const updateContent = vi.fn();
-	const setUsageInfo = vi.fn();
 	const setComplete = vi.fn();
 	const markTranscriptBlockFinalized = vi.fn();
-	const streamingComponent = { updateContent, setUsageInfo, setComplete, markTranscriptBlockFinalized };
+	const streamingComponent = { updateContent, setComplete, markTranscriptBlockFinalized };
 	const requestRender = vi.fn();

 	const ctxBase = {
@@ -82,6 +82,13 @@ function createFixture(opts: {
 }

 describe("EventController #handleMessageEnd abort labeling", () => {
+	beforeEach(async () => {
+		await Settings.init({ inMemory: true, cwd: process.cwd() });
+	});
+	afterEach(() => {
+		resetSettingsForTest();
+	});
+
 	it("C1: SILENT_ABORT_MARKER + aborted -> updateContent stopReason='stop', errorMessage NOT overwritten", async () => {
 		const message = makeAssistantMessage({
 			stopReason: "aborted",
@@ -54,7 +54,6 @@ afterEach(() => {
 function createFixture(streamingMessage?: AssistantMessage) {
 	const streamingComponent = {
 		updateContent: vi.fn(),
-		setUsageInfo: vi.fn(),
 		setComplete: vi.fn(),
 		markTranscriptBlockFinalized: vi.fn(),
 		setErrorPinned: vi.fn(),
@@ -103,7 +103,7 @@ describe("EventController finalizes assistant block when tool-call args stream",
 		expect(finalized).not.toHaveBeenCalled();
 	});

-	it("defers finalization to message_end when the per-turn usage row is enabled", async () => {
+	it("marks the streaming assistant finalized even when the per-turn usage row is enabled", async () => {
 		await Settings.init({ inMemory: true, cwd: process.cwd() });
 		settings.set("display.showTokenUsage", true);
 		const message = makeStreamingMessage([
@@ -111,6 +111,6 @@ describe("EventController finalizes assistant block when tool-call args stream",
 			{ type: "toolCall", id: "tc-2", name: "write", arguments: { file_path: "/tmp/b.ts", content: "y" } },
 		]);
 		const finalized = await dispatchUpdate(message);
-		expect(finalized).not.toHaveBeenCalled();
+		expect(finalized).toHaveBeenCalled();
 	});
 });
@@ -0,0 +1,108 @@
+/**
+ * Regression: when `display.showTokenUsage` is on, the per-turn token-usage row
+ * must render BELOW the turn's tool blocks on the transcript-rebuild path — including
+ * `read` tool groups, which are only materialized when their `toolResult` message is
+ * processed (not in the assistant pass). A naive append in the assistant branch put the
+ * row above the read group, diverging from the live path. The fix defers the row and
+ * flushes it after the turn's tools are placed.
+ */
+import { beforeAll, describe, expect, it, vi } from "bun:test";
+import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
+import { ReadToolGroupComponent } from "@oh-my-pi/pi-coding-agent/modes/components/read-tool-group";
+import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
+import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types";
+import { UiHelpers } from "@oh-my-pi/pi-coding-agent/modes/utils/ui-helpers";
+import type { SessionContext } from "@oh-my-pi/pi-coding-agent/session/session-manager";
+import { Container } from "@oh-my-pi/pi-tui";
+import { formatNumber } from "@oh-my-pi/pi-utils";
+
+// 4242 → "4.2K": distinctive enough not to collide with a read group's render.
+const USAGE_INPUT = 4242;
+const USAGE_LABEL = formatNumber(USAGE_INPUT);
+
+function readTurn(): AgentMessage[] {
+	const assistant = {
+		role: "assistant",
+		content: [{ type: "toolCall", id: "r1", name: "read", arguments: { path: "src/foo.ts" } }],
+		api: "anthropic-messages",
+		provider: "anthropic",
+		model: "claude-sonnet-4-5",
+		stopReason: "stop",
+		usage: {
+			input: USAGE_INPUT,
+			output: 7,
+			cacheRead: 0,
+			cacheWrite: 0,
+			totalTokens: USAGE_INPUT + 7,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		timestamp: Date.now(),
+	} as unknown as AgentMessage;
+	const toolResult = {
+		role: "toolResult",
+		toolCallId: "r1",
+		toolName: "read",
+		content: [{ type: "text", text: "line1\nline2" }],
+		timestamp: Date.now(),
+	} as unknown as AgentMessage;
+	return [assistant, toolResult];
+}
+
+function makeHarness(showTokenUsage: boolean): { ctx: InteractiveModeContext; helpers: UiHelpers } {
+	let helpers: UiHelpers;
+	const ctx = {
+		chatContainer: new Container(),
+		pendingTools: new Map(),
+		ui: { requestRender: vi.fn() },
+		statusLine: { invalidate: vi.fn() },
+		updateEditorBorderColor: vi.fn(),
+		settings: { get: (key: string) => (key === "display.showTokenUsage" ? showTokenUsage : false) },
+		addMessageToChat: (message: AgentMessage) => helpers.addMessageToChat(message),
+		session: {
+			retryAttempt: 0,
+			getToolByName: () => undefined,
+			sessionManager: { getCwd: () => process.cwd() },
+		},
+		get viewSession() {
+			return (this as typeof ctx).session;
+		},
+		toolOutputExpanded: false,
+		hideThinkingBlock: false,
+		clearTransientSessionUi: () => {},
+	} as unknown as InteractiveModeContext;
+	helpers = new UiHelpers(ctx);
+	return { ctx, helpers };
+}
+
+describe("UiHelpers.renderSessionContext token-usage row placement", () => {
+	beforeAll(async () => {
+		await initTheme();
+	});
+
+	it("places the usage row below the read group for a read turn", () => {
+		const { ctx, helpers } = makeHarness(true);
+		helpers.renderSessionContext({ messages: readTurn() } as SessionContext);
+
+		const children = ctx.chatContainer.children;
+		const readIdx = children.findIndex(c => c instanceof ReadToolGroupComponent);
+		expect(readIdx).toBeGreaterThanOrEqual(0);
+
+		// The usage row is the trailing block and renders the turn's input tokens.
+		const last = children[children.length - 1]!;
+		expect(last.render(120).join("\n")).toContain(USAGE_LABEL);
+		// And it sits strictly below the read group (the bug placed it above).
+		expect(children.length - 1).toBeGreaterThan(readIdx);
+		// Exactly one usage row — no duplication.
+		expect(children.filter(c => c.render(120).join("\n").includes(USAGE_LABEL))).toHaveLength(1);
+	});
+
+	it("renders no usage row when showTokenUsage is off", () => {
+		const { ctx, helpers } = makeHarness(false);
+		helpers.renderSessionContext({ messages: readTurn() } as SessionContext);
+
+		const children = ctx.chatContainer.children;
+		expect(children.some(c => c.render(120).join("\n").includes(USAGE_LABEL))).toBe(false);
+		// Last block is the read group, not a usage row.
+		expect(children[children.length - 1]).toBeInstanceOf(ReadToolGroupComponent);
+	});
+});