Skip to content

Commit 5259fa4

Browse files
steipetesteipete
authored
fix(llm): keep OpenAI-compatible reasoning streams active
1 parent 2ffeca1 commit 5259fa4

2 files changed

Lines changed: 162 additions & 9 deletions

File tree

‎src/agents/openai-transport-stream.test.ts‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1685,6 +1685,125 @@ describe("openai transport stream", () => {
16851685
});
16861686
});
16871687

1688+
it("emits reasoning activity for OpenAI-compatible usage-only reasoning chunks", async () => {
1689+
const model = {
1690+
id: "google/gemini-2.5-flash",
1691+
name: "Gemini 2.5 Flash",
1692+
api: "openai-completions",
1693+
provider: "vertex-ai",
1694+
baseUrl: "http://127.0.0.1:8787/v1beta1/projects/test/locations/us/endpoints/openapi",
1695+
reasoning: true,
1696+
input: ["text"],
1697+
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
1698+
contextWindow: 1_000_000,
1699+
maxTokens: 8192,
1700+
} satisfies Model<"openai-completions">;
1701+
const output = createAssistantOutput(model);
1702+
const events: CapturedStreamEvent[] = [];
1703+
1704+
await testing.processOpenAICompletionsStream(
1705+
streamChunks([
1706+
{
1707+
id: "chatcmpl-vertex",
1708+
object: "chat.completion.chunk" as const,
1709+
created: 1775425651,
1710+
model: model.id,
1711+
choices: [],
1712+
usage: {
1713+
prompt_tokens: 8,
1714+
completion_tokens: 23,
1715+
total_tokens: 31,
1716+
completion_tokens_details: { reasoning_tokens: 23 },
1717+
},
1718+
},
1719+
{
1720+
id: "chatcmpl-vertex",
1721+
object: "chat.completion.chunk" as const,
1722+
created: 1775425651,
1723+
model: model.id,
1724+
choices: [
1725+
{
1726+
index: 0,
1727+
delta: { role: "assistant" as const, content: "Hi" },
1728+
logprobs: null,
1729+
finish_reason: "stop" as const,
1730+
},
1731+
],
1732+
},
1733+
]),
1734+
output,
1735+
model,
1736+
{ push: (event) => events.push(event as CapturedStreamEvent) },
1737+
);
1738+
1739+
expect(events.map((event) => event.type)).toEqual([
1740+
"thinking_start",
1741+
"thinking_delta",
1742+
"text_start",
1743+
"text_delta",
1744+
]);
1745+
expect(events[1]).toHaveProperty("delta", "");
1746+
expect(output.content).toEqual([
1747+
{ type: "thinking", thinking: "" },
1748+
{ type: "text", text: "Hi" },
1749+
]);
1750+
});
1751+
1752+
it("does not add trailing reasoning activity after visible OpenAI-compatible text", async () => {
1753+
const model = {
1754+
id: "google/gemini-2.5-flash",
1755+
name: "Gemini 2.5 Flash",
1756+
api: "openai-completions",
1757+
provider: "vertex-ai",
1758+
baseUrl: "http://127.0.0.1:8787/v1beta1/projects/test/locations/us/endpoints/openapi",
1759+
reasoning: true,
1760+
input: ["text"],
1761+
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
1762+
contextWindow: 1_000_000,
1763+
maxTokens: 8192,
1764+
} satisfies Model<"openai-completions">;
1765+
const output = createAssistantOutput(model);
1766+
const events: CapturedStreamEvent[] = [];
1767+
1768+
await testing.processOpenAICompletionsStream(
1769+
streamChunks([
1770+
{
1771+
id: "chatcmpl-vertex",
1772+
object: "chat.completion.chunk" as const,
1773+
created: 1775425651,
1774+
model: model.id,
1775+
choices: [
1776+
{
1777+
index: 0,
1778+
delta: { role: "assistant" as const, content: "Hi" },
1779+
logprobs: null,
1780+
finish_reason: null,
1781+
},
1782+
],
1783+
},
1784+
{
1785+
id: "chatcmpl-vertex",
1786+
object: "chat.completion.chunk" as const,
1787+
created: 1775425651,
1788+
model: model.id,
1789+
choices: [],
1790+
usage: {
1791+
prompt_tokens: 8,
1792+
completion_tokens: 25,
1793+
total_tokens: 33,
1794+
completion_tokens_details: { reasoning_tokens: 23 },
1795+
},
1796+
},
1797+
]),
1798+
output,
1799+
model,
1800+
{ push: (event) => events.push(event as CapturedStreamEvent) },
1801+
);
1802+
1803+
expect(events.map((event) => event.type)).toEqual(["text_start", "text_delta"]);
1804+
expect(output.content).toEqual([{ type: "text", text: "Hi" }]);
1805+
});
1806+
16881807
it("yields to aborts during bursty OpenAI-compatible streams", async () => {
16891808
const model = {
16901809
id: "deepseek-v4-flash",

‎src/agents/openai-transport-stream.ts‎

Lines changed: 43 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2657,6 +2657,11 @@ async function processOpenAICompletionsStream(
26572657
let sawStopFinishReason = false;
26582658
const blockIndex = () => output.content.length - 1;
26592659
const measureUtf8Bytes = (text: string) => Buffer.byteLength(text, "utf8");
2660+
let chunkPushedEvent = false;
2661+
const pushStreamEvent = (event: unknown) => {
2662+
chunkPushedEvent = true;
2663+
stream.push(event);
2664+
};
26602665
const finishCurrentBlock = () => {
26612666
if (!currentBlock) {
26622667
return;
@@ -2697,13 +2702,13 @@ async function processOpenAICompletionsStream(
26972702
currentBlock = {
26982703
type: "thinking",
26992704
thinking: "",
2700-
thinkingSignature: reasoningDelta.signature,
2705+
...(reasoningDelta.signature ? { thinkingSignature: reasoningDelta.signature } : {}),
27012706
};
27022707
output.content.push(currentBlock);
2703-
stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
2708+
pushStreamEvent({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
27042709
}
27052710
currentBlock.thinking += reasoningDelta.text;
2706-
stream.push({
2711+
pushStreamEvent({
27072712
type: "thinking_delta",
27082713
contentIndex: blockIndex(),
27092714
delta: reasoningDelta.text,
@@ -2715,10 +2720,10 @@ async function processOpenAICompletionsStream(
27152720
finishCurrentBlock();
27162721
currentBlock = { type: "text", text: "" };
27172722
output.content.push(currentBlock);
2718-
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
2723+
pushStreamEvent({ type: "text_start", contentIndex: blockIndex(), partial: output });
27192724
}
27202725
currentBlock.text += text;
2721-
stream.push({
2726+
pushStreamEvent({
27222727
type: "text_delta",
27232728
contentIndex: blockIndex(),
27242729
delta: text,
@@ -2782,12 +2787,12 @@ async function processOpenAICompletionsStream(
27822787
};
27832788
currentBlock = block;
27842789
output.content.push(block);
2785-
stream.push({
2790+
pushStreamEvent({
27862791
type: "toolcall_start",
27872792
contentIndex: output.content.indexOf(block),
27882793
partial: output,
27892794
});
2790-
stream.push({
2795+
pushStreamEvent({
27912796
type: "toolcall_delta",
27922797
contentIndex: output.content.indexOf(block),
27932798
delta: toolCall.partialArgs,
@@ -2853,6 +2858,19 @@ async function processOpenAICompletionsStream(
28532858
appendFilteredVisibleTextDelta(delta.text);
28542859
}
28552860
};
2861+
const emitReasoningUsageActivity = (hasReasoningUsageActivity: boolean) => {
2862+
if (!hasReasoningUsageActivity || chunkPushedEvent || !emitReasoning) {
2863+
return;
2864+
}
2865+
const latestBlock = output.content[output.content.length - 1];
2866+
if (currentBlock?.type === "text" || currentBlock?.type === "toolCall") {
2867+
return;
2868+
}
2869+
if (latestBlock?.type === "text" || latestBlock?.type === "toolCall") {
2870+
return;
2871+
}
2872+
appendThinkingDelta({ signature: "", text: "" });
2873+
};
28562874
const flushReasoningTagTextPartitionerAtEnd = () => {
28572875
for (const delta of reasoningTagTextPartitioner.flush()) {
28582876
appendPartitionedVisibleDelta(delta);
@@ -2861,23 +2879,28 @@ async function processOpenAICompletionsStream(
28612879
const cooperativeScheduler = createModelStreamCooperativeScheduler(options?.signal);
28622880
for await (const rawChunk of responseStream as AsyncIterable<unknown>) {
28632881
throwIfModelStreamAborted(options?.signal);
2882+
chunkPushedEvent = false;
28642883
if (!rawChunk || typeof rawChunk !== "object") {
28652884
await cooperativeScheduler.afterEvent();
28662885
continue;
28672886
}
28682887
const chunk = rawChunk as ChatCompletionChunk;
28692888
output.responseId ||= chunk.id;
2889+
let hasReasoningUsageActivity = false;
28702890
if (chunk.usage) {
28712891
output.usage = parseTransportChunkUsage(chunk.usage, model);
2892+
hasReasoningUsageActivity = hasOpenAICompletionsReasoningUsageActivity(chunk.usage);
28722893
}
28732894
const choice = Array.isArray(chunk.choices) ? chunk.choices[0] : undefined;
28742895
if (!choice) {
2896+
emitReasoningUsageActivity(hasReasoningUsageActivity);
28752897
await cooperativeScheduler.afterEvent();
28762898
continue;
28772899
}
28782900
const choiceUsage = (choice as unknown as { usage?: ChatCompletionChunk["usage"] }).usage;
28792901
if (!chunk.usage && choiceUsage) {
28802902
output.usage = parseTransportChunkUsage(choiceUsage, model);
2903+
hasReasoningUsageActivity = hasOpenAICompletionsReasoningUsageActivity(choiceUsage);
28812904
}
28822905
if (choice.finish_reason) {
28832906
const finishReasonResult = mapStopReason(choice.finish_reason);
@@ -2893,6 +2916,7 @@ async function processOpenAICompletionsStream(
28932916
choice.delta ??
28942917
(choice as unknown as { message?: ChatCompletionChunk["choices"][number]["delta"] }).message;
28952918
if (!choiceDelta) {
2919+
emitReasoningUsageActivity(hasReasoningUsageActivity);
28962920
await cooperativeScheduler.afterEvent();
28972921
continue;
28982922
}
@@ -2961,7 +2985,7 @@ async function processOpenAICompletionsStream(
29612985
...(initialSig ? { thoughtSignature: initialSig } : {}),
29622986
};
29632987
output.content.push(block);
2964-
stream.push({
2988+
pushStreamEvent({
29652989
type: "toolcall_start",
29662990
contentIndex: output.content.indexOf(block),
29672991
partial: output,
@@ -2991,7 +3015,7 @@ async function processOpenAICompletionsStream(
29913015
toolCallBlockBytes.set(block, currentBlockArgBytes + nextArgumentBytes);
29923016
block.partialArgs += toolCall.function.arguments;
29933017
block.arguments = parseStreamingJson(block.partialArgs);
2994-
stream.push({
3018+
pushStreamEvent({
29953019
type: "toolcall_delta",
29963020
contentIndex: output.content.indexOf(block),
29973021
delta: toolCall.function.arguments,
@@ -3001,6 +3025,7 @@ async function processOpenAICompletionsStream(
30013025
}
30023026
}
30033027
flushPendingPostToolCallDeltas();
3028+
emitReasoningUsageActivity(hasReasoningUsageActivity);
30043029
await cooperativeScheduler.afterEvent();
30053030
}
30063031
flushReasoningTagTextPartitionerAtEnd();
@@ -4207,6 +4232,15 @@ export function parseTransportChunkUsage(
42074232
return usage;
42084233
}
42094234

4235+
function hasOpenAICompletionsReasoningUsageActivity(
4236+
rawUsage: NonNullable<ChatCompletionChunk["usage"]>,
4237+
) {
4238+
const reasoningTokens = rawUsage.completion_tokens_details?.reasoning_tokens;
4239+
return (
4240+
typeof reasoningTokens === "number" && Number.isFinite(reasoningTokens) && reasoningTokens > 0
4241+
);
4242+
}
4243+
42104244
function mapStopReason(reason: string | null) {
42114245
if (reason === null) {
42124246
return { stopReason: "stop" };

0 commit comments

Comments
 (0)