jb/paseo
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Fix subagent context usage reporting

Browse Source 
Task-level usage reports processed work, not the parent session's active context. Prefer current context data and parent stream usage before falling back to result totals.
This commit is contained in:
Mohamed Boudra
2026-06-20 21:08:50 +07:00
parent 4b45bab7e3
commit fd2fed03a6
2 changed files with 640 additions and 628 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/server/src/server/agent/providers/claude/agent.test.ts
+363 -471
View File
@@ -13,11 +13,11 @@ import {
normalizeClaudeAskUserQuestionUpdatedInput,
toClaudeSdkMcpConfig,
} from "./agent.js";
import type { AgentTimelineItem, AgentUsage, AgentStreamEvent } from "../../agent-sdk-types.js";
import { streamSession } from "../test-utils/session-stream-adapter.js";
import type { AgentSession, AgentTimelineItem, AgentStreamEvent } from "../../agent-sdk-types.js";
interface TestClaudeSession {
translateMessageToEvents(message: SDKMessage): AgentStreamEvent[];
convertUsage(message: SDKMessage): AgentUsage | undefined;
}
afterEach(() => {
@@ -1044,6 +1044,10 @@ describe("ClaudeAgentClient.listImportableSessions", () => {
describe("ClaudeAgentSession context window usage", () => {
const logger = createTestLogger();
interface QueryFactoryForTurnsOptions {
currentContextUsageByTurn?: Array<Record<string, unknown> | undefined>;
}
async function createSessionForTest(): Promise<TestClaudeSession> {
const client = new ClaudeAgentClient({ logger, resolveBinary: async () => "/test/claude/bin" });
const session = await client.createSession({
@@ -1053,11 +1057,38 @@ describe("ClaudeAgentSession context window usage", () => {
return session as unknown as TestClaudeSession;
}
function createQueryFactoryForTurns(turns: Array<Array<Record<string, unknown>>>) {
async function createSessionForTurns(
turns: Array<Array<Record<string, unknown>>>,
options?: QueryFactoryForTurnsOptions,
): Promise<AgentSession> {
const client = new ClaudeAgentClient({
logger,
queryFactory: createQueryFactoryForTurns(turns, options),
resolveBinary: async () => "/test/claude/bin",
});
return await client.createSession({
provider: "claude",
cwd: process.cwd(),
});
}
async function collectStreamEvents(session: AgentSession, prompt = "turn") {
const events: AgentStreamEvent[] = [];
for await (const event of streamSession(session, prompt)) {
events.push(event);
}
return events;
}
function createQueryFactoryForTurns(
turns: Array<Array<Record<string, unknown>>>,
options?: QueryFactoryForTurnsOptions,
) {
return vi.fn(({ prompt }: { prompt: AsyncIterable<unknown> }) => {
const queuedMessages: Array<Record<string, unknown>> = [];
const waiters: Array<() => void> = [];
let turnIndex = 0;
let contextUsageIndex = 0;
const closedRef = { value: false };
function wakeNextWaiter() {
@@ -1106,6 +1137,11 @@ 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;
}),
supportedModels: vi.fn(async () => []),
supportedCommands: vi.fn(async () => []),
rewindFiles: vi.fn(async () => ({ canRewind: true })),
@@ -1116,6 +1152,133 @@ describe("ClaudeAgentSession context window usage", () => {
});
}
function createInitMessage(sessionId = "session-1"): Record<string, unknown> {
return {
type: "system",
subtype: "init",
session_id: sessionId,
permissionMode: "default",
model: "claude-sonnet-4-6",
};
}
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",
subtype: "success",
duration_ms: 100,
duration_api_ms: 75,
is_error: false,
num_turns: 1,
result: "done",
stop_reason: null,
total_cost_usd: 0.25,
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
modelUsage: {
"claude-sonnet-4-6": { contextWindow: 200_000 },
},
permission_denials: [],
uuid: "result-1",
session_id: "session-1",
...overrides,
};
}
function createMessageStartEvent(
usage: Record<string, unknown> = {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
},
): Record<string, unknown> {
return {
type: "stream_event",
event: {
type: "message_start",
message: { usage },
},
session_id: "session-1",
};
}
function createMessageDeltaEvent(outputTokens: number): Record<string, unknown> {
return {
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: outputTokens,
},
},
session_id: "session-1",
};
}
function createAgentToolStartEvent(): Record<string, unknown> {
return {
type: "stream_event",
event: {
type: "content_block_start",
index: 0,
content_block: {
type: "tool_use",
id: "toolu-agent-1",
name: "Agent",
input: {
description: "Check something in a subagent",
prompt: "Return a short answer",
},
},
},
session_id: "session-1",
};
}
function createSubagentTaskNotification(): Record<string, unknown> {
return {
type: "system",
subtype: "task_notification",
uuid: "task-note-1",
task_id: "task-1",
tool_use_id: "toolu-agent-1",
status: "running",
summary: "Subagent is working",
usage: {
total_tokens: 18_876,
tool_uses: 1,
duration_ms: 50,
input_tokens: 12_000,
cache_read_input_tokens: 6_000,
output_tokens: 876,
},
session_id: "session-1",
};
}
test("passes persistSession through to the Claude SDK query options", async () => {
const createResultTurn = (sessionId: string) => [
{
@@ -1389,357 +1552,152 @@ describe("ClaudeAgentSession context window usage", () => {
}
});
test("convertUsage includes contextWindowMaxTokens and derives used tokens from result usage as initial fallback", async () => {
const session = await createSessionForTest();
const usage = session.convertUsage(
test("reports Claude's current context usage after an Agent subagent runs", async () => {
const session = await createSessionForTurns(
[
[
createInitMessage(),
createMessageStartEvent(),
createAgentToolStartEvent(),
createSubagentTaskNotification(),
createMessageDeltaEvent(25),
createSuccessResult({
usage: {
input_tokens: 9_000,
cache_creation_input_tokens: 300,
cache_read_input_tokens: 700,
output_tokens: 400,
},
}),
],
],
{
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
},
{
"claude-sonnet-4-6": { contextWindow: 200_000 },
"claude-opus-4-6": { contextWindow: 1_000_000 },
currentContextUsageByTurn: [createClaudeCurrentContextUsage(12_345, 200_000)],
},
);
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowMaxTokens: 1_000_000,
contextWindowUsedTokens: 22,
});
try {
const result = await session.run("turn");
expect(result.usage).toEqual({
inputTokens: 9_000,
cachedInputTokens: 700,
outputTokens: 400,
totalCostUsd: 0.25,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 12_345,
});
} finally {
await session.close();
}
});
test("contextWindowUsedTokens falls back to result usage when no task_progress was received", async () => {
const session = await createSessionForTest();
test("does not report task notification tokens as parent context usage", async () => {
const session = await createSessionForTurns([
[
createInitMessage(),
createMessageStartEvent(),
createAgentToolStartEvent(),
createSubagentTaskNotification(),
{
type: "system",
subtype: "task_progress",
task_id: "task-1",
description: "Subagent progress",
usage: {
total_tokens: 9_999,
tool_uses: 1,
duration_ms: 50,
},
uuid: "task-progress-1",
session_id: "session-1",
},
createMessageDeltaEvent(25),
createSuccessResult(),
],
]);
const usage = session.convertUsage({
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_creation_input_tokens: 3,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
});
try {
const result = await session.run("turn");
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowUsedTokens: 25,
});
});
test("contextWindowUsedTokens is populated from task_progress usage data", async () => {
const session = await createSessionForTest();
session.translateMessageToEvents({
type: "system",
subtype: "task_progress",
task_id: "task-1",
description: "Processing",
usage: {
total_tokens: 999,
tool_uses: 1,
duration_ms: 50,
input_tokens: 345,
cache_read_input_tokens: 55,
},
uuid: "task-progress-1",
session_id: "session-1",
});
const events = session.translateMessageToEvents({
type: "result",
subtype: "success",
duration_ms: 100,
duration_api_ms: 75,
is_error: false,
num_turns: 1,
result: "done",
stop_reason: null,
total_cost_usd: 0.25,
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
modelUsage: {
"claude-sonnet-4-6": { contextWindow: 200_000 },
},
permission_denials: [],
uuid: "result-1",
session_id: "session-1",
});
expect(events).toContainEqual({
type: "turn_completed",
provider: "claude",
usage: {
expect(result.usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.25,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 999,
},
});
});
test("task_progress emits a usage_updated event", async () => {
const session = await createSessionForTest();
const events = session.translateMessageToEvents({
type: "system",
subtype: "task_progress",
task_id: "task-1",
description: "Processing",
usage: {
total_tokens: 999,
tool_uses: 1,
duration_ms: 50,
},
uuid: "task-progress-1",
session_id: "session-1",
});
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 999,
},
});
});
test("task_notification emits a usage_updated event", async () => {
const session = await createSessionForTest();
const events = session.translateMessageToEvents({
type: "system",
subtype: "task_notification",
uuid: "task-note-1",
task_id: "task-1",
status: "running",
summary: "Background task still running",
usage: {
total_tokens: 777,
tool_uses: 1,
duration_ms: 50,
},
session_id: "session-1",
} as unknown as SDKMessage);
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 777,
},
});
});
test("message_start stream events emit usage_updated with per-request usage", async () => {
const session = await createSessionForTest();
const events = session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
},
},
},
session_id: "session-1",
} as unknown as SDKMessage);
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 150,
},
});
});
test("message_delta stream events update per-request usage", async () => {
const session = await createSessionForTest();
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
},
},
},
session_id: "session-1",
} as unknown as SDKMessage);
const events = session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 25,
},
},
session_id: "session-1",
} as unknown as SDKMessage);
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 175,
},
});
});
} finally {
await session.close();
}
});
test("task_progress usage takes priority over derived result usage", async () => {
const session = await createSessionForTest();
session.translateMessageToEvents({
type: "system",
subtype: "task_progress",
task_id: "task-1",
description: "Processing",
usage: {
total_tokens: 999,
tool_uses: 1,
duration_ms: 50,
input_tokens: 345,
cache_read_input_tokens: 55,
},
uuid: "task-progress-1",
session_id: "session-1",
});
const usage = session.convertUsage({
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_creation_input_tokens: 3,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
});
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowUsedTokens: 999,
});
});
test("contextWindowUsedTokens persists across turns from last task_progress", async () => {
const queryFactory = createQueryFactoryForTurns([
test("falls back to the active result iteration when current and stream usage are unavailable", async () => {
const session = await createSessionForTurns([
[
{
type: "system",
subtype: "init",
session_id: "session-1",
permissionMode: "default",
model: "claude-sonnet-4-6",
},
{
type: "system",
subtype: "task_progress",
task_id: "task-1",
description: "Processing",
createInitMessage(),
createSuccessResult({
usage: {
total_tokens: 999,
tool_uses: 1,
duration_ms: 50,
input_tokens: 345,
cache_read_input_tokens: 55,
input_tokens: 5_000,
cache_creation_input_tokens: 500,
cache_read_input_tokens: 600,
output_tokens: 700,
iterations: [
{
input_tokens: 100,
cache_creation_input_tokens: 10,
cache_read_input_tokens: 20,
output_tokens: 30,
},
{
input_tokens: 2,
cache_creation_input_tokens: 10,
cache_read_input_tokens: 100,
output_tokens: 5,
},
],
},
uuid: "task-progress-1",
session_id: "session-1",
},
{
type: "result",
subtype: "success",
duration_ms: 100,
duration_api_ms: 75,
is_error: false,
num_turns: 1,
result: "done",
stop_reason: null,
total_cost_usd: 0.25,
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
modelUsage: {
"claude-sonnet-4-6": { contextWindow: 200_000 },
},
permission_denials: [],
uuid: "result-1",
session_id: "session-1",
},
}),
],
]);
try {
const result = await session.run("turn");
expect(result.usage).toEqual({
inputTokens: 5_000,
cachedInputTokens: 600,
outputTokens: 700,
totalCostUsd: 0.25,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 117,
});
} finally {
await session.close();
}
});
test("does not use aggregate result totals after the first result turn", async () => {
const session = await createSessionForTurns([
[
createInitMessage(),
createMessageStartEvent(),
createMessageDeltaEvent(25),
createSuccessResult(),
],
[
{
type: "result",
subtype: "success",
duration_ms: 110,
duration_api_ms: 80,
is_error: false,
num_turns: 1,
result: "still done",
stop_reason: null,
createSuccessResult({
total_cost_usd: 0.1,
usage: {
input_tokens: 11,
cache_creation_input_tokens: 3,
cache_read_input_tokens: 6,
output_tokens: 8,
},
modelUsage: {
"claude-sonnet-4-6": { contextWindow: 200_000 },
input_tokens: 1_000,
cache_read_input_tokens: 200,
output_tokens: 300,
},
permission_denials: [],
uuid: "result-2",
session_id: "session-1",
},
}),
],
]);
const client = new ClaudeAgentClient({
logger,
queryFactory,
resolveBinary: async () => "/test/claude/bin",
});
const session = await client.createSession({
provider: "claude",
cwd: process.cwd(),
});
try {
const firstTurn = await session.run("turn 1");
@@ -1751,175 +1709,109 @@ describe("ClaudeAgentSession context window usage", () => {
outputTokens: 7,
totalCostUsd: 0.25,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 999,
contextWindowUsedTokens: 175,
});
// Turn 2 has no task_progress, so contextWindowUsedTokens retains the
// last known value from turn 1 rather than deriving from accumulated
// result.usage (which would be incorrect — those are session-level totals).
expect(secondTurn.usage).toEqual({
inputTokens: 11,
cachedInputTokens: 6,
outputTokens: 8,
inputTokens: 1_000,
cachedInputTokens: 200,
outputTokens: 300,
totalCostUsd: 0.1,
contextWindowMaxTokens: 200_000,
contextWindowUsedTokens: 999,
});
} finally {
await session.close();
}
});
test("convertUsage derives used tokens from result usage as fallback when task_progress is missing", async () => {
const session = await createSessionForTest();
test("message_start stream events emit usage_updated with per-request usage", async () => {
const session = await createSessionForTurns([
[createInitMessage(), createMessageStartEvent(), createSuccessResult()],
]);
const usage = session.convertUsage({
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
});
try {
const events = await collectStreamEvents(session);
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowUsedTokens: 22,
});
expect(events).toContainEqual(
expect.objectContaining({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 150,
},
}),
);
} finally {
await session.close();
}
});
test("convertUsage uses per-request stream usage when no task_progress is available", async () => {
const session = await createSessionForTest();
test("message_delta stream events update per-request usage", async () => {
const session = await createSessionForTurns([
[
createInitMessage(),
createMessageStartEvent(),
createMessageDeltaEvent(25),
createSuccessResult(),
],
]);
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
try {
const events = await collectStreamEvents(session);
expect(events).toContainEqual(
expect.objectContaining({
type: "usage_updated",
provider: "claude",
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
contextWindowUsedTokens: 175,
},
},
},
session_id: "session-1",
} as unknown as SDKMessage);
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 25,
},
},
session_id: "session-1",
} as unknown as SDKMessage);
const usage = session.convertUsage({
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
});
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowUsedTokens: 175,
});
}),
);
} finally {
await session.close();
}
});
test("per-request stream usage is not cumulative across API calls in a turn", async () => {
const session = await createSessionForTest();
const session = await createSessionForTurns([
[
createInitMessage(),
createMessageStartEvent(),
createMessageDeltaEvent(25),
createMessageStartEvent({
input_tokens: 40,
cache_creation_input_tokens: 5,
cache_read_input_tokens: 10,
}),
createMessageDeltaEvent(7),
createSuccessResult(),
],
]);
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
try {
const events = await collectStreamEvents(session);
expect(events).toContainEqual(
expect.objectContaining({
type: "usage_updated",
provider: "claude",
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
contextWindowUsedTokens: 55,
},
},
},
session_id: "session-1",
} as unknown as SDKMessage);
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 25,
},
},
session_id: "session-1",
} as unknown as SDKMessage);
const secondStartEvents = session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
}),
);
expect(events).toContainEqual(
expect.objectContaining({
type: "usage_updated",
provider: "claude",
usage: {
input_tokens: 40,
cache_creation_input_tokens: 5,
cache_read_input_tokens: 10,
contextWindowUsedTokens: 62,
},
},
},
session_id: "session-1",
} as unknown as SDKMessage);
expect(secondStartEvents).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 55,
},
});
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 7,
},
},
session_id: "session-1",
} as unknown as SDKMessage);
const usage = session.convertUsage({
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
});
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowUsedTokens: 62,
});
}),
);
} finally {
await session.close();
}
});
test("result.result is surfaced as an assistant message when no model output was produced", async () => {
packages/server/src/server/agent/providers/claude/agent.ts
+277 -157
View File
@@ -11,9 +11,9 @@ import {
type PermissionResult,
type PermissionUpdate,
type Query,
type SDKControlGetContextUsageResponse,
type SDKMessage,
type SDKPartialAssistantMessage,
type SDKTaskProgressMessage,
type SDKResultMessage,
type SDKSystemMessage,
type SDKUserMessage,
@@ -1603,27 +1603,6 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined {
return maxContextWindow;
}
function readUsageTotalTokens(usage: unknown): number | undefined {
if (!usage || typeof usage !== "object") {
return undefined;
}
const totalTokens = (usage as { total_tokens?: unknown }).total_tokens;
if (typeof totalTokens !== "number" || !Number.isFinite(totalTokens) || totalTokens < 0) {
return undefined;
}
return totalTokens;
}
function readContextWindowUsedTokensFromTaskProgress(
message: SDKTaskProgressMessage,
): number | undefined {
return readUsageTotalTokens(message.usage);
}
function readUsageFromTaskNotification(message: { usage?: unknown }): number | undefined {
return readUsageTotalTokens(message.usage);
}
function readStreamRequestInputTokens(event: Record<string, unknown>): number | undefined {
const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage);
if (!messageUsage) {
@@ -1658,6 +1637,203 @@ function readStreamRequestOutputTokens(event: Record<string, unknown>): number |
return outputTokens;
}
function readLastUsageIteration(usage: unknown): Record<string, unknown> | undefined {
const iterations = toObjectRecord(usage)?.iterations;
if (!Array.isArray(iterations)) {
return undefined;
}
for (let index = iterations.length - 1; index >= 0; index -= 1) {
const candidate = toObjectRecord(iterations[index]);
if (candidate) {
return candidate;
}
}
return undefined;
}
function readUsageTokenTotal(usage: Record<string, unknown>): number | undefined {
const usageWithCacheCreation = usage as typeof usage & {
cache_creation_input_tokens?: unknown;
};
const inputTokens =
typeof usage.input_tokens === "number" && Number.isFinite(usage.input_tokens)
? usage.input_tokens
: 0;
const cacheCreationInputTokens =
typeof usageWithCacheCreation.cache_creation_input_tokens === "number" &&
Number.isFinite(usageWithCacheCreation.cache_creation_input_tokens)
? usageWithCacheCreation.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;
const outputTokens =
typeof usage.output_tokens === "number" && Number.isFinite(usage.output_tokens)
? usage.output_tokens
: 0;
const total = inputTokens + cacheCreationInputTokens + cacheReadInputTokens + outputTokens;
return total > 0 ? total : undefined;
}
function readActiveUsageTokens(usage: unknown): number | undefined {
const activeUsage = readLastUsageIteration(usage);
return activeUsage ? readUsageTokenTotal(activeUsage) : undefined;
}
function readLegacyResultUsageTokens(usage: unknown): number | undefined {
const usageRecord = toObjectRecord(usage);
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";
}
class ClaudeContextUsageState {
private contextWindowMaxTokens: number | undefined;
private streamRequestInputTokens: number | undefined;
private streamRequestOutputTokens: number | undefined;
private completedResultTurns = 0;
beginTurn(): void {
this.streamRequestInputTokens = undefined;
this.streamRequestOutputTokens = undefined;
}
recordModelUsage(modelUsage: unknown): number | undefined {
const contextWindowMaxTokens = extractContextWindowSize(modelUsage);
if (contextWindowMaxTokens !== undefined) {
this.contextWindowMaxTokens = contextWindowMaxTokens;
}
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) {
return null;
}
const eventType = readTrimmedString(streamEvent.type);
if (eventType === "message_start") {
const inputTokens = readStreamRequestInputTokens(streamEvent);
if (typeof inputTokens !== "number") {
return null;
}
this.streamRequestInputTokens = inputTokens;
this.streamRequestOutputTokens = 0;
} else if (eventType === "message_delta") {
const outputTokens = readStreamRequestOutputTokens(streamEvent);
if (typeof outputTokens !== "number") {
return null;
}
this.streamRequestOutputTokens = outputTokens;
} else {
return null;
}
const usedTokens = this.streamUsedTokens();
if (usedTokens === undefined) {
return null;
}
return this.createUsageUpdatedEvent(usedTokens);
}
buildResultUsage(
message: SDKResultMessage,
modelUsage: unknown,
currentContextUsage: ClaudeCurrentContextUsage | undefined,
): AgentUsage | undefined {
try {
if (!message.usage) {
return undefined;
}
const usage: AgentUsage = {
inputTokens: message.usage.input_tokens,
cachedInputTokens: message.usage.cache_read_input_tokens,
outputTokens: message.usage.output_tokens,
totalCostUsd: message.total_cost_usd,
};
const modelContextWindowMaxTokens = this.recordModelUsage(modelUsage ?? message.modelUsage);
this.recordCurrentContextUsage(currentContextUsage);
if (this.contextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.contextWindowMaxTokens;
} else if (modelContextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = modelContextWindowMaxTokens;
}
const activeResultUsageTokens =
readActiveUsageTokens(message.usage) ??
(this.completedResultTurns === 0 ? readLegacyResultUsageTokens(message.usage) : undefined);
const usedTokens =
currentContextUsage?.totalTokens ?? this.streamUsedTokens() ?? activeResultUsageTokens;
if (usedTokens !== undefined) {
usage.contextWindowUsedTokens = usedTokens;
}
return usage;
} finally {
this.completedResultTurns += 1;
}
}
private streamUsedTokens(): number | undefined {
if (
typeof this.streamRequestInputTokens !== "number" ||
typeof this.streamRequestOutputTokens !== "number"
) {
return undefined;
}
return this.streamRequestInputTokens + this.streamRequestOutputTokens;
}
private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent {
const usage: AgentUsage = {
contextWindowUsedTokens,
};
if (this.contextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.contextWindowMaxTokens;
}
return {
type: "usage_updated",
provider: "claude",
usage,
};
}
}
class ClaudeAgentSession implements AgentSession {
readonly provider = "claude" as const;
readonly capabilities = CLAUDE_CAPABILITIES;
@@ -1703,10 +1879,7 @@ class ClaudeAgentSession implements AgentSession {
private pendingInterruptAbort = false;
private foregroundHasVisibleActivity = false;
private activeTurnHasAssistantText = false;
private lastContextWindowUsedTokens: number | undefined;
private lastContextWindowMaxTokens: number | undefined;
private lastStreamRequestInputTokens: number | undefined;
private lastStreamRequestOutputTokens: number | undefined;
private readonly contextUsage = new ClaudeContextUsageState();
private userMessageIds: string[] = [];
private readonly emittedUserMessageIds = new Set<string>();
private readonly rewindTurnAnchors: ClaudeRewindTurnAnchor[] = [];
@@ -1840,6 +2013,7 @@ class ClaudeAgentSession implements AgentSession {
this.activeForegroundTurnId = turnId;
this.foregroundHasVisibleActivity = false;
this.activeTurnHasAssistantText = false;
this.contextUsage.beginTurn();
this.transitionTurnState("foreground", "foreground turn started");
this.clearRecentStderr();
@@ -3005,6 +3179,7 @@ class ClaudeAgentSession implements AgentSession {
id: this.createTurnId("autonomous"),
};
this.activeTurnHasAssistantText = false;
this.contextUsage.beginTurn();
this.notifySubscribers({ type: "turn_started", provider: "claude" });
this.syncTurnState("autonomous turn started");
}
@@ -3098,7 +3273,7 @@ class ClaudeAgentSession implements AgentSession {
if (await this.handleMissingResumedConversation(message, activeQuery)) {
return true;
}
this.routeSdkMessageFromPump(message);
await this.routeSdkMessageFromPump(message, activeQuery);
return false;
};
const drainActiveQuery = async (): Promise<boolean> => {
@@ -3174,7 +3349,7 @@ class ClaudeAgentSession implements AgentSession {
);
}
private routeSdkMessageFromPump(message: SDKMessage): void {
private async routeSdkMessageFromPump(message: SDKMessage, activeQuery: Query): Promise<void> {
if (this.shouldSuppressStaleResult(message)) {
return;
}
@@ -3204,26 +3379,12 @@ class ClaudeAgentSession implements AgentSession {
"provider.claude.parsed_event",
);
const messageEvents = this.translateMessageToEvents(message, {
suppressAssistantText: true,
suppressReasoning: true,
});
const assistantTimelineEvents = this.timelineAssembler
.consume({
message,
runId: turnId,
messageIdHint: identifiers.messageId,
})
.map(
(item) =>
({
type: "timeline",
item,
provider: "claude",
}) satisfies AgentStreamEvent,
);
const events = [...messageEvents, ...assistantTimelineEvents];
const events = await this.buildPumpedMessageEvents(
message,
activeQuery,
identifiers.messageId,
turnId,
);
if (events.length === 0) {
return;
@@ -3258,6 +3419,57 @@ class ClaudeAgentSession implements AgentSession {
this.dispatchEvents(events);
}
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({
message,
runId: turnId,
messageIdHint,
})
.map(
(item) =>
({
type: "timeline",
item,
provider: "claude",
}) satisfies AgentStreamEvent,
);
return [...messageEvents, ...assistantTimelineEvents];
}
private async queryCurrentContextUsage(
activeQuery: Query,
): Promise<ClaudeCurrentContextUsage | undefined> {
const query = activeQuery as Query & {
getContextUsage?: () => Promise<SDKControlGetContextUsageResponse>;
};
if (typeof query.getContextUsage !== "function") {
return undefined;
}
try {
const usage = await withTimeout(query.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,
@@ -3325,6 +3537,7 @@ class ClaudeAgentSession implements AgentSession {
options?: {
suppressAssistantText?: boolean;
suppressReasoning?: boolean;
currentContextUsage?: ClaudeCurrentContextUsage;
},
): AgentStreamEvent[] {
const parentToolUseId =
@@ -3375,7 +3588,9 @@ class ClaudeAgentSession implements AgentSession {
this.appendStreamEventEvents(message, events, options);
break;
case "result":
this.appendResultEvents(message, events);
this.appendResultEvents(message, events, {
currentContextUsage: options?.currentContextUsage,
});
break;
default:
break;
@@ -3456,11 +3671,7 @@ class ClaudeAgentSession implements AgentSession {
return;
}
if (message.subtype === "task_progress") {
this.lastContextWindowUsedTokens =
readContextWindowUsedTokensFromTaskProgress(message) ?? this.lastContextWindowUsedTokens;
if (typeof this.lastContextWindowUsedTokens === "number") {
events.push(this.createUsageUpdatedEvent(this.lastContextWindowUsedTokens));
}
return;
}
}
@@ -3470,12 +3681,12 @@ class ClaudeAgentSession implements AgentSession {
): void {
// TODO: subagent timelines are best-effort. Subagent task_notifications
// arrive without parent_tool_use_id but with tool_use_id pointing at the
// parent's Task call, so they slip past the sidechain router and pollute
// the parent's subagent tool call, so they slip past the sidechain router and pollute
// the parent timeline. Drop them here; eventually thread them into the
// parent Task tool call's sub_agent log instead.
// parent tool call's sub_agent log instead.
const taskUseId = message.tool_use_id;
const cachedTool = taskUseId ? this.toolUseCache.get(taskUseId) : undefined;
if (cachedTool?.name === "Task") {
if (isClaudeSubagentToolName(cachedTool?.name)) {
return;
}
const taskNotificationItem = mapTaskNotificationSystemRecordToToolCall(message);
@@ -3486,11 +3697,6 @@ class ClaudeAgentSession implements AgentSession {
provider: "claude",
});
}
const usage = readUsageFromTaskNotification(message);
if (typeof usage === "number") {
this.lastContextWindowUsedTokens = usage;
events.push(this.createUsageUpdatedEvent(usage));
}
}
private appendUserMessageEvents(
@@ -3569,7 +3775,7 @@ class ClaudeAgentSession implements AgentSession {
events: AgentStreamEvent[],
options: { suppressAssistantText?: boolean; suppressReasoning?: boolean } | undefined,
): void {
const usageUpdatedEvent = this.trackStreamEventUsage(message.event);
const usageUpdatedEvent = this.contextUsage.buildStreamUsageEvent(message.event);
if (usageUpdatedEvent) {
events.push(usageUpdatedEvent);
}
@@ -3585,8 +3791,9 @@ class ClaudeAgentSession implements AgentSession {
private appendResultEvents(
message: Extract<SDKMessage, { type: "result" }>,
events: AgentStreamEvent[],
options?: { currentContextUsage?: ClaudeCurrentContextUsage },
): void {
const usage = this.convertUsage(message, message.modelUsage);
const usage = this.convertUsage(message, message.modelUsage, options?.currentContextUsage);
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
@@ -3753,99 +3960,12 @@ class ClaudeAgentSession implements AgentSession {
return null;
}
private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined {
if (!message.usage) {
return undefined;
}
const usage: AgentUsage = {
inputTokens: message.usage.input_tokens,
cachedInputTokens: message.usage.cache_read_input_tokens,
outputTokens: message.usage.output_tokens,
totalCostUsd: message.total_cost_usd,
};
const contextWindowMaxTokens = extractContextWindowSize(modelUsage ?? message.modelUsage);
if (contextWindowMaxTokens !== undefined) {
this.lastContextWindowMaxTokens = contextWindowMaxTokens;
usage.contextWindowMaxTokens = contextWindowMaxTokens;
} else if (this.lastContextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens;
}
if (typeof this.lastContextWindowUsedTokens === "number") {
// task_progress.total_tokens is the accurate context window fill level.
// Prefer it over result.usage which contains accumulated session totals.
usage.contextWindowUsedTokens = this.lastContextWindowUsedTokens;
} else if (
typeof this.lastStreamRequestInputTokens === "number" &&
typeof this.lastStreamRequestOutputTokens === "number"
) {
usage.contextWindowUsedTokens =
this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens;
} else if (message.usage) {
// Fallback: derive from result.usage when no task_progress has been
// received yet. These values are accumulated across all API calls, but
// for the first turn they equal the per-call values so the estimate is
// reasonable. Once a task_progress arrives it takes over permanently.
const usageWithCacheCreation = message.usage as typeof message.usage & {
cache_creation_input_tokens?: number;
};
const derived =
(message.usage.input_tokens ?? 0) +
(usageWithCacheCreation.cache_creation_input_tokens ?? 0) +
(message.usage.cache_read_input_tokens ?? 0) +
(message.usage.output_tokens ?? 0);
if (Number.isFinite(derived) && derived > 0) {
usage.contextWindowUsedTokens = derived;
}
}
return usage;
}
private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent {
const usage: AgentUsage = {
contextWindowUsedTokens,
};
if (this.lastContextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens;
}
return {
type: "usage_updated",
provider: "claude",
usage,
};
}
private trackStreamEventUsage(event: unknown): AgentStreamEvent | null {
const streamEvent = toObjectRecord(event);
if (!streamEvent) {
return null;
}
const eventType = readTrimmedString(streamEvent.type);
if (eventType === "message_start") {
const inputTokens = readStreamRequestInputTokens(streamEvent);
if (typeof inputTokens !== "number") {
return null;
}
this.lastStreamRequestInputTokens = inputTokens;
this.lastStreamRequestOutputTokens = 0;
} else if (eventType === "message_delta") {
const outputTokens = readStreamRequestOutputTokens(streamEvent);
if (typeof outputTokens !== "number") {
return null;
}
this.lastStreamRequestOutputTokens = outputTokens;
} else {
return null;
}
if (
typeof this.lastStreamRequestInputTokens !== "number" ||
typeof this.lastStreamRequestOutputTokens !== "number"
) {
return null;
}
return this.createUsageUpdatedEvent(
this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens,
);
private convertUsage(
message: SDKResultMessage,
modelUsage?: unknown,
currentContextUsage?: ClaudeCurrentContextUsage,
): AgentUsage | undefined {
return this.contextUsage.buildResultUsage(message, modelUsage, currentContextUsage);
}
private handlePermissionRequest: CanUseTool = async (