Skip to content
This repository was archived by the owner on May 15, 2026. It is now read-only.

Commit 0327f12

Browse files
daniel-lxsdaniel-lxs
authored
feat: implement streaming for native tool calls (#9542)
1 parent cad6145 commit 0327f12

17 files changed

Lines changed: 696 additions & 356 deletions

‎pnpm-lock.yaml‎

Lines changed: 10 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/api/providers/__tests__/minimax.spec.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -372,11 +372,13 @@ describe("MiniMaxHandler", () => {
372372
const firstChunk = await stream.next()
373373

374374
expect(firstChunk.done).toBe(false)
375+
// Provider now yields tool_call_partial chunks, NativeToolCallParser handles reassembly
375376
expect(firstChunk.value).toEqual({
376-
type: "tool_call",
377+
type: "tool_call_partial",
378+
index: 0,
377379
id: "tool-123",
378380
name: "get_weather",
379-
arguments: JSON.stringify({ city: "London" }),
381+
arguments: undefined,
380382
})
381383
})
382384
})

‎src/api/providers/__tests__/openai.spec.ts‎

Lines changed: 48 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -269,13 +269,31 @@ describe("OpenAiHandler", () => {
269269
chunks.push(chunk)
270270
}
271271

272-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
273-
expect(toolCallChunks).toHaveLength(1)
274-
expect(toolCallChunks[0]).toEqual({
275-
type: "tool_call",
272+
// Provider now yields tool_call_partial chunks, NativeToolCallParser handles reassembly
273+
const toolCallPartialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
274+
expect(toolCallPartialChunks).toHaveLength(3)
275+
// First chunk has id and name
276+
expect(toolCallPartialChunks[0]).toEqual({
277+
type: "tool_call_partial",
278+
index: 0,
276279
id: "call_1",
277280
name: "test_tool",
278-
arguments: '{"arg":"value"}',
281+
arguments: "",
282+
})
283+
// Subsequent chunks have arguments
284+
expect(toolCallPartialChunks[1]).toEqual({
285+
type: "tool_call_partial",
286+
index: 0,
287+
id: undefined,
288+
name: undefined,
289+
arguments: '{"arg":',
290+
})
291+
expect(toolCallPartialChunks[2]).toEqual({
292+
type: "tool_call_partial",
293+
index: 0,
294+
id: undefined,
295+
name: undefined,
296+
arguments: '"value"}',
279297
})
280298
})
281299

@@ -318,11 +336,12 @@ describe("OpenAiHandler", () => {
318336
chunks.push(chunk)
319337
}
320338

321-
// Tool calls should still be yielded via the fallback mechanism
322-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
323-
expect(toolCallChunks).toHaveLength(1)
324-
expect(toolCallChunks[0]).toEqual({
325-
type: "tool_call",
339+
// Provider now yields tool_call_partial chunks, NativeToolCallParser handles reassembly
340+
const toolCallPartialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
341+
expect(toolCallPartialChunks).toHaveLength(1)
342+
expect(toolCallPartialChunks[0]).toEqual({
343+
type: "tool_call_partial",
344+
index: 0,
326345
id: "call_fallback",
327346
name: "fallback_tool",
328347
arguments: '{"test":"fallback"}',
@@ -819,12 +838,21 @@ describe("OpenAiHandler", () => {
819838
chunks.push(chunk)
820839
}
821840

822-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
823-
expect(toolCallChunks).toHaveLength(1)
824-
expect(toolCallChunks[0]).toEqual({
825-
type: "tool_call",
841+
// Provider now yields tool_call_partial chunks, NativeToolCallParser handles reassembly
842+
const toolCallPartialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
843+
expect(toolCallPartialChunks).toHaveLength(2)
844+
expect(toolCallPartialChunks[0]).toEqual({
845+
type: "tool_call_partial",
846+
index: 0,
826847
id: "call_1",
827848
name: "test_tool",
849+
arguments: "",
850+
})
851+
expect(toolCallPartialChunks[1]).toEqual({
852+
type: "tool_call_partial",
853+
index: 0,
854+
id: undefined,
855+
name: undefined,
828856
arguments: "{}",
829857
})
830858
})
@@ -870,11 +898,12 @@ describe("OpenAiHandler", () => {
870898
chunks.push(chunk)
871899
}
872900

873-
// Tool calls should still be yielded via the fallback mechanism
874-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
875-
expect(toolCallChunks).toHaveLength(1)
876-
expect(toolCallChunks[0]).toEqual({
877-
type: "tool_call",
901+
// Provider now yields tool_call_partial chunks, NativeToolCallParser handles reassembly
902+
const toolCallPartialChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
903+
expect(toolCallPartialChunks).toHaveLength(1)
904+
expect(toolCallPartialChunks[0]).toEqual({
905+
type: "tool_call_partial",
906+
index: 0,
878907
id: "call_o3_fallback",
879908
name: "o3_fallback_tool",
880909
arguments: '{"o3":"test"}',

‎src/api/providers/__tests__/roo.spec.ts‎

Lines changed: 49 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -636,7 +636,7 @@ describe("RooHandler", () => {
636636
handler = new RooHandler(mockOptions)
637637
})
638638

639-
it("should yield tool calls when finish_reason is tool_calls", async () => {
639+
it("should yield raw tool call chunks when tool_calls present", async () => {
640640
mockCreate.mockResolvedValueOnce({
641641
[Symbol.asyncIterator]: async function* () {
642642
yield {
@@ -689,14 +689,27 @@ describe("RooHandler", () => {
689689
chunks.push(chunk)
690690
}
691691

692-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
693-
expect(toolCallChunks).toHaveLength(1)
694-
expect(toolCallChunks[0].id).toBe("call_123")
695-
expect(toolCallChunks[0].name).toBe("read_file")
696-
expect(toolCallChunks[0].arguments).toBe('{"path":"test.ts"}')
692+
// Verify we get raw tool call chunks
693+
const rawChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
694+
695+
expect(rawChunks).toHaveLength(2)
696+
expect(rawChunks[0]).toEqual({
697+
type: "tool_call_partial",
698+
index: 0,
699+
id: "call_123",
700+
name: "read_file",
701+
arguments: '{"path":"',
702+
})
703+
expect(rawChunks[1]).toEqual({
704+
type: "tool_call_partial",
705+
index: 0,
706+
id: undefined,
707+
name: undefined,
708+
arguments: 'test.ts"}',
709+
})
697710
})
698711

699-
it("should yield tool calls even when finish_reason is not set (fallback behavior)", async () => {
712+
it("should yield raw tool call chunks even when finish_reason is not tool_calls", async () => {
700713
mockCreate.mockResolvedValueOnce({
701714
[Symbol.asyncIterator]: async function* () {
702715
yield {
@@ -718,12 +731,11 @@ describe("RooHandler", () => {
718731
},
719732
],
720733
}
721-
// Stream ends without finish_reason being set to "tool_calls"
722734
yield {
723735
choices: [
724736
{
725737
delta: {},
726-
finish_reason: "stop", // Different finish reason
738+
finish_reason: "stop",
727739
index: 0,
728740
},
729741
],
@@ -738,15 +750,19 @@ describe("RooHandler", () => {
738750
chunks.push(chunk)
739751
}
740752

741-
// Tool calls should still be yielded via the fallback mechanism
742-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
743-
expect(toolCallChunks).toHaveLength(1)
744-
expect(toolCallChunks[0].id).toBe("call_456")
745-
expect(toolCallChunks[0].name).toBe("write_to_file")
746-
expect(toolCallChunks[0].arguments).toBe('{"path":"test.ts","content":"hello"}')
753+
const rawChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
754+
755+
expect(rawChunks).toHaveLength(1)
756+
expect(rawChunks[0]).toEqual({
757+
type: "tool_call_partial",
758+
index: 0,
759+
id: "call_456",
760+
name: "write_to_file",
761+
arguments: '{"path":"test.ts","content":"hello"}',
762+
})
747763
})
748764

749-
it("should handle multiple tool calls", async () => {
765+
it("should handle multiple tool calls with different indices", async () => {
750766
mockCreate.mockResolvedValueOnce({
751767
[Symbol.asyncIterator]: async function* () {
752768
yield {
@@ -800,15 +816,16 @@ describe("RooHandler", () => {
800816
chunks.push(chunk)
801817
}
802818

803-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
804-
expect(toolCallChunks).toHaveLength(2)
805-
expect(toolCallChunks[0].id).toBe("call_1")
806-
expect(toolCallChunks[0].name).toBe("read_file")
807-
expect(toolCallChunks[1].id).toBe("call_2")
808-
expect(toolCallChunks[1].name).toBe("read_file")
819+
const rawChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
820+
821+
expect(rawChunks).toHaveLength(2)
822+
expect(rawChunks[0].index).toBe(0)
823+
expect(rawChunks[0].id).toBe("call_1")
824+
expect(rawChunks[1].index).toBe(1)
825+
expect(rawChunks[1].id).toBe("call_2")
809826
})
810827

811-
it("should accumulate tool call arguments across multiple chunks", async () => {
828+
it("should emit raw chunks for streaming arguments", async () => {
812829
mockCreate.mockResolvedValueOnce({
813830
[Symbol.asyncIterator]: async function* () {
814831
yield {
@@ -876,14 +893,15 @@ describe("RooHandler", () => {
876893
chunks.push(chunk)
877894
}
878895

879-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
880-
expect(toolCallChunks).toHaveLength(1)
881-
expect(toolCallChunks[0].id).toBe("call_789")
882-
expect(toolCallChunks[0].name).toBe("execute_command")
883-
expect(toolCallChunks[0].arguments).toBe('{"command":"npm install"}')
896+
const rawChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
897+
898+
expect(rawChunks).toHaveLength(3)
899+
expect(rawChunks[0].arguments).toBe('{"command":"')
900+
expect(rawChunks[1].arguments).toBe("npm install")
901+
expect(rawChunks[2].arguments).toBe('"}')
884902
})
885903

886-
it("should not yield empty tool calls when no tool calls present", async () => {
904+
it("should not yield tool call chunks when no tool calls present", async () => {
887905
mockCreate.mockResolvedValueOnce({
888906
[Symbol.asyncIterator]: async function* () {
889907
yield {
@@ -902,8 +920,8 @@ describe("RooHandler", () => {
902920
chunks.push(chunk)
903921
}
904922

905-
const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call")
906-
expect(toolCallChunks).toHaveLength(0)
923+
const rawChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
924+
expect(rawChunks).toHaveLength(0)
907925
})
908926
})
909927
})

‎src/api/providers/base-openai-compatible-provider.ts‎

Lines changed: 5 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,6 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
123123
}) as const,
124124
)
125125

126-
const toolCallAccumulator = new Map<number, { id: string; name: string; arguments: string }>()
127-
128126
let lastUsage: OpenAI.CompletionUsage | undefined
129127

130128
for await (const chunk of stream) {
@@ -137,7 +135,6 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
137135
}
138136

139137
const delta = chunk.choices?.[0]?.delta
140-
const finishReason = chunk.choices?.[0]?.finish_reason
141138

142139
if (delta?.content) {
143140
for (const processedChunk of matcher.update(delta.content)) {
@@ -157,56 +154,24 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
157154
}
158155
}
159156

157+
// Emit raw tool call chunks - NativeToolCallParser handles state management
160158
if (delta?.tool_calls) {
161159
for (const toolCall of delta.tool_calls) {
162-
const index = toolCall.index
163-
const existing = toolCallAccumulator.get(index)
164-
165-
if (existing) {
166-
if (toolCall.function?.arguments) {
167-
existing.arguments += toolCall.function.arguments
168-
}
169-
} else {
170-
toolCallAccumulator.set(index, {
171-
id: toolCall.id || "",
172-
name: toolCall.function?.name || "",
173-
arguments: toolCall.function?.arguments || "",
174-
})
175-
}
176-
}
177-
}
178-
179-
if (finishReason === "tool_calls") {
180-
for (const toolCall of toolCallAccumulator.values()) {
181160
yield {
182-
type: "tool_call",
161+
type: "tool_call_partial",
162+
index: toolCall.index,
183163
id: toolCall.id,
184-
name: toolCall.name,
185-
arguments: toolCall.arguments,
164+
name: toolCall.function?.name,
165+
arguments: toolCall.function?.arguments,
186166
}
187167
}
188-
toolCallAccumulator.clear()
189168
}
190169

191170
if (chunk.usage) {
192171
lastUsage = chunk.usage
193172
}
194173
}
195174

196-
// Fallback: If stream ends with accumulated tool calls that weren't yielded
197-
// (e.g., finish_reason was 'stop' or 'length' instead of 'tool_calls')
198-
if (toolCallAccumulator.size > 0) {
199-
for (const toolCall of toolCallAccumulator.values()) {
200-
yield {
201-
type: "tool_call",
202-
id: toolCall.id,
203-
name: toolCall.name,
204-
arguments: toolCall.arguments,
205-
}
206-
}
207-
toolCallAccumulator.clear()
208-
}
209-
210175
if (lastUsage) {
211176
yield this.processUsageMetrics(lastUsage, this.getModel().info)
212177
}

0 commit comments

Comments
 (0)