diff --git a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts index 427e69dcc56aee8335bb95f30c4e464a9b0f8f80..648817534ef7c0b4d9846e6700b14267902aaf1e 100644 --- a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts @@ -2234,6 +2234,12 @@ describe("ClaudeAdapterV2 background wake turns", () => { readonly apiErrorStatus?: number; // null omits the field, as the CLI does on a zero-turn result. readonly terminalReason?: SDKResultMessage["terminal_reason"] | null; + readonly usage?: { + readonly input_tokens: number; + readonly output_tokens: number; + readonly cache_creation_input_tokens: number; + readonly cache_read_input_tokens: number; + }; }) => claudeSdkFrame({ type: "result", @@ -2245,7 +2251,7 @@ describe("ClaudeAdapterV2 background wake turns", () => { result: input.result, stop_reason: "end_turn", total_cost_usd: 0, - usage: { + usage: input.usage ?? { input_tokens: 1, output_tokens: 1, cache_creation_input_tokens: 0, @@ -2993,6 +2999,12 @@ describe("ClaudeAdapterV2 background wake turns", () => { uuid: "00000000-0000-4000-8000-000000000901", result: "", terminalReason, + usage: { + input_tokens: 600, + output_tokens: 40, + cache_creation_input_tokens: 500, + cache_read_input_tokens: 2000, + }, }), ); if (steered) { @@ -3002,6 +3014,12 @@ describe("ClaudeAdapterV2 background wake turns", () => { makeResultFrame({ uuid: "00000000-0000-4000-8000-000000000902", result: "Audit finished after the steer.", + usage: { + input_tokens: 3, + output_tokens: 6, + cache_creation_input_tokens: 2347, + cache_read_input_tokens: 9455, + }, }), ); } @@ -3016,6 +3034,25 @@ describe("ClaudeAdapterV2 background wake turns", () => { event.turnItem.text === WAKE_ASSISTANT_TEXT, ), ); + // Each result's usage covers only its own segment, so the steered + // turn bills both segments. + const completed = harness.events.findLast( + (event) => event.type === "provider_turn.updated", + ); + assert.deepEqual( + completed?.type === "provider_turn.updated" + ? completed.providerTurn.turnTokenUsage + : undefined, + { + usageScope: "main_agent", + usageStatus: "complete", + hasSubagents: false, + inputTokens: 603 + 2847 + 11455, + outputTokens: 46, + cachedInputTokens: 11455, + cacheCreationTokens: 2847, + }, + ); } assert.lengthOf(harness.terminalEvents(), 1); }).pipe( diff --git a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts index aa80b43ad9b8f0cbf72743af82726f373b7cd786..d0b63350f1d4e9d3036458b511195af1e1cf7387 100644 --- a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts @@ -2918,6 +2918,9 @@ interface ActiveClaudeTurnContext { readonly rejectedRateLimitTypes: Set; readonly rateLimitResetTimes: Map; latestAssistantRateLimited: boolean; + // Usage of the results that steers ended, oldest first. The final result's + // usage covers only its own segment, so the turn bills these as well. + readonly steeredSegmentUsage: Array; readonly subagentsByTaskId: Map; readonly subagentsByToolUseId: Map; readonly subagentNodesByTaskId: Map; @@ -5922,6 +5925,7 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( input.context.subagentsByTaskId.size > 0 || input.context.subagentsByToolUseId.size > 0, input.status, + input.context.steeredSegmentUsage, ), }, }), @@ -7660,6 +7664,7 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( const interrupted = (yield* Ref.get(interruptedTurns)).has(context.providerTurnId); const wasSteered = (yield* Ref.get(steeredTurns)).has(context.providerTurnId); if (!interrupted && wasSteered && isClaudeActiveSteeringAbortResult(message)) { + context.steeredSegmentUsage.push(message.usage); return; } yield* Ref.update(steeredTurns, (current) => { @@ -8568,6 +8573,7 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( rejectedRateLimitTypes: new Set(), rateLimitResetTimes: new Map(), latestAssistantRateLimited: false, + steeredSegmentUsage: [], subagentsByTaskId: new Map(), subagentsByToolUseId: new Map(), subagentNodesByTaskId: new Map(), diff --git a/apps/server/src/provider/ClaudeTurnTokenUsage.ts b/apps/server/src/provider/ClaudeTurnTokenUsage.ts index 07fe378692d0dd7eada07608cdde21adcbe938ed..ad69dc1fbc1b57d9b823d8dd184ebccc43d9b769 100644 --- a/apps/server/src/provider/ClaudeTurnTokenUsage.ts +++ b/apps/server/src/provider/ClaudeTurnTokenUsage.ts @@ -7,12 +7,70 @@ function finiteNonNegativeInteger(value: unknown): number | undefined { : undefined; } +const BILLED_USAGE_FIELDS = [ + "input_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "output_tokens", +] as const; + +// Adds earlier segments' counts to a count the final result reports. A count +// the final result lacks stays missing, so status rules see the same fields. +function addKnownCount(finalCount: unknown, earlierCounts: ReadonlyArray): unknown { + const known = finiteNonNegativeInteger(finalCount); + if (known === undefined) return finalCount; + return earlierCounts.reduce( + (sum, count) => sum + (finiteNonNegativeInteger(count) ?? 0), + known, + ); +} + +// A steer ends the running segment with an aborted result, and the SDK scopes +// each result's usage to its own segment. Adds the billed counts of those +// earlier segments onto the final result's usage. Other usage fields describe +// the final result only and are left as they are. +function withEarlierSegmentUsage( + usage: Record, + earlierSegmentUsage: ReadonlyArray, +): Record { + if (earlierSegmentUsage.length === 0) return usage; + const segments = earlierSegmentUsage.map((segment) => (segment ?? {}) as Record); + const total = { ...usage }; + for (const field of BILLED_USAGE_FIELDS) { + total[field] = addKnownCount( + usage[field], + segments.map((segment) => segment[field]), + ); + } + const outputDetails = usage.output_tokens_details as Record | undefined; + if (outputDetails != null) { + total.output_tokens_details = { + ...outputDetails, + thinking_tokens: addKnownCount( + outputDetails.thinking_tokens, + segments.map( + (segment) => + (segment.output_tokens_details as Record | undefined)?.thinking_tokens, + ), + ), + }; + } + return total; +} + +/** + * Normalizes a Claude turn's billed usage from its final SDK result. + * `earlierSegmentUsage` holds the usage of results a steer ended before the + * final one, oldest first. + */ export function normalizeClaudeTurnTokenUsage( result: { readonly subtype: SDKResultMessage["subtype"]; readonly usage?: unknown } | undefined, hasSubagents: boolean, terminalStatus: ProviderRuntimeTurnStatus, + earlierSegmentUsage: ReadonlyArray = [], ): TurnTokenUsage { - const usage = result?.usage as Record | undefined; + const finalUsage = result?.usage as Record | undefined; + const usage = finalUsage && withEarlierSegmentUsage(finalUsage, earlierSegmentUsage); if (!usage) { return { usageStatus: "unavailable",