diff --git a/src/agents/embedded-agent-runner/run.overflow-compaction.harness.ts b/src/agents/embedded-agent-runner/run.overflow-compaction.harness.ts index ec18cb7f87d..7b7f534030a 100644 --- a/src/agents/embedded-agent-runner/run.overflow-compaction.harness.ts +++ b/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 diff --git a/src/agents/embedded-agent-runner/run.ts b/src/agents/embedded-agent-runner/run.ts index 0db7c7de4a3..541bbc4ff52 100644 --- a/src/agents/embedded-agent-runner/run.ts +++ b/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 diff --git a/src/agents/embedded-agent-runner/run/helpers.test.ts b/src/agents/embedded-agent-runner/run/helpers.test.ts index 164cde3a7f2..967e6a9ab7b 100644 --- a/src/agents/embedded-agent-runner/run/helpers.test.ts +++ b/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 diff --git a/src/agents/embedded-agent-runner/run/helpers.ts b/src/agents/embedded-agent-runner/run/helpers.ts index 12119a2eeda..b1d0d1e7e15 100644 --- a/src/agents/embedded-agent-runner/run/helpers.ts +++ b/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, diff --git a/src/agents/embedded-agent-runner/usage-reporting.test.ts b/src/agents/embedded-agent-runner/usage-reporting.test.ts index 39fa68725a8..68339430ea3 100644 --- a/src/agents/embedded-agent-runner/usage-reporting.test.ts +++ b/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: { diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts b/src/agents/embedded-agent-subscribe.handlers.messages.test.ts index da2706a7e87..e230b64e37b 100644 --- a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts +++ b/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({ diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.ts b/src/agents/embedded-agent-subscribe.handlers.messages.ts index 68ec3495dfe..c1c895cf013 100644 --- a/src/agents/embedded-agent-subscribe.handlers.messages.ts +++ b/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) + : 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 | undefined { return value && typeof value === "object" && !Array.isArray(value) ? (value as Record) @@ -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); diff --git a/src/agents/embedded-agent-subscribe.handlers.ts b/src/agents/embedded-agent-subscribe.handlers.ts index 685b3293510..fe4057eeef8 100644 --- a/src/agents/embedded-agent-subscribe.handlers.ts +++ b/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, + ctx.state.pendingAssistantUsage, + ); + } scheduleEvent(evt, () => { return handleMessageEnd(ctx, evt as never); }); diff --git a/src/agents/embedded-agent-subscribe.subscribe-embedded-agent-session.calls-onblockreplyflush-before-tool-execution-start-preserve.test.ts b/src/agents/embedded-agent-subscribe.subscribe-embedded-agent-session.calls-onblockreplyflush-before-tool-execution-start-preserve.test.ts index d3543f5066a..12ef9fd9160 100644 --- a/src/agents/embedded-agent-subscribe.subscribe-embedded-agent-session.calls-onblockreplyflush-before-tool-execution-start-preserve.test.ts +++ b/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((resolve) => { + releaseFirstReply = resolve; + }); + let blockReplyCount = 0; + + subscribeEmbeddedAgentSession({ + session: session as unknown as Parameters[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"); + }); });