With display.showTokenUsage on, the usage row was rendered inside the assistant block above the turn's tool blocks. Finalizing the assistant block was therefore deferred, and the late append recommitted the already-committed tool rows, duplicating them in scrollback (worst with parallel tool calls). 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 event-controller, transcript rebuild, agent-hub). The now-dead setUsageInfo/#usageInfo path on AssistantMessageComponent is removed.
109 lines
4.3 KiB
TypeScript
109 lines
4.3 KiB
TypeScript
/**
|
|
* 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);
|
|
});
|
|
});
|