From 2948a435ed77a6c5edecee8a8857976b2a3445cf Mon Sep 17 00:00:00 2001 From: Slava Goltser Date: Mon, 3 Aug 2026 06:35:23 -0400 Subject: [PATCH] feat(omp): update context usage during active turns (#2503) * feat: add DeepWiki MCP server at project level * feat(omp): poll context window usage during active turns OMP agents previously reported context window usage only at turn completion, unlike Claude and Codex which stream live updates. Add a periodic poll (3s interval) of getSessionStats() that emits usage_updated events while a turn is active, so the client context window meter updates in real-time. * clear context usage poll timer before scheduling next interval * chore: remove .omp/ directory * fix(omp): make live usage polling lifecycle-safe OMP exposes live usage through a stats request rather than stream events. Keep polling scoped to active turns, suppress unchanged snapshots, and invalidate stale completion reads across interrupts, new turns, and session shutdown. --------- Co-authored-by: Mohamed Boudra --- .../server/agent/providers/omp/agent.test.ts | 63 +++++++ .../src/server/agent/providers/omp/agent.ts | 90 ++++----- .../providers/omp/test-utils/fake-omp.ts | 4 + .../providers/omp/test-utils/omp-harness.ts | 7 + .../agent/providers/omp/usage-mapper.test.ts | 39 ---- .../agent/providers/omp/usage-mapper.ts | 24 --- .../agent/providers/omp/usage-poller.test.ts | 174 ++++++++++++++++++ .../agent/providers/omp/usage-poller.ts | 169 +++++++++++++++++ 8 files changed, 452 insertions(+), 118 deletions(-) delete mode 100644 packages/server/src/server/agent/providers/omp/usage-mapper.test.ts delete mode 100644 packages/server/src/server/agent/providers/omp/usage-mapper.ts create mode 100644 packages/server/src/server/agent/providers/omp/usage-poller.test.ts create mode 100644 packages/server/src/server/agent/providers/omp/usage-poller.ts diff --git a/packages/server/src/server/agent/providers/omp/agent.test.ts b/packages/server/src/server/agent/providers/omp/agent.test.ts index d19771447..30876ae9a 100644 --- a/packages/server/src/server/agent/providers/omp/agent.test.ts +++ b/packages/server/src/server/agent/providers/omp/agent.test.ts @@ -1,7 +1,9 @@ import { describe, expect, test } from "vitest"; +import { setImmediate as waitForImmediate } from "node:timers/promises"; import type { PaseoToolCatalog } from "../../tools/types.js"; import type { OmpNoTurnScheduler, OmpProviderIdleScheduler } from "./agent.js"; +import type { OmpUsagePollScheduler } from "./usage-poller.js"; import { OmpHarness } from "./test-utils/omp-harness.js"; class ManualIdleScheduler implements OmpProviderIdleScheduler { @@ -65,6 +67,28 @@ class ManualNoTurnScheduler implements OmpNoTurnScheduler { } } +class ManualUsagePollScheduler implements OmpUsagePollScheduler { + 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("OMP has not scheduled a context usage poll"); + if (poll.active) poll.callback(); + } + + activePollCount(): number { + return this.polls.filter((poll) => poll.active).length; + } +} + function createToolCatalog(): PaseoToolCatalog { return { tools: new Map([ @@ -209,6 +233,45 @@ describe("OMP agent client and session", () => { expect(omp.completedTurnCount()).toBe(1); }); + test("starts and stops context usage polling with the active turn", async () => { + const scheduler = new ManualUsagePollScheduler(); + const omp = new OmpHarness({ usagePollScheduler: scheduler }); + await omp.start(); + omp.runtime().stats = { + contextUsage: { tokens: 130, contextWindow: 200_000 }, + }; + omp.runtime().state.contextUsage = { tokens: 99, contextWindow: 100_000 }; + await omp.requireStartTurn("keep working"); + expect(scheduler.activePollCount()).toBe(1); + scheduler.poll(); + await waitForImmediate(); + expect(omp.usageUpdates()).toEqual([ + { + inputTokens: 0, + cachedInputTokens: 0, + outputTokens: 0, + totalCostUsd: 0, + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 130, + }, + ]); + expect(scheduler.activePollCount()).toBe(1); + omp.runtime().abortError = new Error("abort unavailable"); + await expect(omp.interrupt()).rejects.toThrow("abort unavailable"); + expect(scheduler.activePollCount()).toBe(1); + omp.runtime().abortError = null; + await omp.interrupt(); + expect(scheduler.activePollCount()).toBe(0); + + await omp.runPrompt("finish normally", "done"); + expect(scheduler.activePollCount()).toBe(0); + + await omp.requireStartTurn("close the session"); + expect(scheduler.activePollCount()).toBe(1); + await omp.close(); + expect(scheduler.activePollCount()).toBe(0); + }); + test("does not accept a follow-up until OMP reports stable idle", async () => { const omp = new OmpHarness(); await omp.start(); diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index 01ad77a36..f3f6c7f97 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -27,7 +27,6 @@ import { type AgentSlashCommand, type AgentStreamEvent, type AgentTimelineItem, - type AgentUsage, type FetchCatalogOptions, type ImportableProviderSession, type ImportProviderSessionContext, @@ -79,7 +78,6 @@ import type { OmpImageContent, OmpModel, OmpRuntimeEvent, - OmpSessionStats, OmpSessionState, OmpThinkingLevel, } from "./rpc-types.js"; @@ -102,7 +100,7 @@ import { } from "./host-tools.js"; import { OmpSubagentIndex } from "./subagent-index.js"; import { mapOmpToolDetail } from "./tool-call-mapper.js"; -import { mapOmpUsage } from "./usage-mapper.js"; +import { OmpUsagePoller, type OmpUsagePollScheduler } from "./usage-poller.js"; import { buildOmpRpcUiPermissionResponse, mapOmpRpcUiPermissionRequest, @@ -137,6 +135,7 @@ export interface OmpAgentClientOptions { subagentCardScheduler?: OmpSubagentCardScheduler; providerIdleScheduler?: OmpProviderIdleScheduler; noTurnScheduler?: OmpNoTurnScheduler; + usagePollScheduler?: OmpUsagePollScheduler; } export interface OmpProviderIdleScheduler { @@ -184,6 +183,7 @@ interface OmpAgentSessionOptions { subagentCardScheduler?: OmpSubagentCardScheduler; providerIdleScheduler?: OmpProviderIdleScheduler; noTurnScheduler?: OmpNoTurnScheduler; + usagePollScheduler?: OmpUsagePollScheduler; paseoTools?: PaseoToolCatalog; /** * When false (resumed sessions), replayed session events are dropped until @@ -305,35 +305,6 @@ function parseAutoCompactMode(value: string | undefined): AutoCompactMode { return "unknown"; } -function toAgentUsage(stats: OmpSessionStats): 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 ompModelSupportsImageInput(model: OmpModel | null | undefined): boolean { return model?.input?.includes("image") === true; } @@ -898,6 +869,7 @@ export class OmpAgentSession implements AgentSession { private readonly currentModeId: string | null; private readonly providerIdleScheduler: OmpProviderIdleScheduler; private readonly noTurnScheduler: OmpNoTurnScheduler; + private readonly usagePoller: OmpUsagePoller; private closed = false; private live: boolean; private readonly emittedUserMessageIds = new Set(); @@ -912,6 +884,21 @@ export class OmpAgentSession implements AgentSession { this.live = options.live ?? true; this.providerIdleScheduler = options.providerIdleScheduler ?? createOmpProviderIdleScheduler(); this.noTurnScheduler = options.noTurnScheduler ?? createOmpNoTurnScheduler(); + this.usagePoller = new OmpUsagePoller({ + 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 }, "OMP context usage poll failed"); + }, + }); this.subagentCardTracker = new OmpSubagentCardTracker({ scheduler: options.subagentCardScheduler, }); @@ -976,6 +963,7 @@ export class OmpAgentSession implements AgentSession { this.activePromptRequestId = null; this.clearNoTurnBuffers(); this.activeNoTurnPromptText = payload.text; + this.usagePoller.startTurn(); void (async () => { try { @@ -1000,6 +988,7 @@ export class OmpAgentSession implements AgentSession { if (this.activeTurnId !== turnId) { return; } + this.usagePoller.stopTurn(); this.activeTurnId = null; this.activeClientMessageId = null; this.activeTurnStarted = false; @@ -1129,12 +1118,10 @@ export class OmpAgentSession implements AgentSession { async interrupt(): Promise { const turnId = this.activeTurnId; - try { - await this.runtimeSession.abort(); - } finally { - this.terminalizeActiveWork(); - } + await this.runtimeSession.abort(); if (turnId && this.activeTurnId === turnId) { + this.terminalizeActiveWork(); + this.usagePoller.stopTurn(); this.activeTurnId = null; this.activeClientMessageId = null; this.activeTurnStarted = false; @@ -1169,6 +1156,7 @@ export class OmpAgentSession implements AgentSession { return; } this.closed = true; + this.usagePoller.close(); this.cancelNoTurnPromptCompletion(); try { await this.runtimeSession.close(); @@ -1786,6 +1774,7 @@ export class OmpAgentSession implements AgentSession { } private handleProcessExit(error: string): void { + this.usagePoller.stopTurn(); this.terminalizeActiveWork(); this.subagentIndex.clear(this.runtimeSession); if (!this.activeTurnId) { @@ -2142,6 +2131,7 @@ export class OmpAgentSession implements AgentSession { this.clearNoTurnBuffers(); const errorMessage = latestOmpErrorMessage(messages); if (typeof errorMessage === "string" && errorMessage.length > 0) { + this.usagePoller.stopTurn(); this.emit({ type: "turn_failed", provider: this.provider, @@ -2150,12 +2140,13 @@ export class OmpAgentSession 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 completeTurnAfterProviderIdle( @@ -2181,23 +2172,8 @@ export class OmpAgentSession implements AgentSession { 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) => { - const baseUsage = toAgentUsage(stats); - return mapOmpUsage({ stats, state: this.state, baseUsage }); - }) - .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]); } } @@ -2212,6 +2188,7 @@ export class OmpAgentClient implements AgentClient { private readonly subagentCardScheduler?: OmpSubagentCardScheduler; private readonly providerIdleScheduler?: OmpProviderIdleScheduler; private readonly noTurnScheduler?: OmpNoTurnScheduler; + private readonly usagePollScheduler?: OmpUsagePollScheduler; private readonly runtime: OmpRuntime; constructor(options: OmpAgentClientOptions) { @@ -2234,6 +2211,7 @@ export class OmpAgentClient implements AgentClient { this.subagentCardScheduler = options.subagentCardScheduler; this.providerIdleScheduler = options.providerIdleScheduler; this.noTurnScheduler = options.noTurnScheduler; + this.usagePollScheduler = options.usagePollScheduler; this.runtime = options.runtime ?? createRuntime(options.logger, runtimeSettings); } @@ -2274,6 +2252,7 @@ export class OmpAgentClient implements AgentClient { subagentCardScheduler: this.subagentCardScheduler, providerIdleScheduler: this.providerIdleScheduler, noTurnScheduler: this.noTurnScheduler, + usagePollScheduler: this.usagePollScheduler, paseoTools: launchContext?.paseoTools, }); } catch (error) { @@ -2315,6 +2294,7 @@ export class OmpAgentClient implements AgentClient { subagentCardScheduler: this.subagentCardScheduler, providerIdleScheduler: this.providerIdleScheduler, noTurnScheduler: this.noTurnScheduler, + usagePollScheduler: this.usagePollScheduler, paseoTools: launchContext?.paseoTools, live: false, }); diff --git a/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts b/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts index 31cd6fd0f..6d46c50f1 100644 --- a/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts +++ b/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts @@ -112,6 +112,7 @@ export class FakeOmpSession implements OmpRuntimeSession { readonly hostToolUpdates: OmpRpcHostToolUpdate[] = []; getStateRequestCount = 0; abortRequested = false; + abortError: Error | null = null; readonly canceledExtensionUiRequests: string[] = []; readonly extensionUiResponses: Array<{ id: string; @@ -238,6 +239,9 @@ export class FakeOmpSession implements OmpRuntimeSession { } async abort(): Promise { + if (this.abortError) { + throw this.abortError; + } this.abortRequested = true; } diff --git a/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts b/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts index b3efb6101..30d75bb8f 100644 --- a/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts +++ b/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts @@ -18,6 +18,7 @@ import { type OmpNoTurnScheduler, type OmpProviderIdleScheduler, } from "../agent.js"; +import type { OmpUsagePollScheduler } from "../usage-poller.js"; import type { OmpAgentMessage, OmpRpcSlashCommand } from "../rpc-types.js"; import { FakeOmp } from "./fake-omp.js"; @@ -69,6 +70,7 @@ export class OmpHarness { options: { providerIdleScheduler?: OmpProviderIdleScheduler; noTurnScheduler?: OmpNoTurnScheduler; + usagePollScheduler?: OmpUsagePollScheduler; } = {}, ) { this.client = new OmpAgentClient({ @@ -76,6 +78,7 @@ export class OmpHarness { runtime: this.omp, providerIdleScheduler: options.providerIdleScheduler, noTurnScheduler: options.noTurnScheduler, + usagePollScheduler: options.usagePollScheduler, }); } @@ -407,6 +410,10 @@ export class OmpHarness { return this.events.filter((event) => event.type === "turn_completed").length; } + usageUpdates() { + return this.events.flatMap((event) => (event.type === "usage_updated" ? [event.usage] : [])); + } + requestToolApproval(input: { id: string; tool: "bash" | "edit" | "write"; diff --git a/packages/server/src/server/agent/providers/omp/usage-mapper.test.ts b/packages/server/src/server/agent/providers/omp/usage-mapper.test.ts deleted file mode 100644 index f189d2db3..000000000 --- a/packages/server/src/server/agent/providers/omp/usage-mapper.test.ts +++ /dev/null @@ -1,39 +0,0 @@ -import { describe, expect, test } from "vitest"; - -import { mapOmpUsage } from "./usage-mapper.js"; - -describe("OMP usage mapper", () => { - test("adds OMP context usage to token and cost totals", () => { - const usage = mapOmpUsage({ - stats: { - tokens: { input: 28237, output: 548, cacheRead: 269824, cacheWrite: 0, total: 298609 }, - cost: 0.29253700000000005, - }, - state: { - model: null, - thinkingLevel: "medium", - isStreaming: false, - isCompacting: false, - sessionId: "session", - messageCount: 0, - queuedMessageCount: 0, - contextUsage: { tokens: 23656, contextWindow: 272000, percent: 8.7 }, - }, - baseUsage: { - inputTokens: 28237, - cachedInputTokens: 269824, - outputTokens: 548, - totalCostUsd: 0.29253700000000005, - }, - }); - - expect(usage).toEqual({ - inputTokens: 28237, - cachedInputTokens: 269824, - outputTokens: 548, - totalCostUsd: 0.29253700000000005, - contextWindowMaxTokens: 272000, - contextWindowUsedTokens: 23656, - }); - }); -}); diff --git a/packages/server/src/server/agent/providers/omp/usage-mapper.ts b/packages/server/src/server/agent/providers/omp/usage-mapper.ts deleted file mode 100644 index 8dab93802..000000000 --- a/packages/server/src/server/agent/providers/omp/usage-mapper.ts +++ /dev/null @@ -1,24 +0,0 @@ -import type { AgentUsage } from "../../agent-sdk-types.js"; -import type { OmpSessionState, OmpSessionStats } from "./rpc-types.js"; - -export function mapOmpUsage(input: { - stats: OmpSessionStats; - state: OmpSessionState; - baseUsage: AgentUsage | undefined; -}): AgentUsage | undefined { - const contextWindowUsedTokens = finiteNumber(input.state.contextUsage?.tokens); - const contextWindowMaxTokens = finiteNumber(input.state.contextUsage?.contextWindow); - if (contextWindowUsedTokens === undefined && contextWindowMaxTokens === undefined) { - return input.baseUsage; - } - - return { - ...input.baseUsage, - ...(contextWindowMaxTokens !== undefined ? { contextWindowMaxTokens } : {}), - ...(contextWindowUsedTokens !== undefined ? { contextWindowUsedTokens } : {}), - }; -} - -function finiteNumber(value: number | null | undefined): number | undefined { - return typeof value === "number" && Number.isFinite(value) ? value : undefined; -} diff --git a/packages/server/src/server/agent/providers/omp/usage-poller.test.ts b/packages/server/src/server/agent/providers/omp/usage-poller.test.ts new file mode 100644 index 000000000..9750ec59e --- /dev/null +++ b/packages/server/src/server/agent/providers/omp/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 { OmpSessionStats } from "./rpc-types.js"; +import { OmpUsagePoller, type OmpUsagePollScheduler } from "./usage-poller.js"; + +class ManualPollScheduler implements OmpUsagePollScheduler { + 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("OMP usage poller", () => { + test("emits only changed usage while active", async () => { + const scheduler = new ManualPollScheduler(); + const updates: AgentUsage[] = []; + let stats: OmpSessionStats = { + tokens: { input: 100, cacheRead: 10, output: 20 }, + cost: 0.01, + contextUsage: { contextWindow: 200_000, tokens: 130 }, + }; + const poller = new OmpUsagePoller({ + 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 OmpUsagePoller({ + 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 OmpUsagePoller({ + 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 OmpUsagePoller({ + 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/omp/usage-poller.ts b/packages/server/src/server/agent/providers/omp/usage-poller.ts new file mode 100644 index 000000000..4e41708b5 --- /dev/null +++ b/packages/server/src/server/agent/providers/omp/usage-poller.ts @@ -0,0 +1,169 @@ +import type { AgentUsage } from "../../agent-sdk-types.js"; +import type { OmpSessionStats } from "./rpc-types.js"; + +export interface OmpUsagePollScheduler { + schedulePoll(callback: () => void): () => void; +} + +interface OmpUsagePollerOptions { + readStats(): Promise; + onUsage(usage: AgentUsage, turnId?: string): void; + onPollError(error: unknown): void; + scheduler?: OmpUsagePollScheduler; +} + +function createScheduler(): OmpUsagePollScheduler { + return { + schedulePoll: (callback) => { + const timer = setTimeout(callback, 3_000); + return () => clearTimeout(timer); + }, + }; +} + +function toAgentUsage(stats: OmpSessionStats): 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 OmpUsagePoller { + private readonly scheduler: OmpUsagePollScheduler; + private active = false; + private closed = false; + private generation = 0; + private cancelScheduledPoll: (() => void) | null = null; + private lastUsage: AgentUsage | null = null; + + constructor(private readonly options: OmpUsagePollerOptions) { + 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); + } + } +}