diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 14eb9f6e..c9e731cd 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -737,6 +737,7 @@ async function finalizeExecutedToolCall( ...result, content: afterResult.content ?? result.content, details: afterResult.details ?? result.details, + usage: afterResult.usage ?? result.usage, terminate: afterResult.terminate ?? result.terminate, }; isError = afterResult.isError ?? isError; @@ -780,6 +781,7 @@ function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResul // so the null never enters session history or provider payloads. content: finalized.result.content ?? [], details: finalized.result.details, + usage: finalized.result.usage, ...(finalized.result.addedToolNames?.length ? { addedToolNames: finalized.result.addedToolNames } : {}), isError: finalized.isError, timestamp: Date.now(), diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 1ce57b14..85af6b28 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -32,6 +32,7 @@ import type { AgentHarnessResources, AgentHarnessStreamOptions, AgentHarnessStreamOptionsPatch, + CompactResult, ExecutionEnv, NavigateTreeResult, PendingSessionWrite, @@ -434,9 +435,16 @@ export class AgentHarness< content: result.content, details: result.details, isError, + usage: result.usage, }); return patch - ? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate } + ? { + content: patch.content, + details: patch.details, + isError: patch.isError, + usage: patch.usage, + terminate: patch.terminate, + } : undefined; }, prepareNextTurn: async () => { @@ -690,9 +698,7 @@ export class AgentHarness< } } - async compact( - customInstructions?: string, - ): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> { + async compact(customInstructions?: string): Promise { if (this.phase !== "idle") throw new AgentHarnessError("busy", "compact() requires idle harness"); this.phase = "compaction"; try { @@ -723,6 +729,7 @@ export class AgentHarness< result.tokensBefore, result.details, provided !== undefined, + result.usage, ); const entry = await this.session.getEntry(entryId); if (entry?.type === "compaction") { @@ -764,6 +771,7 @@ export class AgentHarness< let summaryEntry: NavigateTreeResult["summaryEntry"]; let summaryText: string | undefined = hookResult?.summary?.summary; let summaryDetails: unknown = hookResult?.summary?.details; + let summaryUsage = hookResult?.summary?.usage; if (!summaryText && options?.summarize && entries.length > 0) { const model = this.model; if (!model) throw new AgentHarnessError("invalid_state", "No model set for branch summary"); @@ -779,6 +787,7 @@ export class AgentHarness< throw new AgentHarnessError("branch_summary", branchSummary.error.message, branchSummary.error); } summaryText = branchSummary.value.summary; + summaryUsage = branchSummary.value.usage; summaryDetails = { readFiles: branchSummary.value.readFiles, modifiedFiles: branchSummary.value.modifiedFiles, @@ -798,7 +807,12 @@ export class AgentHarness< const summaryId = await this.session.moveTo( newLeafId, summaryText - ? { summary: summaryText, details: summaryDetails, fromHook: hookResult?.summary !== undefined } + ? { + summary: summaryText, + details: summaryDetails, + usage: summaryUsage, + fromHook: hookResult?.summary !== undefined, + } : undefined, ); if (summaryId) { diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts index a8d68673..ab56f3ac 100644 --- a/packages/agent/src/harness/compaction/branch-summarization.ts +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -252,6 +252,7 @@ export async function generateBranchSummary( return ok({ summary: summary || "No summary generated", + usage: response.usage, readFiles, modifiedFiles, }); diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index a7f8ebc0..58938c13 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -101,10 +101,35 @@ export interface CompactionResult { firstKeptEntryId: string; /** Estimated context tokens before compaction. */ tokensBefore: number; + /** Usage from the LLM call(s) that generated this summary, if available. */ + usage?: Usage; /** Optional implementation-specific details stored with the compaction entry. */ details?: T; } +function combineUsage(first: Usage, second: Usage): Usage { + return { + input: first.input + second.input, + output: first.output + second.output, + cacheRead: first.cacheRead + second.cacheRead, + cacheWrite: first.cacheWrite + second.cacheWrite, + ...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined + ? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) } + : {}), + ...(first.reasoning !== undefined || second.reasoning !== undefined + ? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) } + : {}), + totalTokens: first.totalTokens + second.totalTokens, + cost: { + input: first.cost.input + second.cost.input, + output: first.cost.output + second.cost.output, + cacheRead: first.cost.cacheRead + second.cost.cacheRead, + cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite, + total: first.cost.total + second.cost.total, + }, + }; +} + /** Compaction thresholds and retention settings. */ export interface CompactionSettings { /** Enable automatic compaction decisions. */ @@ -474,7 +499,7 @@ export async function generateSummary( customInstructions?: string, previousSummary?: string, thinkingLevel?: ThinkingLevel, -): Promise> { +): Promise> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -523,7 +548,7 @@ export async function generateSummary( const textContent = contentText(response.content); - return ok(textContent); + return ok({ text: textContent, usage: response.usage }); } /** Prepared inputs for a compaction run. */ @@ -656,22 +681,26 @@ export async function compact( } let summary: string; + let summaryUsage: Usage; if (isSplitTurn && turnPrefixMessages.length > 0) { - const historyResult = - messagesToSummarize.length > 0 - ? await generateSummary( - messagesToSummarize, - models, - model, - settings.reserveTokens, - signal, - customInstructions, - previousSummary, - thinkingLevel, - ) - : ok("No prior history."); - if (!historyResult.ok) return err(historyResult.error); + let historyText = "No prior history."; + let historyUsage: Usage | undefined; + if (messagesToSummarize.length > 0) { + const historyResult = await generateSummary( + messagesToSummarize, + models, + model, + settings.reserveTokens, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ); + if (!historyResult.ok) return err(historyResult.error); + historyText = historyResult.value.text; + historyUsage = historyResult.value.usage; + } const turnPrefixResult = await generateTurnPrefixSummary( turnPrefixMessages, models, @@ -681,7 +710,10 @@ export async function compact( thinkingLevel, ); if (!turnPrefixResult.ok) return err(turnPrefixResult.error); - summary = `${historyResult.value}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value}`; + summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`; + summaryUsage = historyUsage + ? combineUsage(historyUsage, turnPrefixResult.value.usage) + : turnPrefixResult.value.usage; } else { const summaryResult = await generateSummary( messagesToSummarize, @@ -694,7 +726,8 @@ export async function compact( thinkingLevel, ); if (!summaryResult.ok) return err(summaryResult.error); - summary = summaryResult.value; + summary = summaryResult.value.text; + summaryUsage = summaryResult.value.usage; } const { readFiles, modifiedFiles } = computeFileLists(fileOps); @@ -704,6 +737,7 @@ export async function compact( summary, firstKeptEntryId, tokensBefore, + usage: summaryUsage, details: { readFiles, modifiedFiles } as CompactionDetails, }); } @@ -714,7 +748,7 @@ async function generateTurnPrefixSummary( reserveTokens: number, signal?: AbortSignal, thinkingLevel?: ThinkingLevel, -): Promise> { +): Promise> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -749,5 +783,8 @@ async function generateTurnPrefixSummary( ); } - return ok(contentText(response.content)); + return ok({ + text: contentText(response.content), + usage: response.usage, + }); } diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index f5224c92..bff41fdd 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/packages/agent/src/harness/session/session.ts @@ -1,4 +1,4 @@ -import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import type { ImageContent, TextContent, Usage } from "@earendil-works/pi-ai"; import type { AgentMessage } from "../../types.ts"; import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.ts"; import type { @@ -247,6 +247,7 @@ export class Session { tokensBefore: number, details?: T, fromHook?: boolean, + usage?: Usage, ): Promise { return this.appendTypedEntry({ type: "compaction", @@ -257,6 +258,7 @@ export class Session { firstKeptEntryId, tokensBefore, details, + usage, fromHook, } satisfies CompactionEntry); } @@ -317,7 +319,7 @@ export class Session { async moveTo( entryId: string | null, - summary?: { summary: string; details?: unknown; fromHook?: boolean }, + summary?: { summary: string; details?: unknown; usage?: Usage; fromHook?: boolean }, ): Promise { if (entryId !== null && !(await this.storage.getEntry(entryId))) { throw new SessionError("not_found", `Entry ${entryId} not found`); @@ -332,6 +334,7 @@ export class Session { fromId: entryId ?? "root", summary: summary.summary, details: summary.details, + usage: summary.usage, fromHook: summary.fromHook, } satisfies BranchSummaryEntry); } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 3f965550..cdc7a542 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -1,4 +1,12 @@ -import type { ImageContent, Model, Models, SimpleStreamOptions, TextContent, Transport } from "@earendil-works/pi-ai"; +import type { + ImageContent, + Model, + Models, + SimpleStreamOptions, + TextContent, + Transport, + Usage, +} from "@earendil-works/pi-ai"; import type { AgentEvent, AgentMessage, AgentTool, QueueMode, ThinkingLevel } from "../index.ts"; import type { Session } from "./session/session.ts"; @@ -365,6 +373,7 @@ export interface CompactionEntry extends SessionTreeEntryBase { firstKeptEntryId: string; tokensBefore: number; details?: T; + usage?: Usage; fromHook?: boolean; } @@ -373,6 +382,7 @@ export interface BranchSummaryEntry extends SessionTreeEntryBase { fromId: string; summary: string; details?: T; + usage?: Usage; fromHook?: boolean; } @@ -572,6 +582,7 @@ export interface ToolResultEvent { content: Array; details: unknown; isError: boolean; + usage?: Usage; } export interface SessionBeforeCompactEvent { @@ -687,6 +698,7 @@ export interface ToolResultPatch { content?: Array; details?: unknown; isError?: boolean; + usage?: Usage; terminate?: boolean; } @@ -697,7 +709,12 @@ export interface SessionBeforeCompactResult { export interface SessionBeforeTreeResult { cancel?: boolean; - summary?: { summary: string; details?: unknown }; + summary?: { + summary: string; + details?: unknown; + /** Usage from the LLM call that generated this summary, if available. */ + usage?: Usage; + }; customInstructions?: string; replaceInstructions?: boolean; label?: string; @@ -738,6 +755,8 @@ export interface CompactResult { summary: string; firstKeptEntryId: string; tokensBefore: number; + /** Usage from the LLM call(s) that generated this summary, if available. */ + usage?: Usage; details?: unknown; } @@ -793,6 +812,7 @@ export interface GenerateBranchSummaryOptions { export interface BranchSummaryResult { summary: string; + usage?: Usage; readFiles: string[]; modifiedFiles: string[]; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 2de50144..52f34891 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -11,6 +11,7 @@ import type { TextContent, Tool, ToolResultMessage, + Usage, } from "@earendil-works/pi-ai"; import type { Static, TSchema } from "typebox"; @@ -69,15 +70,18 @@ export interface BeforeToolCallResult { * - `content`: if provided, replaces the tool result content array in full * - `details`: if provided, replaces the tool result details value in full * - `isError`: if provided, replaces the tool result error flag + * - `usage`: if provided, replaces the tool result usage * - `terminate`: if provided, replaces the early-termination hint * * Omitted fields keep the original executed tool result values. - * There is no deep merge for `content` or `details`. + * There is no deep merge for `content`, `details`, or `usage`. */ export interface AfterToolCallResult { content?: (TextContent | ImageContent)[]; details?: unknown; isError?: boolean; + /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */ + usage?: Usage; /** * Hint that the agent should stop after the current tool batch. * Early termination only happens when every finalized tool result in the batch sets this to true. @@ -273,6 +277,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * - `content` replaces the full content array * - `details` replaces the full details payload * - `isError` replaces the error flag + * - `usage` replaces the tool result usage * - `terminate` replaces the early-termination hint * * Any omitted fields keep their original values. No deep merge is performed. @@ -352,6 +357,8 @@ export interface AgentToolResult { content: (TextContent | ImageContent)[]; /** Arbitrary structured details for logs or UI rendering. */ details: T; + /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */ + usage?: Usage; /** Names of tools introduced by this result and available from this transcript point onward. */ addedToolNames?: string[]; /** diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 03868047..f9641ce7 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -239,6 +239,23 @@ describe("agentLoop with AgentMessage", () => { it("should handle tool calls and results", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; + const toolUsage = { + 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 = { + 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: typeof toolUsage | undefined; const tool: AgentTool = { name: "echo", label: "Echo", @@ -249,6 +266,7 @@ describe("agentLoop with AgentMessage", () => { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, + usage: toolUsage, }; }, }; @@ -264,6 +282,10 @@ describe("agentLoop with AgentMessage", () => { const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, + afterToolCall: async ({ result }) => { + observedToolUsage = result.usage; + return { usage: patchedToolUsage }; + }, }; let callIndex = 0; @@ -305,6 +327,10 @@ describe("agentLoop with AgentMessage", () => { if (toolEnd?.type === "tool_execution_end") { expect(toolEnd.isError).toBe(false); } + expect(observedToolUsage).toEqual(toolUsage); + const messages = await stream.result(); + const toolResult = messages.find((message) => message.role === "toolResult"); + expect(toolResult?.role === "toolResult" ? toolResult.usage : undefined).toEqual(patchedToolUsage); }); it("should not execute tool calls from a length-truncated assistant message", async () => { diff --git a/packages/agent/test/harness/agent-harness.test.ts b/packages/agent/test/harness/agent-harness.test.ts index d13eca84..48f18e24 100644 --- a/packages/agent/test/harness/agent-harness.test.ts +++ b/packages/agent/test/harness/agent-harness.test.ts @@ -5,6 +5,7 @@ import { fauxProvider, fauxToolCall, type RegisterFauxProviderOptions, + type Usage, } from "@earendil-works/pi-ai"; import { getModel } from "@earendil-works/pi-ai/compat"; import { describe, expect, it } from "vitest"; @@ -14,7 +15,7 @@ import { InMemorySessionStorage } from "../../src/harness/session/memory-storage import { Session } from "../../src/harness/session/session.ts"; import type { PromptTemplate, Skill } from "../../src/harness/types.ts"; import type { AgentMessage, AgentTool } from "../../src/types.ts"; -import { calculateTool } from "../utils/calculate.ts"; +import { calculateTool, createCalculateToolWithUsage } from "../utils/calculate.ts"; import { getCurrentTimeTool } from "../utils/get-current-time.ts"; interface AppSkill extends Skill { @@ -60,6 +61,34 @@ function getReasoning(options: unknown): unknown { return options.reasoning; } +function createUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage { + return { + input, + output, + cacheRead, + cacheWrite, + totalTokens: input + output + cacheRead + cacheWrite, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function createUserMessage(text: string): AgentMessage { + return { role: "user", content: [{ type: "text", text }], timestamp: Date.now() }; +} + +function createAssistantMessage(text: string): AgentMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "faux", + provider: "faux", + model: "faux-1", + usage: createUsage(100, 50), + stopReason: "stop", + timestamp: Date.now(), + }; +} + describe("AgentHarness", () => { it("constructs directly and exposes queue modes", () => { const session = new Session(new InMemorySessionStorage()); @@ -427,14 +456,18 @@ describe("AgentHarness", () => { }), ]); const session = new Session(new InMemorySessionStorage()); + const toolUsage = createUsage(1, 2, 3, 4); + const patchedToolUsage = createUsage(5, 6, 7, 8); + const calculateToolWithUsage = createCalculateToolWithUsage(toolUsage); const harness = new AgentHarness({ models, env: new NodeExecutionEnv({ cwd: process.cwd() }), session, model: registration.getModel(), - tools: [calculateTool], + tools: [calculateToolWithUsage], }); const seenToolCalls: Array<{ id: string; name: string; expression: unknown }> = []; + let seenToolUsage: Usage | undefined; harness.on("tool_call", (event) => { seenToolCalls.push({ id: event.toolCallId, name: event.toolName, expression: event.input.expression }); return undefined; @@ -442,9 +475,11 @@ describe("AgentHarness", () => { harness.on("tool_result", (event) => { expect(event.toolCallId).toBe("call-1"); expect(event.toolName).toBe("calculate"); + seenToolUsage = event.usage; return { content: [{ type: "text", text: "patched result" }], details: { patched: true }, + usage: patchedToolUsage, terminate: true, }; }); @@ -455,16 +490,109 @@ describe("AgentHarness", () => { (entry) => entry.type === "message" && entry.message.role === "toolResult", ); expect(seenToolCalls).toEqual([{ id: "call-1", name: "calculate", expression: "2 + 2" }]); + expect(seenToolUsage).toEqual(toolUsage); expect(toolResult).toMatchObject({ type: "message", message: { role: "toolResult", content: [{ type: "text", text: "patched result" }], details: { patched: true }, + usage: patchedToolUsage, }, }); }); + it("persists generated compaction usage", async () => { + const registration = newFaux(); + registration.setResponses([fauxAssistantMessage("## Goal\nTest summary")]); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + }); + + const result = await harness.compact(); + const compaction = (await session.getEntries()).find((entry) => entry.type === "compaction"); + + expect(result.usage?.totalTokens).toBeGreaterThan(0); + expect(compaction?.type === "compaction" ? compaction.usage : undefined).toEqual(result.usage); + }); + + it("persists hook-provided compaction usage", async () => { + const registration = newFaux(); + const usage = createUsage(5, 6, 7, 8); + const session = new Session(new InMemorySessionStorage()); + await session.appendMessage(createUserMessage("one")); + await session.appendMessage(createAssistantMessage("two")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + }); + harness.on("session_before_compact", (event) => ({ + compaction: { + summary: "hook summary", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + usage, + }, + })); + + const result = await harness.compact(); + const compaction = (await session.getEntries()).find((entry) => entry.type === "compaction"); + + expect(result.usage).toEqual(usage); + expect(compaction?.type === "compaction" ? compaction.usage : undefined).toEqual(usage); + }); + + it("persists generated branch summary usage", async () => { + const registration = newFaux(); + registration.setResponses([fauxAssistantMessage("## Goal\nBranch summary")]); + const session = new Session(new InMemorySessionStorage()); + const targetId = await session.appendMessage(createUserMessage("first branch")); + await session.appendMessage(createAssistantMessage("first reply")); + await session.appendMessage(createUserMessage("abandoned work")); + await session.appendMessage(createAssistantMessage("abandoned reply")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + }); + + const result = await harness.navigateTree(targetId, { summarize: true }); + + expect(result.summaryEntry?.usage?.totalTokens).toBeGreaterThan(0); + }); + + it("persists hook-provided branch summary usage", async () => { + const registration = newFaux(); + const usage = createUsage(13, 14, 15, 16); + const session = new Session(new InMemorySessionStorage()); + const targetId = await session.appendMessage(createUserMessage("first branch")); + await session.appendMessage(createAssistantMessage("first reply")); + await session.appendMessage(createUserMessage("abandoned work")); + await session.appendMessage(createAssistantMessage("abandoned reply")); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + }); + harness.on("session_before_tree", () => ({ + summary: { summary: "hook branch summary", usage }, + })); + + const result = await harness.navigateTree(targetId, { summarize: true }); + + expect(result.summaryEntry?.usage).toEqual(usage); + }); + it("preserves app tool types for getters and update events", async () => { const session = new Session(new InMemorySessionStorage()); const env = new NodeExecutionEnv({ cwd: process.cwd() }); diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index 8a228eeb..7cdea0c9 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -6,6 +6,7 @@ import { fauxProvider, type Message, type Model, + type Models, type Usage, } from "@earendil-works/pi-ai"; import { beforeEach, describe, expect, it } from "vitest"; @@ -142,6 +143,17 @@ function createFauxModel(reasoning: boolean, maxTokens = 8192): { faux: FauxProv return { faux, model: faux.getModel() }; } +function createModelsWithSimpleResponses(responses: AssistantMessage[]): Models { + const remaining = [...responses]; + const stub = Object.create(models) as Models; + stub.completeSimple = async () => { + const response = remaining.shift(); + if (!response) throw new Error("No faux completeSimple response queued"); + return response; + }; + return stub; +} + describe("harness compaction", () => { beforeEach(() => { nextId = 0; @@ -501,7 +513,12 @@ describe("harness compaction", () => { await generateSummary(messages, models, model, 2000, undefined, "focus", "old summary"), ); - expect(summary).toContain("Test summary"); + expect(summary.text).toContain("Test summary"); + expect(summary.usage.input).toBeGreaterThan(0); + expect(summary.usage.output).toBeGreaterThan(0); + expect(summary.usage.totalTokens).toBe( + summary.usage.input + summary.usage.output + summary.usage.cacheRead + summary.usage.cacheWrite, + ); expect(promptText).toContain("\nold summary\n"); expect(promptText).toContain("Additional focus: focus"); }); @@ -578,6 +595,30 @@ describe("harness compaction", () => { expect(invalidResult).toMatchObject({ ok: false, error: { code: "invalid_session" } }); }); + it("combines usage for split-turn compaction summaries", async () => { + const messages: AgentMessage[] = [createUserMessage("Summarize this.")]; + const { model } = createFauxModel(false); + const historyUsage = createMockUsage(1, 2, 3, 4); + const turnPrefixUsage = createMockUsage(5, 6, 7, 8); + const usageModels = createModelsWithSimpleResponses([ + { ...fauxAssistantMessage("history summary"), usage: historyUsage }, + { ...fauxAssistantMessage("turn prefix summary"), usage: turnPrefixUsage }, + ]); + const preparation: CompactionPreparation = { + firstKeptEntryId: "entry-keep", + messagesToSummarize: messages, + turnPrefixMessages: messages, + isSplitTurn: true, + tokensBefore: 100, + fileOps: { read: new Set(), written: new Set(), edited: new Set() }, + settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, + }; + + const result = getOrThrow(await compact(preparation, usageModels, model)); + + expect(result.usage).toEqual(createMockUsage(6, 8, 10, 12)); + }); + it("passes reasoning through turn-prefix summaries when enabled", async () => { const messages: AgentMessage[] = [createUserMessage("Summarize this.")]; const seenOptions: Array | undefined> = []; @@ -646,6 +687,7 @@ describe("harness compaction", () => { const result = getOrThrow(await compact(preparation!, models, model)); expect(result.summary.length).toBeGreaterThan(0); expect(result.firstKeptEntryId).toBeTruthy(); + expect(result.usage?.totalTokens).toBeGreaterThan(0); expect(result.details).toBeDefined(); }); }); diff --git a/packages/agent/test/harness/session.test.ts b/packages/agent/test/harness/session.test.ts index c009123d..26af52e2 100644 --- a/packages/agent/test/harness/session.test.ts +++ b/packages/agent/test/harness/session.test.ts @@ -86,6 +86,49 @@ async function runSessionSuite( expect(context.messages[1]?.role).toBe("branchSummary"); }); + it("persists compaction usage", async () => { + const session = new Session(await createStorage()); + const firstKeptEntryId = await session.appendMessage(createUserMessage("one")); + const 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 compactionId = await session.appendCompaction( + "summary", + firstKeptEntryId, + 1234, + undefined, + false, + usage, + ); + + const compactionEntry = await session.getEntry(compactionId); + expect(compactionEntry?.type === "compaction" ? compactionEntry.usage : undefined).toEqual(usage); + }); + + it("persists branch summary usage", async () => { + const session = new Session(await createStorage()); + const user1 = await session.appendMessage(createUserMessage("one")); + const 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 summaryId = await session.moveTo(user1, { summary: "summary text", usage }); + + const summaryEntry = await session.getEntry(summaryId!); + expect(summaryEntry?.type === "branch_summary" ? summaryEntry.usage : undefined).toEqual(usage); + }); + it("supports custom message entries in context", async () => { const session = new Session(await createStorage()); await session.appendMessage(createUserMessage("one")); diff --git a/packages/agent/test/utils/calculate.ts b/packages/agent/test/utils/calculate.ts index 6158efbd..0a4fefc5 100644 --- a/packages/agent/test/utils/calculate.ts +++ b/packages/agent/test/utils/calculate.ts @@ -1,3 +1,4 @@ +import type { Usage } from "@earendil-works/pi-ai"; import { type Static, Type } from "typebox"; import type { AgentTool, AgentToolResult } from "../../src/types.ts"; @@ -30,3 +31,10 @@ export const calculateTool: AgentTool = { return calculate(args.expression); }, }; + +export function createCalculateToolWithUsage(usage: Usage): AgentTool { + return { + ...calculateTool, + execute: async (_toolCallId: string, args: CalculateParams) => ({ ...calculate(args.expression), usage }), + }; +} diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 9fa1ab20..aeaf18c6 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -408,6 +408,8 @@ export interface ToolResultMessage { toolName: string; content: (TextContent | ImageContent)[]; // Supports text and images details?: TDetails; + /** Usage from the tool execution itself, if available. Not part of main LLM context accounting. */ + usage?: Usage; /** * Names from `Context.tools` that became available after this result. * Providers with native deferred tool loading use this as the load point; diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index c4afb9b1..eb4f78e9 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -32,6 +32,7 @@ import type { Model, ProviderHeaders, TextContent, + Usage, } from "@earendil-works/pi-ai/compat"; import { clampThinkingLevel, @@ -104,6 +105,7 @@ import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-promp import { type BashOperations, createLocalBashOperations } from "./tools/bash.ts"; import { createAllToolDefinitions } from "./tools/index.ts"; import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.ts"; +import { addUsageToTotals, createUsageTotals } from "./usage-totals.ts"; // ============================================================================ // Skill Block Parsing @@ -482,6 +484,7 @@ export class AgentSession { content: result.content, details: result.details, isError, + usage: result.usage, }); if (!hookResult) { @@ -492,6 +495,7 @@ export class AgentSession { content: hookResult.content, details: hookResult.details, isError: hookResult.isError ?? isError, + usage: hookResult.usage, }; }; } @@ -1812,6 +1816,7 @@ export class AgentSession { let summary: string; let firstKeptEntryId: string; let tokensBefore: number; + let usage: Usage | undefined; let details: unknown; if (extensionCompaction) { @@ -1819,6 +1824,7 @@ export class AgentSession { summary = extensionCompaction.summary; firstKeptEntryId = extensionCompaction.firstKeptEntryId; tokensBefore = extensionCompaction.tokensBefore; + usage = extensionCompaction.usage; details = extensionCompaction.details; } else { // Generate compaction result @@ -1836,6 +1842,7 @@ export class AgentSession { summary = result.summary; firstKeptEntryId = result.firstKeptEntryId; tokensBefore = result.tokensBefore; + usage = result.usage; details = result.details; } @@ -1843,7 +1850,7 @@ export class AgentSession { throw new Error("Compaction cancelled"); } - this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension); + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage); const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.state.messages = sessionContext.messages; @@ -1869,6 +1876,7 @@ export class AgentSession { firstKeptEntryId, tokensBefore, estimatedTokensAfter, + usage, details, }; this._emit({ @@ -2084,6 +2092,7 @@ export class AgentSession { let summary: string; let firstKeptEntryId: string; let tokensBefore: number; + let usage: Usage | undefined; let details: unknown; if (extensionCompaction) { @@ -2091,6 +2100,7 @@ export class AgentSession { summary = extensionCompaction.summary; firstKeptEntryId = extensionCompaction.firstKeptEntryId; tokensBefore = extensionCompaction.tokensBefore; + usage = extensionCompaction.usage; details = extensionCompaction.details; } else { // Generate compaction result @@ -2108,6 +2118,7 @@ export class AgentSession { summary = compactResult.summary; firstKeptEntryId = compactResult.firstKeptEntryId; tokensBefore = compactResult.tokensBefore; + usage = compactResult.usage; details = compactResult.details; } @@ -2122,7 +2133,7 @@ export class AgentSession { return false; } - this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension); + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage); const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.state.messages = sessionContext.messages; @@ -2148,6 +2159,7 @@ export class AgentSession { firstKeptEntryId, tokensBefore, estimatedTokensAfter, + usage, details, }; this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry }); @@ -2872,7 +2884,7 @@ export class AgentSession { this._branchSummaryAbortController = new AbortController(); try { - let extensionSummary: { summary: string; details?: unknown } | undefined; + let extensionSummary: { summary: string; details?: unknown; usage?: Usage } | undefined; let fromExtension = false; // Emit session_before_tree event @@ -2907,6 +2919,7 @@ export class AgentSession { // Run default summarizer if needed let summaryText: string | undefined; let summaryDetails: unknown; + let summaryUsage: Usage | undefined; if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { const model = this.model!; const { apiKey, headers, env } = await this._getSummarizationRequestAuth(model); @@ -2929,6 +2942,7 @@ export class AgentSession { throw new Error(result.error); } summaryText = result.summary; + summaryUsage = result.usage; summaryDetails = { readFiles: result.readFiles || [], modifiedFiles: result.modifiedFiles || [], @@ -2936,6 +2950,7 @@ export class AgentSession { } else if (extensionSummary) { summaryText = extensionSummary.summary; summaryDetails = extensionSummary.details; + summaryUsage = extensionSummary.usage; } // Determine the new leaf position based on target type @@ -2965,6 +2980,7 @@ export class AgentSession { summaryText, summaryDetails, fromExtension, + summaryUsage, ); summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry; @@ -3037,13 +3053,12 @@ export class AgentSession { let toolResults = 0; let totalMessages = 0; let toolCalls = 0; - let totalInput = 0; - let totalOutput = 0; - let totalCacheRead = 0; - let totalCacheWrite = 0; - let totalCost = 0; + const usageTotals = createUsageTotals(); for (const entry of this.sessionManager.getEntries()) { + if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) { + addUsageToTotals(usageTotals, entry.usage); + } if (entry.type !== "message") continue; totalMessages++; const message = entry.message; @@ -3051,18 +3066,16 @@ export class AgentSession { userMessages++; } else if (message.role === "toolResult") { toolResults++; + if (message.usage) { + addUsageToTotals(usageTotals, message.usage); + } } else if (message.role === "assistant") { assistantMessages++; const assistantMsg = message as AssistantMessage; if (Array.isArray(assistantMsg.content)) { toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length; } - const usage = assistantMsg.usage; - totalInput += usage.input; - totalOutput += usage.output; - totalCacheRead += usage.cacheRead; - totalCacheWrite += usage.cacheWrite; - totalCost += usage.cost.total; + addUsageToTotals(usageTotals, assistantMsg.usage); } } @@ -3075,13 +3088,13 @@ export class AgentSession { toolResults, totalMessages, tokens: { - input: totalInput, - output: totalOutput, - cacheRead: totalCacheRead, - cacheWrite: totalCacheWrite, - total: totalInput + totalOutput + totalCacheRead + totalCacheWrite, + input: usageTotals.input, + output: usageTotals.output, + cacheRead: usageTotals.cacheRead, + cacheWrite: usageTotals.cacheWrite, + total: usageTotals.input + usageTotals.output + usageTotals.cacheRead + usageTotals.cacheWrite, }, - cost: totalCost, + cost: usageTotals.cost, contextUsage: this.getContextUsage(), }; } diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index f2e06876..3366f06a 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -7,7 +7,7 @@ import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; import { contentText } from "@earendil-works/pi-ai"; -import type { Model, SimpleStreamOptions } from "@earendil-works/pi-ai/compat"; +import type { Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; import { completeSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm, @@ -33,6 +33,7 @@ import { export interface BranchSummaryResult { summary?: string; + usage?: Usage; readFiles?: string[]; modifiedFiles?: string[]; aborted?: boolean; @@ -363,6 +364,7 @@ export async function generateBranchSummary( return { summary: summary || "No summary generated", + usage: response.usage, readFiles, modifiedFiles, }; diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 5c71c74d..dafcb09e 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -90,10 +90,35 @@ export interface CompactionResult { firstKeptEntryId: string; tokensBefore: number; estimatedTokensAfter?: number; + /** Usage from the LLM call(s) that generated this summary, if available */ + usage?: Usage; /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ details?: T; } +function combineUsage(first: Usage, second: Usage): Usage { + return { + input: first.input + second.input, + output: first.output + second.output, + cacheRead: first.cacheRead + second.cacheRead, + cacheWrite: first.cacheWrite + second.cacheWrite, + ...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined + ? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) } + : {}), + ...(first.reasoning !== undefined || second.reasoning !== undefined + ? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) } + : {}), + totalTokens: first.totalTokens + second.totalTokens, + cost: { + input: first.cost.input + second.cost.input, + output: first.cost.output + second.cost.output, + cacheRead: first.cost.cacheRead + second.cost.cacheRead, + cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite, + total: first.cost.total + second.cost.total, + }, + }; +} + // ============================================================================ // Types // ============================================================================ @@ -556,7 +581,7 @@ export async function generateSummary( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, -): Promise { +): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -603,7 +628,7 @@ export async function generateSummary( const textContent = contentText(response.content); - return textContent; + return { text: textContent, usage: response.usage }; } // ============================================================================ @@ -759,24 +784,28 @@ export async function compact( // Generate summaries and merge into one let summary: string; + let summaryUsage: Usage; if (isSplitTurn && turnPrefixMessages.length > 0) { - const historyResult = - messagesToSummarize.length > 0 - ? await generateSummary( - messagesToSummarize, - model, - settings.reserveTokens, - apiKey, - headers, - signal, - customInstructions, - previousSummary, - thinkingLevel, - streamFn, - env, - ) - : "No prior history."; + let historyText = "No prior history."; + let historyUsage: Usage | undefined; + if (messagesToSummarize.length > 0) { + const historyResult = await generateSummary( + messagesToSummarize, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + streamFn, + env, + ); + historyText = historyResult.text; + historyUsage = historyResult.usage; + } const turnPrefixResult = await generateTurnPrefixSummary( turnPrefixMessages, model, @@ -789,10 +818,11 @@ export async function compact( streamFn, ); // Merge into single summary - summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; + summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.text}`; + summaryUsage = historyUsage ? combineUsage(historyUsage, turnPrefixResult.usage) : turnPrefixResult.usage; } else { // Just generate history summary - summary = await generateSummary( + const result = await generateSummary( messagesToSummarize, model, settings.reserveTokens, @@ -805,6 +835,8 @@ export async function compact( streamFn, env, ); + summary = result.text; + summaryUsage = result.usage; } // Compute file lists and append to summary @@ -819,6 +851,7 @@ export async function compact( summary, firstKeptEntryId, tokensBefore, + usage: summaryUsage, details: { readFiles, modifiedFiles } as CompactionDetails, }; } @@ -836,7 +869,7 @@ async function generateTurnPrefixSummary( signal?: AbortSignal, thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, -): Promise { +): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, @@ -863,5 +896,8 @@ async function generateTurnPrefixSummary( throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`); } - return contentText(response.content); + return { + text: contentText(response.content), + usage: response.usage, + }; } diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 4bd43dda..3c320e14 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -883,6 +883,10 @@ export class ExtensionRunner { currentEvent.isError = handlerResult.isError; modified = true; } + if (handlerResult.usage !== undefined) { + currentEvent.usage = handlerResult.usage; + modified = true; + } } catch (err) { const message = err instanceof Error ? err.message : String(err); const stack = err instanceof Error ? err.stack : undefined; @@ -904,6 +908,7 @@ export class ExtensionRunner { content: currentEvent.content, details: currentEvent.details, isError: currentEvent.isError, + usage: currentEvent.usage, }; } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 514f7af6..e8bf2627 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -30,6 +30,7 @@ import type { SimpleStreamOptions, TextContent, ToolResultMessage, + Usage, } from "@earendil-works/pi-ai"; import type { AutocompleteItem, @@ -905,6 +906,8 @@ interface ToolResultEventBase { input: Record; content: (TextContent | ImageContent)[]; isError: boolean; + /** Usage from the tool execution itself, if available. */ + usage?: Usage; } export interface BashToolResultEvent extends ToolResultEventBase { @@ -1072,6 +1075,7 @@ export interface ToolResultEventResult { content?: (TextContent | ImageContent)[]; details?: unknown; isError?: boolean; + usage?: Usage; } export interface MessageEndEventResult { @@ -1104,6 +1108,7 @@ export interface SessionBeforeTreeResult { summary?: { summary: string; details?: unknown; + usage?: Usage; }; /** Override custom instructions for summarization */ customInstructions?: string; diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 0ab66595..0dbce18f 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1,5 +1,5 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; -import { type ImageContent, type Message, type TextContent, uuidv7 } from "@earendil-works/pi-ai"; +import { type ImageContent, type Message, type TextContent, type Usage, uuidv7 } from "@earendil-works/pi-ai"; import { randomUUID } from "crypto"; import { appendFileSync, @@ -73,6 +73,8 @@ export interface CompactionEntry extends SessionEntryBase { tokensBefore: number; /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ details?: T; + /** Usage from the LLM call(s) that generated this summary, if available */ + usage?: Usage; /** True if generated by an extension, undefined/false if pi-generated (backward compatible) */ fromHook?: boolean; } @@ -83,6 +85,8 @@ export interface BranchSummaryEntry extends SessionEntryBase { summary: string; /** Extension-specific data (not sent to LLM) */ details?: T; + /** Usage from the LLM call that generated this summary, if available */ + usage?: Usage; /** True if generated by an extension, false if pi-generated */ fromHook?: boolean; } @@ -1096,6 +1100,7 @@ export class SessionManager { tokensBefore: number, details?: T, fromHook?: boolean, + usage?: Usage, ): string { const entry: CompactionEntry = { type: "compaction", @@ -1106,6 +1111,7 @@ export class SessionManager { firstKeptEntryId, tokensBefore, details, + usage, fromHook, }; this._appendEntry(entry); @@ -1372,7 +1378,13 @@ export class SessionManager { * Same as branch(), but also appends a branch_summary entry that captures * context from the abandoned conversation path. */ - branchWithSummary(branchFromId: string | null, summary: string, details?: unknown, fromHook?: boolean): string { + branchWithSummary( + branchFromId: string | null, + summary: string, + details?: unknown, + fromHook?: boolean, + usage?: Usage, + ): string { if (branchFromId !== null && !this.byId.has(branchFromId)) { throw new Error(`Entry ${branchFromId} not found`); } @@ -1385,6 +1397,7 @@ export class SessionManager { fromId: branchFromId ?? "root", summary, details, + usage, fromHook, }; this._appendEntry(entry); diff --git a/packages/coding-agent/src/core/usage-totals.ts b/packages/coding-agent/src/core/usage-totals.ts new file mode 100644 index 00000000..e65e2eb4 --- /dev/null +++ b/packages/coding-agent/src/core/usage-totals.ts @@ -0,0 +1,27 @@ +import type { Usage } from "@earendil-works/pi-ai/compat"; + +export interface UsageTotals { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + cost: number; +} + +export function createUsageTotals(): UsageTotals { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + cost: 0, + }; +} + +export function addUsageToTotals(totals: UsageTotals, usage: Usage): void { + totals.input += usage.input; + totals.output += usage.output; + totals.cacheRead += usage.cacheRead; + totals.cacheWrite += usage.cacheWrite; + totals.cost += usage.cost.total; +} diff --git a/packages/coding-agent/src/modes/interactive/components/footer.ts b/packages/coding-agent/src/modes/interactive/components/footer.ts index fa947480..44cf886b 100644 --- a/packages/coding-agent/src/modes/interactive/components/footer.ts +++ b/packages/coding-agent/src/modes/interactive/components/footer.ts @@ -3,6 +3,7 @@ import { type Component, truncateToWidth, visibleWidth } from "@earendil-works/p import type { AgentSession } from "../../../core/agent-session.ts"; import { areExperimentalFeaturesEnabled } from "../../../core/experimental.ts"; import type { ReadonlyFooterDataProvider } from "../../../core/footer-data-provider.ts"; +import { addUsageToTotals, createUsageTotals } from "../../../core/usage-totals.ts"; import { theme } from "../theme/theme.ts"; /** @@ -84,25 +85,21 @@ export class FooterComponent implements Component { const state = this.session.state; // Calculate cumulative usage from ALL session entries (not just post-compaction messages) - let totalInput = 0; - let totalOutput = 0; - let totalCacheRead = 0; - let totalCacheWrite = 0; - let totalCost = 0; + const usageTotals = createUsageTotals(); let latestCacheHitRate: number | undefined; for (const entry of this.session.sessionManager.getEntries()) { if (entry.type === "message" && entry.message.role === "assistant") { - totalInput += entry.message.usage.input; - totalOutput += entry.message.usage.output; - totalCacheRead += entry.message.usage.cacheRead; - totalCacheWrite += entry.message.usage.cacheWrite; - totalCost += entry.message.usage.cost.total; + addUsageToTotals(usageTotals, entry.message.usage); const latestPromptTokens = entry.message.usage.input + entry.message.usage.cacheRead + entry.message.usage.cacheWrite; latestCacheHitRate = latestPromptTokens > 0 ? (entry.message.usage.cacheRead / latestPromptTokens) * 100 : undefined; + } else if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.usage) { + addUsageToTotals(usageTotals, entry.message.usage); + } else if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) { + addUsageToTotals(usageTotals, entry.usage); } } @@ -130,19 +127,20 @@ export class FooterComponent implements Component { // Build stats line const statsParts = []; - if (totalInput) statsParts.push(`↑${formatTokens(totalInput)}`); - if (totalOutput) statsParts.push(`↓${formatTokens(totalOutput)}`); - if (totalCacheRead) statsParts.push(`R${formatTokens(totalCacheRead)}`); - if (totalCacheWrite) statsParts.push(`W${formatTokens(totalCacheWrite)}`); - if ((totalCacheRead > 0 || totalCacheWrite > 0) && latestCacheHitRate !== undefined) { + if (usageTotals.input) statsParts.push(`↑${formatTokens(usageTotals.input)}`); + if (usageTotals.output) statsParts.push(`↓${formatTokens(usageTotals.output)}`); + if (usageTotals.cacheRead) statsParts.push(`R${formatTokens(usageTotals.cacheRead)}`); + if (usageTotals.cacheWrite) statsParts.push(`W${formatTokens(usageTotals.cacheWrite)}`); + if ((usageTotals.cacheRead > 0 || usageTotals.cacheWrite > 0) && latestCacheHitRate !== undefined) { statsParts.push(`CH${latestCacheHitRate.toFixed(1)}%`); } + // Kimi Coding is subscription-backed despite using API-key authentication. const usingSubscription = state.model ? state.model.provider === "kimi-coding" || this.session.modelRuntime.isUsingOAuth(state.model.provider) : false; - if (totalCost || usingSubscription) { - const costStr = `$${totalCost.toFixed(3)}${usingSubscription ? " (sub)" : ""}`; + if (usageTotals.cost || usingSubscription) { + const costStr = `$${usageTotals.cost.toFixed(3)}${usingSubscription ? " (sub)" : ""}`; statsParts.push(costStr); } diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index 9a442c73..684c5e1c 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/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(); diff --git a/packages/coding-agent/test/branch-summary-extensions.test.ts b/packages/coding-agent/test/branch-summary-extensions.test.ts new file mode 100644 index 00000000..9c47cb81 --- /dev/null +++ b/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); + }); +}); diff --git a/packages/coding-agent/test/compaction-summary-reasoning.test.ts b/packages/coding-agent/test/compaction-summary-reasoning.test.ts index 306a3b1c..eb1d3579 100644 --- a/packages/coding-agent/test/compaction-summary-reasoning.test.ts +++ b/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]); }); }); diff --git a/packages/coding-agent/test/footer-width.test.ts b/packages/coding-agent/test/footer-width.test.ts index 85956d53..310114a4 100644 --- a/packages/coding-agent/test/footer-width.test.ts +++ b/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> = []; + + 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: "", diff --git a/packages/coding-agent/test/session-manager/tree-traversal.test.ts b/packages/coding-agent/test/session-manager/tree-traversal.test.ts index e47ec909..9f447de1 100644 --- a/packages/coding-agent/test/session-manager/tree-traversal.test.ts +++ b/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); } }); diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index ba5ae5c7..ee7bea77 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/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); diff --git a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts index 39a9243c..b334b123 100644 --- a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts +++ b/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 () => { diff --git a/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts b/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts index 121cc006..bfedbaaf 100644 --- a/packages/coding-agent/test/suite/regressions/6324-branch-summary-ambient-auth.test.ts +++ b/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); }); });