fix(agents): preserve embedded completions usage (#96523)

Co-authored-by: Vincent Koc <25068+vincentkoc@users.noreply.github.com>

src/agents/embedded-agent-runner/run.overflow-compaction.harness.ts

@@ -585,6 +585,24 @@ export async function loadRunOverflowCompactionHarness(): Promise<{
     normalizeUsage: vi.fn((usage?: unknown) =>
       usage && typeof usage === "object" ? usage : undefined,
     ),
+    hasNonzeroUsage: vi.fn(
+      (usage?: {
+        total?: number;
+        input?: number;
+        output?: number;
+        cacheRead?: number;
+        cacheWrite?: number;
+        reasoningTokens?: number;
+      }) =>
+        [
+          usage?.total,
+          usage?.input,
+          usage?.output,
+          usage?.cacheRead,
+          usage?.cacheWrite,
+          usage?.reasoningTokens,
+        ].some((value) => (value ?? 0) > 0),
+    ),
     derivePromptTokens: vi.fn(
       (usage?: { input?: number; cacheRead?: number; cacheWrite?: number }) =>
         usage

src/agents/embedded-agent-runner/run.ts

@@ -192,6 +192,7 @@ import {
   resolveActiveErrorContext,
   resolveFinalAssistantRawText,
   resolveFinalAssistantVisibleText,
+  resolveLatestCallUsage,
   resolveMaxRunRetryIterations,
   resolveReportedModelRef,
   MAX_SAME_MODEL_RATE_LIMIT_RETRIES,
@@ -2319,12 +2320,22 @@ async function runEmbeddedAgentInternal(
                 )
               : bootstrapPromptWarningSignaturesSeen);
           const lastAssistantUsage = normalizeUsage(sessionLastAssistant?.usage as UsageLike);
-          const attemptUsage = attempt.attemptUsage ?? lastAssistantUsage;
+          const currentAttemptAssistantUsage = normalizeUsage(
+            currentAttemptAssistant?.usage as UsageLike,
+          );
+          const promptCacheLastCallUsage = normalizeUsage(
+            attempt.promptCache?.lastCallUsage as UsageLike,
+          );
+          const callUsage = resolveLatestCallUsage({
+            currentAttemptCandidates: [currentAttemptAssistantUsage, promptCacheLastCallUsage],
+            carriedCandidates: [lastRunPromptUsage, lastAssistantUsage],
+          });
+          const attemptUsage = attempt.attemptUsage ?? callUsage.currentAttempt;
           mergeUsageIntoAccumulator(usageAccumulator, attemptUsage);
           // Keep prompt size from the latest model call so session totalTokens
           // reflects current context usage, not accumulated tool-loop usage.
-          lastRunPromptUsage = lastAssistantUsage ?? attemptUsage;
-          lastTurnTotal = lastAssistantUsage?.total ?? attemptUsage?.total;
+          lastRunPromptUsage = callUsage.latest;
+          lastTurnTotal = callUsage.latest?.total;
           // Idle-timeout cost-runaway breaker (#76293). Logic lives in the
           // pure helper below so it stays unit-testable; the run loop just
           // feeds it the latest attempt outcome and bails through the

src/agents/embedded-agent-runner/run/helpers.test.ts

@@ -7,6 +7,7 @@ import {
   buildErrorAgentMeta,
   resolveFinalAssistantRawText,
   resolveFinalAssistantVisibleText,
+  resolveLatestCallUsage,
   resolveNextSameModelRateLimitRetryCount,
   resolveSameModelRateLimitRetryDelayMs,
 } from "./helpers.js";
@@ -160,6 +161,36 @@ describe("resolveNextSameModelRateLimitRetryCount", () => {
   });
 });
 
+describe("resolveLatestCallUsage", () => {
+  it("preserves the previous exact call across a zero-usage retry", () => {
+    const previous = { input: 12, output: 3, total: 15 };
+
+    expect(
+      resolveLatestCallUsage({
+        currentAttemptCandidates: [{ input: 0, output: 0, total: 0 }, undefined],
+        carriedCandidates: [previous],
+      }),
+    ).toEqual({
+      currentAttempt: undefined,
+      latest: previous,
+    });
+  });
+
+  it("replaces the previous call when a new nonzero snapshot arrives", () => {
+    const latest = { input: 20, output: 4, total: 24 };
+
+    expect(
+      resolveLatestCallUsage({
+        currentAttemptCandidates: [{ input: 0, output: 0, total: 0 }, latest],
+        carriedCandidates: [{ input: 12, output: 3, total: 15 }],
+      }),
+    ).toEqual({
+      currentAttempt: latest,
+      latest,
+    });
+  });
+});
+
 describe("buildErrorAgentMeta", () => {
   it("preserves active session file for error exits after transcript rotation", () => {
     // Error metadata follows the active session after transcript rotation so

src/agents/embedded-agent-runner/run/helpers.ts

@@ -7,7 +7,12 @@ import type { AssistantMessage } from "../../../llm/types.js";
 import { extractAssistantTextForPhase } from "../../../shared/chat-message-content.js";
 import { resolveAgentConfig } from "../../agent-scope-config.js";
 import { extractAssistantVisibleText } from "../../embedded-agent-utils.js";
-import { derivePromptTokens, normalizeUsage } from "../../usage.js";
+import {
+  derivePromptTokens,
+  hasNonzeroUsage,
+  normalizeUsage,
+  type NormalizedUsage,
+} from "../../usage.js";
 import type { EmbeddedAgentMeta } from "../types.js";
 import { toLastCallUsage, toNormalizedUsage, type UsageAccumulator } from "../usage-accumulator.js";
 
@@ -184,6 +189,20 @@ export function resolveReportedModelRef(params: {
   };
 }
 
+export function resolveLatestCallUsage(params: {
+  currentAttemptCandidates: readonly (NormalizedUsage | undefined)[];
+  carriedCandidates: readonly (NormalizedUsage | undefined)[];
+}): {
+  currentAttempt: NormalizedUsage | undefined;
+  latest: NormalizedUsage | undefined;
+} {
+  const currentAttempt = params.currentAttemptCandidates.find(hasNonzeroUsage);
+  return {
+    currentAttempt,
+    latest: currentAttempt ?? params.carriedCandidates.find(hasNonzeroUsage),
+  };
+}
+
 export function buildUsageAgentMetaFields(params: {
   usageAccumulator: UsageAccumulator;
   lastAssistantUsage?: UsageSnapshot | null;
@@ -194,8 +213,12 @@ export function buildUsageAgentMetaFields(params: {
   if (usage && params.lastTurnTotal && params.lastTurnTotal > 0) {
     usage.total = params.lastTurnTotal;
   }
-  const lastCallUsage =
-    normalizeUsage(params.lastAssistantUsage as never) ?? toLastCallUsage(params.usageAccumulator);
+  const lastAssistantUsage = normalizeUsage(params.lastAssistantUsage as never);
+  const lastCallUsage = hasNonzeroUsage(lastAssistantUsage)
+    ? lastAssistantUsage
+    : hasNonzeroUsage(params.lastRunPromptUsage)
+      ? params.lastRunPromptUsage
+      : toLastCallUsage(params.usageAccumulator);
   const promptTokens = derivePromptTokens(params.lastRunPromptUsage);
   return {
     usage,

src/agents/embedded-agent-runner/usage-reporting.test.ts

@@ -201,6 +201,49 @@ describe("runEmbeddedAgent usage reporting", () => {
     expect(usage?.total).toBe(200);
   });
 
+  it("uses current-attempt usage when the persisted assistant snapshot is zeroed", async () => {
+    mockedRunEmbeddedAttempt.mockResolvedValueOnce(
+      makeAttemptResult({
+        assistantTexts: ["Response 1", "Response 2"],
+        lastAssistant: makeAssistantMessage({
+          usage: {
+            input: 0,
+            output: 0,
+            cacheRead: 0,
+            cacheWrite: 0,
+            totalTokens: 0,
+          } as unknown as AssistantMessage["usage"],
+        }),
+        currentAttemptAssistant: makeAssistantMessage({
+          usage: { input: 150, output: 50, total: 200 } as unknown as AssistantMessage["usage"],
+        }),
+        attemptUsage: { input: 250, output: 100, total: 350 },
+      }),
+    );
+
+    const result = await runEmbeddedAgent({
+      sessionId: "test-session",
+      sessionKey: "test-key",
+      sessionFile: "/tmp/session.json",
+      workspaceDir: "/tmp/workspace",
+      prompt: "hello",
+      timeoutMs: 30000,
+      runId: "run-zeroed-persisted-usage",
+    });
+
+    expect(result.meta.agentMeta?.usage).toMatchObject({
+      input: 250,
+      output: 100,
+      total: 200,
+    });
+    expect(result.meta.agentMeta?.lastCallUsage).toMatchObject({
+      input: 150,
+      output: 50,
+      total: 200,
+    });
+    expect(result.meta.agentMeta?.promptTokens).toBe(150);
+  });
+
   it("reports the resolved model provider when OpenClaw marks the assistant message as the native runtime", async () => {
     mockedResolveModelAsync.mockResolvedValueOnce({
       model: {

src/agents/embedded-agent-subscribe.handlers.messages.test.ts

@@ -1028,6 +1028,77 @@ describe("handleMessageUpdate commentary phase", () => {
 });
 
 describe("handleMessageEnd", () => {
+  it("persists streamed usage when the final assistant snapshot is zeroed", () => {
+    const ctx = createMessageEndContext({
+      state: {
+        pendingAssistantUsage: { input: 7, output: 5, reasoningTokens: 2, total: 12 },
+      },
+    });
+    const message = {
+      role: "assistant",
+      api: "openai-completions",
+      content: [{ type: "text", text: "Done." }],
+      usage: {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+      },
+    };
+
+    void handleMessageEnd(ctx, {
+      type: "message_end",
+      message,
+    } as never);
+
+    expect(firstMockArg(ctx.noteLastAssistant as never, "last assistant")).toMatchObject({
+      usage: {
+        input: 7,
+        output: 5,
+        cacheRead: 0,
+        cacheWrite: 0,
+        reasoningTokens: 2,
+        totalTokens: 12,
+      },
+    });
+    expect(ctx.recordAssistantUsage).toHaveBeenCalledWith(
+      expect.objectContaining({
+        input: 7,
+        output: 5,
+        reasoningTokens: 2,
+        totalTokens: 12,
+      }),
+    );
+  });
+
+  it("keeps authoritative final usage instead of pending stream usage", () => {
+    const ctx = createMessageEndContext({
+      state: {
+        pendingAssistantUsage: { input: 7, output: 5, total: 12 },
+      },
+    });
+    const message = {
+      role: "assistant",
+      content: [{ type: "text", text: "Done." }],
+      usage: {
+        input: 11,
+        output: 3,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 14,
+      },
+    };
+
+    void handleMessageEnd(ctx, {
+      type: "message_end",
+      message,
+    } as never);
+
+    expect(firstMockArg(ctx.noteLastAssistant as never, "last assistant")).toBe(message);
+    expect(ctx.recordAssistantUsage).toHaveBeenCalledWith(message.usage);
+  });
+
   it("warns when assistant text only pretends to call a registered tool", () => {
     const warn = vi.fn();
     const ctx = createMessageEndContext({

src/agents/embedded-agent-subscribe.handlers.messages.ts

@@ -42,6 +42,13 @@ import {
   sanitizeAssistantVisibleStreamText,
 } from "./embedded-agent-utils.js";
 import type { AgentEvent, AgentMessage } from "./runtime/index.js";
+import {
+  hasNonzeroUsage,
+  makeZeroUsageSnapshot,
+  normalizeUsage,
+  type NormalizedUsage,
+  type UsageLike,
+} from "./usage.js";
 
 function shouldSuppressAssistantVisibleOutput(message: AgentMessage | undefined): boolean {
   return resolveAssistantMessagePhase(message) === "commentary";
@@ -80,6 +87,67 @@ function isOpenAiCompletionsAssistantMessage(message: AgentMessage | undefined):
   return api === "openai-completions" || api === "openclaw-openai-completions-transport";
 }
 
+export function preservePendingAssistantUsage(
+  message: AssistantMessage,
+  pendingUsage: NormalizedUsage | undefined,
+): AssistantMessage {
+  if (isTranscriptOnlyOpenClawAssistantMessage(message) || !hasNonzeroUsage(pendingUsage)) {
+    return message;
+  }
+  const messageUsage = normalizeUsage((message as { usage?: UsageLike }).usage);
+  if (hasNonzeroUsage(messageUsage)) {
+    return message;
+  }
+
+  // Pending usage resets at each assistant-message boundary, so it belongs to
+  // this final snapshot. Only replace missing/zero usage; provider totals win.
+  const input = pendingUsage.input ?? 0;
+  const output = pendingUsage.output ?? 0;
+  const cacheRead = pendingUsage.cacheRead ?? 0;
+  const cacheWrite = pendingUsage.cacheWrite ?? 0;
+  message.usage = {
+    ...makeZeroUsageSnapshot(),
+    input,
+    output,
+    cacheRead,
+    cacheWrite,
+    totalTokens: pendingUsage.total ?? input + output + cacheRead + cacheWrite,
+    ...(pendingUsage.reasoningTokens !== undefined
+      ? { reasoningTokens: pendingUsage.reasoningTokens }
+      : {}),
+  };
+  return message;
+}
+
+export function capturePendingAssistantUsage(
+  ctx: EmbeddedAgentSubscribeContext,
+  evt: AgentEvent & { message: AgentMessage; assistantMessageEvent?: unknown },
+): void {
+  const msg = evt.message;
+  if (msg?.role !== "assistant" || isTranscriptOnlyOpenClawAssistantMessage(msg)) {
+    return;
+  }
+  const assistantRecord =
+    evt.assistantMessageEvent && typeof evt.assistantMessageEvent === "object"
+      ? (evt.assistantMessageEvent as Record<string, unknown>)
+      : undefined;
+  const evtType = typeof assistantRecord?.type === "string" ? assistantRecord.type : "";
+  if (evtType === "text_end" || evtType === "done" || evtType === "error") {
+    ctx.recordAssistantUsage(assistantRecord);
+  }
+}
+
+export function resetPendingAssistantUsage(
+  ctx: EmbeddedAgentSubscribeContext,
+  message: AgentMessage,
+): void {
+  if (message?.role !== "assistant" || isTranscriptOnlyOpenClawAssistantMessage(message)) {
+    return;
+  }
+  ctx.state.pendingAssistantUsage = undefined;
+  ctx.state.assistantUsageCommitted = false;
+}
+
 function asRecord(value: unknown): Record<string, unknown> | undefined {
   return value && typeof value === "object" && !Array.isArray(value)
     ? (value as Record<string, unknown>)
@@ -627,7 +695,7 @@ export function handleMessageUpdate(
   const evtType = typeof assistantRecord?.type === "string" ? assistantRecord.type : "";
 
   if (evtType === "text_end" || evtType === "done" || evtType === "error") {
-    ctx.recordAssistantUsage(assistantRecord);
+    capturePendingAssistantUsage(ctx, evt);
     if (evtType === "done" || evtType === "error") {
       ctx.commitAssistantUsage();
     }
@@ -964,7 +1032,7 @@ export function handleMessageEnd(
     return;
   }
 
-  const assistantMessage = msg;
+  const assistantMessage = preservePendingAssistantUsage(msg, ctx.state.pendingAssistantUsage);
   const assistantPhase = resolveAssistantMessagePhase(assistantMessage);
   const suppressVisibleAssistantOutput = shouldSuppressAssistantVisibleOutput(assistantMessage);
   const suppressDeterministicApprovalOutput = shouldSuppressDeterministicApprovalOutput(ctx.state);

src/agents/embedded-agent-subscribe.handlers.ts

@@ -8,9 +8,12 @@ import {
   handleCompactionStart,
 } from "./embedded-agent-subscribe.handlers.lifecycle.js";
 import {
+  capturePendingAssistantUsage,
   handleMessageEnd,
   handleMessageStart,
   handleMessageUpdate,
+  preservePendingAssistantUsage,
+  resetPendingAssistantUsage,
 } from "./embedded-agent-subscribe.handlers.messages.js";
 import {
   handleToolExecutionEnd,
@@ -22,6 +25,7 @@ import type {
   EmbeddedAgentSubscribeEvent,
 } from "./embedded-agent-subscribe.handlers.types.js";
 import { isPromiseLike } from "./embedded-agent-subscribe.promise.js";
+import type { AgentMessage } from "./runtime/index.js";
 
 /** Create the serialized event dispatcher for subscribed embedded-agent sessions. */
 export function createEmbeddedAgentSessionEventHandler(ctx: EmbeddedAgentSubscribeContext) {
@@ -78,16 +82,30 @@ export function createEmbeddedAgentSessionEventHandler(ctx: EmbeddedAgentSubscri
   return (evt: EmbeddedAgentSubscribeEvent) => {
     switch (evt.type) {
       case "message_start":
+        // Delivery from the previous message may still be queued, but usage is
+        // message-scoped. Reset only its accounting boundary synchronously so
+        // this message's streamed usage cannot inherit the prior commit state.
+        resetPendingAssistantUsage(ctx, evt.message as AgentMessage);
         scheduleEvent(evt, () => {
           handleMessageStart(ctx, evt as never);
         });
         return;
       case "message_update":
+        // AgentSession persists message_end after this listener returns, while
+        // delivery handlers may still be queued. Capture usage synchronously so
+        // the following final snapshot can be repaired before persistence.
+        capturePendingAssistantUsage(ctx, evt as never);
         scheduleEvent(evt, () => {
           handleMessageUpdate(ctx, evt as never);
         });
         return;
       case "message_end":
+        if ((evt.message as AgentMessage)?.role === "assistant") {
+          preservePendingAssistantUsage(
+            evt.message as Extract<AgentMessage, { role: "assistant" }>,
+            ctx.state.pendingAssistantUsage,
+          );
+        }
         scheduleEvent(evt, () => {
           return handleMessageEnd(ctx, evt as never);
         });

src/agents/embedded-agent-subscribe.subscribe-embedded-agent-session.calls-onblockreplyflush-before-tool-execution-start-preserve.test.ts

@@ -212,4 +212,103 @@ describe("subscribeEmbeddedAgentSession", () => {
       expect(flushSnapshots).toEqual([["Final reply before lifecycle end."]]);
     });
   });
+
+  it("repairs final usage before persistence when delivery work is queued", async () => {
+    const { session, emit } = createStubSessionHarness();
+    let releaseFirstReply: (() => void) | undefined;
+    const firstReplyPending = new Promise<void>((resolve) => {
+      releaseFirstReply = resolve;
+    });
+    let blockReplyCount = 0;
+
+    subscribeEmbeddedAgentSession({
+      session: session as unknown as Parameters<typeof subscribeEmbeddedAgentSession>[0]["session"],
+      runId: "run-queued-usage-persistence",
+      onBlockReply: () => {
+        blockReplyCount += 1;
+        return blockReplyCount === 1 ? firstReplyPending : undefined;
+      },
+      onBlockReplyFlush: vi.fn(),
+      blockReplyBreak: "message_end",
+    });
+
+    emit({
+      type: "message_start",
+      message: { role: "assistant" },
+    });
+    emit({
+      type: "message_end",
+      message: {
+        role: "assistant",
+        content: [{ type: "text", text: "First reply." }],
+        usage: {
+          input: 3,
+          output: 2,
+          cacheRead: 0,
+          cacheWrite: 0,
+          totalTokens: 5,
+        },
+      },
+    });
+
+    emit({
+      type: "message_start",
+      message: { role: "assistant" },
+    });
+    emit({
+      type: "message_update",
+      message: {
+        role: "assistant",
+        api: "openai-completions",
+        content: [{ type: "text", text: "Second reply." }],
+      },
+      assistantMessageEvent: {
+        type: "text_end",
+        usage: {
+          input: 7,
+          output: 5,
+          cacheRead: 0,
+          cacheWrite: 0,
+          reasoningTokens: 2,
+          totalTokens: 12,
+        },
+      },
+    });
+    const finalMessage = {
+      role: "assistant",
+      api: "openai-completions",
+      content: [{ type: "text", text: "Second reply." }],
+      usage: {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+      },
+    };
+    emit({ type: "message_end", message: finalMessage });
+
+    // AgentSession persists immediately after notifying listeners, so this
+    // mutation must happen before the queued message_end handler executes.
+    expect(finalMessage.usage).toMatchObject({
+      input: 7,
+      output: 5,
+      reasoningTokens: 2,
+      totalTokens: 12,
+    });
+
+    releaseFirstReply?.();
+    await vi.waitFor(() => {
+      expect(blockReplyCount).toBeGreaterThanOrEqual(2);
+    });
+
+    const transcriptOnlyMessage = {
+      role: "assistant",
+      provider: "openclaw",
+      model: "delivery-mirror",
+      content: [{ type: "text", text: "Already delivered." }],
+    };
+    emit({ type: "message_end", message: transcriptOnlyMessage });
+    expect(transcriptOnlyMessage).not.toHaveProperty("usage");
+  });
 });