From 942c53e7f0795b157be41ec06554e0c05fe0e0a6 Mon Sep 17 00:00:00 2001 From: SP <8068616+sp-hk2ldn@users.noreply.github.com> Date: Fri, 6 Mar 2026 23:59:16 +0000 Subject: [PATCH] fix(agents): prevent totalTokens crash when assistant usage is missing (#34977) Merged via squash. Prepared head SHA: 1c14094f3f85f2c87596df2d67fd0ca6af4e9084 Co-authored-by: sp-hk2ldn <8068616+sp-hk2ldn@users.noreply.github.com> Co-authored-by: jalehman <550978+jalehman@users.noreply.github.com> Reviewed-by: @jalehman --- CHANGELOG.md | 1 + ...ed-runner.sanitize-session-history.test.ts | 125 ++++++++++++++++++ src/agents/pi-embedded-runner/google.ts | 117 +++++++++++++++- 3 files changed, 240 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e1ebcf56b75..6f1a55432a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -213,6 +213,7 @@ Docs: https://docs.openclaw.ai - Gateway/probes: keep `/health`, `/healthz`, `/ready`, and `/readyz` reachable when the Control UI is mounted at `/`, preserve plugin-owned route precedence on those paths, and make `/ready` and `/readyz` report channel-backed readiness with startup grace plus `503` on disconnected managed channels, while `/health` and `/healthz` stay shallow liveness probes. (#18446) Thanks @vibecodooor, @mahsumaktas, and @vincentkoc. - Feishu/media downloads: drop invalid timeout fields from SDK method calls now that client-level `httpTimeoutMs` applies to requests. (#38267) Thanks @ant1eicher and @thewilloftheshadow. - PI embedded runner/Feishu docs: propagate sender identity into embedded attempts so Feishu doc auto-grant restores requester access for embedded-runner executions. (#32915) thanks @cszhouwei. +- Agents/usage normalization: normalize missing or partial assistant usage snapshots before compaction accounting so `openclaw agent --json` no longer crashes when provider payloads omit `totalTokens` or related usage fields. (#34977) thanks @sp-hk2ldn. ## 2026.3.2 diff --git a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts index 13884cd904f..e216a45f4f3 100644 --- a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts +++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts @@ -330,6 +330,131 @@ describe("sanitizeSessionHistory", () => { expect(assistants[1]?.usage).toBeDefined(); }); + it("adds a zeroed assistant usage snapshot when usage is missing", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); + + const messages = castAgentMessages([ + { role: "user", content: "question" }, + { + role: "assistant", + content: [{ type: "text", text: "answer without usage" }], + }, + ]); + + const result = await sanitizeOpenAIHistory(messages); + const assistant = result.find((message) => message.role === "assistant") as + | (AgentMessage & { usage?: unknown }) + | undefined; + + expect(assistant?.usage).toEqual(makeZeroUsageSnapshot()); + }); + + it("normalizes mixed partial assistant usage fields to numeric totals", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); + + const messages = castAgentMessages([ + { role: "user", content: "question" }, + { + role: "assistant", + content: [{ type: "text", text: "answer with partial usage" }], + usage: { + output: 3, + cache_read_input_tokens: 9, + }, + }, + ]); + + const result = await sanitizeOpenAIHistory(messages); + const assistant = result.find((message) => message.role === "assistant") as + | (AgentMessage & { usage?: unknown }) + | undefined; + + expect(assistant?.usage).toEqual({ + input: 0, + output: 3, + cacheRead: 9, + cacheWrite: 0, + totalTokens: 12, + }); + }); + + it("preserves existing usage cost while normalizing token fields", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); + + const messages = castAgentMessages([ + { role: "user", content: "question" }, + { + role: "assistant", + content: [{ type: "text", text: "answer with partial usage and cost" }], + usage: { + output: 3, + cache_read_input_tokens: 9, + cost: { + input: 1.25, + output: 2.5, + cacheRead: 0.25, + cacheWrite: 0, + total: 4, + }, + }, + }, + ]); + + const result = await sanitizeOpenAIHistory(messages); + const assistant = result.find((message) => message.role === "assistant") as + | (AgentMessage & { usage?: unknown }) + | undefined; + + expect(assistant?.usage).toEqual({ + ...makeZeroUsageSnapshot(), + input: 0, + output: 3, + cacheRead: 9, + cacheWrite: 0, + totalTokens: 12, + cost: { + input: 1.25, + output: 2.5, + cacheRead: 0.25, + cacheWrite: 0, + total: 4, + }, + }); + }); + + it("preserves unknown cost when token fields already match", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); + + const messages = castAgentMessages([ + { role: "user", content: "question" }, + { + role: "assistant", + content: [{ type: "text", text: "answer with complete numeric usage but no cost" }], + usage: { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + }, + }, + ]); + + const result = await sanitizeOpenAIHistory(messages); + const assistant = result.find((message) => message.role === "assistant") as + | (AgentMessage & { usage?: unknown }) + | undefined; + + expect(assistant?.usage).toEqual({ + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + }); + expect((assistant?.usage as { cost?: unknown } | undefined)?.cost).toBeUndefined(); + }); + it("drops stale usage when compaction summary appears before kept assistant messages", async () => { vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); diff --git a/src/agents/pi-embedded-runner/google.ts b/src/agents/pi-embedded-runner/google.ts index 094aa9142c3..4daf30552b3 100644 --- a/src/agents/pi-embedded-runner/google.ts +++ b/src/agents/pi-embedded-runner/google.ts @@ -25,7 +25,12 @@ import { } from "../session-transcript-repair.js"; import type { TranscriptPolicy } from "../transcript-policy.js"; import { resolveTranscriptPolicy } from "../transcript-policy.js"; -import { makeZeroUsageSnapshot } from "../usage.js"; +import { + makeZeroUsageSnapshot, + normalizeUsage, + type AssistantUsageSnapshot, + type UsageLike, +} from "../usage.js"; import { log } from "./logger.js"; import { dropThinkingBlocks } from "./thinking.js"; import { describeUnknownError } from "./utils.js"; @@ -200,6 +205,111 @@ function stripStaleAssistantUsageBeforeLatestCompaction(messages: AgentMessage[] return touched ? out : messages; } +function normalizeAssistantUsageSnapshot(usage: unknown) { + const normalized = normalizeUsage((usage ?? undefined) as UsageLike | undefined); + if (!normalized) { + return makeZeroUsageSnapshot(); + } + const input = normalized.input ?? 0; + const output = normalized.output ?? 0; + const cacheRead = normalized.cacheRead ?? 0; + const cacheWrite = normalized.cacheWrite ?? 0; + const totalTokens = normalized.total ?? input + output + cacheRead + cacheWrite; + const cost = normalizeAssistantUsageCost(usage); + return { + input, + output, + cacheRead, + cacheWrite, + totalTokens, + ...(cost ? { cost } : {}), + }; +} + +function normalizeAssistantUsageCost(usage: unknown): AssistantUsageSnapshot["cost"] | undefined { + const base = makeZeroUsageSnapshot().cost; + if (!usage || typeof usage !== "object") { + return undefined; + } + const rawCost = (usage as { cost?: unknown }).cost; + if (!rawCost || typeof rawCost !== "object") { + return undefined; + } + const cost = rawCost as Record; + const inputRaw = toFiniteCostNumber(cost.input); + const outputRaw = toFiniteCostNumber(cost.output); + const cacheReadRaw = toFiniteCostNumber(cost.cacheRead); + const cacheWriteRaw = toFiniteCostNumber(cost.cacheWrite); + const totalRaw = toFiniteCostNumber(cost.total); + if ( + inputRaw === undefined && + outputRaw === undefined && + cacheReadRaw === undefined && + cacheWriteRaw === undefined && + totalRaw === undefined + ) { + return undefined; + } + const input = inputRaw ?? base.input; + const output = outputRaw ?? base.output; + const cacheRead = cacheReadRaw ?? base.cacheRead; + const cacheWrite = cacheWriteRaw ?? base.cacheWrite; + const total = totalRaw ?? input + output + cacheRead + cacheWrite; + return { input, output, cacheRead, cacheWrite, total }; +} + +function toFiniteCostNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + +function ensureAssistantUsageSnapshots(messages: AgentMessage[]): AgentMessage[] { + if (messages.length === 0) { + return messages; + } + + let touched = false; + const out = [...messages]; + for (let i = 0; i < out.length; i += 1) { + const message = out[i] as (AgentMessage & { role?: unknown; usage?: unknown }) | undefined; + if (!message || message.role !== "assistant") { + continue; + } + const normalizedUsage = normalizeAssistantUsageSnapshot(message.usage); + const usageCost = + message.usage && typeof message.usage === "object" + ? (message.usage as { cost?: unknown }).cost + : undefined; + const normalizedCost = normalizedUsage.cost; + if ( + message.usage && + typeof message.usage === "object" && + (message.usage as { input?: unknown }).input === normalizedUsage.input && + (message.usage as { output?: unknown }).output === normalizedUsage.output && + (message.usage as { cacheRead?: unknown }).cacheRead === normalizedUsage.cacheRead && + (message.usage as { cacheWrite?: unknown }).cacheWrite === normalizedUsage.cacheWrite && + (message.usage as { totalTokens?: unknown }).totalTokens === normalizedUsage.totalTokens && + ((normalizedCost && + usageCost && + typeof usageCost === "object" && + (usageCost as { input?: unknown }).input === normalizedCost.input && + (usageCost as { output?: unknown }).output === normalizedCost.output && + (usageCost as { cacheRead?: unknown }).cacheRead === normalizedCost.cacheRead && + (usageCost as { cacheWrite?: unknown }).cacheWrite === normalizedCost.cacheWrite && + (usageCost as { total?: unknown }).total === normalizedCost.total) || + (!normalizedCost && usageCost === undefined)) + ) { + continue; + } + out[i] = { + ...(message as unknown as Record), + usage: normalizedUsage, + } as AgentMessage; + touched = true; + } + + return touched ? out : messages; +} + export function findUnsupportedSchemaKeywords(schema: unknown, path: string): string[] { if (!schema || typeof schema !== "object") { return []; @@ -449,8 +559,9 @@ export async function sanitizeSessionHistory(params: { ? sanitizeToolUseResultPairing(sanitizedToolCalls) : sanitizedToolCalls; const sanitizedToolResults = stripToolResultDetails(repairedTools); - const sanitizedCompactionUsage = - stripStaleAssistantUsageBeforeLatestCompaction(sanitizedToolResults); + const sanitizedCompactionUsage = ensureAssistantUsageSnapshots( + stripStaleAssistantUsageBeforeLatestCompaction(sanitizedToolResults), + ); const isOpenAIResponsesApi = params.modelApi === "openai-responses" || params.modelApi === "openai-codex-responses";