diff --git a/docs/.generated/plugin-sdk-api-baseline.sha256 b/docs/.generated/plugin-sdk-api-baseline.sha256 index 73c23704475d..7c79830a7794 100644 --- a/docs/.generated/plugin-sdk-api-baseline.sha256 +++ b/docs/.generated/plugin-sdk-api-baseline.sha256 @@ -1,2 +1,2 @@ -8928bb6ffc8ab7149502410bea3634ad737a574a490e43accced6ce5fe586c0c plugin-sdk-api-baseline.json -5573bf9f5f3d6802a673c2406e2298de06f14128efb2dcdb779fa430b3abe602 plugin-sdk-api-baseline.jsonl +6648558dcec4765bcb21a0837b3c19331129636b55741280ccc2837df9c03579 plugin-sdk-api-baseline.json +c86d4ddbfc2e5842d87796dc06454fbf3ccf61f851417c38553f98dbf6f49279 plugin-sdk-api-baseline.jsonl diff --git a/docs/docs_map.md b/docs/docs_map.md index 18516e7f2594..e3b6881adb67 100644 --- a/docs/docs_map.md +++ b/docs/docs_map.md @@ -1357,6 +1357,7 @@ Do not edit it by hand; run `pnpm docs:map:gen`. - H2: Operations and approval - H2: Setup bootstrap - H2: Model-assisted planner + - H3: CLI harness trust model - H2: Switching to an agent - H2: Message rescue mode - H2: Related diff --git a/packages/agent-core/src/harness/compaction/compaction.test.ts b/packages/agent-core/src/harness/compaction/compaction.test.ts index e8b948697da2..e6ee99f55117 100644 --- a/packages/agent-core/src/harness/compaction/compaction.test.ts +++ b/packages/agent-core/src/harness/compaction/compaction.test.ts @@ -1,9 +1,128 @@ import { describe, expect, it, vi } from "vitest"; import { createAssistantMessageEventStream } from "../../llm.js"; import type { AssistantMessage, Model, StreamFn } from "../../llm.js"; -import { compact, generateSummary } from "./compaction.js"; +import { + calculateContextTokens, + compact, + estimateContextTokens, + generateSummary, +} from "./compaction.js"; import { createFileOps } from "./utils.js"; +describe("calculateContextTokens", () => { + it("prefers the final-iteration context snapshot over aggregate billing usage", () => { + expect( + calculateContextTokens({ + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + totalTokens: 927_907, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }), + ).toBe(163_978); + }); + + it("preserves the numeric compatibility fallback when the snapshot is unavailable", () => { + expect( + calculateContextTokens({ + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + totalTokens: 927_907, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }), + ).toBe(927_907); + }); + + it("estimates the transcript instead of using aggregate billing when context is unavailable", () => { + const estimate = estimateContextTokens([ + { role: "user", content: "hello", timestamp: 0 }, + { + role: "assistant", + content: [{ type: "text", text: "done" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-fable-5", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + totalTokens: 927_907, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 1, + }, + ]); + + expect(estimate.tokens).toBeLessThan(927_907); + expect(estimate.tokens).toBeGreaterThan(0); + expect(estimate.usageTokens).toBe(0); + expect(estimate.lastUsageIndex).toBeNull(); + }); + + it("uses the previous exact snapshot and estimates only the unavailable tail", () => { + const estimate = estimateContextTokens([ + { + role: "assistant", + content: [{ type: "text", text: "previous" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-fable-5", + usage: { + input: 12, + output: 1_000, + cacheRead: 148_862, + cacheWrite: 0, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 149_874, + }, + totalTokens: 149_874, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 0, + }, + { role: "user", content: "next", timestamp: 1 }, + { + role: "assistant", + content: [{ type: "text", text: "done" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-fable-5", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + totalTokens: 927_907, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 2, + }, + ]); + + expect(estimate.usageTokens).toBe(149_874); + expect(estimate.tokens).toBeGreaterThan(149_874); + expect(estimate.tokens).toBeLessThan(927_907); + expect(estimate.lastUsageIndex).toBe(0); + }); +}); + describe("generateSummary thinking options", () => { it("maps explicit Fable off to low effort for compaction", async () => { const model: Model = { diff --git a/packages/agent-core/src/harness/compaction/compaction.ts b/packages/agent-core/src/harness/compaction/compaction.ts index c778727d8f64..39e6ecabd9eb 100644 --- a/packages/agent-core/src/harness/compaction/compaction.ts +++ b/packages/agent-core/src/harness/compaction/compaction.ts @@ -148,6 +148,9 @@ export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = { /** Calculate total context tokens from provider usage. */ export function calculateContextTokens(usage: Usage): number { + if (usage.contextUsage?.state === "available") { + return usage.contextUsage.totalTokens; + } return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite; } function getAssistantUsage(msg: AgentMessage): Usage | undefined { @@ -184,7 +187,7 @@ export interface ContextUsageEstimate { tokens: number; /** Tokens reported by the most recent assistant usage block. */ usageTokens: number; - /** Estimated tokens after the most recent assistant usage block. */ + /** Estimated tokens not covered by usable provider usage. */ trailingTokens: number; /** Index of the message that provided usage, or null when none exists. */ lastUsageIndex: number | null; @@ -195,7 +198,7 @@ function getLastAssistantUsageInfo( ): { usage: Usage; index: number } | undefined { for (let i = messages.length - 1; i >= 0; i--) { const usage = getAssistantUsage(messages[i]); - if (usage) { + if (usage && usage.contextUsage?.state !== "unavailable") { return { usage, index: i }; } } diff --git a/packages/ai/src/internal/anthropic.ts b/packages/ai/src/internal/anthropic.ts index 9e3c0da1ccea..ea98015395eb 100644 --- a/packages/ai/src/internal/anthropic.ts +++ b/packages/ai/src/internal/anthropic.ts @@ -5,3 +5,4 @@ export * from "../providers/anthropic-refusal.js"; export * from "../providers/anthropic-server-fallback.js"; export * from "../providers/anthropic-thinking-replay.js"; export * from "../providers/anthropic-tool-projection.js"; +export * from "../providers/anthropic-usage.js"; diff --git a/packages/ai/src/providers/anthropic-usage.test.ts b/packages/ai/src/providers/anthropic-usage.test.ts new file mode 100644 index 000000000000..a7627c07c0d7 --- /dev/null +++ b/packages/ai/src/providers/anthropic-usage.test.ts @@ -0,0 +1,77 @@ +import { describe, expect, it } from "vitest"; +import { readLastAnthropicIterationUsage } from "./anthropic-usage.js"; + +describe("readLastAnthropicIterationUsage", () => { + it.each(["message", "compaction", "advisor_message"])( + "reads the final %s iteration as the context snapshot", + (type) => { + expect( + readLastAnthropicIterationUsage({ + iterations: [ + { + type: "message", + input_tokens: 1, + output_tokens: 2, + cache_read_input_tokens: 3, + cache_creation_input_tokens: 4, + }, + { + type, + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + ], + }), + ).toEqual({ + state: "valid", + usage: { + contextPromptTokens: 148_874, + totalTokens: 163_978, + }, + }); + }, + ); + + it("reports absent iterations separately from malformed iterations", () => { + expect(readLastAnthropicIterationUsage({ input_tokens: 1 })).toEqual({ state: "absent" }); + }); + + it("does not reuse an earlier iteration when the final iteration is malformed", () => { + expect( + readLastAnthropicIterationUsage({ + iterations: [ + { + type: "message", + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + { + type: "message", + input_tokens: "malformed", + output_tokens: 1, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + ], + }), + ).toEqual({ state: "invalid" }); + }); + + it("rejects a final iteration with incomplete cache usage", () => { + expect( + readLastAnthropicIterationUsage({ + iterations: [ + { + type: "message", + input_tokens: 12, + output_tokens: 15_104, + }, + ], + }), + ).toEqual({ state: "invalid" }); + }); +}); diff --git a/packages/ai/src/providers/anthropic-usage.ts b/packages/ai/src/providers/anthropic-usage.ts new file mode 100644 index 000000000000..98fafc2ab7c2 --- /dev/null +++ b/packages/ai/src/providers/anthropic-usage.ts @@ -0,0 +1,83 @@ +type AnthropicUsagePayload = { + input_tokens?: unknown; + output_tokens?: unknown; + cache_read_input_tokens?: unknown; + cache_creation_input_tokens?: unknown; + iterations?: unknown; +}; + +export type AnthropicPromptUsageSnapshot = { + input: number; + cacheRead: number; + cacheWrite: number; +}; + +export type AnthropicIterationUsageSnapshot = { + contextPromptTokens: number; + totalTokens: number; +}; + +export type AnthropicIterationUsageResult = + | { state: "absent" } + | { state: "invalid" } + | { state: "valid"; usage: AnthropicIterationUsageSnapshot }; + +export function readAnthropicUsageTokenCount(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined; +} + +export function readAnthropicPromptUsageSnapshot( + usage: AnthropicUsagePayload, +): AnthropicPromptUsageSnapshot | undefined { + const input = readAnthropicUsageTokenCount(usage.input_tokens); + const cacheRead = + usage.cache_read_input_tokens == null + ? 0 + : readAnthropicUsageTokenCount(usage.cache_read_input_tokens); + const cacheWrite = + usage.cache_creation_input_tokens == null + ? 0 + : readAnthropicUsageTokenCount(usage.cache_creation_input_tokens); + if (input === undefined || cacheRead === undefined || cacheWrite === undefined) { + return undefined; + } + return { input, cacheRead, cacheWrite }; +} + +export function readLastAnthropicIterationUsage( + usage: AnthropicUsagePayload, +): AnthropicIterationUsageResult { + if (usage.iterations == null) { + return { state: "absent" }; + } + if (!Array.isArray(usage.iterations) || usage.iterations.length === 0) { + return { state: "invalid" }; + } + // Anthropic documents the final iteration as the true context window. + // Top-level cache fields remain cumulative billing totals across iterations. + const iteration = usage.iterations.at(-1); + if (!iteration || typeof iteration !== "object" || Array.isArray(iteration)) { + return { state: "invalid" }; + } + const record = iteration as AnthropicUsagePayload; + const input = readAnthropicUsageTokenCount(record.input_tokens); + const cacheRead = readAnthropicUsageTokenCount(record.cache_read_input_tokens); + const cacheWrite = readAnthropicUsageTokenCount(record.cache_creation_input_tokens); + const outputTokens = readAnthropicUsageTokenCount(record.output_tokens); + if ( + input === undefined || + cacheRead === undefined || + cacheWrite === undefined || + outputTokens === undefined + ) { + return { state: "invalid" }; + } + const contextPromptTokens = input + cacheRead + cacheWrite; + return { + state: "valid", + usage: { + contextPromptTokens, + totalTokens: contextPromptTokens + outputTokens, + }, + }; +} diff --git a/packages/ai/src/providers/anthropic.test.ts b/packages/ai/src/providers/anthropic.test.ts index d1c69cc2e1c9..5833bb261613 100644 --- a/packages/ai/src/providers/anthropic.test.ts +++ b/packages/ai/src/providers/anthropic.test.ts @@ -149,6 +149,339 @@ describe("Anthropic provider", () => { expect(config.authToken).toBeNull(); }); + it("keeps aggregate cache billing buckets out of the context total", async () => { + const client = { + messages: { + create: vi.fn(() => ({ + asResponse: () => + Promise.resolve( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_usage", + model: "claude-fable-5", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: 120_000, + cache_creation_input_tokens: null, + }, + }, + }, + { + type: "content_block_start", + index: 0, + content_block: { type: "text", text: "" }, + }, + { + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text: "Done." }, + }, + { type: "content_block_stop", index: 0 }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 819_661, + cache_creation_input_tokens: 93_130, + iterations: [ + { + type: "compaction", + input_tokens: 12, + output_tokens: 1_000, + cache_read_input_tokens: 819_661, + cache_creation_input_tokens: 93_130, + }, + { + type: "message", + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + ], + }, + }, + { type: "message_stop" }, + ]), + ), + })), + }, + }; + + const result = await streamAnthropic( + makeAnthropicModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello", timestamp: 0 }] }, + { apiKey: "sk-ant-provider", client: client as never }, + ).result(); + + expect(result.usage).toMatchObject({ + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + totalTokens: 927_907, + }); + }); + + it("does not fall back to aggregate usage when the final iteration is malformed", async () => { + const client = { + messages: { + create: vi.fn(() => ({ + asResponse: () => + Promise.resolve( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_invalid_iteration", + model: "claude-fable-5", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: 120_000, + cache_creation_input_tokens: 0, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 819_661, + cache_creation_input_tokens: 93_130, + iterations: [ + { + type: "message", + input_tokens: "malformed", + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + ], + }, + }, + { type: "message_stop" }, + ]), + ), + })), + }, + }; + + const result = await streamAnthropic( + makeAnthropicModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello", timestamp: 0 }] }, + { apiKey: "sk-ant-provider", client: client as never }, + ).result(); + + expect(result.usage.totalTokens).toBe(927_907); + expect(result.usage.contextUsage).toEqual({ state: "unavailable" }); + }); + + it("uses complete final usage when message-start prompt buckets are zero placeholders", async () => { + const client = { + messages: { + create: vi.fn(() => ({ + asResponse: () => + Promise.resolve( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_zero_start", + model: "claude-fable-5", + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + }, + { type: "message_stop" }, + ]), + ), + })), + }, + }; + + const result = await streamAnthropic( + makeAnthropicModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello", timestamp: 0 }] }, + { apiKey: "sk-ant-provider", client: client as never }, + ).result(); + + expect(result.usage.contextUsage).toEqual({ + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }); + }); + + it("does not treat zero start placeholders as complete final prompt usage", async () => { + const client = { + messages: { + create: vi.fn(() => ({ + asResponse: () => + Promise.resolve( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_zero_start_partial_delta", + model: "claude-fable-5", + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { output_tokens: 15_104 }, + }, + { type: "message_stop" }, + ]), + ), + })), + }, + }; + + const result = await streamAnthropic( + makeAnthropicModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello", timestamp: 0 }] }, + { apiKey: "sk-ant-provider", client: client as never }, + ).result(); + + expect(result.usage.contextUsage).toEqual({ state: "unavailable" }); + }); + + it("uses accumulated prompt buckets when the final usage update is partial", async () => { + const client = { + messages: { + create: vi.fn(() => ({ + asResponse: () => + Promise.resolve( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_partial_final_usage", + model: "claude-sonnet-4-6", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: 120_000, + cache_creation_input_tokens: 500, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: null, + }, + }, + { type: "message_stop" }, + ]), + ), + })), + }, + }; + + const result = await streamAnthropic( + makeAnthropicModel(), + { messages: [{ role: "user", content: "hello", timestamp: 0 }] }, + { apiKey: "sk-ant-provider", client: client as never }, + ).result(); + + expect(result.usage.contextUsage).toEqual({ + state: "available", + promptTokens: 149_374, + totalTokens: 164_478, + }); + }); + + it("preserves valid message-start billing buckets when a sibling is malformed", async () => { + const client = { + messages: { + create: vi.fn(() => ({ + asResponse: () => + Promise.resolve( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_malformed_usage", + model: "claude-sonnet-4-6", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: "malformed", + cache_creation_input_tokens: 500, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_creation_input_tokens: null, + }, + }, + { type: "message_stop" }, + ]), + ), + })), + }, + }; + + const result = await streamAnthropic( + makeAnthropicModel(), + { messages: [{ role: "user", content: "hello", timestamp: 0 }] }, + { apiKey: "sk-ant-provider", client: client as never }, + ).result(); + + expect(result.usage).toMatchObject({ + input: 12, + output: 15_104, + cacheRead: 0, + cacheWrite: 500, + totalTokens: 15_616, + }); + expect(result.usage.contextUsage).toEqual({ state: "unavailable" }); + }); + it("preserves provider-signed Anthropic thinking and drops reasoning_content placeholders", async () => { const highSurrogate = String.fromCharCode(0xd83d); const signedThinking = `keep${highSurrogate}signed`; diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index a468fe72dfe1..a0fbf580faa6 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -72,6 +72,12 @@ import { type AnthropicProjectedToolChoice, type AnthropicToolProjection, } from "./anthropic-tool-projection.js"; +import { + readAnthropicPromptUsageSnapshot, + readAnthropicUsageTokenCount, + readLastAnthropicIterationUsage, + type AnthropicPromptUsageSnapshot, +} from "./anthropic-usage.js"; import { resolveCacheRetention } from "./cache-retention.js"; import { resolveCloudflareBaseUrl } from "./cloudflare.js"; import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js"; @@ -516,6 +522,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti // Fallback-served turns bill at the serving model's rates; a boundary // swaps this to the fallback model's cost table. let costModel = model; + let messageStartPromptUsage: AnthropicPromptUsageSnapshot | undefined; try { let client: Anthropic; @@ -590,15 +597,45 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti if (event.type === "message_start") { output.responseId = event.message.id; output.responseModel = event.message.model; - output.usage.input = event.message.usage.input_tokens || 0; - output.usage.output = event.message.usage.output_tokens || 0; - output.usage.cacheRead = event.message.usage.cache_read_input_tokens || 0; - output.usage.cacheWrite = event.message.usage.cache_creation_input_tokens || 0; + const promptUsage = readAnthropicPromptUsageSnapshot(event.message.usage); + const messageStartPromptTokens = promptUsage + ? promptUsage.input + promptUsage.cacheRead + promptUsage.cacheWrite + : 0; + messageStartPromptUsage = messageStartPromptTokens > 0 ? promptUsage : undefined; + const inputTokens = readAnthropicUsageTokenCount(event.message.usage.input_tokens); + if (inputTokens !== undefined) { + output.usage.input = inputTokens; + } + const outputTokens = readAnthropicUsageTokenCount(event.message.usage.output_tokens); + if (outputTokens !== undefined) { + output.usage.output = outputTokens; + } + const cacheReadTokens = + event.message.usage.cache_read_input_tokens == null + ? 0 + : readAnthropicUsageTokenCount(event.message.usage.cache_read_input_tokens); + if (cacheReadTokens !== undefined) { + output.usage.cacheRead = cacheReadTokens; + } + const cacheWriteTokens = + event.message.usage.cache_creation_input_tokens == null + ? 0 + : readAnthropicUsageTokenCount(event.message.usage.cache_creation_input_tokens); + if (cacheWriteTokens !== undefined) { + output.usage.cacheWrite = cacheWriteTokens; + } output.usage.totalTokens = output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite; + if (messageStartPromptUsage && outputTokens !== undefined) { + output.usage.contextUsage = { + state: "available", + promptTokens: messageStartPromptTokens, + totalTokens: messageStartPromptTokens + output.usage.output, + }; + } calculateCost(costModel, output.usage); // Defer start until after message_start so that pre-stream SSE errors // (e.g. invalid thinking signatures) arrive before any non-error event @@ -799,24 +836,56 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti } // Only update usage fields if present (not null). // Preserves input_tokens from message_start when proxies omit it in message_delta. - if (event.usage.input_tokens != null) { - output.usage.input = event.usage.input_tokens; + const inputTokens = readAnthropicUsageTokenCount(event.usage.input_tokens); + if (inputTokens !== undefined) { + output.usage.input = inputTokens; } - if (event.usage.output_tokens != null) { - output.usage.output = event.usage.output_tokens; + const outputTokens = readAnthropicUsageTokenCount(event.usage.output_tokens); + if (outputTokens !== undefined) { + output.usage.output = outputTokens; } - if (event.usage.cache_read_input_tokens != null) { - output.usage.cacheRead = event.usage.cache_read_input_tokens; + // Match the SDK stream accumulator: null means no update, not a zero counter. + const cacheReadTokens = readAnthropicUsageTokenCount(event.usage.cache_read_input_tokens); + if (cacheReadTokens !== undefined) { + output.usage.cacheRead = cacheReadTokens; } - if (event.usage.cache_creation_input_tokens != null) { - output.usage.cacheWrite = event.usage.cache_creation_input_tokens; + const cacheWriteTokens = readAnthropicUsageTokenCount( + event.usage.cache_creation_input_tokens, + ); + if (cacheWriteTokens !== undefined) { + output.usage.cacheWrite = cacheWriteTokens; } - // Anthropic doesn't provide total_tokens, compute from components output.usage.totalTokens = output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite; + const iterationUsage = readLastAnthropicIterationUsage(event.usage); + if (iterationUsage.state === "valid") { + output.usage.contextUsage = { + state: "available", + promptTokens: iterationUsage.usage.contextPromptTokens, + totalTokens: iterationUsage.usage.totalTokens, + }; + } else if (iterationUsage.state === "invalid") { + output.usage.contextUsage = { state: "unavailable" }; + } else if ( + outputTokens !== undefined && + (messageStartPromptUsage !== undefined || + (inputTokens !== undefined && + cacheReadTokens !== undefined && + cacheWriteTokens !== undefined)) + ) { + const promptTokens = + output.usage.input + output.usage.cacheRead + output.usage.cacheWrite; + output.usage.contextUsage = { + state: "available", + promptTokens, + totalTokens: promptTokens + output.usage.output, + }; + } else { + output.usage.contextUsage = { state: "unavailable" }; + } calculateCost(costModel, output.usage); } } diff --git a/packages/ai/src/utils/overflow.test.ts b/packages/ai/src/utils/overflow.test.ts index 976d0ad3089b..54ce3ec1c60d 100644 --- a/packages/ai/src/utils/overflow.test.ts +++ b/packages/ai/src/utils/overflow.test.ts @@ -23,6 +23,25 @@ function errorMessage(message: string): AssistantMessage { }; } +function successfulMessage( + contextUsage?: AssistantMessage["usage"]["contextUsage"], +): AssistantMessage { + return { + ...errorMessage(""), + usage: { + input: 12, + output: 15_104, + cacheRead: 1_100_000, + cacheWrite: 93_130, + ...(contextUsage ? { contextUsage } : {}), + totalTokens: 1_208_246, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + errorMessage: undefined, + }; +} + describe("configured context size overflow", () => { it.each([ "400 Prompt has 256468 tokens, but the configured context size is 256000 tokens", @@ -32,3 +51,22 @@ describe("configured context size overflow", () => { expect(isContextOverflow(errorMessage(text), 256_000)).toBe(true); }); }); + +describe("usage-based overflow", () => { + it("prefers an available context snapshot over aggregate billing usage", () => { + expect( + isContextOverflow( + successfulMessage({ + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }), + 1_000_000, + ), + ).toBe(false); + }); + + it("does not infer overflow from aggregate billing when context is unavailable", () => { + expect(isContextOverflow(successfulMessage({ state: "unavailable" }), 1_000_000)).toBe(false); + }); +}); diff --git a/packages/ai/src/utils/overflow.ts b/packages/ai/src/utils/overflow.ts index 6d5a683d6b8b..12aefe480f8a 100644 --- a/packages/ai/src/utils/overflow.ts +++ b/packages/ai/src/utils/overflow.ts @@ -80,6 +80,16 @@ const NON_OVERFLOW_PATTERNS = [ /too many requests/i, // Generic HTTP 429 style ]; +function resolveContextInputTokens(message: AssistantMessage): number | undefined { + if (message.usage.contextUsage?.state === "available") { + return message.usage.contextUsage.promptTokens; + } + if (message.usage.contextUsage?.state === "unavailable") { + return undefined; + } + return message.usage.input + message.usage.cacheRead; +} + /** * Check if an assistant message represents a context overflow error. * @@ -141,8 +151,8 @@ export function isContextOverflow(message: AssistantMessage, contextWindow?: num // Case 2: Silent overflow (z.ai style) - successful but usage exceeds context if (contextWindow && message.stopReason === "stop") { - const inputTokens = message.usage.input + message.usage.cacheRead; - if (inputTokens > contextWindow) { + const inputTokens = resolveContextInputTokens(message); + if (inputTokens !== undefined && inputTokens > contextWindow) { return true; } } @@ -151,8 +161,8 @@ export function isContextOverflow(message: AssistantMessage, contextWindow?: num // to fit the context window, leaving no room for output. Returns stopReason "length" // with output=0 and input+cacheRead filling the context window. if (contextWindow && message.stopReason === "length" && message.usage.output === 0) { - const inputTokens = message.usage.input + message.usage.cacheRead; - if (inputTokens >= contextWindow * 0.99) { + const inputTokens = resolveContextInputTokens(message); + if (inputTokens !== undefined && inputTokens >= contextWindow * 0.99) { return true; } } diff --git a/packages/llm-core/src/types.ts b/packages/llm-core/src/types.ts index 8cc84072f929..68bfca663d4e 100644 --- a/packages/llm-core/src/types.ts +++ b/packages/llm-core/src/types.ts @@ -263,6 +263,10 @@ export interface Usage { output: number; cacheRead: number; cacheWrite: number; + /** Exact context snapshot for the final provider iteration. */ + contextUsage?: + | { state: "available"; promptTokens: number; totalTokens: number } + | { state: "unavailable" }; totalTokens: number; cost: { input: number; diff --git a/src/agents/anthropic-transport-stream.test.ts b/src/agents/anthropic-transport-stream.test.ts index 1184abbe61b7..ce7459003c1d 100644 --- a/src/agents/anthropic-transport-stream.test.ts +++ b/src/agents/anthropic-transport-stream.test.ts @@ -201,6 +201,297 @@ describe("anthropic transport stream", () => { vi.useRealTimers(); }); + it("keeps aggregate cache billing buckets out of the context total", async () => { + guardedFetchMock.mockResolvedValueOnce( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_usage", + model: "claude-fable-5", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: 120_000, + cache_creation_input_tokens: null, + }, + }, + }, + { + type: "content_block_start", + index: 0, + content_block: { type: "text", text: "" }, + }, + { + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text: "Done." }, + }, + { type: "content_block_stop", index: 0 }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 819_661, + cache_creation_input_tokens: 93_130, + iterations: [ + { + type: "compaction", + input_tokens: 12, + output_tokens: 1_000, + cache_read_input_tokens: 819_661, + cache_creation_input_tokens: 93_130, + }, + { + type: "message", + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + ], + }, + }, + { type: "message_stop" }, + ]), + ); + + const result = await runTransportStream( + makeAnthropicTransportModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello" }] } as AnthropicStreamContext, + { apiKey: "sk-ant-api" } as AnthropicStreamOptions, + ); + + expect(result.usage).toMatchObject({ + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + totalTokens: 927_907, + }); + }); + + it("does not fall back to aggregate usage when the final iteration is malformed", async () => { + guardedFetchMock.mockResolvedValueOnce( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_invalid_iteration", + model: "claude-fable-5", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: 120_000, + cache_creation_input_tokens: 0, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 819_661, + cache_creation_input_tokens: 93_130, + iterations: [ + { + type: "message", + input_tokens: "malformed", + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + ], + }, + }, + { type: "message_stop" }, + ]), + ); + + const result = await runTransportStream( + makeAnthropicTransportModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello" }] } as AnthropicStreamContext, + { apiKey: "sk-ant-api" } as AnthropicStreamOptions, + ); + + expect(result.usage.totalTokens).toBe(927_907); + expect(result.usage.contextUsage).toEqual({ state: "unavailable" }); + }); + + it("uses complete final usage when message-start prompt buckets are zero placeholders", async () => { + guardedFetchMock.mockResolvedValueOnce( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_zero_start", + model: "claude-fable-5", + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: 0, + }, + }, + { type: "message_stop" }, + ]), + ); + + const result = await runTransportStream( + makeAnthropicTransportModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello" }] } as AnthropicStreamContext, + { apiKey: "sk-ant-api" } as AnthropicStreamOptions, + ); + + expect(result.usage.contextUsage).toEqual({ + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }); + }); + + it("does not treat zero start placeholders as complete final prompt usage", async () => { + guardedFetchMock.mockResolvedValueOnce( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_zero_start_partial_delta", + model: "claude-fable-5", + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { output_tokens: 15_104 }, + }, + { type: "message_stop" }, + ]), + ); + + const result = await runTransportStream( + makeAnthropicTransportModel({ id: "claude-fable-5", name: "Claude Fable 5" }), + { messages: [{ role: "user", content: "hello" }] } as AnthropicStreamContext, + { apiKey: "sk-ant-api" } as AnthropicStreamOptions, + ); + + expect(result.usage.contextUsage).toEqual({ state: "unavailable" }); + }); + + it("uses accumulated prompt buckets when the final usage update is partial", async () => { + guardedFetchMock.mockResolvedValueOnce( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_partial_final_usage", + model: "claude-sonnet-4-6", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: 120_000, + cache_creation_input_tokens: 500, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_read_input_tokens: 148_862, + cache_creation_input_tokens: null, + }, + }, + { type: "message_stop" }, + ]), + ); + + const result = await runTransportStream( + makeAnthropicTransportModel(), + { messages: [{ role: "user", content: "hello" }] } as AnthropicStreamContext, + { apiKey: "sk-ant-api" } as AnthropicStreamOptions, + ); + + expect(result.usage.contextUsage).toEqual({ + state: "available", + promptTokens: 149_374, + totalTokens: 164_478, + }); + }); + + it("preserves valid message-start billing buckets when a sibling is malformed", async () => { + guardedFetchMock.mockResolvedValueOnce( + createSseResponse([ + { + type: "message_start", + message: { + id: "msg_malformed_usage", + model: "claude-sonnet-4-6", + usage: { + input_tokens: 12, + output_tokens: 0, + cache_read_input_tokens: "malformed", + cache_creation_input_tokens: 500, + }, + }, + }, + { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { + input_tokens: 12, + output_tokens: 15_104, + cache_creation_input_tokens: null, + }, + }, + { type: "message_stop" }, + ]), + ); + + const result = await runTransportStream( + makeAnthropicTransportModel(), + { messages: [{ role: "user", content: "hello" }] } as AnthropicStreamContext, + { apiKey: "sk-ant-api" } as AnthropicStreamOptions, + ); + + expect(result.usage).toMatchObject({ + input: 12, + output: 15_104, + cacheRead: 0, + cacheWrite: 500, + totalTokens: 15_616, + }); + expect(result.usage.contextUsage).toEqual({ state: "unavailable" }); + }); + it("tags pre-tool narration as commentary when a proxy mislabels stop_reason (pioneer/Bedrock)", async () => { // Bedrock/Vertex-proxied routes (e.g. pioneer; tool ids "toolu_vrtx_…") report // stop_reason "end_turn" on turns that DO carry a tool call. Commentary tagging diff --git a/src/agents/anthropic-transport-stream.ts b/src/agents/anthropic-transport-stream.ts index 5e9763c4e40a..15dc6a66f302 100644 --- a/src/agents/anthropic-transport-stream.ts +++ b/src/agents/anthropic-transport-stream.ts @@ -13,12 +13,16 @@ import { resolveClaudeNativeThinkingLevelMap, resolveOriginalAnthropicToolName, readAnthropicFallbackBoundary, + readAnthropicPromptUsageSnapshot, + readAnthropicUsageTokenCount, + readLastAnthropicIterationUsage, supportsClaudeAdaptiveThinking, supportsClaudeNativeMaxEffort, supportsClaudeNativeXhighEffort, usesClaudeFable5MessagesContract, usesFoundryBearerAuth, type AnthropicOptions, + type AnthropicPromptUsageSnapshot, type AnthropicProjectedToolChoice, type AnthropicThinkingDisplay, type AnthropicToolProjection, @@ -75,6 +79,7 @@ import { sanitizeNonEmptyTransportPayloadText, sanitizeTransportPayloadText, } from "./transport-stream-shared.js"; +import type { ContextUsage } from "./usage.js"; const CLAUDE_CODE_VERSION = "2.1.75"; const ANTHROPIC_MESSAGES_ERROR_BODY_MAX_BYTES = 8 * 1024; @@ -156,6 +161,7 @@ type MutableAssistantOutput = { output: number; cacheRead: number; cacheWrite: number; + contextUsage?: ContextUsage; totalTokens: number; cost: { input: number; output: number; cacheRead: number; cacheWrite: number; total: number }; }; @@ -1150,6 +1156,7 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn { // Fallback-served turns bill at the serving model's rates; a boundary // swaps this to the fallback model's cost table. let costModel = model; + let messageStartPromptUsage: AnthropicPromptUsageSnapshot | undefined; try { const apiKey = options?.apiKey ?? getEnvApiKey(model.provider) ?? ""; if (!apiKey) { @@ -1307,19 +1314,45 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn { const usage = message?.usage ?? {}; output.responseId = typeof message?.id === "string" ? message.id : undefined; output.responseModel = typeof message?.model === "string" ? message.model : undefined; - output.usage.input = typeof usage.input_tokens === "number" ? usage.input_tokens : 0; - output.usage.output = typeof usage.output_tokens === "number" ? usage.output_tokens : 0; - output.usage.cacheRead = - typeof usage.cache_read_input_tokens === "number" ? usage.cache_read_input_tokens : 0; - output.usage.cacheWrite = - typeof usage.cache_creation_input_tokens === "number" - ? usage.cache_creation_input_tokens - : 0; + const promptUsage = readAnthropicPromptUsageSnapshot(usage); + const messageStartPromptTokens = promptUsage + ? promptUsage.input + promptUsage.cacheRead + promptUsage.cacheWrite + : 0; + messageStartPromptUsage = messageStartPromptTokens > 0 ? promptUsage : undefined; + const inputTokens = readAnthropicUsageTokenCount(usage.input_tokens); + if (inputTokens !== undefined) { + output.usage.input = inputTokens; + } + const outputTokens = readAnthropicUsageTokenCount(usage.output_tokens); + if (outputTokens !== undefined) { + output.usage.output = outputTokens; + } + const cacheReadTokens = + usage.cache_read_input_tokens == null + ? 0 + : readAnthropicUsageTokenCount(usage.cache_read_input_tokens); + if (cacheReadTokens !== undefined) { + output.usage.cacheRead = cacheReadTokens; + } + const cacheWriteTokens = + usage.cache_creation_input_tokens == null + ? 0 + : readAnthropicUsageTokenCount(usage.cache_creation_input_tokens); + if (cacheWriteTokens !== undefined) { + output.usage.cacheWrite = cacheWriteTokens; + } output.usage.totalTokens = output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite; + if (messageStartPromptUsage && outputTokens !== undefined) { + output.usage.contextUsage = { + state: "available", + promptTokens: messageStartPromptTokens, + totalTokens: messageStartPromptTokens + output.usage.output, + }; + } calculateCost(costModel, output.usage); // Defer start until after message_start so that pre-stream SSE errors // (e.g. invalid thinking signatures) arrive before any non-error event @@ -1653,23 +1686,56 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn { output.stopReason = mapStopReason(delta.stop_reason); } } - if (typeof usage?.input_tokens === "number") { - output.usage.input = usage.input_tokens; + const inputTokens = readAnthropicUsageTokenCount(usage?.input_tokens); + if (inputTokens !== undefined) { + output.usage.input = inputTokens; } - if (typeof usage?.output_tokens === "number") { - output.usage.output = usage.output_tokens; + const outputTokens = readAnthropicUsageTokenCount(usage?.output_tokens); + if (outputTokens !== undefined) { + output.usage.output = outputTokens; } - if (typeof usage?.cache_read_input_tokens === "number") { - output.usage.cacheRead = usage.cache_read_input_tokens; + // Match the SDK stream accumulator: null means no update, not a zero counter. + const cacheReadTokens = readAnthropicUsageTokenCount(usage?.cache_read_input_tokens); + if (cacheReadTokens !== undefined) { + output.usage.cacheRead = cacheReadTokens; } - if (typeof usage?.cache_creation_input_tokens === "number") { - output.usage.cacheWrite = usage.cache_creation_input_tokens; + const cacheWriteTokens = readAnthropicUsageTokenCount( + usage?.cache_creation_input_tokens, + ); + if (cacheWriteTokens !== undefined) { + output.usage.cacheWrite = cacheWriteTokens; } output.usage.totalTokens = output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite; + const iterationUsage = readLastAnthropicIterationUsage(usage ?? {}); + if (iterationUsage.state === "valid") { + output.usage.contextUsage = { + state: "available", + promptTokens: iterationUsage.usage.contextPromptTokens, + totalTokens: iterationUsage.usage.totalTokens, + }; + } else if (iterationUsage.state === "invalid") { + output.usage.contextUsage = { state: "unavailable" }; + } else if ( + outputTokens !== undefined && + (messageStartPromptUsage !== undefined || + (inputTokens !== undefined && + cacheReadTokens !== undefined && + cacheWriteTokens !== undefined)) + ) { + const promptTokens = + output.usage.input + output.usage.cacheRead + output.usage.cacheWrite; + output.usage.contextUsage = { + state: "available", + promptTokens, + totalTokens: promptTokens + output.usage.output, + }; + } else { + output.usage.contextUsage = { state: "unavailable" }; + } calculateCost(costModel, output.usage); // Gate on the turn CONTAINING a tool call, not the provider's stop_reason // label: Bedrock/Vertex-proxied routes (e.g. pioneer) report "end_turn" on diff --git a/src/agents/command/session-store.test.ts b/src/agents/command/session-store.test.ts index e16748e8ac69..3bcb306fbde2 100644 --- a/src/agents/command/session-store.test.ts +++ b/src/agents/command/session-store.test.ts @@ -1178,6 +1178,62 @@ describe("updateSessionStoreAfterAgentRun", () => { }); }); + it("uses the compaction snapshot when non-CLI last-call context is unavailable", async () => { + await withTempSessionStore(async ({ storePath }) => { + const sessionKey = "agent:main:explicit:test-unavailable-context"; + const sessionId = "test-unavailable-context-session"; + const sessionStore: Record = { + [sessionKey]: { + sessionId, + updatedAt: 1, + totalTokens: 95_000, + totalTokensFresh: true, + }, + }; + await fs.writeFile(storePath, JSON.stringify(sessionStore, null, 2)); + + await updateSessionStoreAfterAgentRun({ + cfg: {} as OpenClawConfig, + sessionId, + sessionKey, + storePath, + sessionStore, + defaultProvider: "anthropic", + defaultModel: "claude-fable-5", + result: { + meta: { + durationMs: 1, + agentMeta: { + sessionId, + provider: "anthropic", + model: "claude-fable-5", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + total: 927_907, + }, + lastCallUsage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + total: 927_907, + }, + compactionTokensAfter: 80_000, + }, + }, + } as EmbeddedAgentRunResult, + }); + + expect(sessionStore[sessionKey]?.totalTokens).toBe(80_000); + expect(sessionStore[sessionKey]?.totalTokensFresh).toBe(true); + expect(sessionStore[sessionKey]?.cacheRead).toBeUndefined(); + }); + }); + it("persists CLI lastCallUsage as the context snapshot (totalTokens)", async () => { await withTempSessionStore(async ({ storePath }) => { const cfg = { diff --git a/src/agents/command/session-store.ts b/src/agents/command/session-store.ts index a290eb75c173..952cbf1abc01 100644 --- a/src/agents/command/session-store.ts +++ b/src/agents/command/session-store.ts @@ -201,10 +201,10 @@ export async function updateSessionStoreAfterAgentRun(params: { const input = usage.input ?? 0; const output = usage.output ?? 0; const usageForContext = isCliProvider(providerUsed, cfg) - ? promptTokens - ? undefined - : lastCallUsage - : usage; + ? lastCallUsage + : lastCallUsage?.contextUsage + ? lastCallUsage + : usage; const totalTokens = deriveSessionTotalTokens({ usage: promptTokens ? undefined : usageForContext, contextTokens, diff --git a/src/agents/embedded-agent-runner.sanitize-session-history.test.ts b/src/agents/embedded-agent-runner.sanitize-session-history.test.ts index 824ccd799d8c..ccd2f02afe83 100644 --- a/src/agents/embedded-agent-runner.sanitize-session-history.test.ts +++ b/src/agents/embedded-agent-runner.sanitize-session-history.test.ts @@ -189,6 +189,70 @@ describe("sanitizeSessionHistory", () => { const getAssistantContentTypes = (messages: AgentMessage[]) => getAssistantMessage(messages).content.map((block: { type: string }) => block.type); + it("preserves a validated context snapshot while normalizing replay usage", async () => { + const assistant = { + role: "assistant", + content: [{ type: "text", text: "done" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-opus-4-6", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + }, + stopReason: "stop", + timestamp: 0, + } as unknown as AgentMessage; + + const out = await sanitizeAnthropicHistory({ + messages: [{ role: "user", content: "hello", timestamp: 0 } as AgentMessage, assistant], + }); + + expect(getAssistantMessage(out).usage).toMatchObject({ + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + totalTokens: 927_907, + }); + }); + + it("preserves an unavailable context snapshot while normalizing replay usage", async () => { + const assistant = { + role: "assistant", + content: [{ type: "text", text: "done" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-opus-4-6", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + }, + stopReason: "stop", + timestamp: 0, + } as unknown as AgentMessage; + + const out = await sanitizeAnthropicHistory({ + messages: [{ role: "user", content: "hello", timestamp: 0 } as AgentMessage, assistant], + }); + + expect(getAssistantMessage(out).usage).toMatchObject({ + contextUsage: { state: "unavailable" }, + totalTokens: 927_907, + }); + }); + const makeThinkingAndTextAssistantMessages = (thinkingSignature = "some_sig"): AgentMessage[] => { const user: UserMessage = { role: "user", diff --git a/src/agents/embedded-agent-runner/replay-history.ts b/src/agents/embedded-agent-runner/replay-history.ts index df1222957168..11ef1ab2990e 100644 --- a/src/agents/embedded-agent-runner/replay-history.ts +++ b/src/agents/embedded-agent-runner/replay-history.ts @@ -431,6 +431,7 @@ function normalizeAssistantUsageSnapshot(usage: unknown) { output, cacheRead, cacheWrite, + ...(normalized.contextUsage ? { contextUsage: { ...normalized.contextUsage } } : {}), totalTokens, ...(cost ? { cost } : {}), }; @@ -489,6 +490,25 @@ function ensureAssistantUsageSnapshots(messages: AgentMessage[]): AgentMessage[] message.usage && typeof message.usage === "object" ? (message.usage as { cost?: unknown }).cost : undefined; + const rawContextUsage = + message.usage && typeof message.usage === "object" + ? (message.usage as { contextUsage?: unknown }).contextUsage + : undefined; + const normalizedContextUsage = normalizedUsage.contextUsage; + const contextUsageMatches = + normalizedContextUsage === undefined + ? rawContextUsage === undefined + : normalizedContextUsage.state === "unavailable" + ? rawContextUsage !== null && + typeof rawContextUsage === "object" && + (rawContextUsage as { state?: unknown }).state === "unavailable" + : rawContextUsage !== null && + typeof rawContextUsage === "object" && + (rawContextUsage as { state?: unknown }).state === "available" && + (rawContextUsage as { promptTokens?: unknown }).promptTokens === + normalizedContextUsage.promptTokens && + (rawContextUsage as { totalTokens?: unknown }).totalTokens === + normalizedContextUsage.totalTokens; const normalizedCost = normalizedUsage.cost; if ( message.usage && @@ -498,6 +518,7 @@ function ensureAssistantUsageSnapshots(messages: AgentMessage[]): AgentMessage[] (message.usage as { cacheRead?: unknown }).cacheRead === normalizedUsage.cacheRead && (message.usage as { cacheWrite?: unknown }).cacheWrite === normalizedUsage.cacheWrite && (message.usage as { totalTokens?: unknown }).totalTokens === normalizedUsage.totalTokens && + contextUsageMatches && ((normalizedCost && usageCost && typeof usageCost === "object" && 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 7b7f534030a7..8680a88d6b70 100644 --- a/src/agents/embedded-agent-runner/run.overflow-compaction.harness.ts +++ b/src/agents/embedded-agent-runner/run.overflow-compaction.harness.ts @@ -612,6 +612,45 @@ export async function loadRunOverflowCompactionHarness(): Promise<{ })() : undefined, ), + deriveContextPromptTokens: vi.fn( + (params: { + lastCallUsage?: { + input?: number; + output?: number; + cacheRead?: number; + cacheWrite?: number; + contextUsage?: + | { state: "available"; promptTokens: number; totalTokens: number } + | { state: "unavailable" }; + total?: number; + }; + promptTokens?: number; + usage?: { input?: number; cacheRead?: number; cacheWrite?: number }; + }) => { + if ( + typeof params.promptTokens === "number" && + Number.isFinite(params.promptTokens) && + params.promptTokens > 0 + ) { + return params.promptTokens; + } + const lastCall = params.lastCallUsage; + if (lastCall?.contextUsage?.state === "available") { + return lastCall.contextUsage.promptTokens; + } + if (lastCall?.contextUsage?.state === "unavailable") { + return undefined; + } + for (const usage of [lastCall, params.usage]) { + const promptTokens = + (usage?.input ?? 0) + (usage?.cacheRead ?? 0) + (usage?.cacheWrite ?? 0); + if (promptTokens > 0) { + return promptTokens; + } + } + return undefined; + }, + ), })); vi.doMock("../cli-backends.js", async () => { diff --git a/src/agents/embedded-agent-runner/run.timeout-triggered-compaction.test.ts b/src/agents/embedded-agent-runner/run.timeout-triggered-compaction.test.ts index aef6f507db01..0619d691347a 100644 --- a/src/agents/embedded-agent-runner/run.timeout-triggered-compaction.test.ts +++ b/src/agents/embedded-agent-runner/run.timeout-triggered-compaction.test.ts @@ -616,22 +616,29 @@ describe("timeout-triggered compaction", () => { expect(result.payloads?.[0]?.text).toContain("timed out"); }); - it("uses prompt/input tokens for ratio, not total tokens", async () => { - // Timeout where total tokens are high (150k) but input/prompt tokens - // are low (20k / 200k = 10%). Should NOT trigger compaction because - // the ratio is based on prompt tokens, not total. + it("uses the explicit context snapshot instead of aggregate billing buckets", async () => { + // Server-side loops can report aggregate cache billing far above the final + // iteration's prompt. Timeout recovery must use the explicit 20k snapshot. mockedRunEmbeddedAttempt.mockResolvedValueOnce( makeAttemptResult({ timedOut: true, lastAssistant: { - usage: { input: 20000, total: 150000 }, + usage: { + input: 20_000, + cacheRead: 150_000, + contextUsage: { + state: "available", + promptTokens: 20_000, + totalTokens: 20_500, + }, + total: 170_500, + }, } as never, }), ); const result = await runEmbeddedAgent(overflowBaseRunParams); - // Despite high total tokens, low prompt tokens mean no compaction expect(mockedCompactDirect).not.toHaveBeenCalled(); expect(result.payloads?.[0]?.isError).toBe(true); expect(result.payloads?.[0]?.text).toContain("timed out"); diff --git a/src/agents/embedded-agent-runner/run.ts b/src/agents/embedded-agent-runner/run.ts index b4ac86751b47..f51c1997133f 100644 --- a/src/agents/embedded-agent-runner/run.ts +++ b/src/agents/embedded-agent-runner/run.ts @@ -141,7 +141,7 @@ import { type SessionSuspensionParams, } from "../session-suspension.js"; import { resolveToolLoopDetectionConfig } from "../tool-loop-detection-config.js"; -import { derivePromptTokens, normalizeUsage, type UsageLike } from "../usage.js"; +import { deriveContextPromptTokens, normalizeUsage, type UsageLike } from "../usage.js"; import { redactRunIdentifier, resolveRunWorkspaceDir } from "../workspace-run.js"; import { runPostCompactionSideEffects } from "./compaction-hooks.js"; import { buildEmbeddedCompactionRuntimeContext } from "./compaction-runtime-context.js"; @@ -2509,7 +2509,9 @@ async function runEmbeddedAgentInternal( // Only consider prompt-side tokens here. API totals include output // tokens, which can make a long generation look like high context // pressure even when the prompt itself was small. - const lastTurnPromptTokens = derivePromptTokens(lastRunPromptUsage); + const lastTurnPromptTokens = deriveContextPromptTokens({ + lastCallUsage: lastRunPromptUsage, + }); const tokenUsedRatio = lastTurnPromptTokens != null && ctxInfo.tokens > 0 ? lastTurnPromptTokens / ctxInfo.tokens diff --git a/src/agents/embedded-agent-runner/run/attempt.prompt-helpers.ts b/src/agents/embedded-agent-runner/run/attempt.prompt-helpers.ts index 8cf27aafc8a7..f2fca05985bb 100644 --- a/src/agents/embedded-agent-runner/run/attempt.prompt-helpers.ts +++ b/src/agents/embedded-agent-runner/run/attempt.prompt-helpers.ts @@ -25,7 +25,7 @@ import { wrapPluginSystemContextSection } from "../../hook-system-context-bounda import { buildActiveImageGenerationTaskPromptContextForSession } from "../../image-generation-task-status.js"; import { buildActiveMusicGenerationTaskPromptContextForSession } from "../../music-generation-task-status.js"; import { resolveEffectiveToolFsWorkspaceOnly } from "../../tool-fs-policy.js"; -import { derivePromptTokens, type NormalizedUsage } from "../../usage.js"; +import { deriveContextPromptTokens, type NormalizedUsage } from "../../usage.js"; import { buildActiveVideoGenerationTaskPromptContextForSession } from "../../video-generation-task-status.js"; import { buildEmbeddedCompactionRuntimeContext } from "../compaction-runtime-context.js"; import { resolveContextEngineCapabilities } from "../context-engine-capabilities.js"; @@ -655,6 +655,6 @@ export function buildAfterTurnRuntimeContextFromUsage( ): ContextEngineRuntimeContext { return buildAfterTurnRuntimeContext({ ...params, - currentTokenCount: derivePromptTokens(params.lastCallUsage), + currentTokenCount: deriveContextPromptTokens({ lastCallUsage: params.lastCallUsage }), }); } diff --git a/src/agents/embedded-agent-runner/run/attempt.test.ts b/src/agents/embedded-agent-runner/run/attempt.test.ts index d3352fc35871..0aed4c75381f 100644 --- a/src/agents/embedded-agent-runner/run/attempt.test.ts +++ b/src/agents/embedded-agent-runner/run/attempt.test.ts @@ -11,6 +11,7 @@ import { addSession, resetProcessRegistryForTests } from "../../bash-process-reg import { createProcessSessionFixture } from "../../bash-process-registry.test-helpers.js"; import { wrapPluginSystemContextSection } from "../../hook-system-context-boundary.js"; import { buildAgentSystemPrompt } from "../../system-prompt.js"; +import type { NormalizedUsage } from "../../usage.js"; import { resetEmbeddedAgentBaseStreamFnCacheForTest, resolveEmbeddedAgentBaseStreamFn, @@ -3419,8 +3420,13 @@ describe("buildAfterTurnRuntimeContext", () => { output: 5, cacheRead: 40, cacheWrite: 2, + contextUsage: { + state: "available", + promptTokens: 23, + totalTokens: 28, + }, total: 57, - }; + } satisfies NormalizedUsage; const promptCache = buildContextEnginePromptCacheInfo({ lastCallUsage }); const legacy = buildAfterTurnRuntimeContextFromUsage({ attempt: { @@ -3445,7 +3451,7 @@ describe("buildAfterTurnRuntimeContext", () => { promptCache, }); - expect(legacy.currentTokenCount).toBe(52); + expect(legacy.currentTokenCount).toBe(23); expect(legacy.promptCache?.lastCallUsage?.total).toBe(57); }); diff --git a/src/agents/embedded-agent-runner/run/helpers.test.ts b/src/agents/embedded-agent-runner/run/helpers.test.ts index 967e6a9ab7b1..5649f75cc2f7 100644 --- a/src/agents/embedded-agent-runner/run/helpers.test.ts +++ b/src/agents/embedded-agent-runner/run/helpers.test.ts @@ -2,8 +2,10 @@ // metadata assembly shared by normal exits and failure paths. import type { AssistantMessage } from "openclaw/plugin-sdk/llm"; import { describe, expect, it } from "vitest"; -import { createUsageAccumulator } from "../usage-accumulator.js"; +import type { NormalizedUsage } from "../../usage.js"; +import { createUsageAccumulator, mergeUsageIntoAccumulator } from "../usage-accumulator.js"; import { + buildUsageAgentMetaFields, buildErrorAgentMeta, resolveFinalAssistantRawText, resolveFinalAssistantVisibleText, @@ -191,6 +193,42 @@ describe("resolveLatestCallUsage", () => { }); }); +describe("buildUsageAgentMetaFields", () => { + it("keeps aggregate billing buckets out of the latest context snapshot", () => { + const usageAccumulator = createUsageAccumulator(); + const latestCallUsage = { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + total: 927_907, + } satisfies NormalizedUsage; + mergeUsageIntoAccumulator(usageAccumulator, latestCallUsage); + + const fields = buildUsageAgentMetaFields({ + usageAccumulator, + lastAssistantUsage: undefined, + lastRunPromptUsage: latestCallUsage, + lastTurnTotal: latestCallUsage.total, + }); + + expect(fields.usage).toMatchObject({ + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + total: 927_907, + }); + expect(fields.lastCallUsage).toEqual(latestCallUsage); + expect(fields.promptTokens).toBe(148_874); + }); +}); + 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 b1d0d1e7e15f..a3c33d4fc1e9 100644 --- a/src/agents/embedded-agent-runner/run/helpers.ts +++ b/src/agents/embedded-agent-runner/run/helpers.ts @@ -8,9 +8,10 @@ import { extractAssistantTextForPhase } from "../../../shared/chat-message-conte import { resolveAgentConfig } from "../../agent-scope-config.js"; import { extractAssistantVisibleText } from "../../embedded-agent-utils.js"; import { - derivePromptTokens, + deriveContextPromptTokens, hasNonzeroUsage, normalizeUsage, + type ContextUsage, type NormalizedUsage, } from "../../usage.js"; import type { EmbeddedAgentMeta } from "../types.js"; @@ -21,6 +22,7 @@ type UsageSnapshot = { output?: number; cacheRead?: number; cacheWrite?: number; + contextUsage?: ContextUsage; total?: number; }; @@ -219,7 +221,9 @@ export function buildUsageAgentMetaFields(params: { : hasNonzeroUsage(params.lastRunPromptUsage) ? params.lastRunPromptUsage : toLastCallUsage(params.usageAccumulator); - const promptTokens = derivePromptTokens(params.lastRunPromptUsage); + const promptTokens = deriveContextPromptTokens({ + lastCallUsage: params.lastRunPromptUsage, + }); return { usage, lastCallUsage, diff --git a/src/agents/embedded-agent-runner/types.ts b/src/agents/embedded-agent-runner/types.ts index 8f54f41dbb4c..0b351f8f4923 100644 --- a/src/agents/embedded-agent-runner/types.ts +++ b/src/agents/embedded-agent-runner/types.ts @@ -15,6 +15,7 @@ import type { } from "../embedded-agent-messaging.types.js"; import type { FallbackAttempt } from "../model-fallback.types.js"; import type { AgentRunTimeoutPhase } from "../run-timeout-attribution.js"; +import type { ContextUsage } from "../usage.js"; export type EmbeddedAgentMeta = { sessionId: string; @@ -59,6 +60,7 @@ export type EmbeddedAgentMeta = { output?: number; cacheRead?: number; cacheWrite?: number; + contextUsage?: ContextUsage; reasoningTokens?: number; total?: number; }; diff --git a/src/agents/embedded-agent-runner/usage-accumulator.test.ts b/src/agents/embedded-agent-runner/usage-accumulator.test.ts index 4b0bd99ce558..ecf1ff10e7bf 100644 --- a/src/agents/embedded-agent-runner/usage-accumulator.test.ts +++ b/src/agents/embedded-agent-runner/usage-accumulator.test.ts @@ -33,6 +33,11 @@ const FINAL_USAGE: UsageInput = { reasoningTokens: 7, cacheRead: 84_000, cacheWrite: 0, + contextUsage: { + state: "available", + promptTokens: 84_150, + totalTokens: 84_190, + }, total: 84_190, }; @@ -72,6 +77,11 @@ describe("usage-accumulator", () => { expect(acc.lastReasoningTokens).toBe(7); expect(acc.lastCacheRead).toBe(84_000); expect(acc.lastCacheWrite).toBe(0); + expect(acc.lastContextUsage).toEqual({ + state: "available", + promptTokens: 84_150, + totalTokens: 84_190, + }); expect(acc.lastTotal).toBe(84_190); }); @@ -158,9 +168,28 @@ describe("usage-accumulator", () => { reasoningTokens: 7, cacheRead: 84_000, cacheWrite: undefined, + contextUsage: { + state: "available", + promptTokens: 84_150, + totalTokens: 84_190, + }, total: 84_190, }); }); - }); + it("preserves an unavailable context snapshot", () => { + const acc = createUsageAccumulator(); + mergeUsageIntoAccumulator(acc, { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + total: 927_907, + }); + + expect(toLastCallUsage(acc)?.contextUsage).toEqual({ state: "unavailable" }); + expect(toNormalizedUsage(acc)?.contextUsage).toBeUndefined(); + }); + }); }); diff --git a/src/agents/embedded-agent-runner/usage-accumulator.ts b/src/agents/embedded-agent-runner/usage-accumulator.ts index 81d4daf95706..df7c8dc2b002 100644 --- a/src/agents/embedded-agent-runner/usage-accumulator.ts +++ b/src/agents/embedded-agent-runner/usage-accumulator.ts @@ -1,7 +1,7 @@ /** * Accumulates and normalizes per-call token usage across embedded runs. */ -import type { NormalizedUsage } from "../usage.js"; +import type { ContextUsage, NormalizedUsage } from "../usage.js"; export type UsageAccumulator = { input: number; @@ -15,6 +15,7 @@ export type UsageAccumulator = { lastOutput: number; lastCacheRead: number; lastCacheWrite: number; + lastContextUsage?: ContextUsage; lastReasoningTokens: number; lastTotal: number; }; @@ -40,14 +41,19 @@ const hasUsageValues = (usage: MaybeUsage): usage is NormalizedUsage => { if (!usage) { return false; } - return [ - usage.input, - usage.output, - usage.cacheRead, - usage.cacheWrite, - usage.reasoningTokens, - usage.total, - ].some((value) => typeof value === "number" && Number.isFinite(value) && value > 0); + return ( + [ + usage.input, + usage.output, + usage.cacheRead, + usage.cacheWrite, + usage.contextUsage?.state === "available" ? usage.contextUsage.promptTokens : undefined, + usage.contextUsage?.state === "available" ? usage.contextUsage.totalTokens : undefined, + usage.reasoningTokens, + usage.total, + ].some((value) => typeof value === "number" && Number.isFinite(value) && value > 0) || + usage.contextUsage?.state === "unavailable" + ); }; export const mergeUsageIntoAccumulator = (target: UsageAccumulator, usage: MaybeUsage) => { @@ -67,6 +73,7 @@ export const mergeUsageIntoAccumulator = (target: UsageAccumulator, usage: Maybe target.lastOutput = usage.output ?? 0; target.lastCacheRead = usage.cacheRead ?? 0; target.lastCacheWrite = usage.cacheWrite ?? 0; + target.lastContextUsage = usage.contextUsage ? { ...usage.contextUsage } : undefined; target.lastReasoningTokens = usage.reasoningTokens ?? 0; target.lastTotal = callTotal; }; @@ -98,6 +105,7 @@ export const toLastCallUsage = (usage: UsageAccumulator): NormalizedUsage | unde usage.lastOutput > 0 || usage.lastCacheRead > 0 || usage.lastCacheWrite > 0 || + usage.lastContextUsage !== undefined || usage.lastReasoningTokens > 0 || usage.lastTotal > 0; if (!hasUsage) { @@ -108,6 +116,7 @@ export const toLastCallUsage = (usage: UsageAccumulator): NormalizedUsage | unde output: usage.lastOutput || undefined, cacheRead: usage.lastCacheRead || undefined, cacheWrite: usage.lastCacheWrite || undefined, + ...(usage.lastContextUsage ? { contextUsage: { ...usage.lastContextUsage } } : {}), ...(usage.lastReasoningTokens > 0 ? { reasoningTokens: usage.lastReasoningTokens } : {}), total: usage.lastTotal || undefined, }; diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.ts b/src/agents/embedded-agent-subscribe.handlers.messages.ts index c1c895cf0130..3145e3f8e02d 100644 --- a/src/agents/embedded-agent-subscribe.handlers.messages.ts +++ b/src/agents/embedded-agent-subscribe.handlers.messages.ts @@ -111,6 +111,7 @@ export function preservePendingAssistantUsage( output, cacheRead, cacheWrite, + ...(pendingUsage.contextUsage ? { contextUsage: { ...pendingUsage.contextUsage } } : {}), totalTokens: pendingUsage.total ?? input + output + cacheRead + cacheWrite, ...(pendingUsage.reasoningTokens !== undefined ? { reasoningTokens: pendingUsage.reasoningTokens } diff --git a/src/agents/sessions/agent-session.context-usage.test.ts b/src/agents/sessions/agent-session.context-usage.test.ts new file mode 100644 index 000000000000..969eb2018cf6 --- /dev/null +++ b/src/agents/sessions/agent-session.context-usage.test.ts @@ -0,0 +1,203 @@ +import { describe, expect, it } from "vitest"; +import type { AgentMessage } from "../runtime/index.js"; +import { AgentSession } from "./agent-session.js"; + +describe("AgentSession context usage", () => { + it("preserves an earlier exact snapshot when unavailable usage precedes any compaction", () => { + const messages = [ + { + role: "assistant", + content: [{ type: "text", text: "large exact response" }], + stopReason: "stop", + usage: { + input: 180_000, + output: 10_000, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 190_000, + contextUsage: { + state: "available" as const, + promptTokens: 180_000, + totalTokens: 190_000, + }, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }, + }, + { role: "user", content: "small follow-up" }, + { + role: "assistant", + content: [{ type: "text", text: "small answer" }], + stopReason: "stop", + usage: { + input: 12, + output: 8, + cacheRead: 180_000, + cacheWrite: 0, + totalTokens: 180_020, + contextUsage: { state: "unavailable" as const }, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }, + }, + ] as unknown as AgentMessage[]; + + const usage = AgentSession.prototype.getContextUsage.call({ + model: { contextWindow: 200_000 }, + messages, + sessionManager: { getBranch: () => [] }, + } as unknown as AgentSession); + + expect(usage?.tokens).toBeGreaterThan(190_000); + }); + + it("uses a content estimate after compaction when provider context usage is unavailable", () => { + const unavailableUsage = { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + totalTokens: 927_907, + contextUsage: { state: "unavailable" as const }, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }; + const messages = [ + { + role: "assistant", + content: [{ type: "text", text: "retained answer" }], + stopReason: "stop", + usage: { + ...unavailableUsage, + contextUsage: { + state: "available" as const, + promptTokens: 120_000, + totalTokens: 125_000, + }, + }, + }, + { role: "user", content: "new prompt" }, + { + role: "assistant", + content: [{ type: "text", text: "new answer" }], + stopReason: "stop", + usage: unavailableUsage, + }, + ] as unknown as AgentMessage[]; + const branchEntries = [ + { + type: "compaction", + id: "compact-1", + parentId: null, + timestamp: "2026-07-05T00:00:00.000Z", + summary: "summary", + firstKeptEntryId: "assistant-old", + tokensBefore: 120_000, + }, + { + type: "message", + id: "assistant-new", + parentId: "compact-1", + timestamp: "2026-07-05T00:00:01.000Z", + message: messages[2], + }, + ]; + + const usage = AgentSession.prototype.getContextUsage.call({ + model: { contextWindow: 200_000 }, + messages, + sessionManager: { getBranch: () => branchEntries }, + } as unknown as AgentSession); + + expect(usage?.tokens).not.toBeNull(); + expect(usage?.tokens).toBeLessThan(1_000); + }); + + it("preserves an earlier exact post-compaction snapshot before an unavailable response", () => { + const exactUsage = { + input: 180_000, + output: 10_000, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 190_000, + contextUsage: { + state: "available" as const, + promptTokens: 180_000, + totalTokens: 190_000, + }, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }; + const messages = [ + { + role: "assistant", + content: [{ type: "text", text: "exact post-compaction answer" }], + stopReason: "stop", + usage: exactUsage, + }, + { role: "user", content: "small follow-up" }, + { + role: "assistant", + content: [{ type: "text", text: "small answer" }], + stopReason: "stop", + usage: { + ...exactUsage, + contextUsage: { state: "unavailable" as const }, + }, + }, + ] as unknown as AgentMessage[]; + const branchEntries = [ + { + type: "compaction", + id: "compact-1", + parentId: null, + timestamp: "2026-07-05T00:00:00.000Z", + summary: "summary", + firstKeptEntryId: "assistant-exact", + tokensBefore: 120_000, + }, + { + type: "message", + id: "assistant-exact", + parentId: "compact-1", + timestamp: "2026-07-05T00:00:01.000Z", + message: messages[0], + }, + { + type: "message", + id: "assistant-unavailable", + parentId: "assistant-exact", + timestamp: "2026-07-05T00:00:02.000Z", + message: messages[2], + }, + ]; + + const usage = AgentSession.prototype.getContextUsage.call({ + model: { contextWindow: 200_000 }, + messages, + sessionManager: { getBranch: () => branchEntries }, + } as unknown as AgentSession); + + expect(usage?.tokens).toBeGreaterThan(190_000); + }); +}); diff --git a/src/agents/sessions/agent-session.ts b/src/agents/sessions/agent-session.ts index 2b9774ce9a02..ffd9bdfbcc98 100644 --- a/src/agents/sessions/agent-session.ts +++ b/src/agents/sessions/agent-session.ts @@ -49,6 +49,7 @@ import { collectEntriesForBranchSummaryFromBranches, compact, estimateContextTokens, + estimateTokens, generateBranchSummary, prepareCompaction, shouldCompact, @@ -329,6 +330,10 @@ type CompactionWorkOutcome = /** Standard thinking levels */ const THINKING_LEVELS: ThinkingLevel[] = ["off", "minimal", "low", "medium", "high"]; +function estimateMessagesFromContent(messages: AgentMessage[]): number { + return messages.reduce((total, message) => total + estimateTokens(message), 0); +} + // ============================================================================ // AgentSession Class // ============================================================================ @@ -2093,6 +2098,12 @@ export class AgentSession { return false; } contextTokens = estimate.tokens; + } else if (assistantMessage.usage.contextUsage?.state === "unavailable") { + const estimatedContextTokens = this.getContextUsage()?.tokens; + if (estimatedContextTokens == null) { + return false; + } + contextTokens = estimatedContextTokens; } else { contextTokens = calculateContextTokens(assistantMessage.usage); } @@ -3135,6 +3146,7 @@ export class AgentSession { // If no such assistant exists, context token count is unknown until the next LLM response. const branchEntries = this.sessionManager.getBranch(); const latestCompaction = getLatestCompactionEntry(branchEntries); + let estimateFromContent = false; if (latestCompaction) { // Check if there's a valid assistant usage after the compaction boundary @@ -3145,25 +3157,32 @@ export class AgentSession { if (entry.type === "message" && entry.message.role === "assistant") { const assistant = entry.message; if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") { + if (assistant.usage.contextUsage?.state === "unavailable") { + estimateFromContent = true; + continue; + } const contextTokens = calculateContextTokens(assistant.usage); if (contextTokens > 0) { hasPostCompactionUsage = true; + estimateFromContent = false; } break; } } } - if (!hasPostCompactionUsage) { + if (!hasPostCompactionUsage && !estimateFromContent) { return { tokens: null, contextWindow, percent: null }; } } - const estimate = estimateContextTokens(this.messages); - const percent = (estimate.tokens / contextWindow) * 100; + const tokens = estimateFromContent + ? estimateMessagesFromContent(this.messages) + : estimateContextTokens(this.messages).tokens; + const percent = (tokens / contextWindow) * 100; return { - tokens: estimate.tokens, + tokens, contextWindow, percent, }; diff --git a/src/agents/transport-stream-shared.ts b/src/agents/transport-stream-shared.ts index dfcccf2632a2..eddbae70433c 100644 --- a/src/agents/transport-stream-shared.ts +++ b/src/agents/transport-stream-shared.ts @@ -6,12 +6,14 @@ import { createAssistantMessageEventStream } from "../llm/utils/event-stream.js"; import { redactSensitiveText } from "../logging/redact.js"; import { truncateErrorDetail } from "./provider-http-errors.js"; +import type { ContextUsage } from "./usage.js"; type TransportUsage = { input: number; output: number; cacheRead: number; cacheWrite: number; + contextUsage?: ContextUsage; totalTokens: number; cost: { input: number; output: number; cacheRead: number; cacheWrite: number; total: number }; }; diff --git a/src/agents/usage.test.ts b/src/agents/usage.test.ts index 0fe285dd882b..25de6d24b181 100644 --- a/src/agents/usage.test.ts +++ b/src/agents/usage.test.ts @@ -13,6 +13,30 @@ import { } from "./usage.js"; describe("normalizeUsage", () => { + it("preserves only complete context snapshots", () => { + expect( + normalizeUsage({ + input: 12, + contextUsage: { state: "available", promptTokens: 148_874, totalTokens: 163_978 }, + }), + ).toMatchObject({ + input: 12, + contextUsage: { state: "available", promptTokens: 148_874, totalTokens: 163_978 }, + }); + expect( + normalizeUsage({ + input: 12, + contextUsage: { state: "available", promptTokens: 163_978, totalTokens: 148_874 }, + }), + ).toEqual({ + input: 12, + output: undefined, + cacheRead: undefined, + cacheWrite: undefined, + total: undefined, + }); + }); + it("normalizes cache fields from provider response", () => { const usage = normalizeUsage({ input: 1000, @@ -374,6 +398,66 @@ describe("deriveContextPromptTokens", () => { ).toBe(81_000); }); + it("prefers explicit prompt buckets over total-minus-output fallback", () => { + expect( + deriveContextPromptTokens({ + lastCallUsage: { input: 20, cacheRead: 100, output: 30, total: 250 }, + }), + ).toBe(120); + }); + + it("prefers an explicit final-iteration context snapshot over aggregate billing usage", () => { + expect( + deriveContextPromptTokens({ + lastCallUsage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + total: 927_907, + }, + }), + ).toBe(148_874); + }); + + it("does not reconstruct context when the provider snapshot is unavailable", () => { + expect( + deriveContextPromptTokens({ + lastCallUsage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + total: 927_907, + }, + }), + ).toBeUndefined(); + }); + + it("does not treat total-only usage as a prompt snapshot", () => { + expect( + deriveContextPromptTokens({ + lastCallUsage: { input: 1_000, total: 1_200 }, + }), + ).toBe(1_000); + expect( + deriveContextPromptTokens({ + lastCallUsage: { total: 1_200 }, + }), + ).toBeUndefined(); + expect( + deriveContextPromptTokens({ + lastCallUsage: { output: 200, total: 1_200 }, + }), + ).toBe(1_000); + }); + it("falls back to accumulated usage when no prompt snapshot exists", () => { expect( deriveContextPromptTokens({ @@ -381,9 +465,51 @@ describe("deriveContextPromptTokens", () => { }), ).toBe(100_000); }); + + it("keeps accumulated usage on its component-based context snapshot", () => { + expect( + deriveContextPromptTokens({ + usage: { input: 10_000, cacheRead: 26_000, output: 1_000, total: 36_000 }, + }), + ).toBe(36_000); + }); }); describe("deriveSessionTotalTokens", () => { + it("prefers the explicit context snapshot over aggregate billing buckets", () => { + expect( + deriveSessionTotalTokens({ + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + total: 927_907, + }, + }), + ).toBe(148_874); + }); + + it("does not store aggregate billing as session context when the snapshot is unavailable", () => { + expect( + deriveSessionTotalTokens({ + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + total: 927_907, + }, + }), + ).toBeUndefined(); + }); + it("includes cache tokens in total calculation", () => { const totalTokens = deriveSessionTotalTokens({ usage: { diff --git a/src/agents/usage.ts b/src/agents/usage.ts index 70c87369c00f..b8b0fa8da7cd 100644 --- a/src/agents/usage.ts +++ b/src/agents/usage.ts @@ -5,12 +5,17 @@ */ import { asFiniteNumber } from "@openclaw/normalization-core/number-coercion"; +export type ContextUsage = + | { state: "available"; promptTokens: number; totalTokens: number } + | { state: "unavailable" }; + /** Provider/SDK usage payload variants accepted by usage normalization. */ export type UsageLike = { input?: number; output?: number; cacheRead?: number; cacheWrite?: number; + contextUsage?: ContextUsage; total?: number; // Common alternates across providers/SDKs. inputTokens?: number; @@ -53,6 +58,7 @@ export type NormalizedUsage = { output?: number; cacheRead?: number; cacheWrite?: number; + contextUsage?: ContextUsage; reasoningTokens?: number; total?: number; }; @@ -72,6 +78,7 @@ export type AssistantUsageSnapshot = { output: number; cacheRead: number; cacheWrite: number; + contextUsage?: ContextUsage; totalTokens: number; cost: { input: number; @@ -105,14 +112,19 @@ export function hasNonzeroUsage(usage?: NormalizedUsage | null): usage is Normal if (!usage) { return false; } - return [ - usage.input, - usage.output, - usage.cacheRead, - usage.cacheWrite, - usage.reasoningTokens, - usage.total, - ].some((v) => typeof v === "number" && Number.isFinite(v) && v > 0); + return ( + [ + usage.input, + usage.output, + usage.cacheRead, + usage.cacheWrite, + usage.contextUsage?.state === "available" ? usage.contextUsage.promptTokens : undefined, + usage.contextUsage?.state === "available" ? usage.contextUsage.totalTokens : undefined, + usage.reasoningTokens, + usage.total, + ].some((v) => typeof v === "number" && Number.isFinite(v) && v > 0) || + usage.contextUsage?.state === "unavailable" + ); } const normalizeTokenCount = (value: unknown): number | undefined => { @@ -178,6 +190,26 @@ export function normalizeUsage(raw?: UsageLike | null): NormalizedUsage | undefi const cacheWrite = normalizeTokenCount( raw.cacheWrite ?? raw.cache_write ?? raw.cache_creation_input_tokens, ); + const contextPromptTokens = + raw.contextUsage?.state === "available" + ? normalizeTokenCount(raw.contextUsage.promptTokens) + : undefined; + const contextTotalTokens = + raw.contextUsage?.state === "available" + ? normalizeTokenCount(raw.contextUsage.totalTokens) + : undefined; + const contextUsage = + raw.contextUsage?.state === "unavailable" + ? ({ state: "unavailable" } as const) + : contextPromptTokens !== undefined && + contextTotalTokens !== undefined && + contextTotalTokens >= contextPromptTokens + ? ({ + state: "available", + promptTokens: contextPromptTokens, + totalTokens: contextTotalTokens, + } as const) + : undefined; const reasoningTokens = normalizeTokenCount( raw.reasoningTokens ?? raw.reasoning_tokens ?? @@ -191,6 +223,7 @@ export function normalizeUsage(raw?: UsageLike | null): NormalizedUsage | undefi output === undefined && cacheRead === undefined && cacheWrite === undefined && + contextUsage === undefined && reasoningTokens === undefined && total === undefined ) { @@ -202,6 +235,7 @@ export function normalizeUsage(raw?: UsageLike | null): NormalizedUsage | undefi output, cacheRead, cacheWrite, + ...(contextUsage ? { contextUsage } : {}), ...(reasoningTokens !== undefined ? { reasoningTokens } : {}), total, }; @@ -266,6 +300,23 @@ export function derivePromptTokens(usage?: { return sum > 0 ? sum : undefined; } +function derivePromptTokensFromTotal(usage?: NormalizedUsage): number | undefined { + const total = usage?.total; + const output = usage?.output; + if ( + typeof total !== "number" || + !Number.isFinite(total) || + total <= 0 || + typeof output !== "number" || + !Number.isFinite(output) || + output < 0 + ) { + return undefined; + } + const promptTokens = total - output; + return promptTokens > 0 ? promptTokens : undefined; +} + /** Resolve context prompt tokens from explicit override, last call, or aggregate usage. */ export function deriveContextPromptTokens(params: { lastCallUsage?: NormalizedUsage; @@ -277,7 +328,24 @@ export function deriveContextPromptTokens(params: { return promptOverride; } - return derivePromptTokens(params.lastCallUsage) ?? derivePromptTokens(params.usage); + if (params.lastCallUsage?.contextUsage?.state === "unavailable") { + return undefined; + } + if (params.lastCallUsage?.contextUsage?.state === "available") { + return params.lastCallUsage.contextUsage.promptTokens; + } + const lastCallPromptTokens = + derivePromptTokens(params.lastCallUsage) ?? derivePromptTokensFromTotal(params.lastCallUsage); + if (lastCallPromptTokens !== undefined) { + return lastCallPromptTokens; + } + if (params.usage?.contextUsage?.state === "unavailable") { + return undefined; + } + if (params.usage?.contextUsage?.state === "available") { + return params.usage.contextUsage.promptTokens; + } + return derivePromptTokens(params.usage); } /** Derive the session prompt-token snapshot stored for context display. */ @@ -288,6 +356,7 @@ export function deriveSessionTotalTokens(params: { total?: number; cacheRead?: number; cacheWrite?: number; + contextUsage?: ContextUsage; }; contextTokens?: number; promptTokens?: number; diff --git a/src/auto-reply/reply/agent-runner-memory.test.ts b/src/auto-reply/reply/agent-runner-memory.test.ts index 2bb18bdddecd..dc75b5e69fa3 100644 --- a/src/auto-reply/reply/agent-runner-memory.test.ts +++ b/src/auto-reply/reply/agent-runner-memory.test.ts @@ -1910,6 +1910,61 @@ describe("runMemoryFlushIfNeeded", () => { expect(compactCall.sessionFile).toContain("active-run-session.jsonl"); }); + it("does not treat unavailable Anthropic context as transcript prompt usage", async () => { + const sessionFile = path.join(rootDir, "unavailable-context-session.jsonl"); + await fs.writeFile( + sessionFile, + `${JSON.stringify({ + message: { + role: "assistant", + content: "small answer", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + totalTokens: 927_907, + }, + }, + })}\n`, + "utf8", + ); + registerMemoryFlushPlanResolverForTest(() => ({ + softThresholdTokens: 4_000, + forceFlushTranscriptBytes: 1_000_000_000, + reserveTokensFloor: 0, + prompt: "Pre-compaction memory flush.\nNO_REPLY", + systemPrompt: "Write memory to memory/YYYY-MM-DD.md.", + relativePath: "memory/2023-11-14.md", + })); + const sessionEntry: SessionEntry = { + sessionId: "session", + sessionFile, + updatedAt: Date.now(), + totalTokensFresh: false, + }; + + await runPreflightCompactionIfNeeded({ + cfg: { agents: { defaults: { compaction: { memoryFlush: {} } } } }, + followupRun: createTestFollowupRun({ + sessionId: "session", + sessionFile, + sessionKey: "main", + }), + defaultModel: "anthropic/claude-opus-4-6", + agentCfgContextTokens: 100_000, + sessionEntry, + sessionStore: { main: sessionEntry }, + sessionKey: "main", + storePath: path.join(rootDir, "sessions.json"), + isHeartbeat: false, + replyOperation: createReplyOperation(), + }); + + expect(compactEmbeddedAgentSessionMock).not.toHaveBeenCalled(); + }); + it("keeps preflight compaction conservative for content appended after latest usage", async () => { const sessionFile = path.join(rootDir, "post-usage-tail-session.jsonl"); await fs.writeFile( diff --git a/src/auto-reply/reply/agent-runner-memory.ts b/src/auto-reply/reply/agent-runner-memory.ts index 5773468cb55b..3fead621b3d8 100644 --- a/src/auto-reply/reply/agent-runner-memory.ts +++ b/src/auto-reply/reply/agent-runner-memory.ts @@ -18,7 +18,7 @@ import { resolveContextConfigProviderForRuntime } from "../../agents/openai-rout import type { AgentMessage } from "../../agents/runtime/index.js"; import { resolveSandboxConfigForAgent, resolveSandboxRuntimeStatus } from "../../agents/sandbox.js"; import { - derivePromptTokens, + deriveContextPromptTokens, hasNonzeroUsage, normalizeUsage, type UsageLike, @@ -34,9 +34,9 @@ import { updateSessionEntry } from "../../config/sessions/session-accessor.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; import { readSessionMessagesAsync } from "../../gateway/session-utils.fs.js"; import { logVerbose } from "../../globals.js"; +import { isAbortError } from "../../infra/abort-signal.js"; import { emitAgentEvent, registerAgentRunContext } from "../../infra/agent-events.js"; import { formatErrorMessage } from "../../infra/errors.js"; -import { isAbortError } from "../../infra/abort-signal.js"; import { resolveMemoryFlushPlan } from "../../plugins/memory-state.js"; import { CommandLane } from "../../process/lanes.js"; import { createLazyImportLoader } from "../../shared/lazy-promise.js"; @@ -447,7 +447,7 @@ function deriveTranscriptUsageSnapshot( if (!usage) { return undefined; } - const promptTokens = derivePromptTokens(usage); + const promptTokens = deriveContextPromptTokens({ lastCallUsage: usage }); const outputRaw = usage.output; const outputTokens = typeof outputRaw === "number" && Number.isFinite(outputRaw) && outputRaw > 0 diff --git a/src/auto-reply/reply/session-fork.runtime.test.ts b/src/auto-reply/reply/session-fork.runtime.test.ts index ebeb213fddd6..c47362f9eb8a 100644 --- a/src/auto-reply/reply/session-fork.runtime.test.ts +++ b/src/auto-reply/reply/session-fork.runtime.test.ts @@ -150,6 +150,9 @@ describe("resolveParentForkTokenCountRuntime", () => { }, }), JSON.stringify({ + type: "message", + id: "active-usage", + parentId: null, message: { role: "assistant", content: "latest", @@ -175,7 +178,140 @@ describe("resolveParentForkTokenCountRuntime", () => { expect(tokens).toBe(78_000); }); - it("keeps parent fork checks conservative for content appended after latest usage", async () => { + it("does not reconstruct parent context from billing buckets when context is unavailable", async () => { + const root = await makeRoot("openclaw-parent-fork-unavailable-context-"); + const sessionsDir = path.join(root, "sessions"); + await fs.mkdir(sessionsDir); + + const sessionId = "parent-unavailable-context"; + const sessionFile = path.join(sessionsDir, "parent.jsonl"); + await fs.writeFile( + sessionFile, + [ + JSON.stringify({ + type: "session", + version: 3, + id: sessionId, + timestamp: new Date().toISOString(), + cwd: process.cwd(), + }), + JSON.stringify({ + message: { + role: "assistant", + content: "latest", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + total: 927_907, + }, + }, + }), + ].join("\n"), + "utf-8", + ); + + const entry: SessionEntry = { + sessionId, + sessionFile, + updatedAt: Date.now(), + totalTokens: 4_567, + totalTokensFresh: false, + }; + + const tokens = await resolveParentForkTokenCountRuntime({ + parentEntry: entry, + storePath: path.join(root, "sessions.json"), + }); + + expect(tokens).toBe(4_567); + }); + + it("uses the exact final-iteration total when context usage is available", async () => { + const root = await makeRoot("openclaw-parent-fork-exact-context-"); + const sessionsDir = path.join(root, "sessions"); + await fs.mkdir(sessionsDir); + + const sessionId = "parent-exact-context"; + const sessionFile = path.join(sessionsDir, "parent.jsonl"); + await fs.writeFile( + sessionFile, + [ + JSON.stringify({ + type: "session", + version: 3, + id: sessionId, + timestamp: new Date().toISOString(), + cwd: process.cwd(), + }), + JSON.stringify({ + type: "message", + id: "active-usage", + parentId: null, + message: { + role: "assistant", + content: "latest", + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { + state: "available", + promptTokens: 148_874, + totalTokens: 163_978, + }, + total: 927_907, + }, + }, + }), + JSON.stringify({ + type: "message", + id: "inactive-side-usage", + parentId: "active-usage", + message: { + role: "assistant", + content: `side branch ${"x".repeat(1_100_000)}`, + usage: { + input: 9_000, + output: 1_000, + contextUsage: { + state: "available", + promptTokens: 9_000, + totalTokens: 10_000, + }, + }, + }, + }), + JSON.stringify({ + type: "leaf", + id: "active-leaf", + parentId: "inactive-side-usage", + targetId: "active-usage", + }), + ].join("\n"), + "utf-8", + ); + + const entry: SessionEntry = { + sessionId, + sessionFile, + updatedAt: Date.now(), + totalTokens: 900_000, + totalTokensFresh: false, + }; + + const tokens = await resolveParentForkTokenCountRuntime({ + parentEntry: entry, + storePath: path.join(root, "sessions.json"), + }); + + expect(tokens).toBe(163_978); + }); + + it("adds only post-usage transcript pressure to an exact context snapshot", async () => { const root = await makeRoot("openclaw-parent-fork-post-usage-tail-"); const sessionsDir = path.join(root, "sessions"); await fs.mkdir(sessionsDir); @@ -196,13 +332,21 @@ describe("resolveParentForkTokenCountRuntime", () => { message: { role: "assistant", content: "latest model call", - usage: { input: 40_000, output: 2_000 }, + usage: { + input: 12, + output: 10_000, + contextUsage: { + state: "available", + promptTokens: 70_000, + totalTokens: 80_000, + }, + }, }, }), JSON.stringify({ message: { role: "tool", - content: `large appended tool result ${"x".repeat(450_000)}`, + content: `large appended tool result ${"x".repeat(100_000)}`, }, }), ].join("\n"), @@ -222,6 +366,7 @@ describe("resolveParentForkTokenCountRuntime", () => { }); expect(tokens).toBeGreaterThan(100_000); + expect(tokens).toBeLessThan(110_000); }); }); diff --git a/src/auto-reply/reply/session-fork.runtime.ts b/src/auto-reply/reply/session-fork.runtime.ts index 660f2f94a0ae..170c266cd13a 100644 --- a/src/auto-reply/reply/session-fork.runtime.ts +++ b/src/auto-reply/reply/session-fork.runtime.ts @@ -95,20 +95,30 @@ export async function resolveParentForkTokenCountRuntime(params: { undefined, 1024 * 1024, ); - const promptTokens = resolvePositiveTokenCount( - derivePromptTokens({ - input: usage?.inputTokens, - cacheRead: usage?.cacheRead, - cacheWrite: usage?.cacheWrite, - }), - ); - const outputTokens = resolvePositiveTokenCount(usage?.outputTokens); - if (typeof promptTokens === "number") { - return maxPositiveTokenCount( - promptTokens + (outputTokens ?? 0), - cachedTokens, - byteEstimateTokens, + let transcriptTokens: number | undefined; + if (usage?.contextUsage?.state === "available") { + const trailingTokens = Math.ceil( + (usage.trailingBytes ?? 0) / FALLBACK_TRANSCRIPT_BYTES_PER_TOKEN, ); + transcriptTokens = resolvePositiveTokenCount(usage.contextUsage.totalTokens + trailingTokens); + if (typeof transcriptTokens === "number") { + return transcriptTokens; + } + } else if (usage?.contextUsage?.state !== "unavailable") { + const promptTokens = resolvePositiveTokenCount( + derivePromptTokens({ + input: usage?.inputTokens, + cacheRead: usage?.cacheRead, + cacheWrite: usage?.cacheWrite, + }), + ); + const outputTokens = resolvePositiveTokenCount(usage?.outputTokens); + if (typeof promptTokens === "number") { + transcriptTokens = promptTokens + (outputTokens ?? 0); + } + } + if (typeof transcriptTokens === "number") { + return maxPositiveTokenCount(transcriptTokens, cachedTokens, byteEstimateTokens); } } catch { // Fall back to cached totals when recent transcript usage cannot be read. diff --git a/src/auto-reply/reply/session-usage.ts b/src/auto-reply/reply/session-usage.ts index b36754584a1d..c3b0b3fc7555 100644 --- a/src/auto-reply/reply/session-usage.ts +++ b/src/auto-reply/reply/session-usage.ts @@ -131,8 +131,12 @@ export async function persistSessionUsageUpdate(params: { typeof params.promptTokens === "number" && Number.isFinite(params.promptTokens) && params.promptTokens > 0; + const hasUsableLastCallUsage = + Boolean(params.lastCallUsage) && params.lastCallUsage?.contextUsage?.state !== "unavailable"; + const hasUsableUsageContextSnapshot = + params.usageIsContextSnapshot === true && params.usage?.contextUsage?.state !== "unavailable"; const hasFreshContextSnapshot = - Boolean(params.lastCallUsage) || hasPromptTokens || params.usageIsContextSnapshot === true; + hasUsableLastCallUsage || hasPromptTokens || hasUsableUsageContextSnapshot; const compactionTokensAfter = resolveNonNegativeTokenCount(params.compactionTokensAfter); const hasCompactionSnapshot = compactionTokensAfter !== undefined; diff --git a/src/auto-reply/reply/session.test.ts b/src/auto-reply/reply/session.test.ts index 521836046ff1..f0488c1024a3 100644 --- a/src/auto-reply/reply/session.test.ts +++ b/src/auto-reply/reply/session.test.ts @@ -3792,9 +3792,7 @@ describe("initSessionState preserves behavior overrides across /new and /reset", // Foreign owners may need the writer lane to finalize before releasing. // The rollover must not hold that lane while it drains them. await runExclusiveSessionStoreWrite(storePath, async () => {}); - expect(readSessionStoreForTest(storePath)[sessionKey]?.sessionId).toBe( - existingSessionId, - ); + expect(readSessionStoreForTest(storePath)[sessionKey]?.sessionId).toBe(existingSessionId); expect(await fs.stat(transcriptPath).catch(() => null)).not.toBeNull(); admission.release(); @@ -3898,9 +3896,7 @@ describe("initSessionState preserves behavior overrides across /new and /reset", const outcomes = await Promise.allSettled([runRollover(0), runRollover(1)]); expect(outcomes.filter((outcome) => outcome.status === "fulfilled")).toHaveLength(1); expect(outcomes.filter((outcome) => outcome.status === "rejected")).toHaveLength(1); - expect(readSessionStoreForTest(storePath)[sessionKey]?.sessionId).not.toBe( - existingSessionId, - ); + expect(readSessionStoreForTest(storePath)[sessionKey]?.sessionId).not.toBe(existingSessionId); }); it.each([ @@ -4620,6 +4616,42 @@ describe("persistSessionUsageUpdate", () => { expect(stored[sessionKey].outputTokens).toBe(10_000); }); + it("keeps the prior total stale when last-call context is unavailable", async () => { + const storePath = await createStorePath("openclaw-usage-unavailable-context-"); + const sessionKey = "main"; + await seedSessionStore({ + storePath, + sessionKey, + entry: { + sessionId: "s1", + updatedAt: Date.now(), + totalTokens: 148_874, + totalTokensFresh: true, + }, + }); + + await persistSessionUsageUpdate({ + storePath, + sessionKey, + usage: { input: 12, output: 15_104, cacheRead: 819_661, cacheWrite: 93_130 }, + lastCallUsage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + total: 927_907, + }, + contextTokensUsed: 200_000, + }); + + const stored = JSON.parse(await fs.readFile(storePath, "utf-8")); + expect(stored[sessionKey].totalTokens).toBe(148_874); + expect(stored[sessionKey].totalTokensFresh).toBe(false); + expect(stored[sessionKey].inputTokens).toBe(12); + expect(stored[sessionKey].cacheRead).toBe(819_661); + }); + it("marks a fresh zero stale when a completed run has no context snapshot", async () => { const storePath = await createStorePath("openclaw-usage-no-snapshot-"); const sessionKey = "main"; diff --git a/src/context-engine/types.ts b/src/context-engine/types.ts index e5dce135dab2..2794d620c823 100644 --- a/src/context-engine/types.ts +++ b/src/context-engine/types.ts @@ -222,6 +222,9 @@ type ContextEnginePromptCacheUsage = { output?: number; cacheRead?: number; cacheWrite?: number; + contextUsage?: + | { state: "available"; promptTokens: number; totalTokens: number } + | { state: "unavailable" }; total?: number; }; diff --git a/src/cron/isolated-agent/run.ts b/src/cron/isolated-agent/run.ts index 053acbd940fc..48d30337c7cf 100644 --- a/src/cron/isolated-agent/run.ts +++ b/src/cron/isolated-agent/run.ts @@ -1135,7 +1135,9 @@ async function finalizeCronRun(params: { const totalTokens = typeof lastCallTotalTokens === "number" && lastCallTotalTokens > 0 ? lastCallTotalTokens - : deriveSessionTotalTokens({ usage, contextTokens, promptTokens }); + : lastCallUsage?.contextUsage?.state === "unavailable" + ? undefined + : deriveSessionTotalTokens({ usage, contextTokens, promptTokens }); const runEstimatedCostUsd = resolveNonNegativeNumber( estimateUsageCost({ usage, diff --git a/src/cron/isolated-agent/run.usage-accounting.test.ts b/src/cron/isolated-agent/run.usage-accounting.test.ts index 195d428765e9..4195def4a79c 100644 --- a/src/cron/isolated-agent/run.usage-accounting.test.ts +++ b/src/cron/isolated-agent/run.usage-accounting.test.ts @@ -148,4 +148,40 @@ describe("runCronIsolatedAgentTurn usage accounting", () => { promptTokens: undefined, }); }); + + it("does not fall back to aggregate billing when final-call context is unavailable", async () => { + const cronSession = makeCronSession(); + resolveCronSessionMock.mockReturnValue(cronSession); + mockRunCronFallbackPassthrough(); + deriveSessionTotalTokensMock.mockReturnValueOnce(undefined); + runEmbeddedAgentMock.mockResolvedValueOnce({ + payloads: [{ text: "done" }], + meta: { + agentMeta: { + usage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + total: 927_907, + }, + lastCallUsage: { + input: 12, + output: 15_104, + cacheRead: 819_661, + cacheWrite: 93_130, + contextUsage: { state: "unavailable" }, + total: 927_907, + }, + }, + }, + }); + + const result = await runCronIsolatedAgentTurn(makeIsolatedAgentParamsFixture()); + + expect(result.status).toBe("ok"); + expect(cronSession.sessionEntry.totalTokens).toBeUndefined(); + expect(cronSession.sessionEntry.totalTokensFresh).toBe(false); + expect(deriveSessionTotalTokensMock).toHaveBeenCalledTimes(1); + }); }); diff --git a/src/gateway/session-transcript-readers.test.ts b/src/gateway/session-transcript-readers.test.ts index 540b9427a2cd..4469701c04ff 100644 --- a/src/gateway/session-transcript-readers.test.ts +++ b/src/gateway/session-transcript-readers.test.ts @@ -110,7 +110,11 @@ describe("session transcript reader facade", () => { content: "metered answer", provider: "openai", model: "gpt-5.5", - usage: { input: 11, output: 7 }, + usage: { + input: 11, + output: 7, + contextUsage: { state: "unavailable" }, + }, }, }, ]); @@ -123,6 +127,7 @@ describe("session transcript reader facade", () => { readLatestRecentSessionUsageFromTranscriptAsync(scope, 4096), ).resolves.toMatchObject({ inputTokens: 11, + contextUsage: { state: "unavailable" }, model: "gpt-5.5", modelProvider: "openai", outputTokens: 7, diff --git a/src/gateway/session-transcript-readers.ts b/src/gateway/session-transcript-readers.ts index 7bbcbe1d5000..e266c3bdaeee 100644 --- a/src/gateway/session-transcript-readers.ts +++ b/src/gateway/session-transcript-readers.ts @@ -3,6 +3,7 @@ import { resolveSessionTranscriptReadTarget } from "../config/sessions/session-a import type { ReadRecentSessionMessagesOptions, ReadSessionMessagesAsyncOptions, + SessionTranscriptUsageSnapshot, } from "./session-utils.fs.js"; import { readFirstUserMessageFromTranscript as readFirstUserMessageFromTranscriptFile, @@ -58,18 +59,6 @@ type ReadSessionMessageByIdResult = { found: boolean; }; -type SessionTranscriptUsageSnapshot = { - modelProvider?: string; - model?: string; - inputTokens?: number; - outputTokens?: number; - cacheRead?: number; - cacheWrite?: number; - totalTokens?: number; - totalTokensFresh?: boolean; - costUsd?: number; -}; - type FileBackedReadScope = { agentId?: string; sessionFile: string; diff --git a/src/gateway/session-utils.fs.test.ts b/src/gateway/session-utils.fs.test.ts index 590b4970933d..d2ace8d25c22 100644 --- a/src/gateway/session-utils.fs.test.ts +++ b/src/gateway/session-utils.fs.test.ts @@ -1437,7 +1437,17 @@ describe("readSessionMessages", () => { writeTranscript(tmpDir, sessionId, [ { type: "session", version: 1, id: sessionId }, { message: { role: "assistant", content: "older", usage: { input: 50, output: 5 } } }, - { message: { role: "assistant", content: "latest", usage: { input: 70, output: 9 } } }, + { + message: { + role: "assistant", + content: "latest", + usage: { + input: 70, + output: 9, + contextUsage: { state: "unavailable" }, + }, + }, + }, ]); const aggregate = await readRecentSessionUsageFromTranscriptAsync( @@ -1455,8 +1465,89 @@ describe("readSessionMessages", () => { 2048, ); + expectUsageFields(aggregate, { + inputTokens: 120, + outputTokens: 14, + contextUsage: { state: "unavailable" }, + }); + expect(aggregate).not.toHaveProperty("totalTokens"); + expect(aggregate).not.toHaveProperty("totalTokensFresh"); + expectUsageFields(latest, { + inputTokens: 70, + outputTokens: 9, + contextUsage: { state: "unavailable" }, + trailingBytes: 0, + }); + }); + + test("counts transcript bytes appended after the latest usage snapshot", async () => { + const sessionId = "test-session-latest-usage-trailing-bytes"; + writeTranscript(tmpDir, sessionId, [ + { type: "session", version: 1, id: sessionId }, + { + message: { + role: "assistant", + content: "latest", + usage: { + input: 70, + output: 9, + contextUsage: { + state: "available", + promptTokens: 70, + totalTokens: 79, + }, + }, + }, + }, + { message: { role: "tool", content: "appended tool result" } }, + ]); + + const latest = await readLatestRecentSessionUsageFromTranscriptAsync( + sessionId, + storePath, + undefined, + undefined, + 2048, + ); + + expectUsageFields(latest, { + contextUsage: { + state: "available", + promptTokens: 70, + totalTokens: 79, + }, + }); + expect(latest?.trailingBytes).toBeGreaterThan(0); + }); + + test("clears an older context marker when aggregate usage has a newer plain snapshot", async () => { + const sessionId = "test-session-aggregate-clears-context-marker"; + writeTranscript(tmpDir, sessionId, [ + { type: "session", version: 1, id: sessionId }, + { + message: { + role: "assistant", + content: "older", + usage: { + input: 50, + output: 5, + contextUsage: { state: "unavailable" }, + }, + }, + }, + { message: { role: "assistant", content: "latest", usage: { input: 70, output: 9 } } }, + ]); + + const aggregate = await readRecentSessionUsageFromTranscriptAsync( + sessionId, + storePath, + undefined, + undefined, + 2048, + ); + expectUsageFields(aggregate, { inputTokens: 120, outputTokens: 14 }); - expectUsageFields(latest, { inputTokens: 70, outputTokens: 9 }); + expect(aggregate).not.toHaveProperty("contextUsage"); }); test("tails transcript lines for manual compaction without loading the whole file", () => { diff --git a/src/gateway/session-utils.fs.ts b/src/gateway/session-utils.fs.ts index 22ba456d8e61..e48c9cdf2461 100644 --- a/src/gateway/session-utils.fs.ts +++ b/src/gateway/session-utils.fs.ts @@ -7,7 +7,12 @@ import { resolveNonNegativeIntegerOption, } from "@openclaw/normalization-core/number-coercion"; import { normalizeLowercaseStringOrEmpty } from "@openclaw/normalization-core/string-coerce"; -import { deriveSessionTotalTokens, hasNonzeroUsage, normalizeUsage } from "../agents/usage.js"; +import { + deriveSessionTotalTokens, + hasNonzeroUsage, + normalizeUsage, + type ContextUsage, +} from "../agents/usage.js"; import { scanSessionTranscriptTree, selectSessionTranscriptTreePathNodes, @@ -1342,13 +1347,15 @@ async function readLastMessagePreviewFromOpenTranscriptAsync(params: { return extractLastMessagePreviewFromTranscriptLines(lines.slice(-LAST_MSG_MAX_LINES)); } -type SessionTranscriptUsageSnapshot = { +export type SessionTranscriptUsageSnapshot = { modelProvider?: string; model?: string; inputTokens?: number; outputTokens?: number; cacheRead?: number; cacheWrite?: number; + contextUsage?: ContextUsage; + trailingBytes?: number; totalTokens?: number; totalTokensFresh?: boolean; costUsd?: number; @@ -1523,6 +1530,9 @@ function extractUsageSnapshotFromTranscriptLine( if (typeof usage?.cacheWrite === "number" && Number.isFinite(usage.cacheWrite)) { snapshot.cacheWrite = usage.cacheWrite; } + if (usage?.contextUsage) { + snapshot.contextUsage = usage.contextUsage; + } if (typeof totalTokens === "number") { snapshot.totalTokens = totalTokens; snapshot.totalTokensFresh = true; @@ -1589,7 +1599,15 @@ function extractAggregateUsageFromTranscriptLines( cacheWrite += current.cacheWrite; sawCacheWrite = true; } - if (typeof current.totalTokens === "number") { + if (current.contextUsage) { + snapshot.contextUsage = current.contextUsage; + } else { + delete snapshot.contextUsage; + } + if (current.contextUsage?.state === "unavailable") { + delete snapshot.totalTokens; + delete snapshot.totalTokensFresh; + } else if (typeof current.totalTokens === "number") { snapshot.totalTokens = current.totalTokens; snapshot.totalTokensFresh = true; } @@ -1631,16 +1649,88 @@ function extractAggregateUsageFromTranscriptLines( return snapshot; } +function hasTranscriptUsage( + snapshot: SessionTranscriptUsageSnapshot | null, +): snapshot is SessionTranscriptUsageSnapshot { + return Boolean( + snapshot && + (snapshot.contextUsage !== undefined || + snapshot.inputTokens !== undefined || + snapshot.outputTokens !== undefined || + snapshot.cacheRead !== undefined || + snapshot.cacheWrite !== undefined || + snapshot.totalTokens !== undefined || + snapshot.costUsd !== undefined), + ); +} + function extractLatestUsageFromTranscriptLines( lines: Iterable, ): SessionTranscriptUsageSnapshot | null { + const parsed = Array.from(lines).flatMap((line) => { + const entry = parseTailTranscriptRecord(line); + return entry ? [{ entry, line }] : []; + }); + const selected = selectBoundedActiveTailRecords( + parsed.map(({ entry }) => entry), + { failClosedOnInvalidLeafControl: true }, + ); + const lineByRecord = new Map(parsed.map(({ entry, line }) => [entry.record, line])); let latest: SessionTranscriptUsageSnapshot | null = null; - for (const line of lines) { - latest = extractUsageSnapshotFromTranscriptLine(line) ?? latest; + let trailingBytes = 0; + for (const entry of selected) { + const line = lineByRecord.get(entry.record); + if (!line) { + continue; + } + const current = extractUsageSnapshotFromTranscriptLine(line); + if (hasTranscriptUsage(current)) { + latest = current; + trailingBytes = 0; + } else if (latest) { + trailingBytes += Buffer.byteLength(line, "utf8") + 1; + } + } + if (latest) { + latest.trailingBytes = trailingBytes; } return latest; } +function hasInvalidLeafControl(lines: Iterable): boolean { + const entries = Array.from(lines).flatMap((line) => { + const entry = parseTailTranscriptRecord(line); + return entry ? [entry.record] : []; + }); + const tree = scanSessionTranscriptTree(entries); + return tree.hasInvalidLeafControl; +} + +async function extractLatestUsageFromTranscriptIndex( + filePath: string, +): Promise { + const index = await readSessionTranscriptIndex(filePath); + if (!index) { + return null; + } + let trailingBytes = 0; + for (let position = index.entries.length - 1; position >= 0; position -= 1) { + const entry = index.entries[position]; + if (!entry) { + continue; + } + if (entry.byteLength <= MAX_TRANSCRIPT_PARSE_LINE_BYTES) { + const current = extractUsageSnapshotFromTranscriptLine(JSON.stringify(entry.record)); + if (hasTranscriptUsage(current)) { + current.trailingBytes = trailingBytes; + return current; + } + } + trailingBytes += entry.byteLength + 1; + } + return null; +} + function extractAggregateUsageFromTranscriptChunk( chunk: string, ): SessionTranscriptUsageSnapshot | null { @@ -1748,6 +1838,9 @@ export async function readLatestRecentSessionUsageFromTranscriptAsync( maxLines: 1000, maxBytes, }); + if (hasInvalidLeafControl(lines)) { + return await extractLatestUsageFromTranscriptIndex(filePath); + } return extractLatestUsageFromTranscriptLines(lines); } catch { return null;