From 5259fa4495ac42408dcb8ab1bffd466507d79c8b Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 2 Jun 2026 08:40:03 -0400 Subject: [PATCH] fix(llm): keep OpenAI-compatible reasoning streams active --- src/agents/openai-transport-stream.test.ts | 119 +++++++++++++++++++++ src/agents/openai-transport-stream.ts | 52 +++++++-- 2 files changed, 162 insertions(+), 9 deletions(-) diff --git a/src/agents/openai-transport-stream.test.ts b/src/agents/openai-transport-stream.test.ts index 77112d820053..cec9127c1772 100644 --- a/src/agents/openai-transport-stream.test.ts +++ b/src/agents/openai-transport-stream.test.ts @@ -1685,6 +1685,125 @@ describe("openai transport stream", () => { }); }); + it("emits reasoning activity for OpenAI-compatible usage-only reasoning chunks", async () => { + const model = { + id: "google/gemini-2.5-flash", + name: "Gemini 2.5 Flash", + api: "openai-completions", + provider: "vertex-ai", + baseUrl: "http://127.0.0.1:8787/v1beta1/projects/test/locations/us/endpoints/openapi", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1_000_000, + maxTokens: 8192, + } satisfies Model<"openai-completions">; + const output = createAssistantOutput(model); + const events: CapturedStreamEvent[] = []; + + await testing.processOpenAICompletionsStream( + streamChunks([ + { + id: "chatcmpl-vertex", + object: "chat.completion.chunk" as const, + created: 1775425651, + model: model.id, + choices: [], + usage: { + prompt_tokens: 8, + completion_tokens: 23, + total_tokens: 31, + completion_tokens_details: { reasoning_tokens: 23 }, + }, + }, + { + id: "chatcmpl-vertex", + object: "chat.completion.chunk" as const, + created: 1775425651, + model: model.id, + choices: [ + { + index: 0, + delta: { role: "assistant" as const, content: "Hi" }, + logprobs: null, + finish_reason: "stop" as const, + }, + ], + }, + ]), + output, + model, + { push: (event) => events.push(event as CapturedStreamEvent) }, + ); + + expect(events.map((event) => event.type)).toEqual([ + "thinking_start", + "thinking_delta", + "text_start", + "text_delta", + ]); + expect(events[1]).toHaveProperty("delta", ""); + expect(output.content).toEqual([ + { type: "thinking", thinking: "" }, + { type: "text", text: "Hi" }, + ]); + }); + + it("does not add trailing reasoning activity after visible OpenAI-compatible text", async () => { + const model = { + id: "google/gemini-2.5-flash", + name: "Gemini 2.5 Flash", + api: "openai-completions", + provider: "vertex-ai", + baseUrl: "http://127.0.0.1:8787/v1beta1/projects/test/locations/us/endpoints/openapi", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1_000_000, + maxTokens: 8192, + } satisfies Model<"openai-completions">; + const output = createAssistantOutput(model); + const events: CapturedStreamEvent[] = []; + + await testing.processOpenAICompletionsStream( + streamChunks([ + { + id: "chatcmpl-vertex", + object: "chat.completion.chunk" as const, + created: 1775425651, + model: model.id, + choices: [ + { + index: 0, + delta: { role: "assistant" as const, content: "Hi" }, + logprobs: null, + finish_reason: null, + }, + ], + }, + { + id: "chatcmpl-vertex", + object: "chat.completion.chunk" as const, + created: 1775425651, + model: model.id, + choices: [], + usage: { + prompt_tokens: 8, + completion_tokens: 25, + total_tokens: 33, + completion_tokens_details: { reasoning_tokens: 23 }, + }, + }, + ]), + output, + model, + { push: (event) => events.push(event as CapturedStreamEvent) }, + ); + + expect(events.map((event) => event.type)).toEqual(["text_start", "text_delta"]); + expect(output.content).toEqual([{ type: "text", text: "Hi" }]); + }); + it("yields to aborts during bursty OpenAI-compatible streams", async () => { const model = { id: "deepseek-v4-flash", diff --git a/src/agents/openai-transport-stream.ts b/src/agents/openai-transport-stream.ts index 90ec22026b77..24d88c3a1ce7 100644 --- a/src/agents/openai-transport-stream.ts +++ b/src/agents/openai-transport-stream.ts @@ -2657,6 +2657,11 @@ async function processOpenAICompletionsStream( let sawStopFinishReason = false; const blockIndex = () => output.content.length - 1; const measureUtf8Bytes = (text: string) => Buffer.byteLength(text, "utf8"); + let chunkPushedEvent = false; + const pushStreamEvent = (event: unknown) => { + chunkPushedEvent = true; + stream.push(event); + }; const finishCurrentBlock = () => { if (!currentBlock) { return; @@ -2697,13 +2702,13 @@ async function processOpenAICompletionsStream( currentBlock = { type: "thinking", thinking: "", - thinkingSignature: reasoningDelta.signature, + ...(reasoningDelta.signature ? { thinkingSignature: reasoningDelta.signature } : {}), }; output.content.push(currentBlock); - stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output }); + pushStreamEvent({ type: "thinking_start", contentIndex: blockIndex(), partial: output }); } currentBlock.thinking += reasoningDelta.text; - stream.push({ + pushStreamEvent({ type: "thinking_delta", contentIndex: blockIndex(), delta: reasoningDelta.text, @@ -2715,10 +2720,10 @@ async function processOpenAICompletionsStream( finishCurrentBlock(); currentBlock = { type: "text", text: "" }; output.content.push(currentBlock); - stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output }); + pushStreamEvent({ type: "text_start", contentIndex: blockIndex(), partial: output }); } currentBlock.text += text; - stream.push({ + pushStreamEvent({ type: "text_delta", contentIndex: blockIndex(), delta: text, @@ -2782,12 +2787,12 @@ async function processOpenAICompletionsStream( }; currentBlock = block; output.content.push(block); - stream.push({ + pushStreamEvent({ type: "toolcall_start", contentIndex: output.content.indexOf(block), partial: output, }); - stream.push({ + pushStreamEvent({ type: "toolcall_delta", contentIndex: output.content.indexOf(block), delta: toolCall.partialArgs, @@ -2853,6 +2858,19 @@ async function processOpenAICompletionsStream( appendFilteredVisibleTextDelta(delta.text); } }; + const emitReasoningUsageActivity = (hasReasoningUsageActivity: boolean) => { + if (!hasReasoningUsageActivity || chunkPushedEvent || !emitReasoning) { + return; + } + const latestBlock = output.content[output.content.length - 1]; + if (currentBlock?.type === "text" || currentBlock?.type === "toolCall") { + return; + } + if (latestBlock?.type === "text" || latestBlock?.type === "toolCall") { + return; + } + appendThinkingDelta({ signature: "", text: "" }); + }; const flushReasoningTagTextPartitionerAtEnd = () => { for (const delta of reasoningTagTextPartitioner.flush()) { appendPartitionedVisibleDelta(delta); @@ -2861,23 +2879,28 @@ async function processOpenAICompletionsStream( const cooperativeScheduler = createModelStreamCooperativeScheduler(options?.signal); for await (const rawChunk of responseStream as AsyncIterable) { throwIfModelStreamAborted(options?.signal); + chunkPushedEvent = false; if (!rawChunk || typeof rawChunk !== "object") { await cooperativeScheduler.afterEvent(); continue; } const chunk = rawChunk as ChatCompletionChunk; output.responseId ||= chunk.id; + let hasReasoningUsageActivity = false; if (chunk.usage) { output.usage = parseTransportChunkUsage(chunk.usage, model); + hasReasoningUsageActivity = hasOpenAICompletionsReasoningUsageActivity(chunk.usage); } const choice = Array.isArray(chunk.choices) ? chunk.choices[0] : undefined; if (!choice) { + emitReasoningUsageActivity(hasReasoningUsageActivity); await cooperativeScheduler.afterEvent(); continue; } const choiceUsage = (choice as unknown as { usage?: ChatCompletionChunk["usage"] }).usage; if (!chunk.usage && choiceUsage) { output.usage = parseTransportChunkUsage(choiceUsage, model); + hasReasoningUsageActivity = hasOpenAICompletionsReasoningUsageActivity(choiceUsage); } if (choice.finish_reason) { const finishReasonResult = mapStopReason(choice.finish_reason); @@ -2893,6 +2916,7 @@ async function processOpenAICompletionsStream( choice.delta ?? (choice as unknown as { message?: ChatCompletionChunk["choices"][number]["delta"] }).message; if (!choiceDelta) { + emitReasoningUsageActivity(hasReasoningUsageActivity); await cooperativeScheduler.afterEvent(); continue; } @@ -2961,7 +2985,7 @@ async function processOpenAICompletionsStream( ...(initialSig ? { thoughtSignature: initialSig } : {}), }; output.content.push(block); - stream.push({ + pushStreamEvent({ type: "toolcall_start", contentIndex: output.content.indexOf(block), partial: output, @@ -2991,7 +3015,7 @@ async function processOpenAICompletionsStream( toolCallBlockBytes.set(block, currentBlockArgBytes + nextArgumentBytes); block.partialArgs += toolCall.function.arguments; block.arguments = parseStreamingJson(block.partialArgs); - stream.push({ + pushStreamEvent({ type: "toolcall_delta", contentIndex: output.content.indexOf(block), delta: toolCall.function.arguments, @@ -3001,6 +3025,7 @@ async function processOpenAICompletionsStream( } } flushPendingPostToolCallDeltas(); + emitReasoningUsageActivity(hasReasoningUsageActivity); await cooperativeScheduler.afterEvent(); } flushReasoningTagTextPartitionerAtEnd(); @@ -4207,6 +4232,15 @@ export function parseTransportChunkUsage( return usage; } +function hasOpenAICompletionsReasoningUsageActivity( + rawUsage: NonNullable, +) { + const reasoningTokens = rawUsage.completion_tokens_details?.reasoning_tokens; + return ( + typeof reasoningTokens === "number" && Number.isFinite(reasoningTokens) && reasoningTokens > 0 + ); +} + function mapStopReason(reason: string | null) { if (reason === null) { return { stopReason: "stop" };