Merge pull request #4247 from badlogic/separate-accumulators

fix(ai): handle mixed chat completion deltas

packages/ai/src/providers/openai-completions.ts

@@ -160,45 +160,104 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
 				partialArgs?: string;
 				streamIndex?: number;
 			}
+			type StreamingBlock = TextContent | ThinkingContent | StreamingToolCallBlock;
+			type StreamingToolCallDelta = NonNullable<ChatCompletionChunk.Choice.Delta["tool_calls"]>[number];
 
-			let currentBlock: TextContent | ThinkingContent | StreamingToolCallBlock | null = null;
+			let textBlock: TextContent | null = null;
-			const blocks = output.content;
+			let thinkingBlock: ThinkingContent | null = null;
-			const getContentIndex = (block: typeof currentBlock) => (block ? blocks.indexOf(block) : -1);
+			const toolCallBlocksByIndex = new Map<number, StreamingToolCallBlock>();
-			const currentContentIndex = () => getContentIndex(currentBlock);
+			const toolCallBlocksById = new Map<string, StreamingToolCallBlock>();
-			const finishCurrentBlock = (block?: typeof currentBlock) => {
+			const blocks = output.content as StreamingBlock[];
-				if (block) {
+			const getContentIndex = (block: StreamingBlock) => blocks.indexOf(block);
-					const contentIndex = getContentIndex(block);
+			const finishBlock = (block: StreamingBlock) => {
-					if (contentIndex === -1) {
+				const contentIndex = getContentIndex(block);
-						return;
+				if (contentIndex === -1) {
-					}
+					return;
-					if (block.type === "text") {
-						stream.push({
-							type: "text_end",
-							contentIndex,
-							content: block.text,
-							partial: output,
-						});
-					} else if (block.type === "thinking") {
-						stream.push({
-							type: "thinking_end",
-							contentIndex,
-							content: block.thinking,
-							partial: output,
-						});
-					} else if (block.type === "toolCall") {
-						block.arguments = parseStreamingJson(block.partialArgs);
-						// 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,
-							toolCall: block,
-							partial: output,
-						});
-					}
 				}
+				if (block.type === "text") {
+					stream.push({
+						type: "text_end",
+						contentIndex,
+						content: block.text,
+						partial: output,
+					});
+				} else if (block.type === "thinking") {
+					stream.push({
+						type: "thinking_end",
+						contentIndex,
+						content: block.thinking,
+						partial: output,
+					});
+				} else if (block.type === "toolCall") {
+					block.arguments = parseStreamingJson(block.partialArgs);
+					// 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,
+						toolCall: block,
+						partial: output,
+					});
+				}
+			};
+			const ensureTextBlock = () => {
+				if (!textBlock) {
+					textBlock = { type: "text", text: "" };
+					blocks.push(textBlock);
+					stream.push({ type: "text_start", contentIndex: getContentIndex(textBlock), partial: output });
+				}
+				return textBlock;
+			};
+			const ensureThinkingBlock = (thinkingSignature: string) => {
+				if (!thinkingBlock) {
+					thinkingBlock = {
+						type: "thinking",
+						thinking: "",
+						thinkingSignature,
+					};
+					blocks.push(thinkingBlock);
+					stream.push({ type: "thinking_start", contentIndex: getContentIndex(thinkingBlock), partial: output });
+				}
+				return thinkingBlock;
+			};
+			const ensureToolCallBlock = (toolCall: StreamingToolCallDelta) => {
+				const streamIndex = typeof toolCall.index === "number" ? toolCall.index : undefined;
+				let block = streamIndex !== undefined ? toolCallBlocksByIndex.get(streamIndex) : undefined;
+				if (!block && toolCall.id) {
+					block = toolCallBlocksById.get(toolCall.id);
+				}
+				if (!block) {
+					block = {
+						type: "toolCall",
+						id: toolCall.id || "",
+						name: toolCall.function?.name || "",
+						arguments: {},
+						partialArgs: "",
+						streamIndex,
+					};
+					if (streamIndex !== undefined) {
+						toolCallBlocksByIndex.set(streamIndex, block);
+					}
+					if (toolCall.id) {
+						toolCallBlocksById.set(toolCall.id, block);
+					}
+					blocks.push(block);
+					stream.push({
+						type: "toolcall_start",
+						contentIndex: getContentIndex(block),
+						partial: output,
+					});
+				}
+				if (streamIndex !== undefined && block.streamIndex === undefined) {
+					block.streamIndex = streamIndex;
+					toolCallBlocksByIndex.set(streamIndex, block);
+				}
+				if (toolCall.id) {
+					toolCallBlocksById.set(toolCall.id, block);
+				}
+				return block;
 			};
 
 			for await (const chunk of openaiStream) {
@@ -237,22 +296,14 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
 						choice.delta.content !== undefined &&
 						choice.delta.content.length > 0
 					) {
-						if (!currentBlock || currentBlock.type !== "text") {
+						const block = ensureTextBlock();
-							finishCurrentBlock(currentBlock);
+						block.text += choice.delta.content;
-							currentBlock = { type: "text", text: "" };
+						stream.push({
-							output.content.push(currentBlock);
+							type: "text_delta",
-							stream.push({ type: "text_start", contentIndex: currentContentIndex(), partial: output });
+							contentIndex: getContentIndex(block),
-						}
+							delta: choice.delta.content,
-
+							partial: output,
-						if (currentBlock.type === "text") {
+						});
-							currentBlock.text += choice.delta.content;
-							stream.push({
-								type: "text_delta",
-								contentIndex: currentContentIndex(),
-								delta: choice.delta.content,
-								partial: output,
-							});
-						}
 					}
 
 					// Some endpoints return reasoning in reasoning_content (llama.cpp),
@@ -260,38 +311,24 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
 					// Use the first non-empty reasoning field to avoid duplication
 					// (e.g., chutes.ai returns both reasoning_content and reasoning with same content)
 					const reasoningFields = ["reasoning_content", "reasoning", "reasoning_text"];
+					const deltaFields = choice.delta as Record<string, unknown>;
 					let foundReasoningField: string | null = null;
 					for (const field of reasoningFields) {
-						if (
+						const value = deltaFields[field];
-							(choice.delta as any)[field] !== null &&
+						if (typeof value === "string" && value.length > 0) {
-							(choice.delta as any)[field] !== undefined &&
+							foundReasoningField = field;
-							(choice.delta as any)[field].length > 0
+							break;
-						) {
-							if (!foundReasoningField) {
-								foundReasoningField = field;
-								break;
-							}
 						}
 					}
 
 					if (foundReasoningField) {
-						if (!currentBlock || currentBlock.type !== "thinking") {
+						const delta = deltaFields[foundReasoningField];
-							finishCurrentBlock(currentBlock);
+						if (typeof delta === "string" && delta.length > 0) {
-							currentBlock = {
+							const block = ensureThinkingBlock(foundReasoningField);
-								type: "thinking",
+							block.thinking += delta;
-								thinking: "",
-								thinkingSignature: foundReasoningField,
-							};
-							output.content.push(currentBlock);
-							stream.push({ type: "thinking_start", contentIndex: currentContentIndex(), partial: output });
-						}
-
-						if (currentBlock.type === "thinking") {
-							const delta = (choice.delta as any)[foundReasoningField];
-							currentBlock.thinking += delta;
 							stream.push({
 								type: "thinking_delta",
-								contentIndex: currentContentIndex(),
+								contentIndex: getContentIndex(block),
 								delta,
 								partial: output,
 							});
@@ -300,52 +337,27 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
 
 					if (choice?.delta?.tool_calls) {
 						for (const toolCall of choice.delta.tool_calls) {
-							const streamIndex = typeof toolCall.index === "number" ? toolCall.index : undefined;
+							const block = ensureToolCallBlock(toolCall);
-							const sameToolCall =
+							if (!block.id && toolCall.id) {
-								currentBlock?.type === "toolCall" &&
+								block.id = toolCall.id;
-								((streamIndex !== undefined && currentBlock.streamIndex === streamIndex) ||
+								toolCallBlocksById.set(toolCall.id, block);
-									(streamIndex === undefined && toolCall.id && currentBlock.id === toolCall.id));
+							}
-
+							if (!block.name && toolCall.function?.name) {
-							if (!sameToolCall) {
+								block.name = toolCall.function.name;
-								finishCurrentBlock(currentBlock);
-								currentBlock = {
-									type: "toolCall",
-									id: toolCall.id || "",
-									name: toolCall.function?.name || "",
-									arguments: {},
-									partialArgs: "",
-									streamIndex,
-								};
-								output.content.push(currentBlock);
-								stream.push({
-									type: "toolcall_start",
-									contentIndex: getContentIndex(currentBlock),
-									partial: output,
-								});
 							}
 
-							const currentToolCallBlock = currentBlock?.type === "toolCall" ? currentBlock : null;
+							let delta = "";
-							if (currentToolCallBlock) {
+							if (toolCall.function?.arguments) {
-								if (!currentToolCallBlock.id && toolCall.id) currentToolCallBlock.id = toolCall.id;
+								delta = toolCall.function.arguments;
-								if (!currentToolCallBlock.name && toolCall.function?.name) {
+								block.partialArgs = (block.partialArgs ?? "") + toolCall.function.arguments;
-									currentToolCallBlock.name = toolCall.function.name;
+								block.arguments = parseStreamingJson(block.partialArgs);
-								}
-								if (currentToolCallBlock.streamIndex === undefined && streamIndex !== undefined) {
-									currentToolCallBlock.streamIndex = streamIndex;
-								}
-								let delta = "";
-								if (toolCall.function?.arguments) {
-									delta = toolCall.function.arguments;
-									currentToolCallBlock.partialArgs += toolCall.function.arguments;
-									currentToolCallBlock.arguments = parseStreamingJson(currentToolCallBlock.partialArgs);
-								}
-								stream.push({
-									type: "toolcall_delta",
-									contentIndex: getContentIndex(currentToolCallBlock),
-									delta,
-									partial: output,
-								});
 							}
+							stream.push({
+								type: "toolcall_delta",
+								contentIndex: getContentIndex(block),
+								delta,
+								partial: output,
+							});
 						}
 					}
 
@@ -365,7 +377,9 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
 				}
 			}
 
-			finishCurrentBlock(currentBlock);
+			for (const block of blocks) {
+				finishBlock(block);
+			}
 			if (options?.signal?.aborted) {
 				throw new Error("Request was aborted");
 			}

packages/ai/test/openai-completions-tool-choice.test.ts

@@ -552,6 +552,238 @@ describe("openai-completions tool_choice", () => {
 		expect(toolCall).not.toHaveProperty("partialArgs");
 	});
 
+	it("accumulates mixed content, reasoning, and parallel tool call deltas independently", async () => {
+		mockState.chunks = [
+			{
+				id: "chatcmpl-mixed-deltas",
+				choices: [
+					{
+						delta: {
+							content: "answer 1",
+							reasoning_content: "think 1",
+							tool_calls: [
+								{
+									index: 0,
+									id: "tc_read_initial",
+									type: "function",
+									function: { name: "read", arguments: '{"path":"README' },
+								},
+								{
+									index: 1,
+									id: "tc_grep_initial",
+									type: "function",
+									function: { name: "grep", arguments: '{"pattern":"TODO' },
+								},
+								{
+									id: "tc_list_no_index",
+									type: "function",
+									function: { name: "list", arguments: '{"path":"packages' },
+								},
+								{
+									id: "tc_write_no_index",
+									type: "function",
+									function: { name: "write", arguments: '{"path":"out' },
+								},
+							],
+						},
+						finish_reason: null,
+					},
+				],
+			},
+			{
+				id: "chatcmpl-mixed-deltas",
+				choices: [
+					{
+						delta: {
+							content: " answer 2",
+							tool_calls: [
+								{
+									index: 1,
+									id: "tc_grep_changed",
+									type: "function",
+									function: { arguments: '","path":"src' },
+								},
+								{
+									id: "tc_write_no_index",
+									type: "function",
+									function: { arguments: '.txt","content":"ok"}' },
+								},
+								{
+									id: "tc_list_no_index",
+									type: "function",
+									function: { arguments: '/ai"}' },
+								},
+							],
+						},
+						finish_reason: null,
+					},
+				],
+			},
+			{
+				id: "chatcmpl-mixed-deltas",
+				choices: [
+					{
+						delta: {
+							content: "\n",
+							reasoning_content: " think 2",
+							tool_calls: [
+								{
+									index: 0,
+									id: "tc_read_changed",
+									type: "function",
+									function: { arguments: '.md"}' },
+								},
+								{
+									index: 1,
+									type: "function",
+									function: { arguments: '"}' },
+								},
+							],
+						},
+						finish_reason: "tool_calls",
+					},
+				],
+				usage: {
+					prompt_tokens: 10,
+					completion_tokens: 8,
+					prompt_tokens_details: { cached_tokens: 0 },
+					completion_tokens_details: { reasoning_tokens: 2 },
+				},
+			},
+		];
+
+		const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
+		const model = { ...baseModel, api: "openai-completions" } as const;
+		const tools: Tool[] = [
+			{
+				name: "read",
+				description: "Read a file",
+				parameters: Type.Object({ path: Type.String() }),
+			},
+			{
+				name: "grep",
+				description: "Search a file",
+				parameters: Type.Object({ pattern: Type.String(), path: Type.String() }),
+			},
+			{
+				name: "list",
+				description: "List a directory",
+				parameters: Type.Object({ path: Type.String() }),
+			},
+			{
+				name: "write",
+				description: "Write a file",
+				parameters: Type.Object({ path: Type.String(), content: Type.String() }),
+			},
+		];
+		const s = streamSimple(
+			model,
+			{
+				messages: [
+					{
+						role: "user",
+						content: "Think, answer, and use tools.",
+						timestamp: Date.now(),
+					},
+				],
+				tools,
+			},
+			{ apiKey: "test" },
+		);
+
+		const eventTypes: string[] = [];
+		const toolEventsByContentIndex = new Map<number, string[]>();
+		for await (const event of s) {
+			eventTypes.push(event.type);
+			if (event.type === "toolcall_start" || event.type === "toolcall_delta" || event.type === "toolcall_end") {
+				const events = toolEventsByContentIndex.get(event.contentIndex) ?? [];
+				events.push(event.type);
+				toolEventsByContentIndex.set(event.contentIndex, events);
+			}
+		}
+
+		const response = await s.result();
+		expect(response.stopReason).toBe("toolUse");
+		expect(eventTypes.filter((type) => type === "text_start")).toHaveLength(1);
+		expect(eventTypes.filter((type) => type === "text_delta")).toHaveLength(3);
+		expect(eventTypes.filter((type) => type === "text_end")).toHaveLength(1);
+		expect(eventTypes.filter((type) => type === "thinking_start")).toHaveLength(1);
+		expect(eventTypes.filter((type) => type === "thinking_delta")).toHaveLength(2);
+		expect(eventTypes.filter((type) => type === "thinking_end")).toHaveLength(1);
+		expect(eventTypes.filter((type) => type === "toolcall_start")).toHaveLength(4);
+		expect(eventTypes.filter((type) => type === "toolcall_delta")).toHaveLength(9);
+		expect(eventTypes.filter((type) => type === "toolcall_end")).toHaveLength(4);
+		expect(toolEventsByContentIndex.get(2)).toEqual([
+			"toolcall_start",
+			"toolcall_delta",
+			"toolcall_delta",
+			"toolcall_end",
+		]);
+		expect(toolEventsByContentIndex.get(3)).toEqual([
+			"toolcall_start",
+			"toolcall_delta",
+			"toolcall_delta",
+			"toolcall_delta",
+			"toolcall_end",
+		]);
+		expect(toolEventsByContentIndex.get(4)).toEqual([
+			"toolcall_start",
+			"toolcall_delta",
+			"toolcall_delta",
+			"toolcall_end",
+		]);
+		expect(toolEventsByContentIndex.get(5)).toEqual([
+			"toolcall_start",
+			"toolcall_delta",
+			"toolcall_delta",
+			"toolcall_end",
+		]);
+
+		expect(response.content).toHaveLength(6);
+		expect(response.content[0]).toEqual({ type: "text", text: "answer 1 answer 2\n" });
+		expect(response.content[1]).toEqual({
+			type: "thinking",
+			thinking: "think 1 think 2",
+			thinkingSignature: "reasoning_content",
+		});
+		const readCall = response.content[2];
+		const grepCall = response.content[3];
+		const listCall = response.content[4];
+		const writeCall = response.content[5];
+		expect(readCall.type).toBe("toolCall");
+		expect(grepCall.type).toBe("toolCall");
+		expect(listCall.type).toBe("toolCall");
+		expect(writeCall.type).toBe("toolCall");
+		if (
+			readCall.type !== "toolCall" ||
+			grepCall.type !== "toolCall" ||
+			listCall.type !== "toolCall" ||
+			writeCall.type !== "toolCall"
+		) {
+			throw new Error("Expected toolCall content");
+		}
+		expect(readCall.id).toBe("tc_read_initial");
+		expect(readCall.name).toBe("read");
+		expect(readCall.arguments).toEqual({ path: "README.md" });
+		expect(readCall).not.toHaveProperty("streamIndex");
+		expect(readCall).not.toHaveProperty("partialArgs");
+		expect(grepCall.id).toBe("tc_grep_initial");
+		expect(grepCall.name).toBe("grep");
+		expect(grepCall.arguments).toEqual({ pattern: "TODO", path: "src" });
+		expect(grepCall).not.toHaveProperty("streamIndex");
+		expect(grepCall).not.toHaveProperty("partialArgs");
+		expect(listCall.id).toBe("tc_list_no_index");
+		expect(listCall.name).toBe("list");
+		expect(listCall.arguments).toEqual({ path: "packages/ai" });
+		expect(listCall).not.toHaveProperty("streamIndex");
+		expect(listCall).not.toHaveProperty("partialArgs");
+		expect(writeCall.id).toBe("tc_write_no_index");
+		expect(writeCall.name).toBe("write");
+		expect(writeCall.arguments).toEqual({ path: "out.txt", content: "ok" });
+		expect(writeCall).not.toHaveProperty("streamIndex");
+		expect(writeCall).not.toHaveProperty("partialArgs");
+	});
+
 	it("does not double-count reasoning tokens in completion usage", async () => {
 		mockState.chunks = [
 			{