From b66c00b552e552d1d4c1f35ccc38e03cbe385bbb Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Tue, 20 Oct 2026 17:33:48 -0400 Subject: [PATCH] feat(pi): report Pi's usage to the Usage plugin Pi agents get a `usage` hook, as Claude, Codex and goose have. For each of the agent's Buzz conversations, it reads the conversation's Pi entries and reports each answer's tokens, model and cost as pi-ai recorded them. The channel and thread come from the conversation's key, the title from its first prompt. Only agents running on this computer are read, as for their sessions, since reading opens their storage. The host no longer drops `usage` from the messages it sends; that was the only place the counts were lost. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Salman Mohammed --- src/bundled/pi/host.live.test.ts | 9 +- src/bundled/pi/host.mjs | 5 +- src/bundled/pi/index.ts | 4 + src/bundled/pi/runtime.ts | 7 ++ src/bundled/pi/transcript.ts | 3 + src/bundled/pi/usage.test.ts | 138 ++++++++++++++++++++++++++++++ src/bundled/pi/usage.ts | 139 +++++++++++++++++++++++++++++++ 7 files changed, 299 insertions(+), 6 deletions(-) create mode 100644 src/bundled/pi/usage.test.ts create mode 100644 src/bundled/pi/usage.ts diff --git a/src/bundled/pi/host.live.test.ts b/src/bundled/pi/host.live.test.ts index 8723f10a01..0ce5faaa7a 100644 --- a/src/bundled/pi/host.live.test.ts +++ b/src/bundled/pi/host.live.test.ts @@ -73,7 +73,7 @@ out.a = { ...(await pi.rpc("wait", { agent, submission: a.submission })), calls: faux.setResponses([fauxAssistantMessage("Plain answer")]); const b = await pi.rpc("deliver", turn("B", "second")); out.b = { ...(await pi.rpc("wait", { agent, submission: b.submission })), conversation: b.conversation }; -// What a transcript reads: the places, and the entries without schemas or usage. +// What a transcript and usage read: the places, and the entries without schemas. out.places = (await pi.rpc("places", { agent })).places; out.entries = (await pi.rpc("entries", { agent, conversation: a.conversation })).entries.map((entry) => [entry.kind, ...(entry.model ?? []).map((message) => message.role + (message.usage ? "+usage" : ""))]); // A request while the run works steers it. @@ -195,12 +195,13 @@ describe.skipIf(!DEPS)("the Pi host on Pi durable", () => { expect(out.places).toEqual([ { key: "chan/root", conversation: out.a.conversation, seen: 100 }, ]); - // Newest first; a system entry's tool schemas are left out. + // Newest first; a system entry's tool schemas are left out, an answer's + // usage kept. expect(out.entries).toEqual([ - ["pi.assistant", "assistant"], + ["pi.assistant", "assistant+usage"], ["pi.user", "user"], ["pi.tool-result", "toolResult"], - ["pi.assistant", "assistant"], + ["pi.assistant", "assistant+usage"], ["pi.system"], ["pi.user", "user"], ]); diff --git a/src/bundled/pi/host.mjs b/src/bundled/pi/host.mjs index 9edfbba691..a76c582958 100644 --- a/src/bundled/pi/host.mjs +++ b/src/bundled/pi/host.mjs @@ -179,8 +179,9 @@ async function piDefaults(dir) { } } -/** A message as a transcript shows it: no usage, image data or signatures. */ -const slim = ({ usage, ...message }) => +/** A message as a transcript and usage show it: no image data or + * signatures. */ +const slim = (message) => Array.isArray(message.content) ? { ...message, diff --git a/src/bundled/pi/index.ts b/src/bundled/pi/index.ts index 5265220f66..7fe53d1460 100644 --- a/src/bundled/pi/index.ts +++ b/src/bundled/pi/index.ts @@ -69,6 +69,10 @@ export const apply: PluginModule["apply"] = (ctx) => { const { model, provider } = config(agent.config); return [provider, model].filter(Boolean).join(" ยท ") || "Pi"; }, + ...(host.spawn && { + usage: (agent, { since, signal }) => + runtime.usage(agent.pubkey, since, signal), + }), sessions: { list: (pubkey) => runtime.sessions(pubkey), subscribe: (listener) => runtime.subscribe(listener), diff --git a/src/bundled/pi/runtime.ts b/src/bundled/pi/runtime.ts index 46d6b30fcd..7396a69732 100644 --- a/src/bundled/pi/runtime.ts +++ b/src/bundled/pi/runtime.ts @@ -54,6 +54,7 @@ import { } from "./pi"; import { DEFAULT_THINKING, PiSetup, runsAt } from "./setup"; import { readTranscript } from "./transcript"; +import { readUsage } from "./usage"; export type Config = Readonly<{ /** A Pi provider; empty uses the one Pi's settings name. */ @@ -274,6 +275,12 @@ export class PiRuntime { }; }); } + /** The usage of the agent's Buzz conversations since `since`. */ + async usage(pubkey: string, since: number, signal: AbortSignal) { + // Reading opens its storage, which only an agent running here may hold. + if (!this.running.has(pubkey)) return []; + return readUsage(this.pi, pubkey, await this.places(pubkey), since, signal); + } /** The conversations a Buzz request of the agent's is working in. */ working(pubkey: string) { let view = this.views.get(pubkey); diff --git a/src/bundled/pi/transcript.ts b/src/bundled/pi/transcript.ts index 607c5cd372..93059b9d0b 100644 --- a/src/bundled/pi/transcript.ts +++ b/src/bundled/pi/transcript.ts @@ -37,6 +37,9 @@ type Message = Readonly<{ stopReason?: string; errorMessage?: string; model?: string; + responseId?: string; + /** pi-ai's token counts and cost of an answer. */ + usage?: unknown; }>; /** A Pi entry as the host reads it: its messages, as pi-ai has them. */ export type PiEntry = Readonly<{ diff --git a/src/bundled/pi/usage.test.ts b/src/bundled/pi/usage.test.ts new file mode 100644 index 0000000000..added27ca7 --- /dev/null +++ b/src/bundled/pi/usage.test.ts @@ -0,0 +1,138 @@ +// @vitest-environment node +import { expect, it, vi } from "vitest"; +import type { PiEntry } from "./transcript"; +import { conversationUsage, readUsage } from "./usage"; + +const PLACE = { conversation: 2, channelId: "chan", root: "root" }; +const ask = (id: number, text: string): PiEntry => ({ + id, + kind: "pi.user", + model: [ + { + role: "user", + content: `\nChannel: general (#chan)\n\n\n\nFrom: Sal (npub1x)\nContent: ${text}\nTags: []\n`, + timestamp: id * 1000, + }, + ], +}); +const answer = (id: number, usage: Record, more = {}) => ({ + id, + kind: "pi.assistant", + model: [ + { + role: "assistant", + content: [{ type: "text", text: "Done." }], + timestamp: id * 1000, + model: "claude-x", + responseId: `resp-${id}`, + usage, + ...more, + }, + ], +}); +const usage = (cost: number, more = {}) => ({ + input: 100, + output: 10, + cacheRead: 50, + cacheWrite: 5, + totalTokens: 165, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: cost }, + ...more, +}); + +it("reports each answer's tokens and cost, where and what was asked", () => { + expect( + conversationUsage( + [ + ask(1, "fix the build"), + answer(2, usage(0.25, { cacheWrite1h: 5 })), + { + id: 3, + kind: "pi.tool-result", + model: [{ role: "toolResult", timestamp: 3000 }], + }, + // A model pi-ai has no rates for is priced at 0: the cost is unknown. + answer(4, usage(0), { model: "local-y", responseId: undefined }), + ], + PLACE, + ), + ).toEqual({ + id: "2", + channelId: "chan", + root: "root", + title: "fix the build", + startedAt: 1000, + lastAt: 4000, + gaps: [], + observations: [ + { + requestId: "resp-2", + at: 2000, + model: "claude-x", + tokens: { + input: 100, + output: 10, + cacheRead: 50, + cacheWrite: 5, + cacheWrite1h: 5, + }, + cost: { usd: 0.25, source: "pi" }, + }, + { + at: 4000, + model: "local-y", + tokens: { input: 100, output: 10, cacheRead: 50, cacheWrite: 5 }, + }, + ], + }); +}); + +it("leaves out counts pi-ai did not record, and a conversation with no messages", () => { + const session = conversationUsage( + [ask(1, "hi"), answer(2, { output: 3, input: -1, cacheRead: "x" })], + { conversation: 5, channelId: "chan" }, + ); + expect(session?.root).toBeUndefined(); + expect(session?.observations).toEqual([ + { requestId: "resp-2", at: 2000, model: "claude-x", tokens: { output: 3 } }, + ]); + expect(conversationUsage([], PLACE)).toBeUndefined(); +}); + +it("reads every page of the conversations used since `since`, newest first", async () => { + const pages: Record = { + "2:": { entries: [answer(4, usage(1)), ask(3, "second")], next: "c" }, + "2:c": { entries: [answer(2, usage(1)), ask(1, "first")] }, + "7:": { entries: [answer(9, usage(1)), ask(8, "later")] }, + }; + const request = vi.fn( + async ( + _method: string, + params: { conversation: number; cursor?: string }, + ) => pages[`${params.conversation}:${params.cursor ?? ""}`], + ); + const sessions = await readUsage( + { request } as never, + "agent", + [ + { conversation: 2, channelId: "chan", root: "root", at: 3000 }, + { conversation: 7, channelId: "other", at: 8000 }, + // Shown nothing since: not read. + { conversation: 9, channelId: "old", at: 500 }, + ], + 1500, + new AbortController().signal, + ); + expect(request).toHaveBeenCalledTimes(3); + expect(request).toHaveBeenCalledWith("entries", { + agent: "agent", + conversation: 2, + cursor: "c", + }); + expect(sessions.map((session) => [session.id, session.title])).toEqual([ + ["7", "later"], + ["2", "first"], + ]); + expect(sessions[1]?.observations).toHaveLength(2); + expect(sessions[0]?.revision).toMatch(/^[0-9a-f]{64}$/); +}); diff --git a/src/bundled/pi/usage.ts b/src/bundled/pi/usage.ts new file mode 100644 index 0000000000..a7dd6b39f9 --- /dev/null +++ b/src/bundled/pi/usage.ts @@ -0,0 +1,139 @@ +// The token use of an agent's Pi conversations, as pi-ai records it on each +// of their answers, read from the agent's Pi durable storage through the host. +// Pi owns the entries; nothing here writes or keeps them. +import { promptSections } from "../../features/agents/activity-transcript"; +import type { + SessionUsage, + UsageObservation, +} from "../../features/agents/service"; +import { withRevision } from "../../features/agents/session-usage"; +import type { Pi } from "./pi"; +import type { PiEntry } from "./transcript"; + +type Json = Record; +/** A Buzz conversation's Pi conversation, as `places` lists it. */ +type Place = Readonly<{ + conversation: number; + channelId: string; + root?: string; + at: number; +}>; + +/** pi-ai's counts, named as Buzz names them. */ +const TOKENS = [ + "input", + "output", + "cacheRead", + "cacheWrite", + "cacheWrite1h", +] as const; + +const object = (value: unknown): Json | undefined => + value !== null && typeof value === "object" && !Array.isArray(value) + ? (value as Json) + : undefined; +const str = (value: unknown) => (typeof value === "string" ? value : ""); +const count = (value: unknown) => + typeof value === "number" && Number.isFinite(value) && value >= 0 + ? value + : undefined; +const text = (content: unknown) => + typeof content === "string" + ? content + : (Array.isArray(content) ? content : []) + .map((part) => str(object(part)?.text)) + .join("\n"); + +/** What was asked, from the conversation's first prompt. */ +function title(entries: readonly PiEntry[]) { + const first = entries.find((entry) => entry.kind === "pi.user")?.model?.[0]; + return promptSections(text(first?.content)) + .find((section) => section.tag === "buzz-event") + ?.body.match(/^Content: (.*)$/m)?.[1] + ?.trim(); +} + +/** A conversation's model requests, its entries oldest first. pi-ai counts a + * request's input without its cached part, and prices the models it knows. */ +export function conversationUsage( + entries: readonly PiEntry[], + place: Omit, +): Omit | undefined { + const observations: UsageObservation[] = []; + const times: number[] = []; + for (const entry of entries) + for (const message of entry.model ?? []) { + if (typeof message.timestamp === "number") times.push(message.timestamp); + const usage = object(message.usage); + if (message.role !== "assistant" || !usage) continue; + const cost = count(object(usage.cost)?.total); + observations.push({ + ...(message.responseId ? { requestId: message.responseId } : {}), + at: message.timestamp ?? 0, + model: message.model || "pi", + tokens: Object.fromEntries( + TOKENS.flatMap((key) => { + const value = count(usage[key]); + return value === undefined ? [] : [[key, value]]; + }), + ), + // pi-ai prices a model it has no rates for at 0: that is unknown. + ...(cost ? { cost: { usd: cost, source: "pi" } } : {}), + }); + } + if (!times.length) return undefined; + const asked = title(entries); + return { + id: String(place.conversation), + channelId: place.channelId, + ...(place.root ? { root: place.root } : {}), + ...(asked ? { title: asked } : {}), + startedAt: Math.min(...times), + lastAt: Math.max(...times), + gaps: [], + observations, + }; +} + +/** All of a conversation's entries, oldest first. */ +async function entries( + pi: Pick, + agent: string, + conversation: number, + signal: AbortSignal, +) { + const found: PiEntry[] = []; + let cursor: unknown; + do { + const page = await pi.request<{ entries: PiEntry[]; next?: unknown }>( + "entries", + { agent, conversation, ...(cursor ? { cursor } : {}) }, + ); + signal.throwIfAborted(); + found.push(...page.entries); + cursor = page.next; + } while (cursor); + return found.reverse(); +} + +/** The agent's Buzz conversations used since `since`. A conversation's + * requests follow the newest message it was shown. */ +export async function readUsage( + pi: Pick, + agent: string, + places: readonly Place[], + since: number, + signal: AbortSignal, +): Promise { + const found: SessionUsage[] = []; + for (const place of places) { + if (place.at < since) continue; + const session = conversationUsage( + await entries(pi, agent, place.conversation, signal), + place, + ); + if (session && session.lastAt >= since) + found.push(await withRevision(session)); + } + return found.sort((a, b) => b.lastAt - a.lastAt); +}