From a3e90b3d95b67dd2f45583bcc39dbbad55f7282d Mon Sep 17 00:00:00 2001 From: joshrzemien Date: Tue, 28 Jul 2026 22:39:16 -0400 Subject: [PATCH] fix(tui): nest usage metrics in read groups (cherry picked from commit 101268adc632e80bb3f8304b95497ac7bd884131) --- .../src/cli/gallery-fixtures/fs.ts | 25 ++++ .../components/chat-transcript-builder.ts | 41 ++++-- .../src/modes/components/read-tool-group.ts | 125 +++++++++++++++++- .../src/modes/components/usage-row.ts | 15 ++- .../src/modes/controllers/event-controller.ts | 22 ++- .../src/modes/utils/ui-helpers.ts | 44 ++++-- .../event-controller-read-grouping.test.ts | 28 +++- .../coding-agent/test/read-tool-group.test.ts | 54 ++++++++ .../test/usage-row-placement.test.ts | 100 ++++++++++---- 9 files changed, 389 insertions(+), 65 deletions(-) diff --git a/packages/coding-agent/src/cli/gallery-fixtures/fs.ts b/packages/coding-agent/src/cli/gallery-fixtures/fs.ts index 23e3baee7..e88d0a6d0 100644 --- a/packages/coding-agent/src/cli/gallery-fixtures/fs.ts +++ b/packages/coding-agent/src/cli/gallery-fixtures/fs.ts @@ -1,5 +1,6 @@ // biome-ignore-all lint/suspicious/noTemplateCurlyInString: sample source-code strings (read fixtures) intentionally contain literal ${...}. // Gallery fixtures for the filesystem tools (read, write, glob). +import type { Usage } from "@oh-my-pi/pi-ai"; import { ReadToolGroupComponent } from "../../modes/components/read-tool-group"; import type { GalleryFixture, GalleryFixtureState, GalleryResult } from "./types"; @@ -47,6 +48,15 @@ const groupedReadDelimitedPath = groupedReadTargets.join(","); const groupedReadRepeatedFile = "packages/coding-agent/src/task/render.ts"; const groupedReadRepeatedRanges = `${groupedReadRepeatedFile}:507-605,1070-1194,1210-1240,1270-1274`; +const GROUPED_READ_USAGE: Usage = { + input: 2400, + output: 113, + cacheRead: 103_000, + cacheWrite: 0, + totalTokens: 105_513, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + function textResult(text: string, details?: unknown, isError?: boolean): GalleryResult { return { content: [{ type: "text", text }], details, isError }; } @@ -81,6 +91,13 @@ function renderReadGroupFixtureState(state: GalleryFixtureState, width: number, false, "read-delimited", ); + component.attachUsage( + ["read-delimited"], + GROUPED_READ_USAGE, + 5300, + 2200, + new Date(2026, 6, 28, 21, 5, 47).getTime(), + ); if (state === "error") { component.updateResult( @@ -88,10 +105,18 @@ function renderReadGroupFixtureState(state: GalleryFixtureState, width: number, false, "read-ranges", ); + component.attachUsage( + ["read-ranges"], + GROUPED_READ_USAGE, + 4700, + 1900, + new Date(2026, 6, 28, 21, 5, 52).getTime(), + ); return component.render(width); } component.updateResult(textResult("Read four render.ts ranges."), false, "read-ranges"); + component.attachUsage(["read-ranges"], GROUPED_READ_USAGE, 4700, 1900, new Date(2026, 6, 28, 21, 5, 52).getTime()); return component.render(width); } diff --git a/packages/coding-agent/src/modes/components/chat-transcript-builder.ts b/packages/coding-agent/src/modes/components/chat-transcript-builder.ts index 0a8466cb1..e33ab4f3c 100644 --- a/packages/coding-agent/src/modes/components/chat-transcript-builder.ts +++ b/packages/coding-agent/src/modes/components/chat-transcript-builder.ts @@ -51,7 +51,7 @@ import { import { CustomMessageComponent } from "./custom-message"; import { EvalExecutionComponent } from "./eval-execution"; import { type LateDiagnosticsFile, LateDiagnosticsMessageComponent } from "./late-diagnostics-message"; -import { ReadToolGroupComponent, readArgsCollapseIntoGroup } from "./read-tool-group"; +import { groupedReadUsageCallIds, ReadToolGroupComponent, readArgsCollapseIntoGroup } from "./read-tool-group"; import { SkillMessageComponent } from "./skill-message"; import { ToolExecutionComponent } from "./tool-execution"; import { TranscriptContainer } from "./transcript-container"; @@ -86,6 +86,7 @@ export class ChatTranscriptBuilder { #pendingUsageDuration: number | undefined; #pendingUsageTtft: number | undefined; #pendingUsageTimestamp: number | undefined; + #pendingReadUsageCallIds: string[] | undefined; #lastAssistantUsage: Usage | undefined; #waitingPoll: ToolExecutionComponent | null = null; #todoSnapshot: ToolExecutionComponent | null = null; @@ -135,6 +136,7 @@ export class ChatTranscriptBuilder { this.#pendingUsageDuration = undefined; this.#pendingUsageTtft = undefined; this.#pendingUsageTimestamp = undefined; + this.#pendingReadUsageCallIds = undefined; this.#lastAssistantUsage = undefined; this.#waitingPoll = null; this.#todoSnapshot = null; @@ -194,30 +196,46 @@ export class ChatTranscriptBuilder { return this.#readGroup; } - // 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, sealing the read run so the row sits under it. + // Defer per-turn metrics until the turn's tool results have materialized. + // Read-only invisible turns attach the metrics to their shared compact + // group; every other turn keeps the standalone row below its tool blocks. #flushPendingUsage(): void { if (!this.#pendingUsage) return; - this.#readGroup?.seal(); - this.#readGroup = null; - this.container.addChild( - createUsageRowBlock( + const usageAttached = + this.#pendingReadUsageCallIds !== undefined && + (this.#readGroup?.attachUsage( + this.#pendingReadUsageCallIds, this.#pendingUsage, this.#pendingUsageDuration, this.#pendingUsageTtft, this.#pendingUsageTimestamp, - ), - ); + ) ?? + false); + if (!usageAttached) { + this.#readGroup?.seal(); + this.#readGroup = null; + this.container.addChild( + createUsageRowBlock( + this.#pendingUsage, + this.#pendingUsageDuration, + this.#pendingUsageTtft, + this.#pendingUsageTimestamp, + ), + ); + } this.#pendingUsage = undefined; this.#pendingUsageDuration = undefined; this.#pendingUsageTtft = undefined; this.#pendingUsageTimestamp = undefined; + this.#pendingReadUsageCallIds = undefined; } #appendChatMessage(message: AgentMessage): void { if (message.role !== "toolResult") this.#flushPendingUsage(); + if (message.role !== "assistant" && message.role !== "toolResult") { + this.#readGroup?.seal(); + this.#readGroup = null; + } switch (message.role) { case "assistant": this.#appendAssistantMessage(message); @@ -403,6 +421,7 @@ export class ChatTranscriptBuilder { this.#pendingUsageDuration = message.duration; this.#pendingUsageTtft = message.ttft; this.#pendingUsageTimestamp = message.timestamp; + this.#pendingReadUsageCallIds = this.#pendingUsage ? groupedReadUsageCallIds(message) : undefined; } #appendToolResult(message: Extract): void { diff --git a/packages/coding-agent/src/modes/components/read-tool-group.ts b/packages/coding-agent/src/modes/components/read-tool-group.ts index de76bd3aa..d973a57a6 100644 --- a/packages/coding-agent/src/modes/components/read-tool-group.ts +++ b/packages/coding-agent/src/modes/components/read-tool-group.ts @@ -1,4 +1,5 @@ import * as path from "node:path"; +import type { AssistantMessage, Usage } from "@oh-my-pi/pi-ai"; import type { Component } from "@oh-my-pi/pi-tui"; import { Container, Text } from "@oh-my-pi/pi-tui"; import { InternalUrlRouter, XD_URL_PREFIX } from "../../internal-urls"; @@ -6,7 +7,9 @@ import { getLanguageFromPath, theme } from "../../modes/theme/theme"; import { parseLineRanges, selectorLineRanges, splitPathAndSel } from "../../tools/path-utils"; import { PREVIEW_LIMITS, shortenPath } from "../../tools/render-utils"; import { fileHyperlink, renderCodeCell, tryResolveInternalUrlSync } from "../../tui"; +import { canonicalizeMessage } from "../../utils/thinking-display"; import type { ToolExecutionHandle } from "./tool-execution"; +import { formatUsageRow } from "./usage-row"; /** * Extract the read call's target path. `path` is the canonical arg; `file_path` @@ -39,6 +42,30 @@ export function readArgsCollapseIntoGroup(args: unknown): boolean { return target.startsWith(XD_URL_PREFIX) || !InternalUrlRouter.instance().canHandle(target); } +/** + * Return the collapsed read calls that can own a turn's usage row. Visible + * assistant content and mixed-tool turns keep their usage as a standalone row + * so the metrics are not misleadingly attributed to one tool. + */ +export function groupedReadUsageCallIds(message: AssistantMessage): string[] | undefined { + const hasVisibleContent = message.content.some( + content => + content.type === "image" || + (content.type === "text" && canonicalizeMessage(content.text)) || + (content.type === "thinking" && canonicalizeMessage(content.thinking)), + ); + if (hasVisibleContent) return undefined; + + const toolCalls = message.content.filter(content => content.type === "toolCall"); + if ( + toolCalls.length === 0 || + toolCalls.some(content => content.name !== "read" || !readArgsCollapseIntoGroup(content.arguments)) + ) { + return undefined; + } + return toolCalls.map(content => content.id); +} + type ReadRenderArgs = { path?: string; file_path?: string; @@ -94,6 +121,13 @@ type ReadEntry = { codeLineNumbers?: Array; }; +type ReadUsageRow = { + usage: Usage; + durationMs?: number; + ttftMs?: number; + timestamp?: number; +}; + /** Number of code lines to show in collapsed preview mode */ const COLLAPSED_PREVIEW_LINES = PREVIEW_LIMITS.OUTPUT_COLLAPSED; @@ -290,6 +324,7 @@ function formatMergedSelectorParts(selectors: string[]): string { export class ReadToolGroupComponent extends Container implements ToolExecutionHandle { #entries = new Map(); + #usageRows = new Map(); #text: Text; #expanded = false; #showContentPreview: boolean; @@ -394,6 +429,31 @@ export class ReadToolGroupComponent extends Container implements ToolExecutionHa this.#updateDisplay(); } + /** + * Nest one request's usage beneath the last visible read call from that + * request. Parallel reads share one row rather than duplicating request totals. + */ + attachUsage( + toolCallIds: readonly string[], + usage: Usage, + durationMs?: number, + ttftMs?: number, + timestamp?: number, + ): boolean { + let anchorId: string | undefined; + for (let index = toolCallIds.length - 1; index >= 0; index--) { + const toolCallId = toolCallIds[index]!; + if (this.#entries.has(toolCallId)) { + anchorId = toolCallId; + break; + } + } + if (!anchorId) return false; + this.#usageRows.set(anchorId, { usage, durationMs, ttftMs, timestamp }); + this.#updateDisplay(); + return true; + } + setArgsComplete(_toolCallId?: string): void { this.#updateDisplay(); } @@ -427,13 +487,15 @@ export class ReadToolGroupComponent extends Container implements ToolExecutionHa if (!this.#shouldRenderPreviewRow(row)) { const statusSymbol = this.#formatStatus(this.#statusForTargets(row.targets)); const pathDisplay = this.#formatRowPath(row); - this.#text.setText( - ` ${statusSymbol} ${theme.fg("toolTitle", theme.bold("Read"))} ${pathDisplay}`.trimEnd(), - ); + const lines = [` ${statusSymbol} ${theme.fg("toolTitle", theme.bold("Read"))} ${pathDisplay}`.trimEnd()]; + const usageRows = this.#usageRowsBySummaryRow(displayRows).get(0) ?? []; + this.#appendUsageRows(lines, usageRows, " "); + this.#text.setText(lines.join("\n")); this.addChild(this.#text); } for (const entry of this.#previewEntriesForRow(row)) { this.#addContentPreview(entry); + this.#addPreviewUsage(entry); } return; } @@ -443,8 +505,9 @@ export class ReadToolGroupComponent extends Container implements ToolExecutionHa const entriesWithoutPreview = entries.filter(entry => !this.#shouldRenderPreview(entry)); const summaryTargets = this.#displayTargetsForEntries(entriesWithoutPreview); const rows = this.#buildSummaryRows(summaryTargets); + const usageRowsBySummaryRow = this.#usageRowsBySummaryRow(rows); for (const [index, row] of rows.entries()) { - this.#appendSummaryRow(lines, row, index, rows.length); + this.#appendSummaryRow(lines, row, index, rows.length, usageRowsBySummaryRow.get(index) ?? []); } this.#text.setText(lines.join("\n")); @@ -453,6 +516,7 @@ export class ReadToolGroupComponent extends Container implements ToolExecutionHa for (const entry of entries) { if (this.#shouldRenderPreview(entry)) { this.#addContentPreview(entry); + this.#addPreviewUsage(entry); } } } @@ -519,9 +583,48 @@ export class ReadToolGroupComponent extends Container implements ToolExecutionHa return rows; } - #appendSummaryRow(lines: string[], row: ReadSummaryRow, index: number, total: number): void { + #appendSummaryRow( + lines: string[], + row: ReadSummaryRow, + index: number, + total: number, + usageRows: ReadUsageRow[], + ): void { const connector = index === total - 1 ? theme.tree.last : theme.tree.branch; lines.push(` ${theme.fg("dim", connector)} ${this.#formatRow(row)}`.trimEnd()); + + const connectorWidth = Bun.stringWidth(connector); + const continuation = + index === total - 1 + ? " ".repeat(connectorWidth) + : `${theme.tree.vertical}${" ".repeat(Math.max(0, connectorWidth - Bun.stringWidth(theme.tree.vertical)))}`; + this.#appendUsageRows(lines, usageRows, ` ${continuation} `); + } + + #usageRowsBySummaryRow(rows: ReadSummaryRow[]): Map { + const usageRowsByIndex = new Map(); + for (const [toolCallId, usageRow] of this.#usageRows) { + for (let index = rows.length - 1; index >= 0; index--) { + const row = rows[index]!; + if (!row.targets.some(target => target.entry.toolCallId === toolCallId)) continue; + const usageRows = usageRowsByIndex.get(index); + if (usageRows) usageRows.push(usageRow); + else usageRowsByIndex.set(index, [usageRow]); + break; + } + } + return usageRowsByIndex; + } + + #appendUsageRows(lines: string[], usageRows: ReadUsageRow[], prefix: string): void { + for (const usageRow of usageRows) { + lines.push( + theme.fg( + "dim", + `${prefix}${formatUsageRow(usageRow.usage, usageRow.durationMs, usageRow.ttftMs, usageRow.timestamp)}`, + ), + ); + } } #formatRow(row: ReadSummaryRow): string { @@ -659,6 +762,18 @@ export class ReadToolGroupComponent extends Container implements ToolExecutionHa this.addChild(component); } + #addPreviewUsage(entry: ReadEntry): void { + const usageRow = this.#usageRows.get(entry.toolCallId); + if (!usageRow) return; + this.addChild( + new Text( + theme.fg("dim", formatUsageRow(usageRow.usage, usageRow.durationMs, usageRow.ttftMs, usageRow.timestamp)), + 3, + 0, + ), + ); + } + #shouldRenderPreview(entry: ReadEntry): boolean { return this.#showContentPreview && entry.contentText !== undefined; } diff --git a/packages/coding-agent/src/modes/components/usage-row.ts b/packages/coding-agent/src/modes/components/usage-row.ts index 8643f145f..f051ca55e 100644 --- a/packages/coding-agent/src/modes/components/usage-row.ts +++ b/packages/coding-agent/src/modes/components/usage-row.ts @@ -15,10 +15,8 @@ function formatUsageTimestamp(ms: number): string { return `${date} ${time}`; } -// `timestamp` is optional and trails the throughput args to preserve the existing -// (usage, durationMs, ttftMs) call contract — this function is part of the package's -// public export surface (./modes/components/*). -export function createUsageRowBlock(usage: Usage, durationMs?: number, ttftMs?: number, timestamp?: number): Container { +/** Format the metrics shared by standalone usage blocks and compact tool groups. */ +export function formatUsageRow(usage: Usage, durationMs?: number, ttftMs?: number, timestamp?: number): string { const totalInput = usage.input + usage.cacheWrite; const parts: string[] = []; // Lead with the turn's local wall-clock time (down to the second), log-line style. @@ -40,8 +38,15 @@ export function createUsageRowBlock(usage: Usage, durationMs?: number, ttftMs?: const tokPerSec = (usage.output / durationMs) * 1000; parts.push(`${theme.icon.throughput} ${tokPerSec.toFixed(1)}/s`); } + return parts.join(" "); +} + +// `timestamp` is optional and trails the throughput args to preserve the existing +// (usage, durationMs, ttftMs) call contract — this function is part of the package's +// public export surface (./modes/components/*). +export function createUsageRowBlock(usage: Usage, durationMs?: number, ttftMs?: number, timestamp?: number): Container { const block = new Container(); block.addChild(new Spacer(1)); - block.addChild(new Text(theme.fg("dim", parts.join(" ")), 1, 0)); + block.addChild(new Text(theme.fg("dim", formatUsageRow(usage, durationMs, ttftMs, timestamp)), 1, 0)); return block; } diff --git a/packages/coding-agent/src/modes/controllers/event-controller.ts b/packages/coding-agent/src/modes/controllers/event-controller.ts index b927a7409..ff3b0b2bf 100644 --- a/packages/coding-agent/src/modes/controllers/event-controller.ts +++ b/packages/coding-agent/src/modes/controllers/event-controller.ts @@ -10,6 +10,7 @@ import { getFileSnapshotStore } from "../../edit/file-snapshot-store"; import { AssistantMessageComponent } from "../../modes/components/assistant-message"; import { detectCacheInvalidation } from "../../modes/components/cache-invalidation-marker"; import { + groupedReadUsageCallIds, ReadToolGroupComponent, readArgsCollapseIntoGroup, readArgsHaveTarget, @@ -954,14 +955,27 @@ export class EventController { } this.#lastAssistantComponent = lastPostToolAssistantComponent ?? this.ctx.streamingComponent; if (settings.get("display.showTokenUsage") && assistantUsageIsBilled(event.message.usage)) { - this.ctx.chatContainer.addChild( - createUsageRowBlock( + const readCallIds = groupedReadUsageCallIds(event.message); + const usageAttached = + readCallIds !== undefined && + (this.#lastReadGroup?.attachUsage( + readCallIds, event.message.usage, event.message.duration, event.message.ttft, event.message.timestamp, - ), - ); + ) ?? + false); + if (!usageAttached) { + this.ctx.chatContainer.addChild( + createUsageRowBlock( + event.message.usage, + event.message.duration, + event.message.ttft, + event.message.timestamp, + ), + ); + } } if (displayMessage === event.message) { this.ctx.transcriptMessageComponents.set(event.message, this.ctx.streamingComponent); diff --git a/packages/coding-agent/src/modes/utils/ui-helpers.ts b/packages/coding-agent/src/modes/utils/ui-helpers.ts index 2a4eb48e9..326c579ac 100644 --- a/packages/coding-agent/src/modes/utils/ui-helpers.ts +++ b/packages/coding-agent/src/modes/utils/ui-helpers.ts @@ -24,7 +24,11 @@ import { type LateDiagnosticsFile, LateDiagnosticsMessageComponent, } from "../../modes/components/late-diagnostics-message"; -import { ReadToolGroupComponent, readArgsCollapseIntoGroup } from "../../modes/components/read-tool-group"; +import { + groupedReadUsageCallIds, + ReadToolGroupComponent, + readArgsCollapseIntoGroup, +} from "../../modes/components/read-tool-group"; import { SkillMessageComponent } from "../../modes/components/skill-message"; import { ToolExecutionComponent } from "../../modes/components/tool-execution"; import { TranscriptBlock } from "../../modes/components/transcript-container"; @@ -307,29 +311,38 @@ export class UiHelpers { let readGroup: ReadToolGroupComponent | null = null; const readToolCallArgs = new Map>(); const readToolCallAssistantComponents = new Map(); - // 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. + // Defer per-turn metrics until the turn's tool results have materialized. + // Read-only invisible turns attach the metrics to their shared compact + // group; every other turn keeps the standalone row below its tool blocks. let pendingUsage: Usage | undefined; let pendingUsageDuration: number | undefined; let pendingUsageTtft: number | undefined; let pendingUsageTimestamp: number | undefined; + let pendingReadUsageCallIds: string[] | undefined; const flushPendingUsage = () => { if (!pendingUsage) return; - readGroup?.seal(); - readGroup = null; - this.ctx.chatContainer.addChild( - createUsageRowBlock(pendingUsage, pendingUsageDuration, pendingUsageTtft, pendingUsageTimestamp), - ); + const usageAttached = + pendingReadUsageCallIds !== undefined && + (readGroup?.attachUsage( + pendingReadUsageCallIds, + pendingUsage, + pendingUsageDuration, + pendingUsageTtft, + pendingUsageTimestamp, + ) ?? + false); + if (!usageAttached) { + readGroup?.seal(); + readGroup = null; + this.ctx.chatContainer.addChild( + createUsageRowBlock(pendingUsage, pendingUsageDuration, pendingUsageTtft, pendingUsageTimestamp), + ); + } pendingUsage = undefined; pendingUsageDuration = undefined; pendingUsageTtft = undefined; pendingUsageTimestamp = undefined; + pendingReadUsageCallIds = undefined; }; // Rebuild-time mirror of the event controller's displaceable-poll // bookkeeping: a `hub` wait that found every watched job still running is @@ -535,6 +548,7 @@ export class UiHelpers { pendingUsageDuration = message.duration; pendingUsageTtft = message.ttft; pendingUsageTimestamp = message.timestamp; + pendingReadUsageCallIds = pendingUsage ? groupedReadUsageCallIds(message) : undefined; } else if (message.role === "toolResult") { if (options.preservedLiveToolCallIds?.has(message.toolCallId)) continue; const pendingReadComponent = this.ctx.pendingTools.get(message.toolCallId); @@ -602,6 +616,8 @@ export class UiHelpers { } } } else { + readGroup?.seal(); + readGroup = null; // A user prompt closes the displacement window, same as the live path. if (message.role === "user") resolveWaitingPoll(); if (message.role === "user") resolveTodoSnapshot(); diff --git a/packages/coding-agent/test/modes/controllers/event-controller-read-grouping.test.ts b/packages/coding-agent/test/modes/controllers/event-controller-read-grouping.test.ts index decdbe368..ad20e4ddf 100644 --- a/packages/coding-agent/test/modes/controllers/event-controller-read-grouping.test.ts +++ b/packages/coding-agent/test/modes/controllers/event-controller-read-grouping.test.ts @@ -14,7 +14,7 @@ */ import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test"; import type { AssistantMessage, ImageContent } from "@oh-my-pi/pi-ai"; -import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { resetSettingsForTest, Settings, settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { AssistantMessageComponent } from "@oh-my-pi/pi-coding-agent/modes/components/assistant-message"; import { ReadToolGroupComponent } from "@oh-my-pi/pi-coding-agent/modes/components/read-tool-group"; import { EventController } from "@oh-my-pi/pi-coding-agent/modes/controllers/event-controller"; @@ -132,6 +132,32 @@ describe("EventController read-group accretion", () => { expect(header(groups[0]!)).toContain("Read (4)"); }); + it("nests a read-only completion's usage inside the active group", async () => { + settings.set("display.showTokenUsage", true); + const { controller, chatContainer } = createFixture(); + const message = assistantMessage([read("usage.ts:1-50")]); + message.usage = { + input: 1234, + output: 7, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 1241, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + message.timestamp = new Date(2026, 0, 2, 3, 4, 5).getTime(); + + await controller.handleEvent({ type: "message_start", message } as AgentSessionEvent); + await controller.handleEvent({ type: "message_update", message } as AgentSessionEvent); + await controller.handleEvent({ type: "message_end", message } as AgentSessionEvent); + + const [group] = readGroups(chatContainer); + expect(group).toBeDefined(); + const usageBlocks = chatContainer.children.filter(component => + Bun.stripANSI(component.render(120).join("\n")).includes("2026-01-02 03:04:05"), + ); + expect(usageBlocks).toEqual([group!]); + }); + it("starts a new group after a completion that renders visible reasoning", async () => { const { controller, chatContainer } = createFixture(); diff --git a/packages/coding-agent/test/read-tool-group.test.ts b/packages/coding-agent/test/read-tool-group.test.ts index 40c421c2e..44c96b3ce 100644 --- a/packages/coding-agent/test/read-tool-group.test.ts +++ b/packages/coding-agent/test/read-tool-group.test.ts @@ -95,6 +95,60 @@ describe("ReadToolGroupComponent", () => { expect(plain).not.toContain(`${themeModule.theme.tree.last} ${themeModule.theme.status.enabled}`); }); + it("nests one usage row beneath the last path from each read-only turn", () => { + const component = new ReadToolGroupComponent(); + const onePath = path.resolve("/tmp/one.ts"); + const twoPath = path.resolve("/tmp/two.ts"); + const threePath = path.resolve("/tmp/three.ts"); + component.updateArgs({ path: onePath }, "read-one"); + component.updateArgs({ path: twoPath }, "read-two"); + component.updateArgs({ path: threePath }, "read-three"); + component.updateResult({ content: [{ type: "text", text: "one" }] }, false, "read-one"); + component.updateResult({ content: [{ type: "text", text: "two" }] }, false, "read-two"); + component.updateResult({ content: [{ type: "text", text: "three" }] }, false, "read-three"); + + const firstUsage = { + input: 1111, + output: 11, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 1122, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + const parallelUsage = { + input: 2222, + output: 22, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2244, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + component.attachUsage(["read-one"], firstUsage, 1000, 500, new Date(2026, 0, 2, 3, 4, 5).getTime()); + component.attachUsage( + ["read-two", "read-three"], + parallelUsage, + 2000, + 600, + new Date(2026, 0, 2, 3, 4, 6).getTime(), + ); + + const lines = Bun.stripANSI(component.render(120).join("\n")).split("\n"); + const onePathIndex = lines.findIndex(line => line.includes(onePath)); + const twoPathIndex = lines.findIndex(line => line.includes(twoPath)); + const threePathIndex = lines.findIndex(line => line.includes(threePath)); + const firstUsageIndex = lines.findIndex(line => line.includes("2026-01-02 03:04:05")); + const parallelUsageIndices = lines + .map((line, index) => (line.includes("2026-01-02 03:04:06") ? index : -1)) + .filter(index => index >= 0); + + expect(firstUsageIndex).toBe(onePathIndex + 1); + expect(lines[firstUsageIndex]?.startsWith(` ${themeModule.theme.tree.vertical} `)).toBe(true); + expect(twoPathIndex).toBeGreaterThan(firstUsageIndex); + expect(threePathIndex).toBeGreaterThan(twoPathIndex); + expect(parallelUsageIndices).toEqual([threePathIndex + 1]); + expect(lines[parallelUsageIndices[0]!]?.startsWith(" ")).toBe(true); + }); + it("splits a single selector-delimited read argument into child rows", () => { const component = new ReadToolGroupComponent(); const onePath = path.resolve("/tmp/one.ts"); diff --git a/packages/coding-agent/test/usage-row-placement.test.ts b/packages/coding-agent/test/usage-row-placement.test.ts index e9a1d4928..001ccd0f5 100644 --- a/packages/coding-agent/test/usage-row-placement.test.ts +++ b/packages/coding-agent/test/usage-row-placement.test.ts @@ -1,10 +1,7 @@ /** - * 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. + * Regression coverage for per-turn usage placement across transcript rebuilds. + * Read-only turns keep their metrics inside the compact read group; other + * assistant turns retain a standalone row below their visible content/tools. */ import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; @@ -25,31 +22,38 @@ const USAGE_LABEL = formatNumber(USAGE_INPUT); // Single-digit month/day/hour/minute/second exercise the formatter's zero-padding. const USAGE_TS = new Date(2026, 0, 2, 3, 4, 5).getTime(); const USAGE_TS_LABEL = "2026-01-02 03:04:05"; +const SECOND_USAGE_TS = new Date(2026, 0, 2, 3, 4, 6).getTime(); +const SECOND_USAGE_TS_LABEL = "2026-01-02 03:04:06"; -function readTurn(): AgentMessage[] { +function readTurn( + toolCallId = "r1", + filePath = "src/foo.ts", + usageInput = USAGE_INPUT, + timestamp = USAGE_TS, +): AgentMessage[] { const assistant = { role: "assistant", - content: [{ type: "toolCall", id: "r1", name: "read", arguments: { path: "src/foo.ts" } }], + content: [{ type: "toolCall", id: toolCallId, name: "read", arguments: { path: filePath } }], api: "anthropic-messages", provider: "anthropic", model: "claude-sonnet-4-5", stopReason: "stop", usage: { - input: USAGE_INPUT, + input: usageInput, output: 7, cacheRead: 0, cacheWrite: 0, - totalTokens: USAGE_INPUT + 7, + totalTokens: usageInput + 7, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, - timestamp: USAGE_TS, + timestamp, } as unknown as AgentMessage; const toolResult = { role: "toolResult", - toolCallId: "r1", + toolCallId, toolName: "read", content: [{ type: "text", text: "line1\nline2" }], - timestamp: Date.now(), + timestamp, } as unknown as AgentMessage; return [assistant, toolResult]; } @@ -86,23 +90,40 @@ describe("UiHelpers.renderSessionContext token-usage row placement", () => { await initTheme(); }); - it("places the usage row below the read group for a read turn", () => { + it("nests the usage row inside the read group for a read-only 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); + const group = children.find( + (component): component is ReadToolGroupComponent => component instanceof ReadToolGroupComponent, + ); + expect(group).toBeDefined(); - // 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); - // The row also carries the message's local timestamp down to the second. - expect(last.render(120).join("\n")).toContain(USAGE_TS_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); + const rendered = group!.render(120).join("\n"); + expect(rendered).toContain(USAGE_LABEL); + expect(rendered).toContain(USAGE_TS_LABEL); + expect(children[children.length - 1]).toBe(group!); + expect(children.filter(component => component.render(120).join("\n").includes(USAGE_LABEL))).toHaveLength(1); + }); + + it("interleaves consecutive read-only turns with their paths in one group", () => { + const { ctx, helpers } = makeHarness(true); + const messages = [...readTurn(), ...readTurn("r2", "src/bar.ts", 2121, SECOND_USAGE_TS)]; + helpers.renderSessionContext({ messages } as SessionContext); + + const groups = ctx.chatContainer.children.filter( + (component): component is ReadToolGroupComponent => component instanceof ReadToolGroupComponent, + ); + expect(groups).toHaveLength(1); + const lines = Bun.stripANSI(groups[0]!.render(120).join("\n")).split("\n"); + const fooIndex = lines.findIndex(line => line.includes("src/foo.ts")); + const firstUsageIndex = lines.findIndex(line => line.includes(USAGE_TS_LABEL)); + const barIndex = lines.findIndex(line => line.includes("src/bar.ts")); + const secondUsageIndex = lines.findIndex(line => line.includes(SECOND_USAGE_TS_LABEL)); + expect(fooIndex).toBeLessThan(firstUsageIndex); + expect(firstUsageIndex).toBeLessThan(barIndex); + expect(barIndex).toBeLessThan(secondUsageIndex); }); it("renders no usage row when showTokenUsage is off", () => { @@ -155,4 +176,33 @@ describe("ChatTranscriptBuilder token-usage row timestamp", () => { expect(rendered).toContain(USAGE_TS_LABEL); expect(rendered).toContain(USAGE_LABEL); }); + + it("keeps grouped read metrics nested on the reusable transcript-builder path", () => { + const builder = new ChatTranscriptBuilder({ + ui: { requestRender: () => {}, requestComponentRender: () => {} } as unknown as TUI, + cwd: process.cwd(), + requestRender: () => {}, + }); + const messages = [...readTurn(), ...readTurn("r2", "src/bar.ts", 2121, SECOND_USAGE_TS)]; + builder.rebuild( + messages.map((message, index) => ({ + type: "message", + id: `m${index}`, + parentId: index === 0 ? null : `m${index - 1}`, + timestamp: new Date(0).toISOString(), + message, + })), + ); + + const groups = builder.container.children.filter( + (component): component is ReadToolGroupComponent => component instanceof ReadToolGroupComponent, + ); + expect(groups).toHaveLength(1); + const rendered = groups[0]!.render(120).join("\n"); + expect(rendered).toContain(USAGE_TS_LABEL); + expect(rendered).toContain(SECOND_USAGE_TS_LABEL); + expect( + builder.container.children.filter(component => component.render(120).join("\n").includes(USAGE_LABEL)), + ).toEqual([groups[0]!]); + }); });