diff --git a/app/api/sessions/[id]/route.ts b/app/api/sessions/[id]/route.ts index 6aadb18d1..586fece33 100644 --- a/app/api/sessions/[id]/route.ts +++ b/app/api/sessions/[id]/route.ts @@ -14,6 +14,8 @@ import { sessionPathKey } from "@/lib/session-path"; import { getRpcSession } from "@/lib/rpc-manager"; import { projectTreeForResponse } from "@/lib/project-tree"; import { computeSessionTotalActiveMs } from "@/lib/session-timing"; +import { computeSessionStats } from "@/lib/session-stats"; +import type { SessionEntry } from "@/lib/types"; export async function GET( req: Request, @@ -38,6 +40,10 @@ export async function GET( const deferToolResultImages = searchParams.has("deferMedia"); const context = buildSessionContext(entries as never, leafId, { deferThinking, deferToolResultImages }); const totalActiveMs = computeSessionTotalActiveMs(entries); + // Cumulative usage over ALL entries, including history compacted away — + // the same aggregation the SDK's getSessionStats() uses. Lets the client + // keep monotonic token/cost counters across compaction and page reloads. + const stats = computeSessionStats(entries as unknown as SessionEntry[]); const header = sm.getHeader(); let modified = header?.timestamp ?? new Date().toISOString(); @@ -72,6 +78,7 @@ export async function GET( tree, context, totalActiveMs, + stats, }); } catch (error) { return NextResponse.json({ error: String(error) }, { status: 500 }); diff --git a/hooks/useAgentSession.ts b/hooks/useAgentSession.ts index 1a5eb10a2..4f18aa81f 100644 --- a/hooks/useAgentSession.ts +++ b/hooks/useAgentSession.ts @@ -18,6 +18,7 @@ import { clearDraft, rekeyDraft, restoreDraftSubmission } from "@/lib/draft-stor import { getPreferredToolPreset, setPreferredToolPreset } from "@/lib/tool-preset-preference"; import { getToolNamesForPreset, type ToolEntry, type ToolPreset } from "@/lib/tool-presets"; import type { SessionStatsInfo } from "@/lib/pi-types"; +import { mergeSessionStats, type SessionFileStats } from "@/lib/session-stats"; import { userMessageKey } from "@/lib/prompt-recovery"; import { AgentEventConnection } from "@/lib/agent-event-connection"; import { getToolExecutionProgress } from "@/lib/tool-execution-progress"; @@ -44,6 +45,8 @@ export interface SessionData { thinkingLevel: string; model: { provider: string; modelId: string } | null; }; + /** Cumulative usage over ALL session-file entries (incl. compacted history). */ + stats?: SessionFileStats; } interface AgentEvent { @@ -422,43 +425,18 @@ export function useAgentSession(opts: UseAgentSessionOptions) { if (sessionStatsOverride) { return { ...sessionStatsOverride, totalActiveMs: data?.totalActiveMs }; } - const tokens = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }; - let cost = 0; - let userMessages = 0; - let assistantMessages = 0; - let toolResults = 0; - let toolCalls = 0; - for (const msg of messages) { - if (msg.role === "user") userMessages += 1; - if (msg.role === "toolResult") toolResults += 1; - if (msg.role !== "assistant") continue; - assistantMessages += 1; - const u = (msg as import("@/lib/types").AssistantMessage).usage; - toolCalls += (msg as import("@/lib/types").AssistantMessage).content.filter((c) => c.type === "toolCall").length; - if (!u) continue; - tokens.input += u.input ?? 0; - tokens.output += u.output ?? 0; - tokens.cacheRead += u.cacheRead ?? 0; - tokens.cacheWrite += u.cacheWrite ?? 0; - cost += u.cost?.total ?? 0; - } - tokens.total = tokens.input + tokens.output + tokens.cacheRead + tokens.cacheWrite; - if (tokens.total === 0 && messages.length === 0) return null; + const fileStats = data?.stats; + const stats = mergeSessionStats(fileStats, data?.context.messages ?? [], messages); + if (stats.tokens.total === 0 && messages.length === 0 && !fileStats) return null; return { sessionFile: data?.filePath || undefined, sessionId: sessionIdRef.current ?? session?.id ?? "", sessionName: session?.name, - userMessages, - assistantMessages, - toolCalls, - toolResults, - totalMessages: messages.length, - tokens, - cost, + ...stats, totalActiveMs: data?.totalActiveMs, ...(contextUsage ? { contextUsage } : {}), } satisfies SessionStatsInfo; - }, [messages, sessionStatsOverride, contextUsage, data?.filePath, data?.totalActiveMs, session?.id, session?.name]); + }, [messages, sessionStatsOverride, contextUsage, data?.context.messages, data?.filePath, data?.totalActiveMs, data?.stats, session?.id, session?.name]); const loadSession = useCallback(async (sid: string, showLoading = false, includeState = false) => { let messagesLoaded = false; diff --git a/lib/session-stats.test.mjs b/lib/session-stats.test.mjs new file mode 100644 index 000000000..1e7fca32e --- /dev/null +++ b/lib/session-stats.test.mjs @@ -0,0 +1,191 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { createJiti } from "jiti"; + +const jiti = createJiti(import.meta.url); +const { computeSessionStats, mergeSessionStats } = await jiti.import("./session-stats.ts"); +const { buildContextEntries } = await jiti.import("@earendil-works/pi-coding-agent"); + +function usage(over = {}) { + return { + input: 10, + output: 5, + cacheRead: 2, + cacheWrite: 1, + cost: { input: 0.1, output: 0.05, cacheRead: 0.01, cacheWrite: 0.02, total: 0.18 }, + ...over, + }; +} + +function userEntry(id, parentId, content = "hi") { + return { type: "message", id, parentId, timestamp: "2026-01-01T00:00:00.000Z", message: { role: "user", content } }; +} + +function assistantEntry(id, parentId, blocks = [{ type: "text", text: "ok" }], u) { + const message = { + role: "assistant", + provider: "test", + model: "test-model", + content: blocks, + }; + if (u) message.usage = u; + return { type: "message", id, parentId, timestamp: "2026-01-01T00:00:00.000Z", message }; +} + +function toolResultEntry(id, parentId, u) { + const message = { role: "toolResult", toolCallId: "tc1", content: [{ type: "text", text: "result" }] }; + if (u) message.usage = u; + return { type: "message", id, parentId, timestamp: "2026-01-01T00:00:00.000Z", message }; +} + +function compactionEntry(id, parentId, summary, u) { + const entry = { + type: "compaction", + id, + parentId, + timestamp: "2026-01-01T00:00:00.000Z", + summary, + firstKeptEntryId: "u2", + tokensBefore: 100, + }; + if (u) entry.usage = u; + return entry; +} + +test("sums usage across ALL entries, including history compacted away", () => { + const entries = [ + userEntry("u1", null), + assistantEntry("a1", "u1", [{ type: "text", text: "old" }], usage()), // pre-compaction + compactionEntry("c1", "a1", "summary of old history", usage({ input: 1, output: 1, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.05 } })), + userEntry("u2", "c1"), + assistantEntry("a2", "u2", [{ type: "text", text: "new" }], usage({ input: 20, cost: { input: 0.2, output: 0.1, cacheRead: 0, cacheWrite: 0, total: 0.3 } })), + ]; + const stats = computeSessionStats(entries); + + assert.equal(stats.tokens.input, 10 + 1 + 20); + assert.equal(stats.tokens.output, 5 + 1 + 5); + assert.equal(stats.tokens.cacheRead, 2 + 2 + 2); + assert.equal(stats.tokens.cacheWrite, 1 + 1 + 1); + assert.equal(stats.tokens.total, stats.tokens.input + stats.tokens.output + stats.tokens.cacheRead + stats.tokens.cacheWrite); + assert.ok(Math.abs(stats.cost - (0.18 + 0.05 + 0.3)) < 1e-9); + assert.equal(stats.userMessages, 2); + assert.equal(stats.assistantMessages, 2); + assert.equal(stats.toolResults, 0); + assert.equal(stats.totalMessages, 4); +}); + +test("counts tool calls and includes tool result usage", () => { + const entries = [ + userEntry("u1", null), + assistantEntry("a1", "u1", [ + { type: "toolCall", toolCallId: "tc1", toolName: "bash", input: {} }, + { type: "toolCall", toolCallId: "tc2", toolName: "read", input: {} }, + { type: "text", text: "done" }, + ], usage()), + toolResultEntry("t1", "a1", usage({ input: 3, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.1 } })), + toolResultEntry("t2", "t1"), // no usage (e.g. local tool) + ]; + const stats = computeSessionStats(entries); + assert.equal(stats.toolCalls, 2); + assert.equal(stats.toolResults, 2); + assert.equal(stats.tokens.input, 10 + 3); + assert.ok(Math.abs(stats.cost - 0.28) < 1e-9); +}); + +test("includes branch summary usage", () => { + const entries = [ + userEntry("u1", null), + assistantEntry("a1", "u1", [{ type: "text", text: "x" }], usage()), + { + type: "branch_summary", + id: "b1", + parentId: "a1", + timestamp: "2026-01-01T00:00:00.000Z", + fromId: "a1", + summary: "side branch", + usage: usage({ input: 7, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.7 } }), + }, + ]; + const stats = computeSessionStats(entries); + assert.equal(stats.tokens.input, 17); + assert.ok(Math.abs(stats.cost - 0.88) < 1e-9); +}); + +test("tolerates entries without usage (older session files)", () => { + const entries = [ + userEntry("u1", null), + compactionEntry("c1", "u1", "old summary without usage"), + userEntry("u2", "c1"), + ]; + const stats = computeSessionStats(entries); + assert.equal(stats.tokens.total, 0); + assert.equal(stats.cost, 0); + assert.equal(stats.totalMessages, 2); +}); + +test("returns zeros for an empty session", () => { + const stats = computeSessionStats([]); + assert.deepEqual(stats, { + userMessages: 0, + assistantMessages: 0, + toolCalls: 0, + toolResults: 0, + totalMessages: 0, + tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + cost: 0, + }); +}); + +test("full-file stats never shrink relative to the post-compaction context", () => { + const entries = [ + userEntry("u1", null), + assistantEntry("a1", "u1", [{ type: "text", text: "old" }], usage()), + assistantEntry("a2", "a1", [{ type: "text", text: "older" }], usage({ input: 40 })), + compactionEntry("c1", "a2", "summary", usage()), + userEntry("u3", "c1"), + assistantEntry("a3", "u3", [{ type: "text", text: "new" }], usage({ input: 5 })), + ]; + const byId = new Map(entries.map((e) => [e.id, e])); + // What the SDK exposes to the UI after compaction: the compaction entry plus + // entries kept/recreated after it — everything before is dropped. + const contextEntries = buildContextEntries(entries, "a3", byId); + + const full = computeSessionStats(entries); + const context = computeSessionStats(contextEntries); + for (const field of ["input", "output", "cacheRead", "cacheWrite", "total"]) { + assert.ok( + full.tokens[field] >= context.tokens[field], + `full.tokens.${field} (${full.tokens[field]}) should be >= context.tokens.${field} (${context.tokens[field]})`, + ); + } + assert.ok(full.cost >= context.cost); +}); + +test("adds messages completed after load to compacted file totals", () => { + const entries = [ + userEntry("u1", null), + assistantEntry("a1", "u1", [{ type: "text", text: "old" }], usage()), + compactionEntry("c1", "a1", "summary", usage()), + userEntry("u2", "c1"), + assistantEntry("a2", "u2", [{ type: "text", text: "loaded" }], usage({ input: 20 })), + ]; + const loadedMessages = entries.slice(-2).map((entry) => entry.message); + const newUser = userEntry("u3", "a2"); + const newAssistant = assistantEntry("a3", "u3", [ + { type: "toolCall", toolCallId: "tc1", toolName: "bash", input: {} }, + ], usage({ input: 40, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.4 } })); + const currentMessages = [...loadedMessages, newUser.message, newAssistant.message]; + + const file = computeSessionStats(entries); + const merged = mergeSessionStats(file, loadedMessages, currentMessages); + + assert.equal(merged.tokens.input, file.tokens.input + 40); + assert.ok(Math.abs(merged.cost - (file.cost + 0.4)) < 1e-9); + assert.equal(merged.userMessages, file.userMessages + 1); + assert.equal(merged.assistantMessages, file.assistantMessages + 1); + assert.equal(merged.toolCalls, file.toolCalls + 1); + assert.equal(merged.totalMessages, file.totalMessages + 2); + + const reloadedFile = computeSessionStats([...entries, newUser, newAssistant]); + assert.deepEqual(mergeSessionStats(reloadedFile, currentMessages, currentMessages), reloadedFile); +}); diff --git a/lib/session-stats.ts b/lib/session-stats.ts new file mode 100644 index 000000000..3cfb75ffb --- /dev/null +++ b/lib/session-stats.ts @@ -0,0 +1,123 @@ +import type { AgentMessage, AgentUsage, SessionEntry } from "./types"; + +export interface SessionFileStats { + userMessages: number; + assistantMessages: number; + toolCalls: number; + toolResults: number; + totalMessages: number; + tokens: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; + }; + cost: number; +} + +function emptyStats(): SessionFileStats { + return { + userMessages: 0, + assistantMessages: 0, + toolCalls: 0, + toolResults: 0, + totalMessages: 0, + tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + cost: 0, + }; +} + +function addUsage(stats: SessionFileStats, usage?: AgentUsage): void { + if (!usage) return; + stats.tokens.input += usage.input ?? 0; + stats.tokens.output += usage.output ?? 0; + stats.tokens.cacheRead += usage.cacheRead ?? 0; + stats.tokens.cacheWrite += usage.cacheWrite ?? 0; + stats.cost += usage.cost?.total ?? 0; +} + +function addMessage(stats: SessionFileStats, message: AgentMessage): void { + stats.totalMessages += 1; + if (message.role === "user") { + stats.userMessages += 1; + } else if (message.role === "toolResult") { + stats.toolResults += 1; + addUsage(stats, message.usage); + } else if (message.role === "assistant") { + stats.assistantMessages += 1; + if (Array.isArray(message.content)) { + stats.toolCalls += message.content.filter((c) => c.type === "toolCall").length; + } + addUsage(stats, message.usage); + } +} + +function finishStats(stats: SessionFileStats): SessionFileStats { + stats.tokens.total = stats.tokens.input + stats.tokens.output + stats.tokens.cacheRead + stats.tokens.cacheWrite; + return stats; +} + +function computeMessageStats(messages: AgentMessage[]): SessionFileStats { + const stats = emptyStats(); + for (const message of messages) { + if (message.role !== "custom") addMessage(stats, message); + } + return finishStats(stats); +} + +export function mergeSessionStats( + fileStats: SessionFileStats | undefined, + loadedMessages: AgentMessage[], + currentMessages: AgentMessage[], +): SessionFileStats { + const current = computeMessageStats(currentMessages); + if (!fileStats) return current; + + const loaded = computeMessageStats(loadedMessages); + const delta = (now: number, before: number) => Math.max(0, now - before); + const tokens = { + input: fileStats.tokens.input + delta(current.tokens.input, loaded.tokens.input), + output: fileStats.tokens.output + delta(current.tokens.output, loaded.tokens.output), + cacheRead: fileStats.tokens.cacheRead + delta(current.tokens.cacheRead, loaded.tokens.cacheRead), + cacheWrite: fileStats.tokens.cacheWrite + delta(current.tokens.cacheWrite, loaded.tokens.cacheWrite), + total: 0, + }; + tokens.total = tokens.input + tokens.output + tokens.cacheRead + tokens.cacheWrite; + return { + userMessages: fileStats.userMessages + delta(current.userMessages, loaded.userMessages), + assistantMessages: fileStats.assistantMessages + delta(current.assistantMessages, loaded.assistantMessages), + toolCalls: fileStats.toolCalls + delta(current.toolCalls, loaded.toolCalls), + toolResults: fileStats.toolResults + delta(current.toolResults, loaded.toolResults), + totalMessages: fileStats.totalMessages + delta(current.totalMessages, loaded.totalMessages), + tokens, + cost: fileStats.cost + delta(current.cost, loaded.cost), + }; +} + +/** + * Aggregate usage across ALL entries in a session file. + * + * Mirrors the SDK's `AgentSession.getSessionStats()`: besides assistant + * (and tool-result) messages, this also counts usage recorded on compaction + * and branch-summary entries. Compaction only appends a summary entry — the + * summarized history stays in the file — so these totals grow monotonically + * for the life of the session. Totals computed over the active context alone + * (the compaction-aware message list) shrink whenever old history is + * summarized away, which is what made the UI token/cost counters appear to be + * reset after compaction. + */ +export function computeSessionStats(entries: SessionEntry[]): SessionFileStats { + const stats = emptyStats(); + + for (const entry of entries) { + if (entry.type === "compaction" || entry.type === "branch_summary") { + addUsage(stats, entry.usage); + continue; + } + if (entry.type !== "message") continue; + addMessage(stats, entry.message); + } + + return finishStats(stats); +} diff --git a/lib/types.ts b/lib/types.ts index 36572df10..c73d1a265 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -55,6 +55,20 @@ export interface UserMessage { timestamp?: number; } +export interface AgentUsage { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + cost: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; + }; +} + export interface AssistantMessage { role: "assistant"; content: AssistantContentBlock[]; @@ -63,19 +77,7 @@ export interface AssistantMessage { stopReason?: string; errorMessage?: string; timestamp?: number; - usage?: { - input: number; - output: number; - cacheRead: number; - cacheWrite: number; - cost: { - input: number; - output: number; - cacheRead: number; - cacheWrite: number; - total: number; - }; - }; + usage?: AgentUsage; } export interface ToolResultMessage { @@ -86,6 +88,7 @@ export interface ToolResultMessage { isError?: boolean; details?: unknown; timestamp?: number; + usage?: AgentUsage; } export interface CustomMessage { @@ -234,6 +237,7 @@ export interface CompactionEntry extends SessionEntryBase { tokensBefore: number; details?: unknown; fromHook?: boolean; + usage?: AgentUsage; } export interface BranchSummaryEntry extends SessionEntryBase { @@ -242,6 +246,7 @@ export interface BranchSummaryEntry extends SessionEntryBase { summary: string; details?: unknown; fromHook?: boolean; + usage?: AgentUsage; } export interface CustomEntry extends SessionEntryBase {