diff --git a/apps/server/src/provider/Layers/ClaudeAdapter.test.ts b/apps/server/src/provider/Layers/ClaudeAdapter.test.ts index a27fc7703331..b1f871595a4c 100644 --- a/apps/server/src/provider/Layers/ClaudeAdapter.test.ts +++ b/apps/server/src/provider/Layers/ClaudeAdapter.test.ts @@ -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, ); @@ -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, ); @@ -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()), @@ -1910,6 +1954,177 @@ 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", () => { @@ -1917,7 +2132,7 @@ describe("ClaudeAdapterLive", () => { 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, ); @@ -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, }, diff --git a/apps/server/src/provider/Layers/ClaudeAdapter.ts b/apps/server/src/provider/Layers/ClaudeAdapter.ts index b291104d1493..6e902b0f4e68 100644 --- a/apps/server/src/provider/Layers/ClaudeAdapter.ts +++ b/apps/server/src/provider/Layers/ClaudeAdapter.ts @@ -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, @@ -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; + } + } + return; + } + if (event.type === "content_block_delta") { if ( (event.delta.type === "text_delta" || event.delta.type === "thinking_delta") && @@ -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", @@ -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",