From 01509156b31c645e83374adf78e5e5bc9489e0ed Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 23 Apr 2026 16:05:21 +0200 Subject: [PATCH] fix(ai): coalesce malformed streamed tool calls closes #3576 --- packages/ai/CHANGELOG.md | 1 + .../ai/src/providers/openai-completions.ts | 71 +++++++---- .../openai-completions-tool-choice.test.ts | 111 ++++++++++++++++++ packages/ai/test/stream.test.ts | 4 +- 4 files changed, 162 insertions(+), 25 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 6fff07ba..9af07a27 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed +- Fixed `openai-completions` streamed tool-call assembly to coalesce deltas by stable tool index when OpenAI-compatible gateways mutate tool call IDs mid-stream, preventing malformed Kimi K2.6/OpenCode tool streams from splitting one call into multiple bogus tool calls ([#3576](https://github.com/badlogic/pi-mono/issues/3576)) - Fixed built-in `kimi-coding` model generation to attach `User-Agent: KimiCLI/1.5` to all generated Kimi models, overriding the Anthropic SDK default UA so direct Kimi Coding requests use the provider's expected client identity ([#3586](https://github.com/badlogic/pi-mono/issues/3586)) ## [0.69.0] - 2026-04-22 diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index bcfb8c27..2c48b195 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -150,33 +150,44 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); - let currentBlock: TextContent | ThinkingContent | (ToolCall & { partialArgs?: string }) | null = null; + interface StreamingToolCallBlock extends ToolCall { + partialArgs?: string; + streamIndex?: number; + } + + let currentBlock: TextContent | ThinkingContent | StreamingToolCallBlock | null = null; const blocks = output.content; - const blockIndex = () => blocks.length - 1; + const getContentIndex = (block: typeof currentBlock) => (block ? blocks.indexOf(block) : -1); + const currentContentIndex = () => getContentIndex(currentBlock); const finishCurrentBlock = (block?: typeof currentBlock) => { if (block) { + const contentIndex = getContentIndex(block); + if (contentIndex === -1) { + return; + } if (block.type === "text") { stream.push({ type: "text_end", - contentIndex: blockIndex(), + contentIndex, content: block.text, partial: output, }); } else if (block.type === "thinking") { stream.push({ type: "thinking_end", - contentIndex: blockIndex(), + contentIndex, content: block.thinking, partial: output, }); } else if (block.type === "toolCall") { block.arguments = parseStreamingJson(block.partialArgs); - // Finalize in-place and strip the scratch buffer so replay only + // Finalize in-place and strip the scratch buffers so replay only // carries parsed arguments. delete block.partialArgs; + delete block.streamIndex; stream.push({ type: "toolcall_end", - contentIndex: blockIndex(), + contentIndex, toolCall: block, partial: output, }); @@ -221,14 +232,14 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA finishCurrentBlock(currentBlock); currentBlock = { type: "text", text: "" }; output.content.push(currentBlock); - stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output }); + stream.push({ type: "text_start", contentIndex: currentContentIndex(), partial: output }); } if (currentBlock.type === "text") { currentBlock.text += choice.delta.content; stream.push({ type: "text_delta", - contentIndex: blockIndex(), + contentIndex: currentContentIndex(), delta: choice.delta.content, partial: output, }); @@ -263,7 +274,7 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA thinkingSignature: foundReasoningField, }; output.content.push(currentBlock); - stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output }); + stream.push({ type: "thinking_start", contentIndex: currentContentIndex(), partial: output }); } if (currentBlock.type === "thinking") { @@ -271,7 +282,7 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA currentBlock.thinking += delta; stream.push({ type: "thinking_delta", - contentIndex: blockIndex(), + contentIndex: currentContentIndex(), delta, partial: output, }); @@ -280,11 +291,13 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA if (choice?.delta?.tool_calls) { for (const toolCall of choice.delta.tool_calls) { - if ( - !currentBlock || - currentBlock.type !== "toolCall" || - (toolCall.id && currentBlock.id !== toolCall.id) - ) { + const streamIndex = typeof toolCall.index === "number" ? toolCall.index : undefined; + const sameToolCall = + currentBlock?.type === "toolCall" && + ((streamIndex !== undefined && currentBlock.streamIndex === streamIndex) || + (streamIndex === undefined && toolCall.id && currentBlock.id === toolCall.id)); + + if (!sameToolCall) { finishCurrentBlock(currentBlock); currentBlock = { type: "toolCall", @@ -292,23 +305,34 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA name: toolCall.function?.name || "", arguments: {}, partialArgs: "", + streamIndex, }; output.content.push(currentBlock); - stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output }); + stream.push({ + type: "toolcall_start", + contentIndex: getContentIndex(currentBlock), + partial: output, + }); } - if (currentBlock.type === "toolCall") { - if (toolCall.id) currentBlock.id = toolCall.id; - if (toolCall.function?.name) currentBlock.name = toolCall.function.name; + const currentToolCallBlock = currentBlock?.type === "toolCall" ? currentBlock : null; + if (currentToolCallBlock) { + if (!currentToolCallBlock.id && toolCall.id) currentToolCallBlock.id = toolCall.id; + if (!currentToolCallBlock.name && toolCall.function?.name) { + currentToolCallBlock.name = toolCall.function.name; + } + if (currentToolCallBlock.streamIndex === undefined && streamIndex !== undefined) { + currentToolCallBlock.streamIndex = streamIndex; + } let delta = ""; if (toolCall.function?.arguments) { delta = toolCall.function.arguments; - currentBlock.partialArgs += toolCall.function.arguments; - currentBlock.arguments = parseStreamingJson(currentBlock.partialArgs); + currentToolCallBlock.partialArgs += toolCall.function.arguments; + currentToolCallBlock.arguments = parseStreamingJson(currentToolCallBlock.partialArgs); } stream.push({ type: "toolcall_delta", - contentIndex: blockIndex(), + contentIndex: getContentIndex(currentToolCallBlock), delta, partial: output, }); @@ -349,8 +373,9 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA } catch (error) { for (const block of output.content) { delete (block as { index?: number }).index; - // partialArgs is only a streaming scratch buffer; never persist it. + // Streaming scratch buffers are only used during parsing; never persist them. delete (block as { partialArgs?: string }).partialArgs; + delete (block as { streamIndex?: number }).streamIndex; } output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error); diff --git a/packages/ai/test/openai-completions-tool-choice.test.ts b/packages/ai/test/openai-completions-tool-choice.test.ts index fdd92d48..4c22b84d 100644 --- a/packages/ai/test/openai-completions-tool-choice.test.ts +++ b/packages/ai/test/openai-completions-tool-choice.test.ts @@ -441,6 +441,117 @@ describe("openai-completions tool_choice", () => { expect(response.content).toEqual([{ type: "text", text: "OK" }]); }); + it("coalesces tool call deltas by stable index when provider mutates ids mid-stream", async () => { + mockState.chunks = [ + { + id: "chatcmpl-kimi-bad-stream", + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: "functions.read:0", + type: "function", + function: { name: "read", arguments: "" }, + }, + ], + }, + finish_reason: null, + }, + ], + }, + { + id: "chatcmpl-kimi-bad-stream", + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: "chatcmpl-tool-a", + type: "function", + function: { name: null, arguments: '{"path":"README' }, + }, + ], + }, + finish_reason: null, + }, + ], + }, + { + id: "chatcmpl-kimi-bad-stream", + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: "chatcmpl-tool-b", + type: "function", + function: { name: null, arguments: '.md"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + prompt_tokens_details: { cached_tokens: 0 }, + completion_tokens_details: { reasoning_tokens: 0 }, + }, + }, + ]; + + const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!; + const model = { ...baseModel, api: "openai-completions" } as const; + const tool: Tool = { + name: "read", + description: "Read a file", + parameters: Type.Object({ + path: Type.String(), + }), + }; + const s = streamSimple( + model, + { + messages: [ + { + role: "user", + content: "Read README.md", + timestamp: Date.now(), + }, + ], + tools: [tool], + }, + { apiKey: "test" }, + ); + + const toolCallContentIndexes: number[] = []; + for await (const event of s) { + if (event.type === "toolcall_start" || event.type === "toolcall_delta" || event.type === "toolcall_end") { + toolCallContentIndexes.push(event.contentIndex); + } + } + + const response = await s.result(); + expect(response.stopReason).toBe("toolUse"); + expect(toolCallContentIndexes).toEqual([0, 0, 0, 0, 0]); + expect(response.content).toHaveLength(1); + const toolCall = response.content[0]; + expect(toolCall.type).toBe("toolCall"); + if (toolCall.type !== "toolCall") { + throw new Error("Expected toolCall content"); + } + expect(toolCall.id).toBe("functions.read:0"); + expect(toolCall.name).toBe("read"); + expect(toolCall.arguments).toEqual({ path: "README.md" }); + expect(toolCall).not.toHaveProperty("streamIndex"); + expect(toolCall).not.toHaveProperty("partialArgs"); + }); + it("preserves prompt_tokens_details.cache_write_tokens from chunk usage", async () => { mockState.chunks = [ { diff --git a/packages/ai/test/stream.test.ts b/packages/ai/test/stream.test.ts index 2066ea69..4bfdedc7 100644 --- a/packages/ai/test/stream.test.ts +++ b/packages/ai/test/stream.test.ts @@ -1092,8 +1092,8 @@ describe("Generate E2E Tests", () => { }); }); - describe("OpenAI Codex Provider (gpt-5.2-codex)", () => { - const llm = getModel("openai-codex", "gpt-5.2-codex"); + describe("OpenAI Codex Provider (gpt-5.4)", () => { + const llm = getModel("openai-codex", "gpt-5.4"); it.skipIf(!openaiCodexToken)("should complete basic text generation", { retry: 3 }, async () => { await basicTextGeneration(llm, { apiKey: openaiCodexToken });