add usage info to branch summary, compaction and tool result entries (#6671)

* add usage info to branch summary entries

* add usage to compaction entries

* allow custom tools to report llm usage in tool results

* allow observing and patching usage in tool_result hooks

* agent-harness: save usage in entries

for compaction, branch summaries and tool results

packages/coding-agent/test/agent-session-stats.test.ts

@@ -1,5 +1,5 @@
 import { Agent } from "@earendil-works/pi-agent-core";
-import { type AssistantMessage, getModel, type Usage } from "@earendil-works/pi-ai/compat";
+import { type AssistantMessage, getModel, type ToolResultMessage, type Usage } from "@earendil-works/pi-ai/compat";
 import { describe, expect, it } from "vitest";
 import { AgentSession } from "../src/core/agent-session.ts";
 import { AuthStorage } from "../src/core/auth-storage.ts";
@@ -48,6 +48,18 @@ function createUserMessage(text: string, timestamp: number) {
 	};
 }
 
+function createToolResultMessage(usage: Usage): ToolResultMessage {
+	return {
+		role: "toolResult",
+		toolCallId: "tool-call-1",
+		toolName: "test_tool",
+		content: [{ type: "text", text: "tool result" }],
+		usage,
+		isError: false,
+		timestamp: 1,
+	};
+}
+
 async function createSession() {
 	const settingsManager = SettingsManager.inMemory();
 	const sessionManager = SessionManager.inMemory();
@@ -143,6 +155,75 @@ describe("AgentSession.getSessionStats", () => {
 		}
 	});
 
+	it("includes branch summary usage in session totals", async () => {
+		const { session, sessionManager } = await createSession();
+
+		try {
+			sessionManager.branchWithSummary(null, "summary", undefined, false, {
+				input: 10,
+				output: 20,
+				cacheRead: 30,
+				cacheWrite: 40,
+				totalTokens: 100,
+				cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+			});
+			syncAgentMessages(session, sessionManager);
+
+			const stats = session.getSessionStats();
+			expect(stats.tokens).toEqual({ input: 10, output: 20, cacheRead: 30, cacheWrite: 40, total: 100 });
+			expect(stats.cost).toBe(1);
+		} finally {
+			session.dispose();
+		}
+	});
+
+	it("includes compaction usage in session totals", async () => {
+		const { session, sessionManager } = await createSession();
+
+		try {
+			const firstKeptEntryId = sessionManager.appendMessage(createUserMessage("hello", 1));
+			sessionManager.appendCompaction("summary", firstKeptEntryId, 100, undefined, false, {
+				input: 10,
+				output: 20,
+				cacheRead: 30,
+				cacheWrite: 40,
+				totalTokens: 100,
+				cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+			});
+			syncAgentMessages(session, sessionManager);
+
+			const stats = session.getSessionStats();
+			expect(stats.tokens).toEqual({ input: 10, output: 20, cacheRead: 30, cacheWrite: 40, total: 100 });
+			expect(stats.cost).toBe(1);
+		} finally {
+			session.dispose();
+		}
+	});
+
+	it("includes tool result usage in session totals", async () => {
+		const { session, sessionManager } = await createSession();
+
+		try {
+			sessionManager.appendMessage(
+				createToolResultMessage({
+					input: 10,
+					output: 20,
+					cacheRead: 30,
+					cacheWrite: 40,
+					totalTokens: 100,
+					cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+				}),
+			);
+			syncAgentMessages(session, sessionManager);
+
+			const stats = session.getSessionStats();
+			expect(stats.tokens).toEqual({ input: 10, output: 20, cacheRead: 30, cacheWrite: 40, total: 100 });
+			expect(stats.cost).toBe(1);
+		} finally {
+			session.dispose();
+		}
+	});
+
 	it("ignores zero-usage messages when checking for post-compaction context usage", async () => {
 		const { session, sessionManager } = await createSession();
 

packages/coding-agent/test/branch-summary-extensions.test.ts

@@ -0,0 +1,55 @@
+import type { Usage } from "@earendil-works/pi-ai/compat";
+import { afterEach, describe, expect, it } from "vitest";
+import { createHarness, type Harness } from "./suite/harness.ts";
+import { assistantMsg, userMsg } from "./utilities.ts";
+
+describe("Branch summary extensions", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+	});
+
+	it("persists extension-provided summary usage in session totals", async () => {
+		const usage: Usage = {
+			input: 10,
+			output: 20,
+			cacheRead: 30,
+			cacheWrite: 40,
+			totalTokens: 100,
+			cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+		};
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_before_tree", () => ({
+						summary: {
+							summary: "Summary provided by extension",
+							usage,
+						},
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+
+		const targetId = harness.sessionManager.appendMessage(userMsg("first branch"));
+		harness.sessionManager.appendMessage(assistantMsg("first reply"));
+		harness.sessionManager.appendMessage(userMsg("abandoned branch work"));
+		harness.sessionManager.appendMessage(assistantMsg("abandoned reply"));
+
+		const result = await harness.session.navigateTree(targetId, { summarize: true });
+		const summaryEntry = result.summaryEntry;
+
+		expect(summaryEntry?.type).toBe("branch_summary");
+		expect(summaryEntry?.fromHook).toBe(true);
+		expect(summaryEntry?.summary).toBe("Summary provided by extension");
+		expect(summaryEntry?.usage).toEqual(usage);
+
+		const stats = harness.session.getSessionStats();
+		expect(stats.tokens).toEqual({ input: 12, output: 22, cacheRead: 30, cacheWrite: 40, total: 104 });
+		expect(stats.cost).toBe(1);
+	});
+});

packages/coding-agent/test/compaction-summary-reasoning.test.ts

@@ -57,7 +57,7 @@ describe("generateSummary reasoning options", () => {
 	});
 
 	it("uses the provided thinking level for reasoning-capable models", async () => {
-		await generateSummary(
+		const result = await generateSummary(
 			messages,
 			createModel(true),
 			2000,
@@ -69,6 +69,9 @@ describe("generateSummary reasoning options", () => {
 			"medium",
 		);
 
+		expect(result.text).toBe("## Goal\nTest summary");
+		expect(result.usage).toEqual(mockSummaryResponse.usage);
+
 		expect(completeSimpleMock).toHaveBeenCalledTimes(1);
 		expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({
 			reasoning: "medium",
@@ -127,8 +130,15 @@ describe("generateSummary reasoning options", () => {
 			settings: { enabled: true, reserveTokens: 500000, keepRecentTokens: 20000 },
 		};
 
-		await compact(preparation, createModel(false, 128000), "test-key");
+		const result = await compact(preparation, createModel(false, 128000), "test-key");
 
+		expect(result.usage).toEqual({
+			...mockSummaryResponse.usage,
+			input: 20,
+			output: 20,
+			totalTokens: 40,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		});
 		expect(completeSimpleMock.mock.calls.map((call) => call[2]?.maxTokens)).toEqual([128000, 128000]);
 	});
 });

packages/coding-agent/test/footer-width.test.ts

@@ -21,20 +21,46 @@ function createSession(options: {
 	reasoning?: boolean;
 	thinkingLevel?: string;
 	usage?: AssistantUsage;
+	branchUsage?: AssistantUsage;
+	compactionUsage?: AssistantUsage;
+	toolUsage?: AssistantUsage;
 }): AgentSession {
 	const usage = options.usage;
-	const entries =
-		usage === undefined
-			? []
-			: [
-					{
-						type: "message",
-						message: {
-							role: "assistant",
-							usage,
-						},
-					},
-				];
+	const entries: Array<Record<string, unknown>> = [];
+
+	if (usage !== undefined) {
+		entries.push({
+			type: "message",
+			message: {
+				role: "assistant",
+				usage,
+			},
+		});
+	}
+
+	if (options.branchUsage !== undefined) {
+		entries.push({
+			type: "branch_summary",
+			usage: options.branchUsage,
+		});
+	}
+
+	if (options.compactionUsage !== undefined) {
+		entries.push({
+			type: "compaction",
+			usage: options.compactionUsage,
+		});
+	}
+
+	if (options.toolUsage !== undefined) {
+		entries.push({
+			type: "message",
+			message: {
+				role: "toolResult",
+				usage: options.toolUsage,
+			},
+		});
+	}
 
 	const session = {
 		state: {
@@ -125,6 +151,44 @@ describe("FooterComponent width handling", () => {
 		}
 	});
 
+	it("includes summary and tool result usage in the total cost", () => {
+		const session = createSession({
+			sessionName: "",
+			usage: {
+				input: 100,
+				output: 10,
+				cacheRead: 0,
+				cacheWrite: 0,
+				cost: { total: 0.5 },
+			},
+			branchUsage: {
+				input: 20,
+				output: 5,
+				cacheRead: 0,
+				cacheWrite: 0,
+				cost: { total: 0.25 },
+			},
+			compactionUsage: {
+				input: 5,
+				output: 2,
+				cacheRead: 0,
+				cacheWrite: 0,
+				cost: { total: 0.125 },
+			},
+			toolUsage: {
+				input: 15,
+				output: 3,
+				cacheRead: 0,
+				cacheWrite: 0,
+				cost: { total: 0.375 },
+			},
+		});
+		const footer = new FooterComponent(session, createFooterData(1));
+
+		const statsLine = stripAnsi(footer.render(120)[1]);
+		expect(statsLine).toContain("$1.250");
+	});
+
 	it("shows the latest cache hit rate when cache usage is present", () => {
 		const session = createSession({
 			sessionName: "",

packages/coding-agent/test/session-manager/tree-traversal.test.ts

@@ -71,7 +71,15 @@ describe("SessionManager append and tree traversal", () => {
 
 			const id1 = session.appendMessage(userMsg("1"));
 			const id2 = session.appendMessage(assistantMsg("2"));
-			const compactionId = session.appendCompaction("summary", id1, 1000);
+			const usage = {
+				input: 10,
+				output: 20,
+				cacheRead: 30,
+				cacheWrite: 40,
+				totalTokens: 100,
+				cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+			};
+			const compactionId = session.appendCompaction("summary", id1, 1000, undefined, false, usage);
 			const _id3 = session.appendMessage(userMsg("3"));
 
 			const entries = session.getEntries();
@@ -83,6 +91,7 @@ describe("SessionManager append and tree traversal", () => {
 				expect(compactionEntry.summary).toBe("summary");
 				expect(compactionEntry.firstKeptEntryId).toBe(id1);
 				expect(compactionEntry.tokensBefore).toBe(1000);
+				expect(compactionEntry.usage).toEqual(usage);
 			}
 
 			expect(entries[3].parentId).toBe(compactionId);
@@ -319,7 +328,15 @@ describe("SessionManager append and tree traversal", () => {
 			const _id2 = session.appendMessage(assistantMsg("2"));
 			const _id3 = session.appendMessage(userMsg("3"));
 
-			const summaryId = session.branchWithSummary(id1, "Summary of abandoned work");
+			const usage = {
+				input: 10,
+				output: 20,
+				cacheRead: 30,
+				cacheWrite: 40,
+				totalTokens: 100,
+				cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+			};
+			const summaryId = session.branchWithSummary(id1, "Summary of abandoned work", undefined, false, usage);
 
 			expect(session.getLeafId()).toBe(summaryId);
 
@@ -329,6 +346,7 @@ describe("SessionManager append and tree traversal", () => {
 			expect(summaryEntry?.parentId).toBe(id1);
 			if (summaryEntry?.type === "branch_summary") {
 				expect(summaryEntry.summary).toBe("Summary of abandoned work");
+				expect(summaryEntry.usage).toEqual(usage);
 			}
 		});
 

packages/coding-agent/test/suite/agent-session-compaction.test.ts

@@ -97,6 +97,14 @@ describe("AgentSession compaction characterization", () => {
 	});
 
 	it("manually compacts using an extension-provided summary", async () => {
+		const summaryUsage = {
+			input: 10,
+			output: 20,
+			cacheRead: 30,
+			cacheWrite: 40,
+			totalTokens: 100,
+			cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+		};
 		const harness = await createHarness({
 			settings: { compaction: { keepRecentTokens: 1 } },
 			extensionFactories: [
@@ -106,6 +114,7 @@ describe("AgentSession compaction characterization", () => {
 							summary: "summary from extension",
 							firstKeptEntryId: event.preparation.firstKeptEntryId,
 							tokensBefore: event.preparation.tokensBefore,
+							usage: summaryUsage,
 							details: { source: "extension" },
 						},
 					}));
@@ -116,14 +125,26 @@ describe("AgentSession compaction characterization", () => {
 
 		await harness.session.prompt("one");
 		await harness.session.prompt("two");
+		const statsBefore = harness.session.getSessionStats();
 
 		const result = await harness.session.compact();
 		const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction");
 		const estimatedTokensAfter = harness.session.messages.reduce((sum, message) => sum + estimateTokens(message), 0);
 
 		expect(result.summary).toBe("summary from extension");
+		expect(result.usage).toEqual(summaryUsage);
 		expect(result.estimatedTokensAfter).toBe(estimatedTokensAfter);
 		expect(compactionEntries).toHaveLength(1);
+		const compactionEntry = compactionEntries[0];
+		if (compactionEntry?.type === "compaction") {
+			expect(compactionEntry.usage).toEqual(summaryUsage);
+		}
+		const statsAfter = harness.session.getSessionStats();
+		expect(statsAfter.tokens.input).toBe(statsBefore.tokens.input + summaryUsage.input);
+		expect(statsAfter.tokens.output).toBe(statsBefore.tokens.output + summaryUsage.output);
+		expect(statsAfter.tokens.cacheRead).toBe(statsBefore.tokens.cacheRead + summaryUsage.cacheRead);
+		expect(statsAfter.tokens.cacheWrite).toBe(statsBefore.tokens.cacheWrite + summaryUsage.cacheWrite);
+		expect(statsAfter.cost).toBe(statsBefore.cost + summaryUsage.cost.total);
 		expect(harness.session.messages[0]?.role).toBe("compactionSummary");
 	});
 
@@ -154,6 +175,22 @@ describe("AgentSession compaction characterization", () => {
 		expect(getStreamCallCount()).toBe(1);
 	});
 
+	it("persists usage from pi-generated manual compaction", async () => {
+		const harness = await createHarness({ withConfiguredAuth: false });
+		harnesses.push(harness);
+		seedCompactableSession(harness);
+		useSummaryStreamFn(harness, "summary from custom stream");
+
+		const result = await harness.session.compact();
+
+		const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction");
+		expect(result.usage).toEqual(createUsage(10));
+		expect(compactionEntries).toHaveLength(1);
+		expect(compactionEntries[0]?.type === "compaction" ? compactionEntries[0].usage : undefined).toEqual(
+			createUsage(10),
+		);
+	});
+
 	it("auto-compacts with a custom streamFn when registry auth is absent", async () => {
 		const harness = await createHarness({ withConfiguredAuth: false });
 		harnesses.push(harness);

packages/coding-agent/test/suite/agent-session-model-extension.test.ts

@@ -1,5 +1,5 @@
 import type { AgentTool, ThinkingLevel } from "@earendil-works/pi-agent-core";
-import { fauxAssistantMessage, fauxToolCall, type Model } from "@earendil-works/pi-ai";
+import { fauxAssistantMessage, fauxToolCall, type Model, type Usage } from "@earendil-works/pi-ai";
 import { Type } from "typebox";
 import { afterEach, describe, expect, it } from "vitest";
 import type { BuildSystemPromptOptions, ExtensionAPI } from "../../src/index.ts";
@@ -156,6 +156,23 @@ describe("AgentSession model and extension characterization", () => {
 	});
 
 	it("allows extension tool_result handlers to modify tool results", async () => {
+		const toolUsage: Usage = {
+			input: 1,
+			output: 2,
+			cacheRead: 3,
+			cacheWrite: 4,
+			totalTokens: 10,
+			cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
+		};
+		const patchedToolUsage: Usage = {
+			input: 5,
+			output: 6,
+			cacheRead: 7,
+			cacheWrite: 8,
+			totalTokens: 26,
+			cost: { input: 0.5, output: 0.6, cacheRead: 0.7, cacheWrite: 0.8, total: 2.6 },
+		};
+		let observedToolUsage: Usage | undefined;
 		const echoTool: AgentTool = {
 			name: "echo",
 			label: "Echo",
@@ -163,17 +180,21 @@ describe("AgentSession model and extension characterization", () => {
 			parameters: Type.Object({ text: Type.String() }),
 			execute: async (_toolCallId, params) => {
 				const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : "";
-				return { content: [{ type: "text", text }], details: { text } };
+				return { content: [{ type: "text", text }], details: { text }, usage: toolUsage };
 			},
 		};
 		const harness = await createHarness({
 			tools: [echoTool],
 			extensionFactories: [
 				(pi) => {
-					pi.on("tool_result", async () => ({
-						content: [{ type: "text", text: "patched result" }],
-						details: { patched: true },
-					}));
+					pi.on("tool_result", async (event) => {
+						observedToolUsage = event.usage;
+						return {
+							content: [{ type: "text", text: "patched result" }],
+							details: { patched: true },
+							usage: patchedToolUsage,
+						};
+					});
 				},
 			],
 		});
@@ -196,9 +217,12 @@ describe("AgentSession model and extension characterization", () => {
 		await harness.session.prompt("hi");
 
 		expect(getAssistantTexts(harness)).toContain("patched result");
-		expect(
-			harness.session.messages.find((message) => message.role === "toolResult" && message.details?.patched === true),
-		).toBeDefined();
+		const toolResult = harness.session.messages.find(
+			(message) => message.role === "toolResult" && message.details?.patched === true,
+		);
+		expect(observedToolUsage).toEqual(toolUsage);
+		expect(toolResult).toBeDefined();
+		expect(toolResult?.role === "toolResult" ? toolResult.usage : undefined).toEqual(patchedToolUsage);
 	});
 
 	it("allows extension context handlers to modify messages before the LLM call", async () => {

packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts

@@ -37,7 +37,7 @@ describe("issue #6324 branch summary ambient auth", () => {
 						cacheRead: 0,
 						cacheWrite: 0,
 						totalTokens: 2,
-						cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+						cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.25 },
 					},
 					stopReason: "stop",
 					timestamp: Date.now(),
@@ -57,5 +57,6 @@ describe("issue #6324 branch summary ambient auth", () => {
 		expect(streamCallCount).toBe(1);
 		expect(result.summaryEntry?.type).toBe("branch_summary");
 		expect(result.summaryEntry?.summary).toContain("branch summary text");
+		expect(result.summaryEntry?.usage?.cost.total).toBe(0.25);
 	});
 });