Stop Claude context meter from doubling requests (#1701)

* fix(claude): stop probing context usage after turns

Use Claude stream/result usage for parent context tracking, and update the meter from compact boundary post tokens. Subagent usage reports stay out of the parent context meter.

* fix(claude): clear compact context fallback between turns
This commit is contained in:
Mohamed Boudra
2026-06-24 17:44:39 +08:00
committed by GitHub
parent 0748149ec9
commit c5442ef0a2
2 changed files with 338 additions and 112 deletions

View File

@@ -1046,7 +1046,7 @@ describe("ClaudeAgentSession context window usage", () => {
const logger = createTestLogger();
interface QueryFactoryForTurnsOptions {
currentContextUsageByTurn?: Array<Record<string, unknown> | undefined>;
getContextUsage?: ReturnType<typeof vi.fn>;
model?: string;
}
@@ -1091,8 +1091,8 @@ describe("ClaudeAgentSession context window usage", () => {
const queuedMessages: Array<Record<string, unknown>> = [];
const waiters: Array<() => void> = [];
let turnIndex = 0;
let contextUsageIndex = 0;
const closedRef = { value: false };
const getContextUsage = options?.getContextUsage ?? vi.fn(async () => undefined);
function wakeNextWaiter() {
const waiter = waiters.shift();
@@ -1140,11 +1140,7 @@ describe("ClaudeAgentSession context window usage", () => {
}),
setPermissionMode: vi.fn(async () => undefined),
setModel: vi.fn(async () => undefined),
getContextUsage: vi.fn(async () => {
const usage = options?.currentContextUsageByTurn?.[contextUsageIndex];
contextUsageIndex += 1;
return usage;
}),
getContextUsage,
supportedModels: vi.fn(async () => []),
supportedCommands: vi.fn(async () => []),
rewindFiles: vi.fn(async () => ({ canRewind: true })),
@@ -1165,26 +1161,6 @@ describe("ClaudeAgentSession context window usage", () => {
};
}
function createClaudeCurrentContextUsage(
totalTokens: number,
maxTokens: number,
): Record<string, unknown> {
return {
categories: [],
totalTokens,
maxTokens,
rawMaxTokens: maxTokens,
percentage: totalTokens / maxTokens,
gridRows: [],
model: "claude-sonnet-4-6",
memoryFiles: [],
mcpTools: [],
agents: [],
isAutoCompactEnabled: true,
apiUsage: null,
};
}
function createSuccessResult(overrides: Record<string, unknown> = {}): Record<string, unknown> {
return {
type: "result",
@@ -1282,6 +1258,21 @@ describe("ClaudeAgentSession context window usage", () => {
};
}
function createCompactBoundary(overrides: Record<string, unknown> = {}): Record<string, unknown> {
return {
type: "system",
subtype: "compact_boundary",
compact_metadata: {
trigger: "manual",
pre_tokens: 14_990,
post_tokens: 704,
},
uuid: "compact-boundary-1",
session_id: "session-1",
...overrides,
};
}
test("passes persistSession through to the Claude SDK query options", async () => {
const createResultTurn = (sessionId: string) => [
{
@@ -1556,7 +1547,10 @@ describe("ClaudeAgentSession context window usage", () => {
}
});
test("reports Claude's current context usage after an Agent subagent runs", async () => {
test("does not probe current context usage after an Agent subagent runs", async () => {
const getContextUsage = vi.fn(async () => {
throw new Error("getContextUsage should not be called during result handling");
});
const session = await createSessionForTurns(
[
[
@@ -1575,21 +1569,20 @@ describe("ClaudeAgentSession context window usage", () => {
}),
],
],
{
currentContextUsageByTurn: [createClaudeCurrentContextUsage(12_345, 200_000)],
},
{ getContextUsage },
);
try {
const result = await session.run("turn");
expect(getContextUsage).not.toHaveBeenCalled();
expect(result.usage).toEqual({
inputTokens: 9_000,
cachedInputTokens: 700,
outputTokens: 400,
totalCostUsd: 0.25,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 12_345,
contextWindowUsedTokens: 175,
});
} finally {
await session.close();
@@ -1637,6 +1630,121 @@ describe("ClaudeAgentSession context window usage", () => {
}
});
test("uses parent request usage after a real subagent tool result", async () => {
const getContextUsage = vi.fn(async () => {
throw new Error("getContextUsage should not be called during result handling");
});
const session = await createSessionForTurns(
[
[
createInitMessage(),
createMessageStartEvent({
input_tokens: 3,
cache_creation_input_tokens: 16_999,
cache_read_input_tokens: 0,
}),
createAgentToolStartEvent(),
createMessageDeltaEvent(163),
{
type: "assistant",
parent_tool_use_id: "toolu-agent-1",
message: {
role: "assistant",
content: [{ type: "text", text: "SUBAGENT_OK" }],
usage: {
input_tokens: 3,
cache_creation_input_tokens: 1_182,
cache_read_input_tokens: 0,
output_tokens: 8,
},
},
uuid: "subagent-assistant-1",
session_id: "session-1",
},
{
...createSubagentTaskNotification(),
status: "completed",
summary: "Probe subagent test",
usage: {
total_tokens: 1_193,
tool_uses: 0,
},
},
{
type: "user",
parent_tool_use_id: null,
message: {
role: "user",
content: [
{
type: "tool_result",
tool_use_id: "toolu-agent-1",
content: [
{ type: "text", text: "SUBAGENT_OK" },
{
type: "text",
text: "agentId: subagent-1\n<usage>subagent_tokens: 1194\ntool_uses: 0</usage>",
},
],
},
],
},
uuid: "subagent-tool-result-1",
session_id: "session-1",
},
createMessageStartEvent({
input_tokens: 1,
cache_creation_input_tokens: 253,
cache_read_input_tokens: 16_999,
}),
createMessageDeltaEvent(8),
createSuccessResult({
usage: {
input_tokens: 4,
cache_creation_input_tokens: 17_252,
cache_read_input_tokens: 16_999,
output_tokens: 171,
iterations: [
{
input_tokens: 1,
cache_creation_input_tokens: 253,
cache_read_input_tokens: 16_999,
output_tokens: 8,
},
],
},
modelUsage: {
"claude-sonnet-4-6": {
inputTokens: 7,
outputTokens: 180,
cacheReadInputTokens: 16_999,
cacheCreationInputTokens: 18_434,
contextWindow: 200_000,
},
},
}),
],
],
{ getContextUsage },
);
try {
const result = await session.run("turn");
expect(getContextUsage).not.toHaveBeenCalled();
expect(result.usage).toEqual({
inputTokens: 4,
cachedInputTokens: 16_999,
outputTokens: 171,
totalCostUsd: 0.25,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 17_261,
});
} finally {
await session.close();
}
});
test("falls back to the active result iteration when current and stream usage are unavailable", async () => {
const session = await createSessionForTurns([
[
@@ -1842,6 +1950,165 @@ describe("ClaudeAgentSession context window usage", () => {
}
});
test("manual compact boundary updates context usage from post tokens", async () => {
const session = await createSessionForTurns([
[
createInitMessage(),
createMessageStartEvent(),
createMessageDeltaEvent(25),
createCompactBoundary(),
createSuccessResult({
total_cost_usd: 0.04,
usage: {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
output_tokens: 0,
iterations: [],
},
}),
],
]);
try {
const events = await collectStreamEvents(session, "/compact");
expect(events).toContainEqual(
expect.objectContaining({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 704,
},
}),
);
expect(events).toContainEqual(
expect.objectContaining({
type: "turn_completed",
provider: "claude",
usage: {
inputTokens: 0,
cachedInputTokens: 0,
outputTokens: 0,
totalCostUsd: 0.04,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 704,
},
}),
);
} finally {
await session.close();
}
});
test("zero-token stream events after compact keep post-token usage", async () => {
const session = await createSessionForTurns([
[
createInitMessage(),
createMessageStartEvent(),
createMessageDeltaEvent(25),
createCompactBoundary(),
createMessageStartEvent({
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
}),
createMessageDeltaEvent(0),
createSuccessResult({
total_cost_usd: 0.04,
usage: {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
output_tokens: 0,
iterations: [],
},
}),
],
]);
try {
const events = await collectStreamEvents(session, "/compact");
expect(
events.filter(
(event) => event.type === "usage_updated" && event.usage.contextWindowUsedTokens === 0,
),
).toEqual([]);
expect(events).toContainEqual(
expect.objectContaining({
type: "turn_completed",
provider: "claude",
usage: {
inputTokens: 0,
cachedInputTokens: 0,
outputTokens: 0,
totalCostUsd: 0.04,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 704,
},
}),
);
} finally {
await session.close();
}
});
test("starting a new turn clears interrupted compact usage", async () => {
const session = await createSessionForTurns([
[
createSuccessResult({
total_cost_usd: 0.04,
usage: {
input_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
output_tokens: 0,
iterations: [],
},
}),
],
]);
try {
const compactEvents = (session as unknown as TestClaudeSession).translateMessageToEvents(
createCompactBoundary(),
);
expect(compactEvents).toContainEqual(
expect.objectContaining({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 704,
},
}),
);
const events = await collectStreamEvents(session, "next turn");
expect(events).toContainEqual(
expect.objectContaining({
type: "turn_completed",
provider: "claude",
usage: expect.objectContaining({
inputTokens: 0,
cachedInputTokens: 0,
outputTokens: 0,
totalCostUsd: 0.04,
}),
}),
);
expect(
events.some(
(event) =>
event.type === "turn_completed" && event.usage.contextWindowUsedTokens !== undefined,
),
).toBe(false);
} finally {
await session.close();
}
});
test("result.result is surfaced as an assistant message when no model output was produced", async () => {
const session = await createSessionForTest();

View File

@@ -11,7 +11,6 @@ import {
type PermissionResult,
type PermissionUpdate,
type Query,
type SDKControlGetContextUsageResponse,
type SDKMessage,
type SDKPartialAssistantMessage,
type SDKResultMessage,
@@ -1682,31 +1681,6 @@ function readLegacyResultUsageTokens(usage: unknown): number | undefined {
return usageRecord ? readUsageTokenTotal(usageRecord) : undefined;
}
interface ClaudeCurrentContextUsage {
totalTokens: number;
maxTokens?: number;
}
function readCurrentContextUsage(
value: SDKControlGetContextUsageResponse | unknown,
): ClaudeCurrentContextUsage | undefined {
const record = toObjectRecord(value);
if (!record) {
return undefined;
}
const totalTokens = record.totalTokens;
if (typeof totalTokens !== "number" || !Number.isFinite(totalTokens) || totalTokens < 0) {
return undefined;
}
const maxTokens = record.maxTokens;
return {
totalTokens,
...(typeof maxTokens === "number" && Number.isFinite(maxTokens) && maxTokens > 0
? { maxTokens }
: {}),
};
}
function isClaudeSubagentToolName(name: string | undefined): boolean {
return name === "Task" || name === "Agent";
}
@@ -1715,6 +1689,7 @@ class ClaudeContextUsageState {
private contextWindowMaxTokens: number | undefined;
private streamRequestInputTokens: number | undefined;
private streamRequestOutputTokens: number | undefined;
private compactedContextWindowUsedTokens: number | undefined;
private completedResultTurns = 0;
constructor(initialContextWindowMaxTokens?: number) {
@@ -1724,6 +1699,7 @@ class ClaudeContextUsageState {
beginTurn(): void {
this.streamRequestInputTokens = undefined;
this.streamRequestOutputTokens = undefined;
this.compactedContextWindowUsedTokens = undefined;
}
setInitialContextWindowMaxTokens(contextWindowMaxTokens: number | undefined): void {
@@ -1738,12 +1714,6 @@ class ClaudeContextUsageState {
return this.contextWindowMaxTokens;
}
recordCurrentContextUsage(usage: ClaudeCurrentContextUsage | undefined): void {
if (usage?.maxTokens !== undefined) {
this.contextWindowMaxTokens = usage.maxTokens;
}
}
buildStreamUsageEvent(event: unknown): AgentStreamEvent | null {
const streamEvent = toObjectRecord(event);
if (!streamEvent) {
@@ -1774,11 +1744,7 @@ class ClaudeContextUsageState {
return this.createUsageUpdatedEvent(usedTokens);
}
buildResultUsage(
message: SDKResultMessage,
modelUsage: unknown,
currentContextUsage: ClaudeCurrentContextUsage | undefined,
): AgentUsage | undefined {
buildResultUsage(message: SDKResultMessage, modelUsage: unknown): AgentUsage | undefined {
try {
if (!message.usage) {
return undefined;
@@ -1791,7 +1757,6 @@ class ClaudeContextUsageState {
};
const modelContextWindowMaxTokens = this.recordModelUsage(modelUsage ?? message.modelUsage);
this.recordCurrentContextUsage(currentContextUsage);
if (this.contextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.contextWindowMaxTokens;
} else if (modelContextWindowMaxTokens !== undefined) {
@@ -1802,12 +1767,13 @@ class ClaudeContextUsageState {
readActiveUsageTokens(message.usage) ??
(this.completedResultTurns === 0 ? readLegacyResultUsageTokens(message.usage) : undefined);
const usedTokens =
currentContextUsage?.totalTokens ?? this.streamUsedTokens() ?? activeResultUsageTokens;
this.streamUsedTokens() ?? activeResultUsageTokens ?? this.compactedContextWindowUsedTokens;
if (usedTokens !== undefined) {
usage.contextWindowUsedTokens = usedTokens;
}
return usage;
} finally {
this.compactedContextWindowUsedTokens = undefined;
this.completedResultTurns += 1;
}
}
@@ -1819,7 +1785,8 @@ class ClaudeContextUsageState {
) {
return undefined;
}
return this.streamRequestInputTokens + this.streamRequestOutputTokens;
const usedTokens = this.streamRequestInputTokens + this.streamRequestOutputTokens;
return usedTokens > 0 ? usedTokens : undefined;
}
private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent {
@@ -1835,6 +1802,24 @@ class ClaudeContextUsageState {
usage,
};
}
buildCompactionUsageEvent(postTokens: number | undefined): AgentStreamEvent {
this.streamRequestInputTokens = undefined;
this.streamRequestOutputTokens = undefined;
this.compactedContextWindowUsedTokens = postTokens;
const usage: AgentUsage = {};
if (this.contextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.contextWindowMaxTokens;
}
if (postTokens !== undefined) {
usage.contextWindowUsedTokens = postTokens;
}
return {
type: "usage_updated",
provider: "claude",
usage,
};
}
}
class ClaudeAgentSession implements AgentSession {
@@ -3282,7 +3267,7 @@ class ClaudeAgentSession implements AgentSession {
if (await this.handleMissingResumedConversation(message, activeQuery)) {
return true;
}
await this.routeSdkMessageFromPump(message, activeQuery);
await this.routeSdkMessageFromPump(message);
return false;
};
const drainActiveQuery = async (): Promise<boolean> => {
@@ -3358,7 +3343,7 @@ class ClaudeAgentSession implements AgentSession {
);
}
private async routeSdkMessageFromPump(message: SDKMessage, activeQuery: Query): Promise<void> {
private async routeSdkMessageFromPump(message: SDKMessage): Promise<void> {
if (this.shouldSuppressStaleResult(message)) {
return;
}
@@ -3388,12 +3373,7 @@ class ClaudeAgentSession implements AgentSession {
"provider.claude.parsed_event",
);
const events = await this.buildPumpedMessageEvents(
message,
activeQuery,
identifiers.messageId,
turnId,
);
const events = await this.buildPumpedMessageEvents(message, identifiers.messageId, turnId);
if (events.length === 0) {
return;
@@ -3430,18 +3410,12 @@ class ClaudeAgentSession implements AgentSession {
private async buildPumpedMessageEvents(
message: SDKMessage,
activeQuery: Query,
messageIdHint: string | null,
turnId: string | null,
): Promise<AgentStreamEvent[]> {
const currentContextUsage =
message.type === "result" && message.subtype === "success"
? await this.queryCurrentContextUsage(activeQuery)
: undefined;
const messageEvents = this.translateMessageToEvents(message, {
suppressAssistantText: true,
suppressReasoning: true,
currentContextUsage,
});
const assistantTimelineEvents = this.timelineAssembler
.consume({
@@ -3461,18 +3435,6 @@ class ClaudeAgentSession implements AgentSession {
return [...messageEvents, ...assistantTimelineEvents];
}
private async queryCurrentContextUsage(
activeQuery: Query,
): Promise<ClaudeCurrentContextUsage | undefined> {
try {
const usage = await withTimeout(activeQuery.getContextUsage(), 3_000, "timeout");
return readCurrentContextUsage(usage);
} catch (error) {
this.logger.debug({ err: error }, "Claude context usage query failed");
return undefined;
}
}
private async handleMissingResumedConversation(
message: SDKMessage,
activeQuery: Query,
@@ -3540,7 +3502,6 @@ class ClaudeAgentSession implements AgentSession {
options?: {
suppressAssistantText?: boolean;
suppressReasoning?: boolean;
currentContextUsage?: ClaudeCurrentContextUsage;
},
): AgentStreamEvent[] {
const parentToolUseId =
@@ -3591,9 +3552,7 @@ class ClaudeAgentSession implements AgentSession {
this.appendStreamEventEvents(message, events, options);
break;
case "result":
this.appendResultEvents(message, events, {
currentContextUsage: options?.currentContextUsage,
});
this.appendResultEvents(message, events);
break;
default:
break;
@@ -3667,6 +3626,7 @@ class ClaudeAgentSession implements AgentSession {
},
provider: "claude",
});
events.push(this.contextUsage.buildCompactionUsageEvent(compactMetadata?.postTokens));
return;
}
if (message.subtype === "task_notification") {
@@ -3794,9 +3754,8 @@ class ClaudeAgentSession implements AgentSession {
private appendResultEvents(
message: Extract<SDKMessage, { type: "result" }>,
events: AgentStreamEvent[],
options?: { currentContextUsage?: ClaudeCurrentContextUsage },
): void {
const usage = this.convertUsage(message, message.modelUsage, options?.currentContextUsage);
const usage = this.convertUsage(message, message.modelUsage);
if (message.subtype === "success") {
// Built-in slash commands (e.g. /voice, /usage, "Unknown command: …")
// run client-side in the Claude CLI with no model turn — output_tokens
@@ -3963,12 +3922,8 @@ class ClaudeAgentSession implements AgentSession {
return null;
}
private convertUsage(
message: SDKResultMessage,
modelUsage?: unknown,
currentContextUsage?: ClaudeCurrentContextUsage,
): AgentUsage | undefined {
return this.contextUsage.buildResultUsage(message, modelUsage, currentContextUsage);
private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined {
return this.contextUsage.buildResultUsage(message, modelUsage);
}
private handlePermissionRequest: CanUseTool = async (
@@ -4869,7 +4824,9 @@ function hasToolLikeBlock(block?: ClaudeContentChunk | null): boolean {
return type.includes("tool");
}
function readCompactionMetadata(source: unknown): { trigger?: string; preTokens?: number } | null {
function readCompactionMetadata(
source: unknown,
): { trigger?: string; preTokens?: number; postTokens?: number } | null {
const sourceRecord = toObjectRecord(source);
if (!sourceRecord) {
return null;
@@ -4887,7 +4844,9 @@ function readCompactionMetadata(source: unknown): { trigger?: string; preTokens?
const trigger = typeof metadata.trigger === "string" ? metadata.trigger : undefined;
const preTokensRaw = metadata.preTokens ?? metadata.pre_tokens;
const preTokens = typeof preTokensRaw === "number" ? preTokensRaw : undefined;
return { trigger, preTokens };
const postTokensRaw = metadata.postTokens ?? metadata.post_tokens;
const postTokens = typeof postTokensRaw === "number" ? postTokensRaw : undefined;
return { trigger, preTokens, postTokens };
}
return null;
}