From bbd6510e9abc222adb156888172c59ae1d646fc9 Mon Sep 17 00:00:00 2001 From: MrShitFox Date: Sun, 9 Aug 2026 00:51:35 +0300 Subject: [PATCH] fix(executors): preserve Command Code usage in Responses streams --- open-sse/executors/commandCode.ts | 129 +++++++++++-- open-sse/transformer/responsesTransformer.ts | 103 +++++++++- tests/unit/command-code-executor.test.ts | 189 +++++++++++++++++++ tests/unit/responses-transformer.test.ts | 12 +- 4 files changed, 411 insertions(+), 22 deletions(-) diff --git a/open-sse/executors/commandCode.ts b/open-sse/executors/commandCode.ts index cc2642de5c..08576fe051 100644 --- a/open-sse/executors/commandCode.ts +++ b/open-sse/executors/commandCode.ts @@ -420,7 +420,61 @@ type AggregateState = { usage: JsonRecord | null; }; +function firstRecord(record: JsonRecord, keys: readonly string[]): JsonRecord { + for (const key of keys) { + const value = record[key]; + if (isRecord(value)) return value; + } + return {}; +} + +function firstNumber(record: JsonRecord, keys: readonly string[]): number | undefined { + for (const key of keys) { + const value = numberValue(record[key]); + if (value !== undefined) return value; + } + return undefined; +} + +/** Keep earlier finish-step usage when the terminal finish event omits it. */ +function mergeCommandCodeUsage(previous: JsonRecord | null, next: unknown): JsonRecord | null { + if (!isRecord(next)) return previous; + + const merged: JsonRecord = { ...(previous || {}), ...next }; + for (const key of [ + "inputTokenDetails", + "input_token_details", + "input_tokens_details", + "prompt_tokens_details", + "outputTokenDetails", + "output_token_details", + "output_tokens_details", + "completion_tokens_details", + "reasoningTokenDetails", + "reasoning_token_details", + ]) { + const before = isRecord(previous?.[key]) ? previous[key] : {}; + const after = isRecord(next[key]) ? next[key] : {}; + if (Object.keys(before).length > 0 || Object.keys(after).length > 0) { + merged[key] = { ...before, ...after }; + } + } + return merged; +} + +function rememberCommandCodeUsage(state: AggregateState, event: JsonRecord): void { + const usage = + event.type === "finish-step" + ? (event.usage ?? event.totalUsage) + : (event.totalUsage ?? event.usage); + state.usage = mergeCommandCodeUsage(state.usage, usage); +} + function applyEventToAggregate(event: JsonRecord, state: AggregateState): void { + // Some Command Code protocol revisions attach usage to the terminal payload + // without preserving the event type. Capture it before event-specific handling. + rememberCommandCodeUsage(state, event); + switch (event.type) { case "text-delta": state.content += stringValue(event.text) || ""; @@ -440,9 +494,10 @@ function applyEventToAggregate(event: JsonRecord, state: AggregateState): void { }); break; } + case "finish-step": + break; case "finish": state.finishReason = mapFinishReason(event.finishReason); - state.usage = isRecord(event.totalUsage) ? event.totalUsage : null; break; } } @@ -460,30 +515,72 @@ function applyEventToAggregateOrThrow(event: JsonRecord, state: AggregateState): function usageFromCommandCode(usage: JsonRecord | null) { if (!usage) return undefined; - const details = isRecord(usage.inputTokenDetails) ? usage.inputTokenDetails : {}; - const cacheRead = numberValue(details.cacheReadTokens) || 0; - const noCache = numberValue(details.noCacheTokens) || 0; + const inputDetails = firstRecord(usage, [ + "inputTokenDetails", + "input_token_details", + "input_tokens_details", + "prompt_tokens_details", + ]); + const outputDetails = firstRecord(usage, [ + "outputTokenDetails", + "output_token_details", + "output_tokens_details", + "completion_tokens_details", + ]); + const reasoningDetails = firstRecord(usage, [ + "reasoningTokenDetails", + "reasoning_token_details", + "reasoning_tokens_details", + ]); + const cacheRead = + firstNumber(usage, [ + "cachedInputTokens", + "cached_input_tokens", + "cacheReadInputTokens", + "cache_read_input_tokens", + "cacheReadTokens", + "cache_read_tokens", + "cached_tokens", + ]) ?? + firstNumber(inputDetails, [ + "cachedTokens", + "cached_tokens", + "cacheReadTokens", + "cache_read_tokens", + ]); + const noCache = firstNumber(inputDetails, ["noCacheTokens", "no_cache_tokens"]); // Command Code's totalUsage.inputTokens is the FULL prompt total and already // includes the cached portion (noCacheTokens + cacheReadTokens = inputTokens), // so we must NOT add cacheRead back — that would double-count. There is no // cache-write field in the upstream payload, so cache creation stays unset. - const inputTokens = numberValue(usage.inputTokens) || 0; - const prompt = inputTokens; - const completion = numberValue(usage.outputTokens) || 0; + const prompt = + firstNumber(usage, ["inputTokens", "input_tokens", "promptTokens", "prompt_tokens"]) ?? + (noCache ?? 0) + (cacheRead ?? 0); + const reasoning = + firstNumber(usage, ["reasoningTokens", "reasoning_tokens"]) ?? + firstNumber(outputDetails, ["reasoningTokens", "reasoning_tokens"]) ?? + firstNumber(reasoningDetails, ["reasoningTokens", "reasoning_tokens"]); + const textOutput = firstNumber(outputDetails, ["textTokens", "text_tokens"]); + const completion = + firstNumber(usage, [ + "outputTokens", + "output_tokens", + "completionTokens", + "completion_tokens", + ]) ?? (textOutput ?? 0) + (reasoning ?? 0); + const total = firstNumber(usage, ["totalTokens", "total_tokens"]) ?? prompt + completion; const result: JsonRecord = { prompt_tokens: prompt, + prompt_tokens_details: { cached_tokens: cacheRead ?? 0 }, completion_tokens: completion, - total_tokens: prompt + completion, + completion_tokens_details: { reasoning_tokens: reasoning ?? 0 }, + total_tokens: total, }; // Surface the cache breakdown as informational fields so logUsage prints // `| cache_read=X | no_cache=Y` and appendRequestLog persists them. These are // NOT added to prompt_tokens (already included) — metering stays accurate. - if (cacheRead > 0) result.cache_read_input_tokens = cacheRead; - if (noCache > 0) result.no_cache_tokens = noCache; - // Keep reasoning_token_details (reasoningTokens) when present so stream.ts's - // extractUsage can surface it as reasoning_tokens. - const reasoningDetails = isRecord(usage.reasoningTokenDetails) ? usage.reasoningTokenDetails : {}; - const reasoning = numberValue(reasoningDetails.reasoningTokens); + if (cacheRead !== undefined && cacheRead > 0) result.cache_read_input_tokens = cacheRead; + if (noCache !== undefined && noCache > 0) result.no_cache_tokens = noCache; if (reasoning !== undefined && reasoning > 0) result.reasoning_tokens = reasoning; return result; } @@ -523,6 +620,7 @@ function createStreamResponse( const emitEvent = (event: unknown) => { if (!isRecord(event) || closed) return; + rememberCommandCodeUsage(state, event); if (!sentRole) { sentRole = true; controller.enqueue(sse(chatCompletionChunk(id, model, { role: "assistant" }))); @@ -562,9 +660,10 @@ function createStreamResponse( } case "reasoning-end": break; + case "finish-step": + break; case "finish": { state.finishReason = mapFinishReason(event.finishReason); - state.usage = isRecord(event.totalUsage) ? event.totalUsage : null; controller.enqueue(sse(chatCompletionChunk(id, model, {}, state.finishReason))); // Emit a standards-compliant usage-only chunk (choices: []) before // [DONE] when upstream reported usage. stream.ts's extractUsage diff --git a/open-sse/transformer/responsesTransformer.ts b/open-sse/transformer/responsesTransformer.ts index dca61a31fa..e98c4d98db 100644 --- a/open-sse/transformer/responsesTransformer.ts +++ b/open-sse/transformer/responsesTransformer.ts @@ -37,6 +37,102 @@ async function getPath() { return _path || null; } +type UsageRecord = Record; + +function usageRecord(value: unknown): UsageRecord { + return typeof value === "object" && value !== null && !Array.isArray(value) + ? (value as UsageRecord) + : {}; +} + +function usageNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + +function usageDetails(record: UsageRecord, ...keys: string[]): UsageRecord { + for (const key of keys) { + const value = usageRecord(record[key]); + if (Object.keys(value).length > 0) return value; + } + return {}; +} + +/** Normalize Chat Completions and Responses usage into the Responses API shape. */ +function normalizeResponsesUsage(previous: unknown, raw: unknown): UsageRecord | null { + const source = usageRecord(raw); + if (Object.keys(source).length === 0) return usageRecord(previous); + + const before = usageRecord(previous); + const beforeInputDetails = usageDetails(before, "input_tokens_details", "prompt_tokens_details"); + const beforeOutputDetails = usageDetails( + before, + "output_tokens_details", + "completion_tokens_details" + ); + const inputDetails = usageDetails( + source, + "input_tokens_details", + "prompt_tokens_details", + "inputTokenDetails", + "input_token_details" + ); + const outputDetails = usageDetails( + source, + "output_tokens_details", + "completion_tokens_details", + "outputTokenDetails", + "output_token_details", + "reasoningTokenDetails", + "reasoning_token_details" + ); + + const inputTokens = + usageNumber(source.input_tokens) ?? + usageNumber(source.prompt_tokens) ?? + usageNumber(source.inputTokens) ?? + usageNumber(source.promptTokens) ?? + usageNumber(before.input_tokens) ?? + usageNumber(before.prompt_tokens) ?? + 0; + const cachedTokens = + usageNumber(source.cache_read_input_tokens) ?? + usageNumber(source.cached_input_tokens) ?? + usageNumber(source.cachedInputTokens) ?? + usageNumber(source.cached_tokens) ?? + usageNumber(inputDetails.cached_tokens) ?? + usageNumber(inputDetails.cachedTokens) ?? + usageNumber(inputDetails.cacheReadTokens) ?? + usageNumber(beforeInputDetails.cached_tokens) ?? + 0; + const outputTokens = + usageNumber(source.output_tokens) ?? + usageNumber(source.completion_tokens) ?? + usageNumber(source.outputTokens) ?? + usageNumber(source.completionTokens) ?? + usageNumber(before.output_tokens) ?? + usageNumber(before.completion_tokens) ?? + 0; + const reasoningTokens = + usageNumber(source.reasoning_tokens) ?? + usageNumber(source.reasoningTokens) ?? + usageNumber(outputDetails.reasoning_tokens) ?? + usageNumber(outputDetails.reasoningTokens) ?? + usageNumber(beforeOutputDetails.reasoning_tokens) ?? + 0; + const totalTokens = + usageNumber(source.total_tokens) ?? + usageNumber(source.totalTokens) ?? + inputTokens + outputTokens; + + return { + input_tokens: inputTokens, + input_tokens_details: { cached_tokens: cachedTokens }, + output_tokens: outputTokens, + output_tokens_details: { reasoning_tokens: reasoningTokens }, + total_tokens: totalTokens, + }; +} + // Create log directory for responses (Node.js only) export function createResponsesLogger(model, logsDir = null) { // Skip logging in worker environment (no fs) @@ -477,10 +573,11 @@ export function createResponsesApiTransformStream( continue; } + if (parsed.usage) { + state.usage = normalizeResponsesUsage(state.usage, parsed.usage); + } + if (!parsed.choices?.length) { - if (parsed.usage) { - state.usage = parsed.usage; - } // #6906: trailing usage-only chunk after finish_reason already deferred // completion — send it now with the usage just captured above. if (state.awaitingTrailingUsage && !state.completedSent) { diff --git a/tests/unit/command-code-executor.test.ts b/tests/unit/command-code-executor.test.ts index 3fd1dacac2..1dcef8e77b 100644 --- a/tests/unit/command-code-executor.test.ts +++ b/tests/unit/command-code-executor.test.ts @@ -11,10 +11,18 @@ const { REGISTRY, getRegistryEntry } = await import("../../open-sse/config/provi const { CommandCodeExecutor, COMMAND_CODE_VERSION } = await import("../../open-sse/executors/commandCode.ts"); const { getExecutor, hasSpecializedExecutor } = await import("../../open-sse/executors/index.ts"); +const { createResponsesApiTransformStream } = + await import("../../open-sse/transformer/responsesTransformer.ts"); const core = await import("../../src/lib/db/core.ts"); const originalFetch = globalThis.fetch; +type JsonRecord = Record; +type ResponsesEvent = { + event: string; + data: { response: JsonRecord & { usage?: unknown; output?: JsonRecord[] } }; +}; + const PINNED_COMMAND_CODE_MODELS = [ "claude-opus-4-7", "claude-opus-4-6", @@ -60,6 +68,27 @@ function parseSsePayloads(sse: string) { .map((line) => JSON.parse(line)); } +async function responsesFromChatSse(sse: string): Promise { + const input = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(sse)); + controller.close(); + }, + }); + const transformed = await new Response( + input.pipeThrough(createResponsesApiTransformStream(null, 60_000)) + ).text(); + + return transformed + .split("\n\n") + .map((part) => { + const event = part.match(/^event:\s*(.+)$/m)?.[1]; + const data = part.match(/^data:\s*(.+)$/m)?.[1]; + return event && data ? ({ event, data: JSON.parse(data) } as ResponsesEvent) : null; + }) + .filter((entry): entry is ResponsesEvent => entry !== null); +} + test.afterEach(() => { globalThis.fetch = originalFetch; }); @@ -270,7 +299,9 @@ test("Command Code data: SSE lines aggregate into non-stream ChatCompletion JSON assert.equal(json.choices[0].finish_reason, "length"); assert.deepEqual(json.usage, { prompt_tokens: 3, + prompt_tokens_details: { cached_tokens: 2 }, completion_tokens: 5, + completion_tokens_details: { reasoning_tokens: 0 }, total_tokens: 8, cache_read_input_tokens: 2, }); @@ -466,7 +497,9 @@ test("Command Code stream emits a usage-only chunk with actual tokens before [DO assert.ok(usageChunk, "expected a usage-only chunk (choices: []) in the stream"); assert.deepEqual(usageChunk.usage, { prompt_tokens: 10, + prompt_tokens_details: { cached_tokens: 4 }, completion_tokens: 6, + completion_tokens_details: { reasoning_tokens: 1 }, total_tokens: 16, cache_read_input_tokens: 4, reasoning_tokens: 1, @@ -506,9 +539,165 @@ test("Command Code non-stream usage keeps inputTokens as prompt_tokens and repor const json = await response.json(); assert.deepEqual(json.usage, { prompt_tokens: 5, + prompt_tokens_details: { cached_tokens: 3 }, completion_tokens: 2, + completion_tokens_details: { reasoning_tokens: 0 }, total_tokens: 7, cache_read_input_tokens: 3, no_cache_tokens: 2, }); }); + +test("Command Code preserves finish-step usage through a finish without totalUsage", async () => { + globalThis.fetch = async () => + commandCodeStream([ + { type: "text-delta", text: "Hi" }, + { + type: "finish-step", + usage: { + inputTokens: 7308, + inputTokenDetails: { noCacheTokens: 27, cacheReadTokens: 7281 }, + outputTokens: 177, + outputTokenDetails: { textTokens: 12, reasoningTokens: 165 }, + totalTokens: 7485, + }, + }, + { type: "finish", finishReason: "stop", totalUsage: null }, + ]); + + const { response } = await getExecutor("command-code").execute({ + model: "gpt-5.4-mini", + stream: true, + credentials: { apiKey: "cc_test_key" }, + body: { messages: [{ role: "user", content: "Hi" }] }, + }); + const sse = await response.text(); + const chunks = parseSsePayloads(sse); + const usageChunk = chunks.find( + (chunk) => Array.isArray(chunk.choices) && chunk.choices.length === 0 + ); + + assert.deepEqual(usageChunk?.usage, { + prompt_tokens: 7308, + prompt_tokens_details: { cached_tokens: 7281 }, + completion_tokens: 177, + completion_tokens_details: { reasoning_tokens: 165 }, + total_tokens: 7485, + cache_read_input_tokens: 7281, + no_cache_tokens: 27, + reasoning_tokens: 165, + }); + + const completed = (await responsesFromChatSse(sse)).find( + (event) => event.event === "response.completed" + ); + assert.deepEqual(completed?.data.response.usage, { + input_tokens: 7308, + input_tokens_details: { cached_tokens: 7281 }, + output_tokens: 177, + output_tokens_details: { reasoning_tokens: 165 }, + total_tokens: 7485, + }); +}); + +test("Command Code accepts OpenAI-style usage aliases with absent optional details", async () => { + globalThis.fetch = async () => + commandCodeStream([ + { + type: "finish-step", + usage: { + prompt_tokens: 11, + prompt_tokens_details: { cached_tokens: 4 }, + completion_tokens: 5, + completion_tokens_details: { reasoning_tokens: 2 }, + total_tokens: 16, + }, + }, + { type: "finish", finishReason: "stop" }, + ]); + + const { response } = await getExecutor("command-code").execute({ + model: "gpt-5.4-mini", + stream: true, + credentials: { apiKey: "cc_test_key" }, + body: { messages: [{ role: "user", content: "Hi" }] }, + }); + const sse = await response.text(); + const usageChunk = parseSsePayloads(sse).find( + (chunk) => Array.isArray(chunk.choices) && chunk.choices.length === 0 + ); + assert.deepEqual(usageChunk?.usage, { + prompt_tokens: 11, + prompt_tokens_details: { cached_tokens: 4 }, + completion_tokens: 5, + completion_tokens_details: { reasoning_tokens: 2 }, + total_tokens: 16, + cache_read_input_tokens: 4, + reasoning_tokens: 2, + }); + + globalThis.fetch = async () => + commandCodeStream([ + { type: "finish-step", usage: { inputTokens: 4, outputTokens: 3, totalTokens: 7 } }, + { type: "finish", finishReason: "stop", totalUsage: null }, + ]); + const fallback = await getExecutor("command-code").execute({ + model: "gpt-5.4-mini", + stream: true, + credentials: { apiKey: "cc_test_key" }, + body: { messages: [{ role: "user", content: "Hi" }] }, + }); + const fallbackSse = await fallback.response.text(); + const completed = (await responsesFromChatSse(fallbackSse)).find( + (event) => event.event === "response.completed" + ); + assert.deepEqual(completed?.data.response.usage, { + input_tokens: 4, + input_tokens_details: { cached_tokens: 0 }, + output_tokens: 3, + output_tokens_details: { reasoning_tokens: 0 }, + total_tokens: 7, + }); +}); + +test("Command Code preserves tool-call streaming while finalizing finish-step usage", async () => { + globalThis.fetch = async () => + commandCodeStream([ + { + type: "tool-call", + toolCallId: "call_1", + toolName: "lookup", + input: { query: "hello" }, + }, + { type: "finish-step", usage: { inputTokens: 3, outputTokens: 2, totalTokens: 5 } }, + { type: "finish", finishReason: "tool-calls" }, + ]); + + const { response } = await getExecutor("command-code").execute({ + model: "gpt-5.4-mini", + stream: true, + credentials: { apiKey: "cc_test_key" }, + body: { messages: [{ role: "user", content: "Hi" }] }, + }); + const sse = await response.text(); + assert.ok( + parseSsePayloads(sse).some( + (chunk) => chunk.choices?.[0]?.delta?.tool_calls?.[0]?.id === "call_1" + ) + ); + + const completed = (await responsesFromChatSse(sse)).find( + (event) => event.event === "response.completed" + ); + assert.equal( + completed?.data.response.output?.some((item) => item.type === "function_call"), + true + ); + assert.deepEqual(completed?.data.response.usage, { + input_tokens: 3, + input_tokens_details: { cached_tokens: 0 }, + output_tokens: 2, + output_tokens_details: { reasoning_tokens: 0 }, + total_tokens: 5, + }); +}); diff --git a/tests/unit/responses-transformer.test.ts b/tests/unit/responses-transformer.test.ts index 6ea4bd6bfe..f8b1c5a0b1 100644 --- a/tests/unit/responses-transformer.test.ts +++ b/tests/unit/responses-transformer.test.ts @@ -76,8 +76,10 @@ test("createResponsesApiTransformStream converts plain chat deltas into Response assert.ok(types.includes("response.output_text.done")); assert.equal(completed.output[0].content[0].text, "Hello"); assert.deepEqual(completed.usage, { - prompt_tokens: 1, - completion_tokens: 2, + input_tokens: 1, + input_tokens_details: { cached_tokens: 0 }, + output_tokens: 2, + output_tokens_details: { reasoning_tokens: 0 }, total_tokens: 3, }); assert.equal(doneMarker.data, "[DONE]"); @@ -374,8 +376,10 @@ test("createResponsesApiTransformStream ignores malformed events and preserves u assert.equal(completed.id, "resp_chatcmpl_edge"); assert.equal(completed.output[0].content[0].text, "ok"); assert.deepEqual(completed.usage, { - prompt_tokens: 2, - completion_tokens: 1, + input_tokens: 2, + input_tokens_details: { cached_tokens: 0 }, + output_tokens: 1, + output_tokens_details: { reasoning_tokens: 0 }, total_tokens: 3, }); });