feat(pi): emit live usage updates during turns (#3532)

Port the OMP usage poller into the Pi provider so context/token usage updates during active turns instead of only at completion. - Add PiUsagePoller with 3s throttle, dedup, and generation-based cancellation - Wire poller into PiRpcAgentSession turn lifecycle - Add unit tests for polling, dedup, error handling, and turn completion Co-authored-by: Slava Goltser <slavikg@gmail.com>
2026-08-19 07:42:19 +02:00
parent d10edbc5f8
commit 9c9a85a0cf
5 changed files with 565 additions and 51 deletions
@@ -19,14 +19,19 @@ import { describe, expect, onTestFinished, test } from "vitest";
 import type { AgentSession, AgentSessionConfig, AgentStreamEvent } from "../../agent-sdk-types.js";
 import { PiRpcAgentClient, PiRpcAgentSession, transformPiModels } from "./agent.js";
 import { FakePi } from "./test-utils/fake-pi.js";
+import type { PiUsagePollScheduler } from "./usage-poller.js";

 const ONE_BY_ONE_PNG_BASE64 =
  "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII=";

-function createClient(pi = new FakePi()): PiRpcAgentClient {
+function createClient(
+  pi = new FakePi(),
+  usagePollScheduler?: PiUsagePollScheduler,
+): PiRpcAgentClient {
  return new PiRpcAgentClient({
    logger: pino({ level: "silent" }),
    runtime: pi,
+    ...(usagePollScheduler ? { usagePollScheduler } : {}),
  });
 }

@@ -46,6 +51,28 @@ function createConfig(overrides: Partial<AgentSessionConfig> = {}): AgentSession
  };
 }

+class ManualUsagePollScheduler implements PiUsagePollScheduler {
+  private readonly polls: Array<{ active: boolean; callback: () => void }> = [];
+
+  schedulePoll(callback: () => void): () => void {
+    const poll = { active: true, callback };
+    this.polls.push(poll);
+    return () => {
+      poll.active = false;
+    };
+  }
+
+  poll(): void {
+    const poll = this.polls.shift();
+    if (!poll) throw new Error("Pi has not scheduled a context usage poll");
+    if (poll.active) poll.callback();
+  }
+
+  activePollCount(): number {
+    return this.polls.filter((poll) => poll.active).length;
+  }
+}
+
 function readUtf8File(pathname: string): string {
  const fd = openSync(pathname, "r");
  try {
@@ -89,12 +116,15 @@ async function flushTurnScheduling(): Promise<void> {
  await waitForImmediate();
 }

-async function createSession(pi = new FakePi()): Promise<{
+async function createSession(
+  pi = new FakePi(),
+  usagePollScheduler?: PiUsagePollScheduler,
+): Promise<{
  pi: FakePi;
  session: PiRpcAgentSession;
  events: SessionEvents;
 }> {
-  const client = createClient(pi);
+  const client = createClient(pi, usagePollScheduler);
  const session = (await client.createSession(createConfig())) as PiRpcAgentSession;
  const events = new SessionEvents(session);
  return { pi, session, events };
@@ -192,6 +222,13 @@ class SessionEvents {
    );
  }

+  usageUpdatedEvents() {
+    return this.events.filter(
+      (event): event is Extract<AgentStreamEvent, { type: "usage_updated" }> =>
+        event.type === "usage_updated",
+    );
+  }
+
  nextTurnCompletion(): Promise<Extract<AgentStreamEvent, { type: "turn_completed" }>> {
    return this.nextEvent(
      (event): event is Extract<AgentStreamEvent, { type: "turn_completed" }> =>
@@ -1313,6 +1350,145 @@ describe("PiRpcAgentSession", () => {
    expect(completion).toMatchObject({ type: "turn_completed", turnId });
    expect(events.turnCompletedEvents()).toHaveLength(1);
  });
+
+  test("emits usage_updated during an active turn and with the turn id at completion", async () => {
+    const scheduler = new ManualUsagePollScheduler();
+    const { pi, session, events } = await createSession(new FakePi(), scheduler);
+    const fakeSession = pi.latestSession();
+    fakeSession.stats = {
+      tokens: { input: 100, cacheRead: 10, output: 20 },
+      cost: 0.01,
+      contextUsage: { contextWindow: 200_000, tokens: 130 },
+    };
+
+    const { turnId } = await session.startTurn("hello");
+    scheduler.poll();
+    await flushTurnScheduling();
+
+    expect(events.usageUpdatedEvents()).toHaveLength(1);
+    expect(events.usageUpdatedEvents()[0]).not.toHaveProperty("turnId");
+    expect(events.usageUpdatedEvents()[0]).toMatchObject({
+      type: "usage_updated",
+      provider: "pi",
+      usage: {
+        inputTokens: 100,
+        cachedInputTokens: 10,
+        outputTokens: 20,
+        totalCostUsd: 0.01,
+        contextWindowMaxTokens: 200_000,
+        contextWindowUsedTokens: 130,
+      },
+    });
+
+    fakeSession.stats = {
+      ...fakeSession.stats,
+      contextUsage: { contextWindow: 200_000, tokens: 150 },
+    };
+    fakeSession.finishTurn();
+    await flushTurnScheduling();
+
+    expect(events.usageUpdatedEvents()).toHaveLength(2);
+    expect(events.usageUpdatedEvents()[1]).toMatchObject({
+      type: "usage_updated",
+      provider: "pi",
+      turnId,
+      usage: expect.objectContaining({ contextWindowUsedTokens: 150 }),
+    });
+    expect(events.turnCompletedEvents()).toHaveLength(1);
+  });
+
+  test("does not re-emit unchanged usage during a turn or at completion", async () => {
+    const scheduler = new ManualUsagePollScheduler();
+    const { pi, session, events } = await createSession(new FakePi(), scheduler);
+    const fakeSession = pi.latestSession();
+    fakeSession.stats = { contextUsage: { contextWindow: 200_000, tokens: 130 } };
+
+    await session.startTurn("hello");
+    scheduler.poll();
+    await flushTurnScheduling();
+    scheduler.poll();
+    await flushTurnScheduling();
+    expect(events.usageUpdatedEvents()).toHaveLength(1);
+
+    fakeSession.finishTurn();
+    await flushTurnScheduling();
+    expect(events.usageUpdatedEvents()).toHaveLength(1);
+    expect(events.turnCompletedEvents()).toHaveLength(1);
+  });
+
+  test("poll errors do not fail the turn and final usage still emits", async () => {
+    const scheduler = new ManualUsagePollScheduler();
+    const { pi, session, events } = await createSession(new FakePi(), scheduler);
+    const fakeSession = pi.latestSession();
+    fakeSession.stats = { contextUsage: { contextWindow: 200_000, tokens: 130 } };
+    fakeSession.getSessionStatsError = new Error("stats unavailable");
+
+    const { turnId } = await session.startTurn("hello");
+    scheduler.poll();
+    await flushTurnScheduling();
+    expect(events.turnCompletedEvents()).toHaveLength(0);
+    expect(events.usageUpdatedEvents()).toHaveLength(0);
+
+    fakeSession.getSessionStatsError = null;
+    fakeSession.finishTurn();
+    await flushTurnScheduling();
+    expect(events.turnCompletedEvents()).toHaveLength(1);
+    expect(events.usageUpdatedEvents()).toHaveLength(1);
+    expect(events.usageUpdatedEvents()[0]).toMatchObject({
+      turnId,
+      usage: { contextWindowUsedTokens: 130 },
+    });
+  });
+
+  test("stops scheduling polls after turn completion and close", async () => {
+    const scheduler = new ManualUsagePollScheduler();
+    const { pi, session } = await createSession(new FakePi(), scheduler);
+    const fakeSession = pi.latestSession();
+    fakeSession.stats = { contextUsage: { contextWindow: 200_000, tokens: 130 } };
+
+    await session.startTurn("hello");
+    expect(scheduler.activePollCount()).toBe(1);
+    fakeSession.finishTurn();
+    await flushTurnScheduling();
+    expect(scheduler.activePollCount()).toBe(0);
+
+    await session.startTurn("second");
+    expect(scheduler.activePollCount()).toBe(1);
+    await session.close();
+    expect(scheduler.activePollCount()).toBe(0);
+  });
+
+  test("dedupes unchanged usage across turns and emits when it changes", async () => {
+    const scheduler = new ManualUsagePollScheduler();
+    const { pi, session, events } = await createSession(new FakePi(), scheduler);
+    const fakeSession = pi.latestSession();
+    fakeSession.stats = { contextUsage: { contextWindow: 200_000, tokens: 130 } };
+
+    const first = await session.startTurn("first");
+    fakeSession.finishTurn();
+    await flushTurnScheduling();
+    expect(events.usageUpdatedEvents()).toHaveLength(1);
+    expect(events.usageUpdatedEvents()[0]).toMatchObject({
+      turnId: first.turnId,
+      usage: { contextWindowUsedTokens: 130 },
+    });
+
+    await session.startTurn("second");
+    fakeSession.finishTurn();
+    await flushTurnScheduling();
+    expect(events.usageUpdatedEvents()).toHaveLength(1);
+    expect(events.turnCompletedEvents()).toHaveLength(2);
+
+    fakeSession.stats = { contextUsage: { contextWindow: 200_000, tokens: 160 } };
+    const third = await session.startTurn("third");
+    fakeSession.finishTurn();
+    await flushTurnScheduling();
+    expect(events.usageUpdatedEvents()).toHaveLength(2);
+    expect(events.usageUpdatedEvents()[1]).toMatchObject({
+      turnId: third.turnId,
+      usage: { contextWindowUsedTokens: 160 },
+    });
+  });
 });

 describe("PiRpcAgentClient", () => {
@@ -29,7 +29,6 @@ import {
  type AgentSlashCommand,
  type AgentSlashCommandKind,
  type AgentStreamEvent,
-  type AgentUsage,
  type FetchCatalogOptions,
  type ImportableProviderSession,
  type ImportProviderSessionContext,
@@ -74,10 +73,10 @@ import type {
  PiModel,
  PiRpcSlashCommand,
  PiRuntimeEvent,
-  PiSessionStats,
  PiSessionState,
  PiThinkingLevel,
 } from "./rpc-types.js";
+import { PiUsagePoller, type PiUsagePollScheduler } from "./usage-poller.js";
 import {
  mapToolDetail,
  parseToolArgs,
@@ -184,6 +183,7 @@ export interface PiRpcAgentClientOptions {
  runtimeSettings?: ProviderRuntimeSettings;
  providerParams?: unknown;
  runtime?: PiRuntime;
+  usagePollScheduler?: PiUsagePollScheduler;
 }

 interface PiPromptPayload {
@@ -224,6 +224,8 @@ interface PiRpcAgentSessionOptions {
  currentModeId?: string | null;
  cleanup?: () => void;
  extensionTimeoutMs?: number;
+  logger: Logger;
+  usagePollScheduler?: PiUsagePollScheduler;
 }

 interface PiResumeConfig {
@@ -367,35 +369,6 @@ function mapThinkingOption(option: (typeof PI_THINKING_OPTIONS)[number]) {
  return mappedOption;
 }

-function toAgentUsage(stats: PiSessionStats): AgentUsage | undefined {
-  const inputTokens = stats.tokens?.input ?? 0;
-  const cachedInputTokens = stats.tokens?.cacheRead ?? 0;
-  const outputTokens = stats.tokens?.output ?? 0;
-  const totalCostUsd = stats.cost ?? 0;
-  const contextWindowMaxTokens = stats.contextUsage?.contextWindow ?? undefined;
-  const contextWindowUsedTokens = stats.contextUsage?.tokens ?? undefined;
-
-  if (
-    inputTokens === 0 &&
-    cachedInputTokens === 0 &&
-    outputTokens === 0 &&
-    totalCostUsd === 0 &&
-    contextWindowMaxTokens === undefined &&
-    contextWindowUsedTokens === undefined
-  ) {
-    return undefined;
-  }
-
-  return {
-    inputTokens,
-    cachedInputTokens,
-    outputTokens,
-    totalCostUsd,
-    ...(typeof contextWindowMaxTokens === "number" ? { contextWindowMaxTokens } : {}),
-    ...(typeof contextWindowUsedTokens === "number" ? { contextWindowUsedTokens } : {}),
-  };
-}
-
 function piModelSupportsImageInput(model: PiModel | null | undefined): boolean {
  return model?.input?.includes("image") === true;
 }
@@ -1246,6 +1219,8 @@ export class PiRpcAgentSession implements AgentSession {
  private commandCache: AgentSlashCommand[] | null = null;
  private state: PiSessionState;
  private readonly currentModeId: string | null;
+  private readonly logger: Logger;
+  private readonly usagePoller: PiUsagePoller;
  private closed = false;
  // Pi reports an aborted OpenAI Responses stream before the abort RPC resolves.
  // Keep the turn active until that RPC acknowledges the user-requested cancellation.
@@ -1266,6 +1241,22 @@ export class PiRpcAgentSession implements AgentSession {
      this.state.thinkingLevel ??
      null;
    this.extensionTimeoutMs = options.extensionTimeoutMs ?? DEFAULT_PI_EXTENSION_RESULT_TIMEOUT_MS;
+    this.logger = options.logger;
+    this.usagePoller = new PiUsagePoller({
+      scheduler: options.usagePollScheduler,
+      readStats: () => this.runtimeSession.getSessionStats(),
+      onUsage: (usage, turnId) => {
+        this.emit({
+          type: "usage_updated",
+          provider: this.provider,
+          usage,
+          ...(turnId === undefined ? {} : { turnId }),
+        });
+      },
+      onPollError: (error) => {
+        this.logger.debug({ err: error }, "Pi context usage poll failed");
+      },
+    });

    this.runtimeSession.onEvent((event) => {
      this.handleRuntimeEvent(event);
@@ -1301,6 +1292,7 @@ export class PiRpcAgentSession implements AgentSession {
    const payload = convertPromptInput(prompt, { model: this.state.model });
    const turnId = randomUUID();
    this.activeTurnId = turnId;
+    this.usagePoller.startTurn();
    this.lastInterruptedTurnId = null;
    this.activeClientMessageId = options?.clientMessageId ?? null;
    this.activeAssistantMessageId = null;
@@ -1332,6 +1324,7 @@ export class PiRpcAgentSession implements AgentSession {
        if (this.activeTurnId !== turnId) {
          return;
        }
+        this.usagePoller.stopTurn();
        this.activeTurnId = null;
        this.activeClientMessageId = null;
        this.activeTurnStarted = false;
@@ -1459,6 +1452,7 @@ export class PiRpcAgentSession implements AgentSession {
      if (this.interruptedTerminalError?.turnId === turnId) {
        const terminalError = this.interruptedTerminalError;
        this.interruptedTerminalError = null;
+        this.usagePoller.stopTurn();
        this.activeTurnId = null;
        this.activeClientMessageId = null;
        this.activeTurnStarted = false;
@@ -1474,6 +1468,7 @@ export class PiRpcAgentSession implements AgentSession {
      throw error;
    }
    if (turnId && this.activeTurnId === turnId) {
+      this.usagePoller.stopTurn();
      this.activeTurnId = null;
      this.activeClientMessageId = null;
      this.activeTurnStarted = false;
@@ -1527,6 +1522,7 @@ export class PiRpcAgentSession implements AgentSession {
      return;
    }
    this.closed = true;
+    this.usagePoller.close();
    try {
      await this.runtimeSession.close();
    } finally {
@@ -2047,6 +2043,7 @@ export class PiRpcAgentSession implements AgentSession {
      return;
    }
    const turnId = this.activeTurnId;
+    this.usagePoller.stopTurn();
    this.activeTurnId = null;
    this.activeClientMessageId = null;
    this.activeTurnStarted = false;
@@ -2311,6 +2308,7 @@ export class PiRpcAgentSession implements AgentSession {
    this.clearNoTurnBuffers();
    const errorMessage = latestPiErrorMessage(messages);
    if (typeof errorMessage === "string" && errorMessage.length > 0) {
+      this.usagePoller.stopTurn();
      this.emit({
        type: "turn_failed",
        provider: this.provider,
@@ -2319,34 +2317,21 @@ export class PiRpcAgentSession implements AgentSession {
      });
      return;
    }
+    const finalUsage = this.usagePoller.completeTurn(turnId);
    this.emit({
      type: "turn_completed",
      provider: this.provider,
      turnId,
    });
-    void this.refreshAfterTurn(turnId);
+    void this.refreshAfterTurn(finalUsage);
  }

  private async refreshState(): Promise<void> {
    this.state = await this.runtimeSession.getState();
  }

-  private async refreshAfterTurn(turnId: string | undefined): Promise<void> {
-    await this.refreshState().catch(() => undefined);
-    const usage = await this.runtimeSession
-      .getSessionStats()
-      .then((stats) => {
-        return toAgentUsage(stats);
-      })
-      .catch(() => undefined);
-    if (usage) {
-      this.emit({
-        type: "usage_updated",
-        provider: this.provider,
-        turnId,
-        usage,
-      });
-    }
+  private async refreshAfterTurn(finalUsage: Promise<void>): Promise<void> {
+    await Promise.all([this.refreshState().catch(() => undefined), finalUsage]);
  }
 }

@@ -2358,6 +2343,7 @@ export class PiRpcAgentClient implements AgentClient {
  private readonly runtimeSettings?: ProviderRuntimeSettings;
  private readonly providerParams: PiProviderParams;
  private readonly runtime: PiRuntime;
+  private readonly usagePollScheduler?: PiUsagePollScheduler;

  constructor(options: PiRpcAgentClientOptions) {
    this.provider = PI_PROVIDER;
@@ -2366,6 +2352,7 @@ export class PiRpcAgentClient implements AgentClient {
    this.runtimeSettings = options.runtimeSettings;
    this.providerParams = PiProviderParamsSchema.parse(options.providerParams ?? {});
    this.runtime = options.runtime ?? createRuntime(options.logger, options.runtimeSettings);
+    this.usagePollScheduler = options.usagePollScheduler;
  }

  async createSession(
@@ -2405,6 +2392,8 @@ export class PiRpcAgentClient implements AgentClient {
        capabilities: capabilitiesForSession(mcpConfig !== null),
        cleanup: combineCleanup([mcpConfig?.cleanup, paseoExtension?.cleanup]),
        extensionTimeoutMs: this.providerParams.extensionTimeoutMs,
+        logger: this.logger,
+        usagePollScheduler: this.usagePollScheduler,
      });
    } catch (error) {
      await runtimeSession.close().catch(() => undefined);
@@ -2466,6 +2455,8 @@ export class PiRpcAgentClient implements AgentClient {
        capabilities: capabilitiesForSession(mcpConfig !== null),
        cleanup: combineCleanup([mcpConfig?.cleanup, paseoExtension?.cleanup]),
        extensionTimeoutMs: this.providerParams.extensionTimeoutMs,
+        logger: this.logger,
+        usagePollScheduler: this.usagePollScheduler,
      });
    } catch (error) {
      await runtimeSession.close().catch(() => undefined);
@@ -125,6 +125,7 @@ export class FakePiSession implements PiRuntimeSession {
  compactError: Error | null = null;
  emitCompactEnd = true;
  getStateError: Error | null = null;
+  getSessionStatsError: Error | null = null;
  promptAck: PiPromptAck = {};
  branchResponse: { text?: string; cancelled?: boolean } = { text: "" };
  readonly branchRequests: string[] = [];
@@ -250,6 +251,9 @@ export class FakePiSession implements PiRuntimeSession {
  }

  async getSessionStats(): Promise<PiSessionStats> {
+    if (this.getSessionStatsError) {
+      throw this.getSessionStatsError;
+    }
    return this.stats;
  }

@@ -0,0 +1,174 @@
+import { setImmediate as waitForImmediate } from "node:timers/promises";
+import { describe, expect, test } from "vitest";
+
+import type { AgentUsage } from "../../agent-sdk-types.js";
+import type { PiSessionStats } from "./rpc-types.js";
+import { PiUsagePoller, type PiUsagePollScheduler } from "./usage-poller.js";
+
+class ManualPollScheduler implements PiUsagePollScheduler {
+  private readonly polls: Array<{ active: boolean; callback: () => void }> = [];
+
+  schedulePoll(callback: () => void): () => void {
+    const poll = { active: true, callback };
+    this.polls.push(poll);
+    return () => {
+      poll.active = false;
+    };
+  }
+
+  poll(): void {
+    const poll = this.polls.shift();
+    if (!poll) throw new Error("No context usage poll is scheduled");
+    if (poll.active) poll.callback();
+  }
+
+  activePollCount(): number {
+    return this.polls.filter((poll) => poll.active).length;
+  }
+}
+
+function deferred<T>() {
+  let resolve!: (value: T) => void;
+  const promise = new Promise<T>((resolvePromise) => {
+    resolve = resolvePromise;
+  });
+  return { promise, resolve };
+}
+
+describe("Pi usage poller", () => {
+  test("emits only changed usage while active", async () => {
+    const scheduler = new ManualPollScheduler();
+    const updates: AgentUsage[] = [];
+    let stats: PiSessionStats = {
+      tokens: { input: 100, cacheRead: 10, output: 20 },
+      cost: 0.01,
+      contextUsage: { contextWindow: 200_000, tokens: 130 },
+    };
+    const poller = new PiUsagePoller({
+      scheduler,
+      readStats: async () => stats,
+      onUsage: (update) => updates.push(update),
+      onPollError: (error) => {
+        throw error;
+      },
+    });
+
+    poller.startTurn();
+    scheduler.poll();
+    await waitForImmediate();
+    scheduler.poll();
+    await waitForImmediate();
+    stats = { ...stats, contextUsage: { contextWindow: 200_000, tokens: 150 } };
+    scheduler.poll();
+    await waitForImmediate();
+
+    expect(updates).toEqual([
+      {
+        inputTokens: 100,
+        cachedInputTokens: 10,
+        outputTokens: 20,
+        totalCostUsd: 0.01,
+        contextWindowMaxTokens: 200_000,
+        contextWindowUsedTokens: 130,
+      },
+      {
+        inputTokens: 100,
+        cachedInputTokens: 10,
+        outputTokens: 20,
+        totalCostUsd: 0.01,
+        contextWindowMaxTokens: 200_000,
+        contextWindowUsedTokens: 150,
+      },
+    ]);
+    expect(scheduler.activePollCount()).toBe(1);
+    poller.stopTurn();
+    expect(scheduler.activePollCount()).toBe(0);
+  });
+
+  test("drops an in-flight poll without hiding the final refresh", async () => {
+    const scheduler = new ManualPollScheduler();
+    const inFlightStats = deferred<PiSessionStats>();
+    const finalStats = { contextUsage: { contextWindow: 200_000, tokens: 150 } };
+    const finalUsage = {
+      inputTokens: 0,
+      cachedInputTokens: 0,
+      outputTokens: 0,
+      totalCostUsd: 0,
+      contextWindowMaxTokens: 200_000,
+      contextWindowUsedTokens: 150,
+    };
+    const updates: Array<{ usage: AgentUsage; turnId?: string }> = [];
+    let readCount = 0;
+    const poller = new PiUsagePoller({
+      scheduler,
+      readStats: async () => {
+        readCount += 1;
+        return readCount === 1 ? inFlightStats.promise : finalStats;
+      },
+      onUsage: (usage, turnId) => updates.push({ usage, turnId }),
+      onPollError: (error) => {
+        throw error;
+      },
+    });
+
+    poller.startTurn();
+    scheduler.poll();
+    await expect(poller.completeTurn("turn-1")).resolves.toBeUndefined();
+    inFlightStats.resolve({
+      contextUsage: { contextWindow: 200_000, tokens: 130 },
+    });
+    await waitForImmediate();
+
+    expect(updates).toEqual([{ usage: finalUsage, turnId: "turn-1" }]);
+    await expect(poller.completeTurn()).resolves.toBeUndefined();
+    expect(updates).toEqual([{ usage: finalUsage, turnId: "turn-1" }]);
+  });
+
+  test("drops a final refresh when a newer turn starts", async () => {
+    const scheduler = new ManualPollScheduler();
+    const finalStats = deferred<PiSessionStats>();
+    const updates: AgentUsage[] = [];
+    const poller = new PiUsagePoller({
+      scheduler,
+      readStats: () => finalStats.promise,
+      onUsage: (usage) => updates.push(usage),
+      onPollError: (error) => {
+        throw error;
+      },
+    });
+
+    poller.startTurn();
+    const completion = poller.completeTurn();
+    poller.startTurn();
+    finalStats.resolve({ contextUsage: { contextWindow: 200_000, tokens: 150 } });
+
+    await expect(completion).resolves.toBeUndefined();
+    expect(updates).toEqual([]);
+    expect(scheduler.activePollCount()).toBe(1);
+    poller.stopTurn();
+  });
+
+  test("close permanently suppresses pending usage", async () => {
+    const scheduler = new ManualPollScheduler();
+    const finalStats = deferred<PiSessionStats>();
+    const updates: AgentUsage[] = [];
+    const poller = new PiUsagePoller({
+      scheduler,
+      readStats: () => finalStats.promise,
+      onUsage: (usage) => updates.push(usage),
+      onPollError: (error) => {
+        throw error;
+      },
+    });
+
+    poller.startTurn();
+    const completion = poller.completeTurn();
+    poller.close();
+    poller.startTurn();
+    finalStats.resolve({ contextUsage: { contextWindow: 200_000, tokens: 150 } });
+
+    await expect(completion).resolves.toBeUndefined();
+    expect(updates).toEqual([]);
+    expect(scheduler.activePollCount()).toBe(0);
+  });
+});
@@ -0,0 +1,169 @@
+import type { AgentUsage } from "../../agent-sdk-types.js";
+import type { PiSessionStats } from "./rpc-types.js";
+
+export interface PiUsagePollScheduler {
+  schedulePoll(callback: () => void): () => void;
+}
+
+interface PiUsagePollerOptions {
+  readStats(): Promise<PiSessionStats>;
+  onUsage(usage: AgentUsage, turnId?: string): void;
+  onPollError(error: unknown): void;
+  scheduler?: PiUsagePollScheduler;
+}
+
+function createScheduler(): PiUsagePollScheduler {
+  return {
+    schedulePoll: (callback) => {
+      const timer = setTimeout(callback, 3_000);
+      return () => clearTimeout(timer);
+    },
+  };
+}
+
+function toAgentUsage(stats: PiSessionStats): AgentUsage | undefined {
+  const inputTokens = stats.tokens?.input ?? 0;
+  const cachedInputTokens = stats.tokens?.cacheRead ?? 0;
+  const outputTokens = stats.tokens?.output ?? 0;
+  const totalCostUsd = stats.cost ?? 0;
+  const contextWindowMaxTokens = stats.contextUsage?.contextWindow ?? undefined;
+  const contextWindowUsedTokens = stats.contextUsage?.tokens ?? undefined;
+
+  if (
+    inputTokens === 0 &&
+    cachedInputTokens === 0 &&
+    outputTokens === 0 &&
+    totalCostUsd === 0 &&
+    contextWindowMaxTokens === undefined &&
+    contextWindowUsedTokens === undefined
+  ) {
+    return undefined;
+  }
+
+  return {
+    inputTokens,
+    cachedInputTokens,
+    outputTokens,
+    totalCostUsd,
+    ...(typeof contextWindowMaxTokens === "number" ? { contextWindowMaxTokens } : {}),
+    ...(typeof contextWindowUsedTokens === "number" ? { contextWindowUsedTokens } : {}),
+  };
+}
+
+function isSameUsage(left: AgentUsage, right: AgentUsage): boolean {
+  return (
+    left.inputTokens === right.inputTokens &&
+    left.cachedInputTokens === right.cachedInputTokens &&
+    left.outputTokens === right.outputTokens &&
+    left.totalCostUsd === right.totalCostUsd &&
+    left.contextWindowMaxTokens === right.contextWindowMaxTokens &&
+    left.contextWindowUsedTokens === right.contextWindowUsedTokens
+  );
+}
+
+export class PiUsagePoller {
+  private readonly scheduler: PiUsagePollScheduler;
+  private active = false;
+  private closed = false;
+  private generation = 0;
+  private cancelScheduledPoll: (() => void) | null = null;
+  private lastUsage: AgentUsage | null = null;
+
+  constructor(private readonly options: PiUsagePollerOptions) {
+    this.scheduler = options.scheduler ?? createScheduler();
+  }
+
+  startTurn(): void {
+    if (this.closed || this.active) {
+      return;
+    }
+    this.active = true;
+    this.generation += 1;
+    this.schedule(this.generation);
+  }
+
+  stopTurn(): void {
+    this.active = false;
+    this.invalidatePendingWork();
+  }
+
+  async completeTurn(turnId?: string): Promise<void> {
+    if (this.closed) {
+      return;
+    }
+    this.active = false;
+    const completionGeneration = this.invalidatePendingWork();
+    let usage: AgentUsage | undefined;
+    try {
+      usage = toAgentUsage(await this.options.readStats());
+    } catch {
+      return;
+    }
+    if (this.closed || this.generation !== completionGeneration) {
+      return;
+    }
+    this.publishUsage(usage, turnId);
+  }
+
+  close(): void {
+    this.closed = true;
+    this.active = false;
+    this.invalidatePendingWork();
+  }
+
+  private invalidatePendingWork(): number {
+    this.generation += 1;
+    this.cancelScheduledPoll?.();
+    this.cancelScheduledPoll = null;
+    return this.generation;
+  }
+
+  private schedule(generation: number): void {
+    this.cancelScheduledPoll = this.scheduler.schedulePoll(() => {
+      this.cancelScheduledPoll = null;
+      void this.poll(generation);
+    });
+  }
+
+  private async poll(generation: number): Promise<void> {
+    let usage: AgentUsage | undefined;
+    try {
+      usage = toAgentUsage(await this.options.readStats());
+    } catch (error) {
+      try {
+        if (this.active && !this.closed && this.generation === generation) {
+          this.options.onPollError(error);
+        }
+      } finally {
+        if (this.active && !this.closed && this.generation === generation) {
+          this.schedule(generation);
+        }
+      }
+      return;
+    }
+    try {
+      if (this.active && !this.closed && this.generation === generation) {
+        this.publishUsage(usage);
+      }
+    } finally {
+      if (this.active && !this.closed && this.generation === generation) {
+        this.schedule(generation);
+      }
+    }
+  }
+
+  private takeChangedUsage(usage: AgentUsage | undefined): AgentUsage | undefined {
+    if (!usage || (this.lastUsage && isSameUsage(this.lastUsage, usage))) {
+      return undefined;
+    }
+    this.lastUsage = usage;
+    return usage;
+  }
+
+  private publishUsage(usage: AgentUsage | undefined, turnId?: string): void {
+    const changedUsage = this.takeChangedUsage(usage);
+    if (changedUsage) {
+      this.options.onUsage(changedUsage, turnId);
+    }
+  }
+}