fix(cli): avoid zero usage cost during cache refresh (#103998)

* fix(cli): wait for usage-cost cache refresh * test(cli): type usage cost handler response * fix(cli): honor usage-cost timeout budget * fix(cli): give usage audit a complete settle budget * fix(cli): wait through stale usage caches --------- Co-authored-by: Peter Steinberger <steipete@gmail.com>
2026-08-28 05:16:23 -06:00 · 2026-07-11 17:07:37 -07:00
parent c749d714e9
commit dbe53d6212
2 changed files with 174 additions and 6 deletions
@@ -3,16 +3,21 @@ import fs from "node:fs";
 import os from "node:os";
 import path from "node:path";
 import { Command } from "commander";
+import { SessionManager } from "openclaw/plugin-sdk/agent-sessions";
 import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import { withEnvOverride } from "../config/test-helpers.js";
+import type { OpenClawConfig } from "../config/types.openclaw.js";
+import { testApi as usageTestApi, usageHandlers } from "../gateway/server-methods/usage.js";
 import { GatewayLockError } from "../infra/gateway-lock.js";
 import { registerGatewayCli } from "./gateway-cli.js";

 type DiscoveredBeacon = Awaited<
  ReturnType<typeof import("../infra/bonjour-discovery.js").discoverGatewayBeacons>
 >[number];
+type UsageCostHandlerArgs = Parameters<(typeof usageHandlers)["usage.cost"]>[0];

-const callGateway = vi.fn<(opts: unknown) => Promise<{ ok: true }>>(async () => ({ ok: true }));
+const defaultCallGateway = async (): Promise<unknown> => ({ ok: true });
+const callGateway = vi.fn<(opts: unknown) => Promise<unknown>>(defaultCallGateway);
 const formatGatewayClientRequestErrorJson = vi.fn();
 const formatGatewayTransportErrorJson = vi.fn();
 const startGatewayServer = vi.fn<
@@ -156,6 +161,8 @@ describe("gateway-cli coverage", () => {

  beforeEach(() => {
    gatewayProgram = createGatewayProgram();
+    callGateway.mockReset();
+    callGateway.mockImplementation(defaultCallGateway);
    runtimeLogs.length = 0;
    runtimeErrors.length = 0;
    defaultRuntime.log.mockClear();
@@ -253,6 +260,118 @@ describe("gateway-cli coverage", () => {
    expect(costCall?.params).toEqual({ days: 7, agentScope: "all" });
  });

+  it("waits for real all-agent usage caches before printing totals", async () => {
+    const stateDir = fs.mkdtempSync(path.join(os.tmpdir(), "openclaw-usage-cost-cli-"));
+    const config = {
+      agents: {
+        list: [{ id: "main", default: true }, { id: "dev" }],
+      },
+      session: {},
+    } as OpenClawConfig;
+    const seedUsage = (agentId: string, totalTokens: number, totalCost: number) => {
+      const sessionsDir = path.join(stateDir, "agents", agentId, "sessions");
+      fs.mkdirSync(sessionsDir, { recursive: true });
+      const session = SessionManager.create(sessionsDir, sessionsDir);
+      session.appendMessage({
+        role: "assistant",
+        content: [{ type: "text", text: "done" }],
+        api: "openai-responses",
+        provider: "openai",
+        model: "gpt-5.4",
+        usage: {
+          input: totalTokens,
+          output: 0,
+          cacheRead: 0,
+          cacheWrite: 0,
+          totalTokens,
+          cost: {
+            input: totalCost,
+            output: 0,
+            cacheRead: 0,
+            cacheWrite: 0,
+            total: totalCost,
+          },
+        },
+        stopReason: "stop",
+        timestamp: Date.now(),
+      });
+    };
+
+    try {
+      await withEnvOverride({ OPENCLAW_STATE_DIR: stateDir }, async () => {
+        seedUsage("main", 30, 0.03);
+        seedUsage("dev", 70, 0.07);
+        usageTestApi.costUsageCache.clear();
+        const observedStatuses: Array<string | undefined> = [];
+        callGateway.mockImplementation(async (raw) => {
+          const request = raw as { method?: string; params?: Record<string, unknown> };
+          if (request.method !== "usage.cost") {
+            return { ok: true };
+          }
+          return await new Promise((resolve, reject) => {
+            const respond: UsageCostHandlerArgs["respond"] = (ok, payload, error) => {
+              if (!ok) {
+                reject(new Error(error?.message ?? "usage.cost failed"));
+                return;
+              }
+              const summary = payload as { cacheStatus?: { status?: string } };
+              observedStatuses.push(summary.cacheStatus?.status);
+              resolve(payload);
+            };
+            const result = usageHandlers["usage.cost"]({
+              respond,
+              params: request.params ?? {},
+              context: { getRuntimeConfig: () => config },
+            } as unknown as UsageCostHandlerArgs);
+            Promise.resolve(result).catch(reject);
+          });
+        });
+
+        await runGatewayCommand(["gateway", "usage-cost", "--all-agents", "--days", "7", "--json"]);
+
+        expect(observedStatuses[0]).toBe("refreshing");
+        expect(observedStatuses.at(-1)).toBe("fresh");
+        expect(callGateway.mock.calls.length).toBeGreaterThanOrEqual(2);
+        const firstCostCall = firstMockArg(callGateway) as { timeoutMs?: number };
+        expect(firstCostCall.timeoutMs).toBeGreaterThan(290_000);
+        expect(defaultRuntime.writeJson).toHaveBeenCalledWith(
+          expect.objectContaining({
+            totals: expect.objectContaining({ totalTokens: 100, totalCost: 0.1 }),
+            cacheStatus: expect.objectContaining({ status: "fresh" }),
+          }),
+        );
+      });
+    } finally {
+      usageTestApi.costUsageCache.clear();
+      fs.rmSync(stateDir, { recursive: true, force: true });
+    }
+  });
+
+  it.each(["refreshing", "partial", "stale"] as const)(
+    "uses --timeout as the command-wide usage-cost settle budget for %s caches",
+    async (status) => {
+      callGateway.mockResolvedValue({
+        cacheStatus: { status, cachedFiles: 0, pendingFiles: 1 },
+      });
+
+      await expectGatewayExit([
+        "gateway",
+        "usage-cost",
+        "--all-agents",
+        "--timeout",
+        "50",
+        "--json",
+      ]);
+
+      expect(callGateway).toHaveBeenCalledTimes(1);
+      const costCall = firstMockArg(callGateway) as { method?: string; timeoutMs?: number };
+      expect(costCall.method).toBe("usage.cost");
+      expect(costCall.timeoutMs).toBeGreaterThan(0);
+      expect(costCall.timeoutMs).toBeLessThanOrEqual(50);
+      expect(runtimeErrors.join("\n")).toContain("Timed out waiting for usage cost cache refresh");
+    },
+  );
+
  it("rejects combining --agent with --all-agents for usage-cost", async () => {
    callGateway.mockClear();

@@ -17,10 +17,12 @@ import type {
 import type { WriteDiagnosticSupportExportResult } from "../../logging/diagnostic-support-export.js";
 import { defaultRuntime } from "../../runtime.js";
 import { createLazyImportLoader } from "../../shared/lazy-promise.js";
+import { sleep } from "../../utils/sleep.js";
 import { inheritOptionFromParent } from "../command-options.js";
 import { addGatewayServiceCommands } from "../daemon-cli/register-service-commands.js";
 import { parseGatewayPortOption } from "../gateway-port-option.js";
 import { formatHelpExamples } from "../help-format.js";
+import { parseTimeoutMsWithFallback } from "../parse-timeout.js";
 import type { GatewayRpcOpts } from "./call.js";
 import type { GatewayDiscoverOpts } from "./discover.js";
 import { addGatewayRunCommand } from "./run-command.js";
@@ -50,6 +52,11 @@ const daemonStatusGatherModuleLoader = createLazyImportLoader(
  () => import("../daemon-cli/status.gather.js"),
 );

+const DEFAULT_GATEWAY_RPC_TIMEOUT_MS = 10_000;
+const DEFAULT_USAGE_COST_TIMEOUT_MS = 5 * 60_000;
+const USAGE_COST_SETTLE_INITIAL_POLL_MS = 250;
+const USAGE_COST_SETTLE_MAX_POLL_MS = 5_000;
+
 function loadConfigModule() {
  return configModuleLoader.load();
 }
@@ -90,12 +97,12 @@ function loadDaemonStatusGatherModule() {
  return daemonStatusGatherModuleLoader.load();
 }

-function gatewayCallOpts(cmd: Command): Command {
+function gatewayCallOpts(cmd: Command, defaultTimeoutMs = DEFAULT_GATEWAY_RPC_TIMEOUT_MS): Command {
  return cmd
    .option("--url <url>", "Gateway WebSocket URL (defaults to gateway.remote.url when configured)")
    .option("--token <token>", "Gateway token (if required)")
    .option("--password <password>", "Gateway password (password auth)")
-    .option("--timeout <ms>", "Timeout in ms", "10000")
+    .option("--timeout <ms>", "Timeout in ms", String(defaultTimeoutMs))
    .option("--expect-final", "Wait for final response (agent)", false)
    .option("--json", "Output JSON", false);
 }
@@ -105,6 +112,48 @@ async function callGatewayCli(method: string, opts: GatewayRpcOpts, params?: unk
  return mod.callGatewayCli(method, opts, params);
 }

+async function loadSettledCostUsageSummary(
+  rpcOpts: GatewayRpcOpts,
+  params: { days: number; agentId?: string; agentScope?: "all" },
+): Promise<CostUsageSummary> {
+  const timeoutMs = parseTimeoutMsWithFallback(rpcOpts.timeout, DEFAULT_USAGE_COST_TIMEOUT_MS, {
+    invalidType: "error",
+  });
+  const deadline = Date.now() + timeoutMs;
+  let lastSummary: CostUsageSummary | undefined;
+  let pollMs = USAGE_COST_SETTLE_INITIAL_POLL_MS;
+  for (;;) {
+    const remainingBeforeCallMs = deadline - Date.now();
+    if (remainingBeforeCallMs <= 0) {
+      throw createUsageCostSettleTimeoutError(lastSummary);
+    }
+    const callOpts = { ...rpcOpts, timeout: String(remainingBeforeCallMs) };
+    const summary = (await callGatewayCli("usage.cost", callOpts, params)) as CostUsageSummary;
+    lastSummary = summary;
+    const status = summary.cacheStatus?.status;
+    if (!status || status === "fresh") {
+      return summary;
+    }
+
+    const remainingMs = deadline - Date.now();
+    if (remainingMs <= 0) {
+      throw createUsageCostSettleTimeoutError(summary);
+    }
+    // The existing RPC timeout is the whole command budget. Giving each retry a
+    // fresh timeout would let a short bounded audit keep running for minutes.
+    await sleep(Math.min(pollMs, remainingMs));
+    pollMs = Math.min(pollMs * 2, USAGE_COST_SETTLE_MAX_POLL_MS);
+  }
+}
+
+function createUsageCostSettleTimeoutError(summary?: CostUsageSummary): Error {
+  const cachedFiles = summary?.cacheStatus?.cachedFiles ?? 0;
+  const pendingFiles = summary?.cacheStatus?.pendingFiles ?? 0;
+  return new Error(
+    `Timed out waiting for usage cost cache refresh (${cachedFiles} cached, ${pendingFiles} pending)`,
+  );
+}
+
 async function runGatewayCommand(
  action: () => Promise<void>,
  label?: string,
@@ -567,17 +616,16 @@ export function registerGatewayCli(program: Command) {
            if (agentId && opts.allAgents) {
              throw new Error("Use --agent or --all-agents, not both");
            }
-            const result = await callGatewayCli("usage.cost", rpcOpts, {
+            const summary = await loadSettledCostUsageSummary(rpcOpts, {
              days,
              ...(agentId ? { agentId } : {}),
              ...(opts.allAgents ? { agentScope: "all" } : {}),
            });
            if (rpcOpts.json) {
-              defaultRuntime.writeJson(result);
+              defaultRuntime.writeJson(summary);
              return;
            }
            const rich = isRich();
-            const summary = result as CostUsageSummary;
            for (const line of await renderCostUsageSummaryAsync(summary, days, rich)) {
              defaultRuntime.log(line);
            }
@@ -586,6 +634,7 @@ export function registerGatewayCli(program: Command) {
          { json: Boolean(opts.json) },
        );
      }),
+    DEFAULT_USAGE_COST_TIMEOUT_MS,
  );

  gatewayCallOpts(