Adds back support for partial tool calls during streaming + tests

Author: JoaquinGimenez1

.changeset/ten-sloths-cross.md

@@ -0,0 +1,5 @@
+---
+"workers-ai-provider": minor
+---
+
+Adds support for new tool call format during streaming

packages/workers-ai-provider/package.json

@@ -41,6 +41,7 @@
 		"@ai-sdk/provider": "^1.1.3"
 	},
 	"devDependencies": {
-		"@cloudflare/workers-types": "^4.20250525.0"
+		"@cloudflare/workers-types": "^4.20250525.0",
+		"zod": "^3.25.28"
 	}
 }

packages/workers-ai-provider/src/streaming.ts

@@ -2,10 +2,12 @@ import { events } from "fetch-event-stream";
 
 import type { LanguageModelV1StreamPart } from "@ai-sdk/provider";
 import { mapWorkersAIUsage } from "./map-workersai-usage";
+import { processPartialToolCalls } from "./utils";
 
 export function getMappedStream(response: Response) {
 	const chunkEvent = events(response);
 	let usage = { promptTokens: 0, completionTokens: 0 };
+	const partialToolCalls: any[] = [];
 
 	return new ReadableStream<LanguageModelV1StreamPart>({
 		async start(controller) {
@@ -20,12 +22,26 @@ export function getMappedStream(response: Response) {
 				if (chunk.usage) {
 					usage = mapWorkersAIUsage(chunk);
 				}
+				if (chunk.tool_calls) {
+					partialToolCalls.push(...chunk.tool_calls);
+				}
 				chunk.response?.length &&
 					controller.enqueue({
 						type: "text-delta",
 						textDelta: chunk.response,
 					});
 			}
+
+			if (partialToolCalls.length > 0) {
+				const toolCalls = processPartialToolCalls(partialToolCalls);
+				toolCalls.map((toolCall) => {
+					controller.enqueue({
+						type: "tool-call",
+						...toolCall,
+					});
+				});
+			}
+
 			controller.enqueue({
 				type: "finish",
 				finishReason: "stop",

packages/workers-ai-provider/src/utils.ts

@@ -1,4 +1,4 @@
-import type { LanguageModelV1 } from "@ai-sdk/provider";
+import type { LanguageModelV1, LanguageModelV1FunctionToolCall } from "@ai-sdk/provider";
 
 /**
  * General AI run interface with overloads to handle distinct return types.
@@ -182,33 +182,79 @@ export function lastMessageWasUser<T extends { role: string }>(messages: T[]) {
 	return messages.length > 0 && messages[messages.length - 1]!.role === "user";
 }
 
-export function processToolCalls(output: any) {
+function mergePartialToolCalls(partialCalls: any[]) {
+	const mergedCallsByIndex: any = {};
+
+	for (const partialCall of partialCalls) {
+		const index = partialCall.index;
+
+		if (!mergedCallsByIndex[index]) {
+			mergedCallsByIndex[index] = {
+				id: partialCall.id || "",
+				type: partialCall.type || "",
+				function: {
+					name: partialCall.function?.name || "",
+					arguments: "",
+				},
+			};
+		} else {
+			if (partialCall.id) {
+				mergedCallsByIndex[index].id = partialCall.id;
+			}
+			if (partialCall.type) {
+				mergedCallsByIndex[index].type = partialCall.type;
+			}
+
+			if (partialCall.function?.name) {
+				mergedCallsByIndex[index].function.name = partialCall.function.name;
+			}
+		}
+
+		// Append arguments if available, this assumes arguments come in the right order
+		if (partialCall.function?.arguments) {
+			mergedCallsByIndex[index].function.arguments += partialCall.function.arguments;
+		}
+	}
+
+	return Object.values(mergedCallsByIndex);
+}
+
+function processToolCall(toolCall: any): LanguageModelV1FunctionToolCall {
 	// Check for OpenAI format tool calls first
+	if (toolCall.function && toolCall.id) {
+		return {
+			toolCallType: "function",
+			toolCallId: toolCall.id,
+			toolName: toolCall.function.name,
+			args:
+				typeof toolCall.function.arguments === "string"
+					? toolCall.function.arguments
+					: JSON.stringify(toolCall.function.arguments || {}),
+		};
+	}
+	return {
+		toolCallType: "function",
+		toolCallId: toolCall.name,
+		toolName: toolCall.name,
+		args:
+			typeof toolCall.arguments === "string"
+				? toolCall.arguments
+				: JSON.stringify(toolCall.arguments || {}),
+	};
+}
+
+export function processToolCalls(output: any): LanguageModelV1FunctionToolCall[] {
 	if (output.tool_calls && Array.isArray(output.tool_calls)) {
 		return output.tool_calls.map((toolCall: any) => {
-			// Handle new format
-			if (toolCall.function && toolCall.id) {
-				return {
-					toolCallType: "function",
-					toolCallId: toolCall.id,
-					toolName: toolCall.function.name,
-					args:
-						typeof toolCall.function.arguments === "string"
-							? toolCall.function.arguments
-							: JSON.stringify(toolCall.function.arguments || {}),
-				};
-			}
-			return {
-				toolCallType: "function",
-				toolCallId: toolCall.name,
-				toolName: toolCall.name,
-				args:
-					typeof toolCall.arguments === "string"
-						? toolCall.arguments
-						: JSON.stringify(toolCall.arguments || {}),
-			};
+			const processedToolCall = processToolCall(toolCall);
+			return processedToolCall;
 		});
 	}
 
 	return [];
 }
+
+export function processPartialToolCalls(partialToolCalls: any[]) {
+	const mergedToolCalls = mergePartialToolCalls(partialToolCalls);
+	return processToolCalls({ tool_calls: mergedToolCalls });
+}