Skip to content

fix: per-request stream usage for context window occupancy #126

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
harrryyd merged 1 commit into from
Jun 13, 2026
Merged

fix: per-request stream usage for context window occupancy #126

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
251 changes: 233 additions & 18 deletions apps/server/src/provider/Layers/ClaudeAdapter.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1684,7 +1684,7 @@ describe("ClaudeAdapterLive", () => {
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;

const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 6).pipe(
const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 5).pipe(
Stream.runCollect,
Effect.forkChild,
);
Expand Down Expand Up @@ -1726,12 +1726,12 @@ describe("ClaudeAdapterLive", () => {
);
});

it.effect("emits thread token usage updates from Claude task progress", () => {
it.effect("does not emit token usage updates from task progress", () => {
const harness = makeHarness();
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;

const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 6).pipe(
const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 5).pipe(
Stream.runCollect,
Effect.forkChild,
);
Expand Down Expand Up @@ -1759,20 +1759,64 @@ describe("ClaudeAdapterLive", () => {
const runtimeEvents = Array.from(yield* Fiber.join(runtimeEventsFiber));
const usageEvent = runtimeEvents.find((event) => event.type === "thread.token-usage.updated");
const progressEvent = runtimeEvents.find((event) => event.type === "task.progress");
assert.equal(usageEvent?.type, "thread.token-usage.updated");
if (usageEvent?.type === "thread.token-usage.updated") {
assert.deepEqual(usageEvent.payload, {
usage: {
usedTokens: 321,
lastUsedTokens: 321,
toolUses: 2,
durationMs: 654,
},
assert.equal(usageEvent, undefined);
assert.equal(progressEvent?.type, "task.progress");
if (progressEvent?.type === "task.progress") {
assert.equal(progressEvent.payload.description, "Thinking through the patch");
assert.deepEqual(progressEvent.payload.usage, {
total_tokens: 321,
tool_uses: 2,
duration_ms: 654,
});
}
assert.equal(progressEvent?.type, "task.progress");
if (usageEvent && progressEvent) {
assert.notStrictEqual(usageEvent.eventId, progressEvent.eventId);
}).pipe(
Effect.provideService(Random.Random, makeDeterministicRandomService()),
Effect.provide(harness.layer),
);
});

it.effect("does not emit token usage updates from task notification", () => {
const harness = makeHarness();
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;

const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 5).pipe(
Stream.runCollect,
Effect.forkChild,
);

yield* adapter.startSession({
threadId: THREAD_ID,
provider: ProviderDriverKind.make("claudeAgent"),
runtimeMode: "full-access",
});

harness.query.emit({
type: "system",
subtype: "task_notification",
task_id: "task-notif-1",
status: "completed",
description: "Background task finished",
usage: {
total_tokens: 888,
tool_uses: 5,
},
session_id: "sdk-session-task-notif",
uuid: "task-notif-1",
} as unknown as SDKMessage);

const runtimeEvents = Array.from(yield* Fiber.join(runtimeEventsFiber));
const usageEvent = runtimeEvents.find((event) => event.type === "thread.token-usage.updated");
const completedEvent = runtimeEvents.find((event) => event.type === "task.completed");
assert.equal(usageEvent, undefined);
assert.equal(completedEvent?.type, "task.completed");
if (completedEvent?.type === "task.completed") {
assert.equal(completedEvent.payload.taskId, "task-notif-1");
assert.equal(completedEvent.payload.status, "completed");
assert.deepEqual(completedEvent.payload.usage, {
total_tokens: 888,
tool_uses: 5,
});
}
}).pipe(
Effect.provideService(Random.Random, makeDeterministicRandomService()),
Expand Down Expand Up @@ -1910,14 +1954,185 @@ describe("ClaudeAdapterLive", () => {
);
});

it.effect(
"uses per-request usage from message_start stream events for context window occupancy",
() => {
const harness = makeHarness();
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;

const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 7).pipe(
Stream.runCollect,
Effect.forkChild,
);

yield* adapter.startSession({
threadId: THREAD_ID,
provider: ProviderDriverKind.make("claudeAgent"),
runtimeMode: "full-access",
});

yield* adapter.sendTurn({
threadId: THREAD_ID,
input: "hello",
attachments: [],
});

harness.query.emit({
type: "stream_event",
session_id: "sdk-session-msg-start",
uuid: "stream-msg-start-1",
parent_tool_use_id: null,
event: {
type: "message_start",
message: {
id: "msg_1",
type: "message",
role: "assistant",
model: "claude-sonnet-4-5",
stop_reason: null,
stop_sequence: null,
content: [],
usage: {
input_tokens: 400,
output_tokens: 100,
},
},
},
} as unknown as SDKMessage);

harness.query.emit({
type: "result",
subtype: "success",
is_error: false,
duration_ms: 1234,
duration_api_ms: 1200,
num_turns: 1,
result: "done",
stop_reason: "end_turn",
session_id: "sdk-session-result-msg-start",
usage: {
input_tokens: 4000,
output_tokens: 1000,
},
modelUsage: {
"claude-sonnet-4-5": {
contextWindow: 200000,
maxOutputTokens: 64000,
},
},
} as unknown as SDKMessage);
harness.query.finish();

const runtimeEvents = Array.from(yield* Fiber.join(runtimeEventsFiber));
const usageEvent = runtimeEvents.find(
(event) => event.type === "thread.token-usage.updated",
);
assert.equal(usageEvent?.type, "thread.token-usage.updated");
if (usageEvent?.type === "thread.token-usage.updated") {
// usedTokens should be the per-request value from message_start (500),
// NOT the accumulated value from result.usage (5000).
assert.equal(usageEvent.payload.usage.usedTokens, 500);
assert.equal(usageEvent.payload.usage.lastUsedTokens, 500);
// totalProcessedTokens reflects the accumulated total for cost tracking.
assert.equal(usageEvent.payload.usage.totalProcessedTokens, 5000);
assert.equal(usageEvent.payload.usage.maxTokens, 200000);
}
}).pipe(
Effect.provideService(Random.Random, makeDeterministicRandomService()),
Effect.provide(harness.layer),
);
},
);

it.effect(
"uses per-request usage from message_delta stream events for context window occupancy",
() => {
const harness = makeHarness();
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;

const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 7).pipe(
Stream.runCollect,
Effect.forkChild,
);

yield* adapter.startSession({
threadId: THREAD_ID,
provider: ProviderDriverKind.make("claudeAgent"),
runtimeMode: "full-access",
});

yield* adapter.sendTurn({
threadId: THREAD_ID,
input: "hello",
attachments: [],
});

harness.query.emit({
type: "stream_event",
session_id: "sdk-session-msg-delta",
uuid: "stream-msg-delta-1",
parent_tool_use_id: null,
event: {
type: "message_delta",
delta: { stop_reason: "end_turn", stop_sequence: null },
usage: {
input_tokens: 300,
output_tokens: 200,
},
},
} as unknown as SDKMessage);

harness.query.emit({
type: "result",
subtype: "success",
is_error: false,
duration_ms: 1234,
duration_api_ms: 1200,
num_turns: 1,
result: "done",
stop_reason: "end_turn",
session_id: "sdk-session-result-msg-delta",
usage: {
input_tokens: 4000,
output_tokens: 1000,
},
modelUsage: {
"claude-sonnet-4-5": {
contextWindow: 200000,
maxOutputTokens: 64000,
},
},
} as unknown as SDKMessage);
harness.query.finish();

const runtimeEvents = Array.from(yield* Fiber.join(runtimeEventsFiber));
const usageEvent = runtimeEvents.find(
(event) => event.type === "thread.token-usage.updated",
);
assert.equal(usageEvent?.type, "thread.token-usage.updated");
if (usageEvent?.type === "thread.token-usage.updated") {
assert.equal(usageEvent.payload.usage.usedTokens, 500);
assert.equal(usageEvent.payload.usage.lastUsedTokens, 500);
assert.equal(usageEvent.payload.usage.totalProcessedTokens, 5000);
assert.equal(usageEvent.payload.usage.maxTokens, 200000);
}
}).pipe(
Effect.provideService(Random.Random, makeDeterministicRandomService()),
Effect.provide(harness.layer),
);
},
);

it.effect(
"preserves oversized Claude result totals after task progress snapshots are recorded",
() => {
const harness = makeHarness();
return Effect.gen(function* () {
const adapter = yield* ClaudeAdapter;

const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 9).pipe(
const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 8).pipe(
Stream.runCollect,
Effect.forkChild,
);
Expand Down Expand Up @@ -1977,8 +2192,8 @@ describe("ClaudeAdapterLive", () => {
if (finalUsageEvent?.type === "thread.token-usage.updated") {
assert.deepEqual(finalUsageEvent.payload, {
usage: {
usedTokens: 190000,
lastUsedTokens: 190000,
usedTokens: 200000,
lastUsedTokens: 200000,
totalProcessedTokens: 535000,
maxTokens: 200000,
},
Expand Down
65 changes: 25 additions & 40 deletions apps/server/src/provider/Layers/ClaudeAdapter.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1506,8 +1506,9 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* (
// The SDK result.usage contains *accumulated* totals across all API calls
// (input_tokens, cache_read_input_tokens, etc. summed over every request).
// This does NOT represent the current context window size.
// Instead, use the last known context-window-accurate usage from task_progress
// events and treat the accumulated total as totalProcessedTokens.
// Instead, use per-request usage captured from message_start/message_delta
// stream events. If none were captured (e.g. no streaming), fall back to
// the accumulated total, treating it as totalProcessedTokens.
const accumulatedSnapshot = normalizeClaudeTokenUsage(
result?.usage,
resultContextWindow ?? context.lastKnownContextWindow,
Expand Down Expand Up @@ -1677,6 +1678,28 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* (

const { event } = message;

if (event.type === "message_start") {
const usage = event.message?.usage;
if (usage && typeof usage === "object") {
const normalized = normalizeClaudeTokenUsage(usage, context.lastKnownContextWindow);
if (normalized) {
context.lastKnownTokenUsage = normalized;
}
}
return;
}

if (event.type === "message_delta") {
const usage = event.usage;
if (usage && typeof usage === "object") {
const normalized = normalizeClaudeTokenUsage(usage, context.lastKnownContextWindow);
if (normalized) {
context.lastKnownTokenUsage = normalized;
Comment on lines +1695 to +1697

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Merge message_delta usage with the start snapshot

When the SDK sends the documented streamed shape where message_delta.usage contains only cumulative output_tokens, normalizing and assigning it here overwrites the earlier message_start snapshot that had the input/cache tokens. A normal streamed turn can therefore complete with usedTokens equal to just the output tokens (for example 15 instead of 25 input + 15 output), so the context-window occupancy remains severely underreported. Preserve the input/cache counts from message_start when applying a delta, or ignore input-less deltas instead of replacing the whole snapshot.

Useful? React with 👍 / 👎.

}
}
return;
}

if (event.type === "content_block_delta") {
if (
(event.delta.type === "text_delta" || event.delta.type === "thinking_delta") &&
Expand Down Expand Up @@ -2236,25 +2259,6 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* (
});
return;
case "task_progress":
if (message.usage) {
const normalizedUsage = normalizeClaudeTokenUsage(
message.usage,
context.lastKnownContextWindow,
);
if (normalizedUsage) {
context.lastKnownTokenUsage = normalizedUsage;
const usageStamp = yield* makeEventStamp();
yield* offerRuntimeEvent({
...base,
eventId: usageStamp.eventId,
createdAt: usageStamp.createdAt,
type: "thread.token-usage.updated",
payload: {
usage: normalizedUsage,
},
});
}
}
yield* offerRuntimeEvent({
...base,
type: "task.progress",
Expand All @@ -2268,25 +2272,6 @@ export const makeClaudeAdapter = Effect.fn("makeClaudeAdapter")(function* (
});
return;
case "task_notification":
if (message.usage) {
const normalizedUsage = normalizeClaudeTokenUsage(
message.usage,
context.lastKnownContextWindow,
);
if (normalizedUsage) {
context.lastKnownTokenUsage = normalizedUsage;
const usageStamp = yield* makeEventStamp();
yield* offerRuntimeEvent({
...base,
eventId: usageStamp.eventId,
createdAt: usageStamp.createdAt,
type: "thread.token-usage.updated",
payload: {
usage: normalizedUsage,
},
});
}
}
yield* offerRuntimeEvent({
...base,
type: "task.completed",
Expand Down
Loading