From 9c9a85a0cf37e78b44d4d6a0ef81ea7f6aca266b Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Wed, 19 Aug 2026 07:42:19 +0200 Subject: [PATCH] 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 --- .../server/agent/providers/pi/agent.test.ts | 182 +++++++++++++++++- .../src/server/agent/providers/pi/agent.ts | 87 ++++----- .../agent/providers/pi/test-utils/fake-pi.ts | 4 + .../agent/providers/pi/usage-poller.test.ts | 174 +++++++++++++++++ .../server/agent/providers/pi/usage-poller.ts | 169 ++++++++++++++++ 5 files changed, 565 insertions(+), 51 deletions(-) create mode 100644 packages/server/src/server/agent/providers/pi/usage-poller.test.ts create mode 100644 packages/server/src/server/agent/providers/pi/usage-poller.ts diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index c367b57ab..378b556ff 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -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 = {}): 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 { 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 => + event.type === "usage_updated", + ); + } + nextTurnCompletion(): Promise> { return this.nextEvent( (event): event is Extract => @@ -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", () => { diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 82bd0a823..a30e46c8b 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -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 { this.state = await this.runtimeSession.getState(); } - private async refreshAfterTurn(turnId: string | undefined): Promise { - 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): Promise { + 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); diff --git a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts index bf0f5907d..c931a853e 100644 --- a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts +++ b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts @@ -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 { + if (this.getSessionStatsError) { + throw this.getSessionStatsError; + } return this.stats; } diff --git a/packages/server/src/server/agent/providers/pi/usage-poller.test.ts b/packages/server/src/server/agent/providers/pi/usage-poller.test.ts new file mode 100644 index 000000000..e59b0d1f0 --- /dev/null +++ b/packages/server/src/server/agent/providers/pi/usage-poller.test.ts @@ -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() { + let resolve!: (value: T) => void; + const promise = new Promise((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(); + 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(); + 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(); + 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); + }); +}); diff --git a/packages/server/src/server/agent/providers/pi/usage-poller.ts b/packages/server/src/server/agent/providers/pi/usage-poller.ts new file mode 100644 index 000000000..5fcbe2f02 --- /dev/null +++ b/packages/server/src/server/agent/providers/pi/usage-poller.ts @@ -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; + 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 { + 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 { + 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); + } + } +}