Skip to content

Commit cc4b2c0

Browse files
author
Mihir Gada
committed
Fix Workers AI tool call ID collisions
1 parent 630180b commit cc4b2c0

7 files changed

Lines changed: 289 additions & 38 deletions

File tree

‎.changeset/fix-workers-ai-toolcallid-collisions.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"workers-ai-provider": patch
3+
---
4+
5+
Rewrite Workers AI tool call IDs exposed to the AI SDK and restore the original provider IDs when building follow-up prompts to avoid collisions across chat turns.

‎packages/workers-ai-provider/src/convert-to-workersai-chat-messages.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { LanguageModelV3DataContent, LanguageModelV3Prompt } from "@ai-sdk/provider";
2+
import { toWorkersAIToolCallId } from "./utils";
23
import type { WorkersAIContentPart, WorkersAIChatPrompt } from "./workersai-chat-prompt";
34

45
/**
@@ -144,7 +145,7 @@ export function convertToWorkersAIChatMessages(prompt: LanguageModelV3Prompt): {
144145
arguments: JSON.stringify(part.input),
145146
name: part.toolName,
146147
},
147-
id: part.toolCallId,
148+
id: toWorkersAIToolCallId(part.toolCallId),
148149
type: "function",
149150
});
150151
break;
@@ -216,7 +217,7 @@ export function convertToWorkersAIChatMessages(prompt: LanguageModelV3Prompt): {
216217
messages.push({
217218
content,
218219
name: toolResponse.toolName,
219-
tool_call_id: toolResponse.toolCallId,
220+
tool_call_id: toWorkersAIToolCallId(toolResponse.toolCallId),
220221
role: "tool",
221222
});
222223
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import type {
66
import { generateId } from "ai";
77
import { mapWorkersAIFinishReason } from "./map-workersai-finish-reason";
88
import { mapWorkersAIUsage } from "./map-workersai-usage";
9+
import { createAISDKToolCallId } from "./utils";
910

1011
/**
1112
* Prepend a stream-start event to an existing LanguageModelV3 stream.
@@ -317,7 +318,7 @@ export function getMappedStream(
317318
closeToolCall(lastActiveToolIndex, controller);
318319
}
319320

320-
const id = tcId || generateId();
321+
const id = createAISDKToolCallId(tcId);
321322
const toolName = tcName || "";
322323
activeToolCalls.set(tcIndex, { id, toolName, args: "" });
323324
lastActiveToolIndex = tcIndex;

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

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -323,6 +323,23 @@ export function prepareToolsAndToolChoice(
323323
// Tool call processing
324324
// ---------------------------------------------------------------------------
325325

326+
const TOOL_CALL_ID_MARKER = "::cf-wai-tool-call::";
327+
328+
export function createAISDKToolCallId(toolCallId: string | null | undefined): string {
329+
const originalId = toolCallId || generateId();
330+
return `${originalId}${TOOL_CALL_ID_MARKER}${generateId()}`;
331+
}
332+
333+
export function toWorkersAIToolCallId(toolCallId: string): string {
334+
const markerIndex = toolCallId.lastIndexOf(TOOL_CALL_ID_MARKER);
335+
if (markerIndex === -1) return toolCallId;
336+
337+
const suffixIndex = markerIndex + TOOL_CALL_ID_MARKER.length;
338+
if (suffixIndex >= toolCallId.length) return toolCallId;
339+
340+
return toolCallId.slice(0, markerIndex);
341+
}
342+
326343
/** Workers AI flat tool call format (non-streaming, native) */
327344
interface FlatToolCall {
328345
name: string;
@@ -406,7 +423,7 @@ function processToolCall(toolCall: FlatToolCall | OpenAIToolCall): LanguageModel
406423
typeof fn.arguments === "string"
407424
? fn.arguments
408425
: JSON.stringify(fn.arguments || {}),
409-
toolCallId: toolCall.id || generateId(),
426+
toolCallId: createAISDKToolCallId(toolCall.id),
410427
type: "tool-call",
411428
toolName: fn.name,
412429
};
@@ -419,7 +436,7 @@ function processToolCall(toolCall: FlatToolCall | OpenAIToolCall): LanguageModel
419436
typeof flat.arguments === "string"
420437
? flat.arguments
421438
: JSON.stringify(flat.arguments || {}),
422-
toolCallId: flat.id || generateId(),
439+
toolCallId: createAISDKToolCallId(flat.id),
423440
type: "tool-call",
424441
toolName: flat.name,
425442
};

‎packages/workers-ai-provider/test/convert-to-workersai-chat-messages.test.ts‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { describe, it, expect } from "vitest";
22
import { convertToWorkersAIChatMessages } from "../src/convert-to-workersai-chat-messages";
3+
import { createAISDKToolCallId } from "../src/utils";
34

45
describe("convertToWorkersAIChatMessages", () => {
56
describe("tool call ID preservation", () => {
@@ -152,6 +153,82 @@ describe("convertToWorkersAIChatMessages", () => {
152153
expect(messages[0].tool_calls).toBeDefined();
153154
expect(messages[0].tool_calls![0].id).toBe(originalId);
154155
});
156+
157+
it("should restore rewritten assistant tool call IDs before sending to Workers AI", () => {
158+
const originalId = "functions.list_toolbox_tools:0";
159+
const rewrittenId = createAISDKToolCallId(originalId);
160+
161+
const { messages } = convertToWorkersAIChatMessages([
162+
{
163+
role: "assistant" as const,
164+
content: [
165+
{
166+
type: "tool-call" as const,
167+
toolCallId: rewrittenId,
168+
toolName: "list_toolbox_tools",
169+
input: {},
170+
},
171+
],
172+
},
173+
]);
174+
175+
expect(messages[0].tool_calls![0].id).toBe(originalId);
176+
});
177+
178+
it("should restore rewritten tool result IDs for tool error outputs", () => {
179+
const originalId = "functions.invoke_toolbox_tool:1";
180+
const rewrittenId = createAISDKToolCallId(originalId);
181+
182+
const { messages } = convertToWorkersAIChatMessages([
183+
{
184+
role: "tool" as const,
185+
content: [
186+
{
187+
type: "tool-result" as const,
188+
toolCallId: rewrittenId,
189+
toolName: "invoke_toolbox_tool",
190+
output: { type: "error-text", value: "tool failed" } as any,
191+
},
192+
],
193+
},
194+
]);
195+
196+
expect(messages[0].tool_call_id).toBe(originalId);
197+
expect(messages[0].content).toBe("tool failed");
198+
});
199+
200+
it("should round-trip already-unique GLM-style IDs back to their exact original", () => {
201+
const originalId = "chatcmpl-tool-8a89c35582d60474";
202+
const rewrittenId = createAISDKToolCallId(originalId);
203+
204+
const { messages } = convertToWorkersAIChatMessages([
205+
{
206+
role: "assistant" as const,
207+
content: [
208+
{
209+
type: "tool-call" as const,
210+
toolCallId: rewrittenId,
211+
toolName: "get_weather",
212+
input: { city: "London" },
213+
},
214+
],
215+
},
216+
{
217+
role: "tool" as const,
218+
content: [
219+
{
220+
type: "tool-result" as const,
221+
toolCallId: rewrittenId,
222+
toolName: "get_weather",
223+
output: { type: "json", value: { weather: "Raining" } } as any,
224+
},
225+
],
226+
},
227+
]);
228+
229+
expect(messages[0].tool_calls![0].id).toBe(originalId);
230+
expect(messages[1].tool_call_id).toBe(originalId);
231+
});
155232
});
156233

157234
describe("basic message conversion", () => {

‎packages/workers-ai-provider/test/stream-text.test.ts‎

Lines changed: 128 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { setupServer } from "msw/node";
55
import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest";
66
import { z } from "zod/v4";
77
import { createWorkersAI } from "../src/index";
8+
import { toWorkersAIToolCallId } from "../src/utils";
89

910
const TEST_ACCOUNT_ID = "test-account-id";
1011
const TEST_API_KEY = "test-api-key";
@@ -206,7 +207,8 @@ describe("REST API - Streaming Text Tests", () => {
206207

207208
expect(toolCalls).toHaveLength(1);
208209
expect(toolCalls[0].toolName).toBe("get_weather");
209-
expect(toolCalls[0].toolCallId).toBe("call123");
210+
expect(toolCalls[0].toolCallId).not.toBe("call123");
211+
expect(toWorkersAIToolCallId(toolCalls[0].toolCallId)).toBe("call123");
210212
expect(await result.finishReason).toBe("tool-calls");
211213
});
212214

@@ -263,7 +265,8 @@ describe("REST API - Streaming Text Tests", () => {
263265

264266
expect(toolCalls).toHaveLength(1);
265267
expect(toolCalls[0].toolName).toBe("get_weather");
266-
expect(toolCalls[0].toolCallId).toBe("chatcmpl-tool-abc");
268+
expect(toolCalls[0].toolCallId).not.toBe("chatcmpl-tool-abc");
269+
expect(toWorkersAIToolCallId(toolCalls[0].toolCallId)).toBe("chatcmpl-tool-abc");
267270
expect(await result.finishReason).toBe("tool-calls");
268271
});
269272

@@ -496,7 +499,8 @@ describe("Binding - Streaming Text Tests", () => {
496499

497500
expect(toolCalls).toHaveLength(1);
498501
expect(toolCalls[0].toolName).toBe("get_weather");
499-
expect(toolCalls[0].toolCallId).toBe("call_abc");
502+
expect(toolCalls[0].toolCallId).not.toBe("call_abc");
503+
expect(toWorkersAIToolCallId(toolCalls[0].toolCallId)).toBe("call_abc");
500504
});
501505

502506
it("should handle streamed multiple tool calls via binding", async () => {
@@ -575,6 +579,98 @@ describe("Binding - Streaming Text Tests", () => {
575579
expect(toolCalls[1].toolName).toBe("get_temperature");
576580
});
577581

582+
it("should rewrite repeated Kimi-style tool call IDs across turns and restore originals in prompts", async () => {
583+
const capturedInputs: any[] = [];
584+
const workersai = createWorkersAI({
585+
binding: {
586+
run: async (_modelName: string, inputs: any) => {
587+
capturedInputs.push(inputs);
588+
return mockStream([
589+
{
590+
tool_calls: [
591+
{
592+
id: "functions.list_toolbox_tools:0",
593+
type: "function",
594+
index: 0,
595+
function: { name: "list_toolbox_tools", arguments: "{}" },
596+
},
597+
],
598+
},
599+
{ finish_reason: "tool_calls" },
600+
"[DONE]",
601+
]);
602+
},
603+
},
604+
});
605+
606+
const model = workersai(TEST_MODEL);
607+
const tools = {
608+
list_toolbox_tools: {
609+
description: "List tools",
610+
inputSchema: z.object({}),
611+
},
612+
};
613+
614+
const first = streamText({
615+
model,
616+
messages: [{ role: "user", content: "first" }],
617+
tools,
618+
});
619+
const firstToolCalls: any[] = [];
620+
for await (const chunk of first.fullStream) {
621+
if (chunk.type === "tool-call") firstToolCalls.push(chunk);
622+
}
623+
624+
const second = streamText({
625+
model,
626+
messages: [
627+
{ role: "user", content: "first" },
628+
{
629+
role: "assistant",
630+
content: [
631+
{
632+
type: "tool-call",
633+
toolCallId: firstToolCalls[0].toolCallId,
634+
toolName: "list_toolbox_tools",
635+
input: {},
636+
},
637+
],
638+
},
639+
{
640+
role: "tool",
641+
content: [
642+
{
643+
type: "tool-result",
644+
toolCallId: firstToolCalls[0].toolCallId,
645+
toolName: "list_toolbox_tools",
646+
output: { type: "text", value: "[]" },
647+
},
648+
],
649+
},
650+
{ role: "user", content: "second" },
651+
] as any,
652+
tools,
653+
});
654+
const secondToolCalls: any[] = [];
655+
for await (const chunk of second.fullStream) {
656+
if (chunk.type === "tool-call") secondToolCalls.push(chunk);
657+
}
658+
659+
expect(firstToolCalls).toHaveLength(1);
660+
expect(secondToolCalls).toHaveLength(1);
661+
expect(firstToolCalls[0].toolCallId).not.toBe(secondToolCalls[0].toolCallId);
662+
expect(toWorkersAIToolCallId(firstToolCalls[0].toolCallId)).toBe(
663+
"functions.list_toolbox_tools:0",
664+
);
665+
expect(toWorkersAIToolCallId(secondToolCalls[0].toolCallId)).toBe(
666+
"functions.list_toolbox_tools:0",
667+
);
668+
expect(capturedInputs[1].messages[1].tool_calls[0].id).toBe(
669+
"functions.list_toolbox_tools:0",
670+
);
671+
expect(capturedInputs[1].messages[2].tool_call_id).toBe("functions.list_toolbox_tools:0");
672+
});
673+
578674
it("should handle streamed OpenAI-format tool calls with reasoning via binding", async () => {
579675
const workersai = createWorkersAI({
580676
binding: {
@@ -649,7 +745,8 @@ describe("Binding - Streaming Text Tests", () => {
649745
expect(reasoning).toBe("Let me check the weather.");
650746
expect(toolCalls).toHaveLength(1);
651747
expect(toolCalls[0].toolName).toBe("get_weather");
652-
expect(toolCalls[0].toolCallId).toBe("chatcmpl-tool-abc");
748+
expect(toolCalls[0].toolCallId).not.toBe("chatcmpl-tool-abc");
749+
expect(toWorkersAIToolCallId(toolCalls[0].toolCallId)).toBe("chatcmpl-tool-abc");
653750
expect(await result.finishReason).toBe("tool-calls");
654751
});
655752

@@ -1583,7 +1680,17 @@ describe("Incremental Tool Call Streaming", () => {
15831680
const toolCall = parts.find((p) => p.type === "tool-call");
15841681
expect(toolCall).toBeDefined();
15851682
expect(toolCall.toolName).toBe("get_weather");
1586-
expect(toolCall.toolCallId).toBe("call1");
1683+
expect(toolCall.toolCallId).not.toBe("call1");
1684+
expect(toWorkersAIToolCallId(toolCall.toolCallId)).toBe("call1");
1685+
1686+
const toolEventIds = parts
1687+
.filter((p) =>
1688+
["tool-input-start", "tool-input-delta", "tool-input-end", "tool-call"].includes(
1689+
p.type,
1690+
),
1691+
)
1692+
.map((p) => p.toolCallId ?? p.id);
1693+
expect(new Set(toolEventIds)).toEqual(new Set([toolCall.toolCallId]));
15871694

15881695
// The AI SDK assembles and parses the full arguments from incremental events
15891696
const args = toolCall.args ?? toolCall.input;
@@ -2242,7 +2349,8 @@ describe("Eager tool-input-end streaming (issue #488)", () => {
22422349
"tool-input-end",
22432350
"tool-call",
22442351
]);
2245-
expect(toolCallData[0].toolCallId).toBe("solo");
2352+
expect(toolCallData[0].toolCallId).not.toBe("solo");
2353+
expect(toWorkersAIToolCallId(toolCallData[0].toolCallId)).toBe("solo");
22462354
expect(toolCallData[0].toolName).toBe("get_weather");
22472355
});
22482356

@@ -2498,8 +2606,10 @@ describe("Eager tool-input-end streaming (issue #488)", () => {
24982606
"tool-call",
24992607
]);
25002608

2501-
expect(toolCalls[0].toolCallId).toBe("call_1");
2502-
expect(toolCalls[1].toolCallId).toBe("call_2");
2609+
expect(toolCalls[0].toolCallId).not.toBe("call_1");
2610+
expect(toolCalls[1].toolCallId).not.toBe("call_2");
2611+
expect(toWorkersAIToolCallId(toolCalls[0].toolCallId)).toBe("call_1");
2612+
expect(toWorkersAIToolCallId(toolCalls[1].toolCallId)).toBe("call_2");
25032613
});
25042614

25052615
it("should not double-close a tool call (finalization + new index)", async () => {
@@ -2639,8 +2749,10 @@ describe("Eager tool-input-end streaming (issue #488)", () => {
26392749
"tool-input-end",
26402750
"tool-call",
26412751
]);
2642-
expect(toolCalls[0].toolCallId).toBe("call_1");
2643-
expect(toolCalls[1].toolCallId).toBe("call_2");
2752+
expect(toolCalls[0].toolCallId).not.toBe("call_1");
2753+
expect(toolCalls[1].toolCallId).not.toBe("call_2");
2754+
expect(toWorkersAIToolCallId(toolCalls[0].toolCallId)).toBe("call_1");
2755+
expect(toWorkersAIToolCallId(toolCalls[1].toolCallId)).toBe("call_2");
26442756
});
26452757

26462758
it("should handle three sequential OpenAI-format tool calls with eager close", async () => {
@@ -2736,9 +2848,12 @@ describe("Eager tool-input-end streaming (issue #488)", () => {
27362848
}
27372849

27382850
expect(toolCalls).toHaveLength(3);
2739-
expect(toolCalls[0].toolCallId).toBe("call_1");
2740-
expect(toolCalls[1].toolCallId).toBe("call_2");
2741-
expect(toolCalls[2].toolCallId).toBe("call_3");
2851+
expect(toolCalls[0].toolCallId).not.toBe("call_1");
2852+
expect(toolCalls[1].toolCallId).not.toBe("call_2");
2853+
expect(toolCalls[2].toolCallId).not.toBe("call_3");
2854+
expect(toWorkersAIToolCallId(toolCalls[0].toolCallId)).toBe("call_1");
2855+
expect(toWorkersAIToolCallId(toolCalls[1].toolCallId)).toBe("call_2");
2856+
expect(toWorkersAIToolCallId(toolCalls[2].toolCallId)).toBe("call_3");
27422857

27432858
const toolEvents = events.filter((e) =>
27442859
["tool-input-start", "tool-input-end", "tool-call"].includes(e),

0 commit comments

Comments
 (0)