Merge pull request #1094 from fabro-sh/codex/live-usage-refresh

Refresh live Usage totals with bounded stream batching

apps/fabro-web/app/lib/cross-tab-sse.ts

@@ -116,6 +116,8 @@ interface SubscribeOptions<TPayload extends EventPayload> {
   resyncKeys: () => KeyOrMatcher[];
   fallbackSubscribe: () => () => void;
   debounceMs?: number;
+  /** Flush pending keys within this interval, even while events keep arriving. */
+  maxWaitMs?: number;
 }
 
 export interface SubscribeToCrossTabSseOptions<TPayload extends EventPayload>
@@ -136,6 +138,8 @@ interface LocalSubscription {
   pendingKeys: Map<string, KeyOrMatcher>;
   debounceTimer: ReturnType<typeof setTimeout> | null;
   debounceMs: number;
+  maxWaitMs?: number;
+  maxWaitTimer: ReturnType<typeof setTimeout> | null;
   resolveInvalidation: (payload: EventPayload) => EventInvalidation;
   resyncKeys: () => KeyOrMatcher[];
 }
@@ -295,6 +299,8 @@ export class CrossTabSseCoordinator {
         pendingKeys: new Map(),
         debounceTimer: null,
         debounceMs: options.debounceMs ?? 300,
+        maxWaitMs: options.maxWaitMs,
+        maxWaitTimer: null,
         resolveInvalidation: options.resolveInvalidation as (payload: EventPayload) => EventInvalidation,
         resyncKeys: options.resyncKeys,
       };
@@ -304,6 +310,7 @@ export class CrossTabSseCoordinator {
         options.resolveInvalidation as (payload: EventPayload) => EventInvalidation;
       subscription.resyncKeys = options.resyncKeys;
       subscription.debounceMs = options.debounceMs ?? subscription.debounceMs;
+      subscription.maxWaitMs = options.maxWaitMs ?? subscription.maxWaitMs;
     }
 
     subscription.refcount += 1;
@@ -702,6 +709,12 @@ export class CrossTabSseCoordinator {
       return;
     }
 
+    if (subscription.maxWaitMs !== undefined && subscription.maxWaitTimer === null) {
+      subscription.maxWaitTimer = setTimeout(
+        () => this.flushInvalidations(subscription),
+        subscription.maxWaitMs,
+      );
+    }
     if (subscription.debounceTimer) {
       clearTimeout(subscription.debounceTimer);
     }
@@ -712,6 +725,7 @@ export class CrossTabSseCoordinator {
   }
 
   private flushInvalidations(subscription: LocalSubscription) {
+    this.clearSubscriptionTimer(subscription);
     if (subscription.pendingKeys.size === 0) return;
     const keys = [...subscription.pendingKeys.values()];
     subscription.pendingKeys.clear();
@@ -940,9 +954,10 @@ export class CrossTabSseCoordinator {
   }
 
   private clearSubscriptionTimer(subscription: LocalSubscription) {
-    if (!subscription.debounceTimer) return;
-    clearTimeout(subscription.debounceTimer);
+    if (subscription.debounceTimer !== null) clearTimeout(subscription.debounceTimer);
+    if (subscription.maxWaitTimer !== null) clearTimeout(subscription.maxWaitTimer);
     subscription.debounceTimer = null;
+    subscription.maxWaitTimer = null;
   }
 
   private resetIdleState() {

apps/fabro-web/app/lib/run-events.test.tsx

@@ -1,5 +1,5 @@
-import { describe, expect, test } from "bun:test";
-import type { Key } from "swr";
+import { describe, expect, jest, test } from "bun:test";
+import { unstable_serialize, type Key } from "swr";
 
 import { loadPetriFixture } from "./petri-fixtures";
 import { streamItemName } from "./petri-stream";
@@ -431,3 +431,94 @@ async function waitFor(condition: () => boolean, timeoutMs = 200) {
   }
   throw new Error("condition did not become true before timeout");
 }
+
+for (const mode of ["coordinated", "fallback"] as const) {
+  describe(`${mode} live usage batching`, () => {
+    function subscribe() {
+      const source = new FakeEventSource();
+      const coordinator = mode === "coordinated"
+        ? createCoordinator(() => source)
+        : createFallbackCoordinator();
+      const keys: Key[] = [];
+      const usageKey = unstable_serialize(queryKeys.runs.usage("run-live"));
+      const cleanup = subscribeToRunEvents("run-live", (key) => {
+        keys.push(key);
+        return Promise.resolve();
+      }, () => source, { coordinator });
+      jest.advanceTimersByTime(100);
+      expect(source.onmessage).not.toBeNull();
+      keys.length = 0;
+      let seq = 0;
+      return {
+        source, keys, cleanup,
+        close: () => { cleanup(); coordinator.close(); },
+        usageRefreshes: () => keys.filter((key) => unstable_serialize(key) === usageKey).length,
+        progress: () => source.emit(makePetriItem(++seq, { event: "step.progress.recorded" }, {
+          run_id: "run-live", stage: { name: "agent" },
+        })),
+      };
+    }
+
+    test("refreshes bursts after 300ms and continuous progress at least once per second", () => {
+      jest.useFakeTimers();
+      const subscription = subscribe();
+      try {
+        subscription.progress();
+        subscription.progress();
+        jest.advanceTimersByTime(299);
+        expect(subscription.usageRefreshes()).toBe(0);
+        jest.advanceTimersByTime(1);
+        expect(subscription.usageRefreshes()).toBe(1);
+        // No new event means no polling or leftover maximum-wait flush.
+        jest.advanceTimersByTime(1_000);
+        expect(subscription.usageRefreshes()).toBe(1);
+
+        for (let batch = 0; batch < 2; batch += 1) {
+          for (let event = 0; event < 10; event += 1) {
+            subscription.progress();
+            expect(subscription.usageRefreshes()).toBe(batch + 1);
+            jest.advanceTimersByTime(100);
+          }
+          expect(subscription.usageRefreshes()).toBe(batch + 2);
+        }
+      } finally {
+        subscription.close();
+        jest.useRealTimers();
+      }
+    });
+
+    test("flushes terminal usage immediately without a later duplicate refresh", () => {
+      jest.useFakeTimers();
+      const subscription = subscribe();
+      try {
+        subscription.progress();
+        subscription.source.emit(makePlatformItem(2, {
+          kind: "run.lifecycle", transition: "succeeded", status: { kind: "succeeded" },
+        }, { run_id: "run-live" }));
+        expect(subscription.usageRefreshes()).toBe(1);
+        const count = subscription.keys.length;
+        jest.advanceTimersByTime(1_000);
+        expect(subscription.keys).toHaveLength(count);
+        expect(subscription.source.closed).toBe(mode === "fallback");
+      } finally {
+        subscription.close();
+        jest.useRealTimers();
+      }
+    });
+
+    test("unsubscribing discards both pending timers", () => {
+      jest.useFakeTimers();
+      const subscription = subscribe();
+      try {
+        subscription.progress();
+        subscription.cleanup();
+        jest.advanceTimersByTime(2_000);
+        expect(subscription.keys).toEqual([]);
+        expect(subscription.source.closed).toBe(true);
+      } finally {
+        subscription.close();
+        jest.useRealTimers();
+      }
+    });
+  });
+}

apps/fabro-web/app/lib/run-events.ts

@@ -40,6 +40,8 @@ export interface RunEventPayload extends EventPayload {
 
 interface RunEventOptions {
   debounceMs?: number;
+  /** Keep live queries updating even when progress never pauses. */
+  maxWaitMs?: number;
   coordinator?: CrossTabSseCoordinator;
   /** Called with every well-formed stream item of the run, before invalidation. */
   onItem?: (item: RunStreamItem) => void;
@@ -131,7 +133,7 @@ export function queryKeysForStreamItem(
           immediate: false,
         };
       }
-      return { keys: [queryKeys.runs.state(runId), stream, ...stageKeys], immediate: false };
+      return { keys: [queryKeys.runs.state(runId), queryKeys.runs.usage(runId), stream, ...stageKeys], immediate: false };
     }
     case "control.requested":
       return {
@@ -196,13 +198,14 @@ export function subscribeToRunEvents(
   runId: string,
   mutate: MutateFn,
   eventSourceFactory: (url: string) => EventSourceLike = createBrowserEventSource,
-  { debounceMs = 300, coordinator, onItem }: RunEventOptions = {},
+  { debounceMs = 300, maxWaitMs = 1_000, coordinator, onItem }: RunEventOptions = {},
 ): () => void {
   return subscribeToCrossTabSse<RunEventPayload>({
     coordinator,
     subscriptionKey: `run:${runId}`,
     mutate,
     debounceMs,
+    maxWaitMs,
     resyncKeys: () => resyncKeysForRun(runId),
     resolveInvalidation: (payload) => {
       if (payload.run_id !== runId) return { keys: [] };
@@ -216,6 +219,7 @@ export function subscribeToRunEvents(
         mutate,
         eventSourceFactory,
         debounceMs,
+        maxWaitMs,
         resolveInvalidation: (payload) => {
           const result = runInvalidation(runId, payload, onItem);
           return { ...result, close: result.immediate };

apps/fabro-web/app/lib/sse.ts

@@ -29,6 +29,7 @@ export interface SharedEventSubscription {
   resolvers: Map<symbol, EventResolver>;
   pendingKeys: Map<string, KeyOrMatcher>;
   debounceTimer: ReturnType<typeof setTimeout> | null;
+  maxWaitTimer: ReturnType<typeof setTimeout> | null;
 }
 
 export function sseKeyDedupeId(key: KeyOrMatcher): string {
@@ -47,6 +48,7 @@ export function subscribeToSharedEventSource<TPayload extends EventPayload>({
   resolveInvalidation,
   eventSourceFactory = createBrowserEventSource,
   debounceMs = 300,
+  maxWaitMs,
 }: {
   subscriptions: Map<string, SharedEventSubscription>;
   subscriptionKey: string;
@@ -55,6 +57,8 @@ export function subscribeToSharedEventSource<TPayload extends EventPayload>({
   resolveInvalidation: (payload: TPayload) => EventInvalidation;
   eventSourceFactory?: (url: string) => EventSourceLike;
   debounceMs?: number;
+  /** Flush pending keys within this interval, even while events keep arriving. */
+  maxWaitMs?: number;
 }): () => void {
   let subscription = subscriptions.get(subscriptionKey);
   if (!subscription) {
@@ -66,6 +70,7 @@ export function subscribeToSharedEventSource<TPayload extends EventPayload>({
       resolvers: new Map(),
       pendingKeys: new Map(),
       debounceTimer: null,
+      maxWaitTimer: null,
     };
     subscriptions.set(subscriptionKey, subscription);
 
@@ -92,7 +97,7 @@ export function subscribeToSharedEventSource<TPayload extends EventPayload>({
         immediate ||= Boolean(invalidation.immediate);
       }
 
-      queueInvalidations(current, [...keys.values()], { debounceMs, immediate });
+      queueInvalidations(current, [...keys.values()], { debounceMs, maxWaitMs, immediate });
 
       if (close) {
         closeSharedEventSource(subscriptions, subscriptionKey, { flushPending: true });
@@ -133,9 +138,11 @@ function queueInvalidations(
   keys: KeyOrMatcher[],
   {
     debounceMs,
+    maxWaitMs,
     immediate,
   }: {
     debounceMs: number;
+    maxWaitMs?: number;
     immediate?: boolean;
   },
 ) {
@@ -149,6 +156,9 @@ function queueInvalidations(
     return;
   }
 
+  if (maxWaitMs !== undefined && subscription.maxWaitTimer === null) {
+    subscription.maxWaitTimer = setTimeout(() => flushInvalidations(subscription), maxWaitMs);
+  }
   if (subscription.debounceTimer) {
     clearTimeout(subscription.debounceTimer);
   }
@@ -159,6 +169,7 @@ function queueInvalidations(
 }
 
 function flushInvalidations(subscription: SharedEventSubscription) {
+  clearInvalidationTimers(subscription);
   if (subscription.pendingKeys.size === 0) return;
   const keys = [...subscription.pendingKeys.values()];
   subscription.pendingKeys.clear();
@@ -181,13 +192,18 @@ function closeSharedEventSource(
   if (flushPending) {
     flushInvalidations(subscription);
   }
-  if (subscription.debounceTimer) {
-    clearTimeout(subscription.debounceTimer);
-  }
+  clearInvalidationTimers(subscription);
   subscription.source.close();
   subscriptions.delete(subscriptionKey);
 }
 
+function clearInvalidationTimers(subscription: SharedEventSubscription) {
+  if (subscription.debounceTimer !== null) clearTimeout(subscription.debounceTimer);
+  if (subscription.maxWaitTimer !== null) clearTimeout(subscription.maxWaitTimer);
+  subscription.debounceTimer = null;
+  subscription.maxWaitTimer = null;
+}
+
 function stringifyKeyValue(value: unknown): string {
   if (Array.isArray(value)) {
     return `[${value.map((item) => stringifyKeyValue(item)).join(",")}]`;

apps/fabro-web/app/routes/run-usage.live.test.tsx

@@ -0,0 +1,102 @@
+import { expect, jest, spyOn, test } from "bun:test";
+import TestRenderer, { act } from "react-test-renderer";
+import { SWRConfig, useSWRConfig } from "swr";
+import { AxiosHeaders } from "axios";
+import type { RunUsage } from "@qltysh/fabro-api-client";
+
+import { runOutputsApi } from "../lib/api-client";
+import { createCrossTabSseCoordinator } from "../lib/cross-tab-sse";
+import { subscribeToRunEvents } from "../lib/run-events";
+import type { EventSourceLike, MutateFn } from "../lib/sse";
+import { makeUsage } from "../lib/test-fixtures";
+import { makePebbleItem, setupReactTestEnv, textContent } from "../lib/test-utils";
+import RunUsageRoute from "./run-usage";
+
+function usage(responses: number): RunUsage {
+  const model = { provider: "anthropic", model_id: "claude-sonnet-5" };
+  const totals = makeUsage({ input: responses * 10_000, output: responses * 8_000 }, responses * 108_500);
+  const timing = { wall_time_ms: 0, inference_time_ms: 0, tool_time_ms: 0, active_time_ms: 0 };
+  return {
+    stages: [{
+      stage: { id: "agent", name: "agent" }, model, usage: totals, timing,
+      state: "running", started_at: "2026-10-08T16:57:29Z",
+    }],
+    by_model: [{ model, stages: 1, usage: totals }],
+    totals: { usage: totals, timing },
+  };
+}
+
+for (const mode of ["coordinated", "fallback"] as const) {
+  test(`${mode} progress updates the mounted Usage page through its real query cache`, async () => {
+    jest.useFakeTimers();
+    const teardown = setupReactTestEnv();
+    let remote = usage(1);
+    const retrieve = spyOn(runOutputsApi, "retrieveRunUsage").mockImplementation(async () => ({
+      data: remote, status: 200, statusText: "OK", headers: new AxiosHeaders(),
+      config: { headers: new AxiosHeaders() },
+    }));
+    const source: EventSourceLike = { onmessage: null, close() {} };
+    const coordinator = createCrossTabSseCoordinator({
+      channelFactory: () => {
+        if (mode === "fallback") throw new Error("BroadcastChannel unavailable");
+        return { onmessage: null, postMessage() {}, close() {} };
+      },
+      eventSourceFactory: () => source,
+      addVisibilityChangeListener: () => () => {},
+      addPagehideListener: () => () => {},
+      timing: { electionJitterMs: 0 },
+    });
+    let mutate!: MutateFn;
+    function CacheAccess() {
+      mutate = useSWRConfig().mutate as MutateFn;
+      return null;
+    }
+    let renderer: TestRenderer.ReactTestRenderer | undefined;
+    let cleanup: (() => void) | undefined;
+    try {
+      await act(async () => {
+        renderer = TestRenderer.create(
+          <SWRConfig value={{
+            provider: () => new Map(), dedupingInterval: 0,
+            revalidateOnFocus: false, revalidateOnReconnect: false, shouldRetryOnError: false,
+          }}>
+            <CacheAccess />
+            <RunUsageRoute params={{ id: "run-1" }} />
+          </SWRConfig>,
+        );
+      });
+      expect(retrieve).toHaveBeenCalledTimes(1);
+      expect(textContent(renderer!.root)).toContain("$0.11");
+      await act(async () => {
+        cleanup = subscribeToRunEvents("run-1", mutate, () => source, { coordinator });
+        jest.advanceTimersByTime(100);
+      });
+      expect(source.onmessage).not.toBeNull();
+      const fetchesBeforeProgress = retrieve.mock.calls.length;
+      remote = usage(2);
+      await act(async () => {
+        source.onmessage!({ data: JSON.stringify(makePebbleItem(
+          1, "2026-10-08T16:57:43Z", { name: "agent" }, "AssistantMessage", {},
+        )) });
+        jest.advanceTimersByTime(300);
+      });
+      expect(retrieve).toHaveBeenCalledTimes(fetchesBeforeProgress + 1);
+      // Stage row, page total, model row and model total all use the new snapshot.
+      const rows = renderer!.root.findAllByType("tr").filter((row) => textContent(row).includes("$0.22"));
+      expect(rows).toHaveLength(4);
+      for (const row of rows) {
+        expect(textContent(row)).toContain("20.0k");
+        expect(textContent(row)).toContain("16.0k");
+      }
+      expect(textContent(renderer!.root)).not.toContain("$0.11");
+      expect(remote.stages[0].state).toBe("running");
+    } finally {
+      cleanup?.();
+      coordinator.close();
+      act(() => renderer?.unmount());
+      retrieve.mockRestore();
+      teardown();
+      jest.useRealTimers();
+    }
+  });
+}