mirror of
https://github.com/getpaseo/paseo.git
synced 2026-10-14 13:22:37 +00:00
fix(claude): count input reported late in the context meter (#6703)
* fix(claude): read cumulative input usage from message_delta The context meter took input_tokens only from message_start. Backends behind a translating gateway report 0 there and the real cumulative counts on message_delta, so contextWindowUsedTokens collapsed to output_tokens. Keep the message_start counters per field and let a counter present on message_delta replace it, the same rule the Anthropic SDK MessageStream uses. A missing counter keeps the message_start value, so a delta without cache fields does not drop cache_read_input_tokens. Closes #4294 * test(claude): ignore a message_delta that precedes its message_start A canceled request can leave its terminal message_delta queued ahead of the next turn. Pin that such a straggler neither emits usage_updated nor leaks into the next request's counters. * test(claude): take the message_delta usage object as a single parameter Keep createMessageDeltaEvent(outputTokens) for existing call sites and add createMessageDeltaUsageEvent(usage) for the tests that set input counters on the delta, instead of a second positional parameter. * test(claude): verify late input usage through the real provider process Compare start-side and delta-side counters through two real provider turns against a local Messages API fixture, including omitted cache counters. The delta cases fail on the baseline with output-only context totals. Co-authored-by: kywoo26 <kywoo26@naver.com> * test(claude): always close the context usage fixture --------- Co-authored-by: kywoo26 <kywoo26@naver.com> Co-authored-by: paseo-bot[bot] <266920839+paseo-bot[bot]@users.noreply.github.com>
This commit is contained in:
paseo-bot[bot]
co-authored by
kywoo26
kywoo26
parent
99e4b53aa0
commit
a2b040869d
+133
@@ -0,0 +1,133 @@
|
||||
import { createServer } from "node:http";
|
||||
import { mkdir, mkdtemp, rm } from "node:fs/promises";
|
||||
import path from "node:path";
|
||||
import pino from "pino";
|
||||
import { expect, test } from "vitest";
|
||||
import { collectSessionTurnEvents } from "../test-utils/session-stream-adapter.js";
|
||||
import { ClaudeAgentClient } from "./agent.js";
|
||||
|
||||
// Exercise the real Claude Code process against a local Messages API fixture.
|
||||
// The input counter's location is the only difference between the two backends.
|
||||
test.each([
|
||||
{ placement: "start", cacheTokens: 0 },
|
||||
{ placement: "delta", cacheTokens: 0 },
|
||||
{ placement: "start", cacheTokens: 1_000 },
|
||||
{ placement: "delta", cacheTokens: 1_000 },
|
||||
])(
|
||||
"counts $placement input and $cacheTokens cached tokens through Claude Code",
|
||||
async ({ placement, cacheTokens }) => {
|
||||
const cwd = await mkdtemp(path.resolve(".claude-context-usage-"));
|
||||
const configDir = path.join(cwd, "config");
|
||||
await mkdir(configDir);
|
||||
let requestCount = 0;
|
||||
const server = createServer(async (request, response) => {
|
||||
for await (const chunk of request) {
|
||||
void chunk;
|
||||
}
|
||||
if (!request.url?.startsWith("/v1/messages")) {
|
||||
response.writeHead(200, { "content-type": "application/json" });
|
||||
response.end("{}");
|
||||
return;
|
||||
}
|
||||
if (request.url.includes("count_tokens")) {
|
||||
response.writeHead(200, { "content-type": "application/json" });
|
||||
response.end(JSON.stringify({ input_tokens: 20_236 }));
|
||||
return;
|
||||
}
|
||||
const inputTokens = 20_236 + requestCount++;
|
||||
response.writeHead(200, { "content-type": "text/event-stream" });
|
||||
const emit = (data: { type: string; [key: string]: unknown }) => {
|
||||
response.write(`event: ${data.type}\ndata: ${JSON.stringify(data)}\n\n`);
|
||||
};
|
||||
emit({
|
||||
type: "message_start",
|
||||
message: {
|
||||
id: `msg_context_usage_${requestCount}`,
|
||||
type: "message",
|
||||
role: "assistant",
|
||||
model: "claude-sonnet-4-6",
|
||||
content: [],
|
||||
stop_reason: null,
|
||||
stop_sequence: null,
|
||||
usage: {
|
||||
input_tokens: placement === "start" ? inputTokens : 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: cacheTokens,
|
||||
},
|
||||
},
|
||||
});
|
||||
emit({ type: "content_block_start", index: 0, content_block: { type: "text", text: "" } });
|
||||
emit({ type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "hello" } });
|
||||
emit({ type: "content_block_stop", index: 0 });
|
||||
emit({
|
||||
type: "message_delta",
|
||||
delta: { stop_reason: "end_turn", stop_sequence: null },
|
||||
usage: {
|
||||
output_tokens: 64,
|
||||
...(placement === "delta" ? { input_tokens: inputTokens } : {}),
|
||||
},
|
||||
});
|
||||
emit({ type: "message_stop" });
|
||||
response.end();
|
||||
});
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
||||
const address = server.address();
|
||||
if (!address || typeof address === "string") throw new Error("Fixture has no TCP address");
|
||||
const client = new ClaudeAgentClient({
|
||||
logger: pino({ level: "silent" }),
|
||||
runtimeSettings: {
|
||||
env: {
|
||||
ANTHROPIC_BASE_URL: `http://127.0.0.1:${address.port}`,
|
||||
ANTHROPIC_AUTH_TOKEN: "fixture-token",
|
||||
ANTHROPIC_API_KEY: "",
|
||||
CLAUDE_CONFIG_DIR: configDir,
|
||||
CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC: "1",
|
||||
},
|
||||
},
|
||||
});
|
||||
const session = await client.createSession(
|
||||
{ provider: "claude", cwd, model: "claude-sonnet-4-6" },
|
||||
undefined,
|
||||
{ persistSession: false },
|
||||
);
|
||||
try {
|
||||
for (let turn = 0; turn < 2; turn++) {
|
||||
const events = await collectSessionTurnEvents(session, "say hello", {
|
||||
signal: AbortSignal.timeout(20_000),
|
||||
});
|
||||
const expectedTokens = 20_236 + turn + cacheTokens + 64;
|
||||
expect(events).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: "turn_completed",
|
||||
usage: expect.objectContaining({
|
||||
inputTokens: 20_236 + turn,
|
||||
outputTokens: 64,
|
||||
cachedInputTokens: cacheTokens,
|
||||
contextWindowUsedTokens: expectedTokens,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expect(events).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: "usage_updated",
|
||||
usage: expect.objectContaining({ contextWindowUsedTokens: expectedTokens }),
|
||||
}),
|
||||
);
|
||||
}
|
||||
expect(requestCount).toBe(2);
|
||||
} finally {
|
||||
try {
|
||||
await session.close();
|
||||
} finally {
|
||||
try {
|
||||
server.closeAllConnections();
|
||||
await new Promise<void>((resolve) => server.close(() => resolve()));
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
60_000,
|
||||
);
|
||||
@@ -2135,13 +2135,15 @@ describe("ClaudeAgentSession context window usage", () => {
|
||||
}
|
||||
|
||||
function createMessageDeltaEvent(outputTokens: number): Record<string, unknown> {
|
||||
return createMessageDeltaUsageEvent({ output_tokens: outputTokens });
|
||||
}
|
||||
|
||||
function createMessageDeltaUsageEvent(usage: Record<string, unknown>): Record<string, unknown> {
|
||||
return {
|
||||
type: "stream_event",
|
||||
event: {
|
||||
type: "message_delta",
|
||||
usage: {
|
||||
output_tokens: outputTokens,
|
||||
},
|
||||
usage,
|
||||
},
|
||||
session_id: "session-1",
|
||||
};
|
||||
@@ -2887,6 +2889,119 @@ describe("ClaudeAgentSession context window usage", () => {
|
||||
}
|
||||
});
|
||||
|
||||
test("message_delta input usage replaces an empty message_start count", async () => {
|
||||
// Translating gateways (LiteLLM in front of an OpenAI-compatible backend) do not know the
|
||||
// prompt size when message_start is emitted, so they report input_tokens: 0 there and send
|
||||
// the cumulative counts on message_delta, as the Anthropic Messages API allows.
|
||||
const session = await createSessionForTurns([
|
||||
[
|
||||
createInitMessage(),
|
||||
createMessageStartEvent({
|
||||
input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
}),
|
||||
createMessageDeltaUsageEvent({ input_tokens: 20_236, output_tokens: 64 }),
|
||||
createSuccessResult({
|
||||
usage: {
|
||||
input_tokens: 20_236,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
output_tokens: 64,
|
||||
iterations: [],
|
||||
},
|
||||
}),
|
||||
],
|
||||
]);
|
||||
|
||||
try {
|
||||
const events = await collectStreamEvents(session);
|
||||
|
||||
expect(events).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: "usage_updated",
|
||||
provider: "claude",
|
||||
usage: {
|
||||
contextWindowUsedTokens: 20_300,
|
||||
},
|
||||
}),
|
||||
);
|
||||
expect(events).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: "turn_completed",
|
||||
usage: expect.objectContaining({
|
||||
inputTokens: 20_236,
|
||||
outputTokens: 64,
|
||||
contextWindowUsedTokens: 20_300,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
} finally {
|
||||
await session.close();
|
||||
}
|
||||
});
|
||||
|
||||
test("message_delta without cache counters keeps the message_start cache usage", async () => {
|
||||
const session = await createSessionForTurns([
|
||||
[
|
||||
createInitMessage(),
|
||||
createMessageStartEvent({
|
||||
input_tokens: 5,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 1_000,
|
||||
}),
|
||||
createMessageDeltaUsageEvent({ input_tokens: 5, output_tokens: 64 }),
|
||||
createSuccessResult(),
|
||||
],
|
||||
]);
|
||||
|
||||
try {
|
||||
const events = await collectStreamEvents(session);
|
||||
|
||||
expect(events).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: "usage_updated",
|
||||
provider: "claude",
|
||||
usage: {
|
||||
contextWindowUsedTokens: 1_069,
|
||||
},
|
||||
}),
|
||||
);
|
||||
} finally {
|
||||
await session.close();
|
||||
}
|
||||
});
|
||||
|
||||
test("message_delta before this request's message_start does not seed usage", async () => {
|
||||
// A canceled request can leave its terminal message_delta queued behind the next turn.
|
||||
// Nothing is known about the new request until its message_start, so that straggler
|
||||
// must not produce a usage_updated event or leak into the next request's counters.
|
||||
const session = await createSessionForTurns([
|
||||
[
|
||||
createInitMessage(),
|
||||
createMessageDeltaUsageEvent({ input_tokens: 20_236, output_tokens: 64 }),
|
||||
createMessageStartEvent({
|
||||
input_tokens: 40,
|
||||
cache_creation_input_tokens: 5,
|
||||
cache_read_input_tokens: 10,
|
||||
}),
|
||||
createMessageDeltaEvent(7),
|
||||
createSuccessResult(),
|
||||
],
|
||||
]);
|
||||
|
||||
try {
|
||||
const events = await collectStreamEvents(session);
|
||||
|
||||
const usedTokens = events
|
||||
.filter((event) => event.type === "usage_updated")
|
||||
.map((event) => event.usage.contextWindowUsedTokens);
|
||||
expect(usedTokens).toEqual([55, 62]);
|
||||
} finally {
|
||||
await session.close();
|
||||
}
|
||||
});
|
||||
|
||||
test("per-request stream usage is not cumulative across API calls in a turn", async () => {
|
||||
const session = await createSessionForTurns([
|
||||
[
|
||||
|
||||
@@ -1855,30 +1855,59 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined {
|
||||
return maxContextWindow;
|
||||
}
|
||||
|
||||
function readStreamRequestInputTokens(event: Record<string, unknown>): number | undefined {
|
||||
const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage);
|
||||
if (!messageUsage) {
|
||||
interface StreamRequestInputUsage {
|
||||
inputTokens: number;
|
||||
cacheCreationInputTokens: number;
|
||||
cacheReadInputTokens: number;
|
||||
}
|
||||
|
||||
function readUsageCount(usage: Record<string, unknown>, key: string): number | undefined {
|
||||
const value = usage[key];
|
||||
if (typeof value !== "number" || !Number.isFinite(value) || value < 0) {
|
||||
return undefined;
|
||||
}
|
||||
const usage = messageUsage;
|
||||
const inputTokens =
|
||||
typeof usage.input_tokens === "number" && Number.isFinite(usage.input_tokens)
|
||||
? usage.input_tokens
|
||||
: undefined;
|
||||
const cacheCreationInputTokens =
|
||||
typeof usage.cache_creation_input_tokens === "number" &&
|
||||
Number.isFinite(usage.cache_creation_input_tokens)
|
||||
? usage.cache_creation_input_tokens
|
||||
: 0;
|
||||
const cacheReadInputTokens =
|
||||
typeof usage.cache_read_input_tokens === "number" &&
|
||||
Number.isFinite(usage.cache_read_input_tokens)
|
||||
? usage.cache_read_input_tokens
|
||||
: 0;
|
||||
if (typeof inputTokens !== "number" || inputTokens < 0) {
|
||||
return value;
|
||||
}
|
||||
|
||||
function readStreamRequestInputUsage(
|
||||
event: Record<string, unknown>,
|
||||
): StreamRequestInputUsage | undefined {
|
||||
const usage = toObjectRecord(toObjectRecord(event.message)?.usage);
|
||||
if (!usage) {
|
||||
return undefined;
|
||||
}
|
||||
return inputTokens + cacheCreationInputTokens + cacheReadInputTokens;
|
||||
const inputTokens = readUsageCount(usage, "input_tokens");
|
||||
if (inputTokens === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
return {
|
||||
inputTokens,
|
||||
cacheCreationInputTokens: readUsageCount(usage, "cache_creation_input_tokens") ?? 0,
|
||||
cacheReadInputTokens: readUsageCount(usage, "cache_read_input_tokens") ?? 0,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The input counters on `message_delta.usage` are cumulative whole-request totals and are
|
||||
* optional. Mirror the Anthropic SDK's MessageStream: a counter present on the delta replaces
|
||||
* the `message_start` value, a missing counter keeps it. Translating gateways report
|
||||
* `input_tokens: 0` on `message_start` and only know the real count by `message_delta`.
|
||||
*/
|
||||
function mergeStreamDeltaInputUsage(
|
||||
current: StreamRequestInputUsage,
|
||||
event: Record<string, unknown>,
|
||||
): StreamRequestInputUsage {
|
||||
const usage = toObjectRecord(event.usage);
|
||||
if (!usage) {
|
||||
return current;
|
||||
}
|
||||
return {
|
||||
inputTokens: readUsageCount(usage, "input_tokens") ?? current.inputTokens,
|
||||
cacheCreationInputTokens:
|
||||
readUsageCount(usage, "cache_creation_input_tokens") ?? current.cacheCreationInputTokens,
|
||||
cacheReadInputTokens:
|
||||
readUsageCount(usage, "cache_read_input_tokens") ?? current.cacheReadInputTokens,
|
||||
};
|
||||
}
|
||||
|
||||
function readStreamRequestOutputTokens(event: Record<string, unknown>): number | undefined {
|
||||
@@ -1958,7 +1987,7 @@ function readClaudeParentToolUseId(message: SDKMessage): string | null {
|
||||
|
||||
class ClaudeContextUsageState {
|
||||
private contextWindowMaxTokens: number | undefined;
|
||||
private streamRequestInputTokens: number | undefined;
|
||||
private streamRequestInputUsage: StreamRequestInputUsage | undefined;
|
||||
private streamRequestOutputTokens: number | undefined;
|
||||
private compactedContextWindowUsedTokens: number | undefined;
|
||||
private completedResultTurns = 0;
|
||||
@@ -1968,7 +1997,7 @@ class ClaudeContextUsageState {
|
||||
}
|
||||
|
||||
beginTurn(): void {
|
||||
this.streamRequestInputTokens = undefined;
|
||||
this.streamRequestInputUsage = undefined;
|
||||
this.streamRequestOutputTokens = undefined;
|
||||
this.compactedContextWindowUsedTokens = undefined;
|
||||
}
|
||||
@@ -1992,11 +2021,11 @@ class ClaudeContextUsageState {
|
||||
}
|
||||
const eventType = readTrimmedString(streamEvent.type);
|
||||
if (eventType === "message_start") {
|
||||
const inputTokens = readStreamRequestInputTokens(streamEvent);
|
||||
if (typeof inputTokens !== "number") {
|
||||
const inputUsage = readStreamRequestInputUsage(streamEvent);
|
||||
if (!inputUsage) {
|
||||
return null;
|
||||
}
|
||||
this.streamRequestInputTokens = inputTokens;
|
||||
this.streamRequestInputUsage = inputUsage;
|
||||
this.streamRequestOutputTokens = 0;
|
||||
} else if (eventType === "message_delta") {
|
||||
const outputTokens = readStreamRequestOutputTokens(streamEvent);
|
||||
@@ -2004,6 +2033,12 @@ class ClaudeContextUsageState {
|
||||
return null;
|
||||
}
|
||||
this.streamRequestOutputTokens = outputTokens;
|
||||
if (this.streamRequestInputUsage) {
|
||||
this.streamRequestInputUsage = mergeStreamDeltaInputUsage(
|
||||
this.streamRequestInputUsage,
|
||||
streamEvent,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
return null;
|
||||
}
|
||||
@@ -2050,13 +2085,15 @@ class ClaudeContextUsageState {
|
||||
}
|
||||
|
||||
private streamUsedTokens(): number | undefined {
|
||||
if (
|
||||
typeof this.streamRequestInputTokens !== "number" ||
|
||||
typeof this.streamRequestOutputTokens !== "number"
|
||||
) {
|
||||
const inputUsage = this.streamRequestInputUsage;
|
||||
if (!inputUsage || typeof this.streamRequestOutputTokens !== "number") {
|
||||
return undefined;
|
||||
}
|
||||
const usedTokens = this.streamRequestInputTokens + this.streamRequestOutputTokens;
|
||||
const usedTokens =
|
||||
inputUsage.inputTokens +
|
||||
inputUsage.cacheCreationInputTokens +
|
||||
inputUsage.cacheReadInputTokens +
|
||||
this.streamRequestOutputTokens;
|
||||
return usedTokens > 0 ? usedTokens : undefined;
|
||||
}
|
||||
|
||||
@@ -2075,7 +2112,7 @@ class ClaudeContextUsageState {
|
||||
}
|
||||
|
||||
buildCompactionUsageEvent(postTokens: number | undefined): AgentStreamEvent {
|
||||
this.streamRequestInputTokens = undefined;
|
||||
this.streamRequestInputUsage = undefined;
|
||||
this.streamRequestOutputTokens = undefined;
|
||||
this.compactedContextWindowUsedTokens = postTokens;
|
||||
const usage: AgentUsage = {};
|
||||
|
||||
Reference in New Issue
Block a user