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
2026-07-20 16:41:43 +02:00
parent c179395218
commit 2fd3868401
29 changed files with 844 additions and 122 deletions
@@ -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);
+	});
+});