fix(ui): avoid usage reloads on websocket reconnect (#108726)

* fix(ui): reuse usage data across reconnects * test(ui): use valid profile cost summaries * test(ui): prove usage reconnect lifecycle --------- Co-authored-by: Peter Steinberger <steipete@gmail.com>
2026-07-24 20:11:14 +00:00 · 2026-07-16 19:40:13 +08:00
parent 3f25f4b5fc
commit dd58667be8
5 changed files with 451 additions and 9 deletions
--- a/ui/src/e2e/usage-reconnect.e2e.test.ts
+++ b/ui/src/e2e/usage-reconnect.e2e.test.ts
@@ -0,0 +1,256 @@
+// Control UI tests cover proxy-style same-client reconnects through the real browser lifecycle.
+import { mkdir } from "node:fs/promises";
+import path from "node:path";
+import { chromium, type Browser, type BrowserContext, type Page } from "playwright";
+import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import {
+  canRunPlaywrightChromium,
+  installMockGateway,
+  resolvePlaywrightChromiumExecutablePath,
+  startControlUiE2eServer,
+  type ControlUiE2eServer,
+  type MockGatewayControls,
+} from "../test-helpers/control-ui-e2e.ts";
+
+const chromiumExecutablePath = resolvePlaywrightChromiumExecutablePath(chromium.executablePath());
+const chromiumAvailable = canRunPlaywrightChromium(chromiumExecutablePath);
+const allowMissingChromium = process.env.OPENCLAW_UI_E2E_ALLOW_MISSING_CHROMIUM === "1";
+const describeControlUiE2e = chromiumAvailable || !allowMissingChromium ? describe : describe.skip;
+const proofDir = process.env.OPENCLAW_UI_E2E_ARTIFACT_DIR?.trim();
+
+let browser: Browser;
+let server: ControlUiE2eServer;
+
+const totals = {
+  input: 100,
+  output: 20,
+  cacheRead: 0,
+  cacheWrite: 0,
+  totalTokens: 120,
+  totalCost: 0.01,
+  inputCost: 0.008,
+  outputCost: 0.002,
+  cacheReadCost: 0,
+  cacheWriteCost: 0,
+  missingCostEntries: 0,
+};
+
+function today(): string {
+  const date = new Date();
+  return `${date.getFullYear()}-${String(date.getMonth() + 1).padStart(2, "0")}-${String(date.getDate()).padStart(2, "0")}`;
+}
+
+function costSummary(cacheStatus?: {
+  status: "refreshing" | "partial" | "ready";
+  cachedFiles: number;
+  pendingFiles: number;
+  staleFiles: number;
+}) {
+  return {
+    updatedAt: Date.now(),
+    days: 1,
+    daily: [{ date: today(), ...totals }],
+    totals,
+    ...(cacheStatus ? { cacheStatus } : {}),
+  };
+}
+
+function sessionsUsage(cacheStatus?: ReturnType<typeof costSummary>["cacheStatus"]) {
+  return {
+    updatedAt: Date.now(),
+    startDate: today(),
+    endDate: today(),
+    sessions: [
+      {
+        key: "agent:main:proxy-proof",
+        label: "Proxy proof",
+        agentId: "main",
+        modelProvider: "openai",
+        model: "gpt-5.5",
+        updatedAt: Date.now(),
+        usage: {
+          ...totals,
+          activityDates: [today()],
+          dailyBreakdown: [{ date: today(), tokens: totals.totalTokens, cost: totals.totalCost }],
+          messageCounts: {
+            total: 2,
+            user: 1,
+            assistant: 1,
+            toolCalls: 0,
+            toolResults: 0,
+            errors: 0,
+          },
+        },
+      },
+    ],
+    totals,
+    aggregates: {
+      messages: { total: 2, user: 1, assistant: 1, toolCalls: 0, toolResults: 0, errors: 0 },
+      tools: { totalCalls: 0, uniqueTools: 0, tools: [] },
+      byModel: [],
+      byProvider: [],
+      byAgent: [{ agentId: "main", totals }],
+      byChannel: [],
+      daily: [
+        {
+          date: today(),
+          tokens: totals.totalTokens,
+          cost: totals.totalCost,
+          messages: 2,
+          toolCalls: 0,
+          errors: 0,
+        },
+      ],
+    },
+    ...(cacheStatus ? { cacheStatus } : {}),
+  };
+}
+
+async function createContext(): Promise<BrowserContext> {
+  if (proofDir) {
+    await mkdir(proofDir, { recursive: true });
+  }
+  return browser.newContext({
+    locale: "en-US",
+    serviceWorkers: "block",
+    viewport: { height: 900, width: 1440 },
+  });
+}
+
+async function requestCount(gateway: MockGatewayControls, method: string): Promise<number> {
+  return (await gateway.getRequests(method)).length;
+}
+
+async function waitForRequestCount(
+  gateway: MockGatewayControls,
+  method: string,
+  count: number,
+): Promise<void> {
+  await expect.poll(() => requestCount(gateway, method), { timeout: 10_000 }).toBe(count);
+}
+
+async function proxyReconnect(
+  page: Page,
+  gateway: MockGatewayControls,
+  expectedSocketCount: number,
+): Promise<void> {
+  await gateway.closeLatest(1001, "proxy idle timeout");
+  await page.locator("openclaw-connection-banner").waitFor({ state: "visible" });
+  await expect.poll(() => gateway.getSocketCount(), { timeout: 10_000 }).toBe(expectedSocketCount);
+  await page.locator("openclaw-connection-banner").waitFor({ state: "hidden" });
+}
+
+async function captureProof(page: Page, name: string): Promise<void> {
+  if (!proofDir) {
+    return;
+  }
+  await page.screenshot({ fullPage: true, path: path.join(proofDir, name) });
+}
+
+async function usageBadges(page: Page): Promise<string[]> {
+  return (await page.locator(".usage-metric-badge").allTextContents()).map((value) =>
+    value.replace(/\s+/gu, " ").trim(),
+  );
+}
+
+describeControlUiE2e("Control UI usage proxy reconnect lifecycle", () => {
+  beforeAll(async () => {
+    if (!chromiumAvailable) {
+      throw new Error(`Playwright Chromium is not available at ${chromiumExecutablePath}`);
+    }
+    server = await startControlUiE2eServer();
+    browser = await chromium.launch({ executablePath: chromiumExecutablePath });
+  });
+
+  afterAll(async () => {
+    await browser?.close();
+    await server?.close();
+  });
+
+  it("avoids a reload storm but retries Usage work interrupted by a proxy drop", async () => {
+    const context = await createContext();
+    const page = await context.newPage();
+    const gateway = await installMockGateway(page, {
+      methodResponses: {
+        "sessions.usage": sessionsUsage(),
+        "usage.cost": costSummary(),
+        "usage.status": { updatedAt: Date.now(), providers: [] },
+      },
+    });
+
+    try {
+      const response = await page.goto(`${server.baseUrl}usage`);
+      expect(response?.status()).toBe(200);
+      await waitForRequestCount(gateway, "sessions.usage", 1);
+      await waitForRequestCount(gateway, "usage.cost", 1);
+      await page.locator(".daily-chart-compact").waitFor({ timeout: 10_000 });
+      await expect.poll(() => usageBadges(page)).toEqual(["120 Tokens", "$0.01 Cost", "1 session"]);
+
+      for (const socketCount of [2, 3, 4]) {
+        await proxyReconnect(page, gateway, socketCount);
+        expect(await requestCount(gateway, "sessions.usage")).toBe(1);
+        expect(await requestCount(gateway, "usage.cost")).toBe(1);
+      }
+
+      await gateway.deferNext("sessions.usage");
+      await gateway.deferNext("usage.cost");
+      await page.getByRole("button", { name: "Refresh", exact: true }).click();
+      await waitForRequestCount(gateway, "sessions.usage", 2);
+      await waitForRequestCount(gateway, "usage.cost", 2);
+
+      await proxyReconnect(page, gateway, 5);
+      await waitForRequestCount(gateway, "sessions.usage", 3);
+      await waitForRequestCount(gateway, "usage.cost", 3);
+      await page.locator(".daily-chart-compact").waitFor({ timeout: 10_000 });
+      await expect.poll(() => usageBadges(page)).toEqual(["120 Tokens", "$0.01 Cost", "1 session"]);
+      await captureProof(page, "usage-after-interrupted-retry.png");
+    } finally {
+      await context.close();
+    }
+  });
+
+  it("resumes Profile cache settlement once and then preserves the settled result", async () => {
+    const context = await createContext();
+    const page = await context.newPage();
+    const refreshing = {
+      status: "refreshing" as const,
+      cachedFiles: 1,
+      pendingFiles: 1,
+      staleFiles: 0,
+    };
+    const gateway = await installMockGateway(page, {
+      methodResponses: {
+        "sessions.usage": sessionsUsage(refreshing),
+        "usage.cost": costSummary(refreshing),
+      },
+    });
+
+    try {
+      const response = await page.goto(`${server.baseUrl}settings/profile`);
+      expect(response?.status()).toBe(200);
+      await waitForRequestCount(gateway, "sessions.usage", 1);
+      await waitForRequestCount(gateway, "usage.cost", 1);
+      await page.locator(".profile-stats").waitFor({ timeout: 10_000 });
+
+      await gateway.setMethodResponse("sessions.usage", sessionsUsage());
+      await gateway.setMethodResponse("usage.cost", costSummary());
+      await proxyReconnect(page, gateway, 2);
+      await waitForRequestCount(gateway, "sessions.usage", 2);
+      await waitForRequestCount(gateway, "usage.cost", 2);
+      await page.locator(".profile-stats").waitFor();
+      await expect
+        .poll(() => page.locator(".profile-stats__value").first().textContent())
+        .toBe("120");
+
+      await proxyReconnect(page, gateway, 3);
+      expect(await requestCount(gateway, "sessions.usage")).toBe(2);
+      expect(await requestCount(gateway, "usage.cost")).toBe(2);
+      await expect
+        .poll(() => page.locator(".profile-stats__value").first().textContent())
+        .toBe("120");
+      await captureProof(page, "profile-after-cache-settlement.png");
+    } finally {
+      await context.close();
+    }
+  });
+});
--- a/ui/src/pages/gateway-source-replacement.test.ts
+++ b/ui/src/pages/gateway-source-replacement.test.ts
@@ -210,6 +210,78 @@ describe("gateway source replacement across reconnect with a reused client", ()
    expect(page.usageResult).not.toBe(staleResult);
  });

+  it("keeps loaded usage data across a same-client reconnect", async () => {
+    const request = vi.fn();
+    const client = { request } as unknown as GatewayBrowserClient;
+    const context = contextWithClient(client, { connected: true });
+    const result = { sessions: [{ key: "cached" }] } as unknown as UsageRouteData["result"];
+    const page = createPage("openclaw-usage-page", context) as TestPage & {
+      routeData: UsageRouteData;
+      applyGatewaySnapshot: (snapshot: ApplicationGatewaySnapshot) => void;
+    };
+    page.routeData = {
+      gateway: context.gateway,
+      gatewaySnapshot: context.gateway.snapshot,
+      query: {
+        startDate: "2026-07-08",
+        endDate: "2026-07-08",
+        scope: "family",
+        timeZone: "local",
+        agentId: null,
+      },
+      result,
+      costSummary: null,
+      providerUsageSummary: null,
+      error: null,
+    };
+
+    document.body.append(page);
+    await page.updateComplete;
+    page.applyGatewaySnapshot({ ...context.gateway.snapshot, connected: false });
+    page.applyGatewaySnapshot(context.gateway.snapshot);
+    await Promise.resolve();
+
+    expect(request).not.toHaveBeenCalled();
+  });
+
+  it("retries a usage load interrupted by a same-client disconnect", async () => {
+    const request = vi.fn(async (method: string) =>
+      method === "sessions.usage" ? { sessions: [] } : {},
+    );
+    const client = { request } as unknown as GatewayBrowserClient;
+    const context = contextWithClient(client, { connected: true });
+    const page = createPage("openclaw-usage-page", context) as TestPage & {
+      routeData: UsageRouteData;
+      usageLoading: boolean;
+      applyGatewaySnapshot: (snapshot: ApplicationGatewaySnapshot) => void;
+    };
+    page.routeData = {
+      gateway: context.gateway,
+      gatewaySnapshot: context.gateway.snapshot,
+      query: {
+        startDate: "2026-07-08",
+        endDate: "2026-07-08",
+        scope: "family",
+        timeZone: "local",
+        agentId: null,
+      },
+      result: { sessions: [{ key: "cached" }] } as unknown as UsageRouteData["result"],
+      costSummary: null,
+      providerUsageSummary: null,
+      error: null,
+    };
+
+    document.body.append(page);
+    await page.updateComplete;
+    page.usageLoading = true;
+    page.applyGatewaySnapshot({ ...context.gateway.snapshot, connected: false });
+    page.applyGatewaySnapshot(context.gateway.snapshot);
+
+    await vi.waitFor(() =>
+      expect(request).toHaveBeenCalledWith("sessions.usage", expect.any(Object)),
+    );
+  });
+
  it("preserves matching skills route data on the first bind", async () => {
    const request = vi.fn();
    const client = { request } as unknown as GatewayBrowserClient;
--- a/ui/src/pages/profile/profile-page.test.ts
+++ b/ui/src/pages/profile/profile-page.test.ts
@@ -1,6 +1,8 @@
 /* @vitest-environment jsdom */

-import { afterEach, beforeEach, expect, it } from "vitest";
+import { afterEach, beforeEach, expect, it, vi } from "vitest";
+import type { GatewayBrowserClient } from "../../api/gateway.ts";
+import type { CostUsageSummary, SessionsUsageResult } from "../../api/types.ts";
 import type { RouteId } from "../../app-route-paths.ts";
 import type { ApplicationContext, ApplicationGatewaySnapshot } from "../../app/context.ts";
 import { i18n, t } from "../../i18n/index.ts";
@@ -17,10 +19,22 @@ type ProfilePageElement = HTMLElement & {
  updateComplete: Promise<boolean>;
 };

-function createContext(): ApplicationContext<RouteId> {
+type ProfilePageLifecycle = HTMLElement & {
+  context: ApplicationContext<RouteId>;
+  client: GatewayBrowserClient | null;
+  connected: boolean;
+  costSummary: CostUsageSummary | null;
+  sessionsResult: SessionsUsageResult | null;
+  applyGatewaySnapshot: (snapshot: ApplicationGatewaySnapshot) => void;
+};
+
+function createContext(
+  client: GatewayBrowserClient | null = null,
+  connected = false,
+): ApplicationContext<RouteId> {
  const snapshot: ApplicationGatewaySnapshot = {
-    client: null,
-    connected: false,
+    client,
+    connected,
    reconnecting: false,
    hello: null,
    assistantAgentId: "main",
@@ -31,17 +45,60 @@ function createContext(): ApplicationContext<RouteId> {
  const subscribe = () => () => undefined;
  return {
    gateway: { snapshot, subscribe },
-    agents: { subscribe },
-    agentIdentity: { subscribe },
+    agents: { subscribe, ensureList: vi.fn(async () => null) },
+    agentIdentity: { subscribe, ensure: vi.fn(async () => undefined) },
  } as unknown as ApplicationContext<RouteId>;
 }

+function createCostSummary(cacheStatus?: CostUsageSummary["cacheStatus"]): CostUsageSummary {
+  return {
+    updatedAt: 0,
+    days: 0,
+    daily: [],
+    totals: {
+      input: 0,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: 0,
+      totalCost: 1,
+      inputCost: 0,
+      outputCost: 0,
+      cacheReadCost: 0,
+      cacheWriteCost: 0,
+      missingCostEntries: 0,
+    },
+    ...(cacheStatus ? { cacheStatus } : {}),
+  };
+}
+
+function createSessionsResult(): SessionsUsageResult {
+  const totals = createCostSummary().totals;
+  return {
+    updatedAt: 0,
+    startDate: "2026-07-08",
+    endDate: "2026-07-08",
+    sessions: [],
+    totals,
+    aggregates: {
+      messages: { total: 0, user: 0, assistant: 0, toolCalls: 0, toolResults: 0, errors: 0 },
+      tools: { totalCalls: 0, uniqueTools: 0, tools: [] },
+      byModel: [],
+      byProvider: [],
+      byAgent: [],
+      byChannel: [],
+      daily: [],
+    },
+  };
+}
+
 beforeEach(async () => {
  await i18n.setLocale("en");
 });

 afterEach(async () => {
  document.body.replaceChildren();
+  vi.restoreAllMocks();
  await i18n.setLocale("en");
 });

@@ -61,3 +118,45 @@ it("refreshes translated copy when the locale changes while mounted", async () =
  expect(note?.textContent?.trim()).toBe(t("profilePage.offline"));
  expect(note?.textContent?.trim()).not.toBe(englishNote);
 });
+
+it("keeps settled profile usage across a same-client reconnect", async () => {
+  const request = vi.fn();
+  const client = { request } as unknown as GatewayBrowserClient;
+  const context = createContext(client, true);
+  const page = new ProfilePage() as unknown as ProfilePageLifecycle;
+  page.context = context;
+  page.client = client;
+  page.connected = true;
+  page.costSummary = createCostSummary();
+  page.sessionsResult = createSessionsResult();
+
+  page.applyGatewaySnapshot({ ...context.gateway.snapshot, connected: false });
+  page.applyGatewaySnapshot(context.gateway.snapshot);
+  await Promise.resolve();
+
+  expect(request).not.toHaveBeenCalled();
+});
+
+it("resumes a settling profile cache after a same-client reconnect", async () => {
+  const request = vi.fn(async (method: string) =>
+    method === "sessions.usage" ? { sessions: [] } : createCostSummary(),
+  );
+  const client = { request } as unknown as GatewayBrowserClient;
+  const context = createContext(client, true);
+  const page = new ProfilePage() as unknown as ProfilePageLifecycle;
+  page.context = context;
+  page.client = client;
+  page.connected = true;
+  page.costSummary = createCostSummary({
+    status: "refreshing",
+    cachedFiles: 0,
+    pendingFiles: 1,
+    staleFiles: 0,
+  });
+  page.sessionsResult = createSessionsResult();
+
+  page.applyGatewaySnapshot({ ...context.gateway.snapshot, connected: false });
+  page.applyGatewaySnapshot(context.gateway.snapshot);
+
+  await vi.waitFor(() => expect(request).toHaveBeenCalledWith("usage.cost", expect.any(Object)));
+});
--- a/ui/src/pages/profile/profile-page.ts
+++ b/ui/src/pages/profile/profile-page.ts
@@ -156,7 +156,11 @@ export class ProfilePage extends OpenClawLightDomElement {
        void this.context.agentIdentity.ensure([list.defaultId]);
      }
    });
-    if (clientChanged || becameConnected || (!this.costSummary && !this.loading && !this.error)) {
+    if (
+      clientChanged ||
+      (becameConnected && (!this.costSummary || this.isCacheSettling())) ||
+      (!this.costSummary && !this.loading && !this.error)
+    ) {
      void this.loadProfile();
    }
  }
--- a/ui/src/pages/usage/usage-page.ts
+++ b/ui/src/pages/usage/usage-page.ts
@@ -115,6 +115,7 @@ class UsagePage extends OpenClawLightDomElement {
  private queryDebounceTimer: number | null = null;
  private routeDataInitialized = false;
  private routeDataEnabled = true;
+  private usageReloadPending = false;
  private hasBoundGatewaySource = false;
  private observedAgentScopeId: string | null | undefined;
  private readonly subscriptions = new SubscriptionsController(this)
@@ -178,12 +179,16 @@ class UsagePage extends OpenClawLightDomElement {
      this.resetForClientChange();
    }
    if (!snapshot.connected || !snapshot.client) {
+      this.usageReloadPending ||= this.usageLoading;
      this.invalidateRequests();
      return;
    }

    void this.context.agents.ensureList();
-    if (this.routeDataInitialized && (clientChanged || becameConnected)) {
+    if (
+      this.routeDataInitialized &&
+      (clientChanged || (becameConnected && (this.usageReloadPending || this.usageResult === null)))
+    ) {
      void this.loadUsage();
    }
  }
@@ -251,6 +256,7 @@ class UsagePage extends OpenClawLightDomElement {
    this.usageCostSummary = null;
    this.providerUsageSummary = null;
    this.usageError = null;
+    this.usageReloadPending = false;
    this.usageAgentId = this.context.agentSelection.state.scopeId;
    this.clearSelectionsAndDetails();
  }
@@ -300,10 +306,15 @@ class UsagePage extends OpenClawLightDomElement {

  private async loadUsage() {
    const client = this.client;
-    if (!client || !this.connected || this.usageLoading) {
+    if (!client || !this.connected) {
+      this.usageReloadPending = true;
+      return;
+    }
+    if (this.usageLoading) {
      return;
    }

+    this.usageReloadPending = false;
    this.routeDataEnabled = false;
    const requestId = ++this.usageRequestId;
    const startDate = this.usageStartDate;