feat(pi): show Pi agents' usage on the Usage page

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"],
     ]);

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,

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),

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);

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<{

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: `<context>\nChannel: general (#chan)\n</context>\n\n<buzz-event type="@mention">\nFrom: Sal (npub1x)\nContent: ${text}\nTags: []\n</buzz-event>`,
+      timestamp: id * 1000,
+    },
+  ],
+});
+const answer = (id: number, usage: Record<string, unknown>, 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<string, { entries: PiEntry[]; next?: string }> = {
+    "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}$/);
+});

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<string, unknown>;
+/** 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<Place, "at">,
+): Omit<SessionUsage, "revision"> | 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<Pi, "request">,
+  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<Pi, "request">,
+  agent: string,
+  places: readonly Place[],
+  since: number,
+  signal: AbortSignal,
+): Promise<SessionUsage[]> {
+  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);
+}