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 <boudra.moha@gmail.com>

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

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<string>();
@@ -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<void> {
     const turnId = this.activeTurnId;
-    try {
+    await this.runtimeSession.abort();
-      await this.runtimeSession.abort();
-    } finally {
-      this.terminalizeActiveWork();
-    }
     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<void> {
+  private async refreshAfterTurn(finalUsage: Promise<void>): Promise<void> {
-    await this.refreshState().catch(() => undefined);
+    await Promise.all([this.refreshState().catch(() => undefined), finalUsage]);
-    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,
-      });
-    }
   }
 }
 
@@ -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,
       });

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<void> {
+    if (this.abortError) {
+      throw this.abortError;
+    }
     this.abortRequested = true;
   }
 

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

packages/server/src/server/agent/providers/omp/usage-mapper.test.ts

@@ -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,
-    });
-  });
-});

packages/server/src/server/agent/providers/omp/usage-mapper.ts

@@ -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;
-}

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<T>() {
+  let resolve!: (value: T) => void;
+  const promise = new Promise<T>((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<OmpSessionStats>();
+    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<OmpSessionStats>();
+    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<OmpSessionStats>();
+    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);
+  });
+});

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<OmpSessionStats>;
+  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<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);
+    }
+  }
+}