Skip to content

Commit dfd6f81

Browse files
roomoteroomote
committed
fix: isolate native tool parser streams (#1468)
1 parent dfed27d commit dfd6f81

6 files changed

Lines changed: 181 additions & 74 deletions

File tree

‎src/api/providers/lm-studio.ts‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import {
1111

1212
import type { ApiHandlerOptions } from "../../shared/api"
1313

14-
import { NativeToolCallParser } from "../../core/assistant-message/NativeToolCallParser"
1514
import { TagMatcher } from "../../utils/tag-matcher"
1615

1716
import { convertToOpenAiMessages } from "../transform/openai-format"
@@ -118,6 +117,7 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
118117
text: chunk.data,
119118
}) as const,
120119
)
120+
const activeToolCallIds = new Set<string>()
121121

122122
for await (const chunk of results) {
123123
const delta = chunk.choices[0]?.delta
@@ -142,6 +142,9 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
142142
// Handle tool calls in stream - emit partial chunks for NativeToolCallParser
143143
if (delta?.tool_calls) {
144144
for (const toolCall of delta.tool_calls) {
145+
if (toolCall.id) {
146+
activeToolCallIds.add(toolCall.id)
147+
}
145148
yield {
146149
type: "tool_call_partial",
147150
index: toolCall.index,
@@ -153,11 +156,11 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
153156
}
154157

155158
// Process finish_reason to emit tool_call_end events
156-
if (finishReason) {
157-
const endEvents = NativeToolCallParser.processFinishReason(finishReason)
158-
for (const event of endEvents) {
159-
yield event
159+
if (finishReason === "tool_calls") {
160+
for (const id of activeToolCallIds) {
161+
yield { type: "tool_call_end", id }
160162
}
163+
activeToolCallIds.clear()
161164
}
162165
}
163166

‎src/api/providers/openrouter.ts‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,6 @@ import {
1414
} from "@roo-code/types"
1515
import { TelemetryService } from "@roo-code/telemetry"
1616

17-
import { NativeToolCallParser } from "../../core/assistant-message/NativeToolCallParser"
18-
1917
import type { ApiHandlerOptions } from "../../shared/api"
2018

2119
import {
@@ -401,6 +399,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
401399
// When reasoning_details has displayable content (reasoning.text or reasoning.summary),
402400
// we skip yielding the top-level reasoning field to avoid duplicate display.
403401
let hasYieldedReasoningFromDetails = false
402+
const activeToolCallIds = new Set<string>()
404403

405404
for await (const chunk of stream) {
406405
// OpenRouter returns an error object instead of the OpenAI SDK throwing an error.
@@ -491,6 +490,9 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
491490
// Emit raw tool call chunks - NativeToolCallParser handles state management
492491
if ("tool_calls" in delta && Array.isArray(delta.tool_calls)) {
493492
for (const toolCall of delta.tool_calls) {
493+
if (toolCall.id) {
494+
activeToolCallIds.add(toolCall.id)
495+
}
494496
yield {
495497
type: "tool_call_partial",
496498
index: toolCall.index,
@@ -508,11 +510,11 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
508510

509511
// Process finish_reason to emit tool_call_end events
510512
// This ensures tool calls are finalized even if the stream doesn't properly close
511-
if (finishReason) {
512-
const endEvents = NativeToolCallParser.processFinishReason(finishReason)
513-
for (const event of endEvents) {
514-
yield event
513+
if (finishReason === "tool_calls") {
514+
for (const id of activeToolCallIds) {
515+
yield { type: "tool_call_end", id }
515516
}
517+
activeToolCallIds.clear()
516518
}
517519

518520
if (chunk.usage) {

‎src/api/providers/qwen-code.ts‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@ import { type ModelInfo, type QwenCodeModelId, qwenCodeModels, qwenCodeDefaultMo
88

99
import type { ApiHandlerOptions } from "../../shared/api"
1010

11-
import { NativeToolCallParser } from "../../core/assistant-message/NativeToolCallParser"
12-
1311
import { convertToOpenAiMessages } from "../transform/openai-format"
1412
import { ApiStream } from "../transform/stream"
1513

@@ -243,6 +241,7 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan
243241

244242
let fullContent = ""
245243

244+
const activeToolCallIds = new Set<string>()
246245
for await (const apiChunk of stream) {
247246
const delta = apiChunk.choices[0]?.delta ?? {}
248247
const finishReason = apiChunk.choices[0]?.finish_reason
@@ -293,6 +292,9 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan
293292
// Handle tool calls in stream - emit partial chunks for NativeToolCallParser
294293
if (delta.tool_calls) {
295294
for (const toolCall of delta.tool_calls) {
295+
if (toolCall.id) {
296+
activeToolCallIds.add(toolCall.id)
297+
}
296298
yield {
297299
type: "tool_call_partial",
298300
index: toolCall.index,
@@ -304,11 +306,11 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan
304306
}
305307

306308
// Process finish_reason to emit tool_call_end events
307-
if (finishReason) {
308-
const endEvents = NativeToolCallParser.processFinishReason(finishReason)
309-
for (const event of endEvents) {
310-
yield event
309+
if (finishReason === "tool_calls") {
310+
for (const id of activeToolCallIds) {
311+
yield { type: "tool_call_end", id }
311312
}
313+
activeToolCallIds.clear()
312314
}
313315

314316
if (apiChunk.usage) {

‎src/core/assistant-message/NativeToolCallParser.ts‎

Lines changed: 72 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -51,28 +51,43 @@ export type ToolCallStreamEvent = ApiStreamToolCallStartChunk | ApiStreamToolCal
5151
* provider-level raw chunks into start/delta/end events.
5252
*/
5353
export class NativeToolCallParser {
54+
private static readonly defaultScope = {}
55+
5456
// Streaming state management for argument accumulation (keyed by tool call id)
5557
// Note: name is string to accommodate dynamic MCP tools (mcp--serverName--toolName)
56-
private static streamingToolCalls = new Map<
57-
string,
58-
{
59-
id: string
60-
name: string
61-
argumentsAccumulator: string
62-
}
58+
private static streamingToolCallsByScope = new WeakMap<
59+
object,
60+
Map<string, { id: string; name: string; argumentsAccumulator: string }>
6361
>()
6462

65-
// Raw chunk tracking state (keyed by index from API stream)
66-
private static rawChunkTracker = new Map<
67-
number,
68-
{
69-
id: string
70-
name: string
71-
hasStarted: boolean
72-
deltaBuffer: string[]
73-
}
63+
// Raw chunk tracking state (keyed by index from one API stream)
64+
private static rawChunkTrackersByScope = new WeakMap<
65+
object,
66+
Map<number, { id: string; name: string; hasStarted: boolean; deltaBuffer: string[] }>
7467
>()
7568

69+
public static createScope(): object {
70+
return {}
71+
}
72+
73+
private static getStreamingToolCalls(scope = this.defaultScope) {
74+
let streamingToolCalls = this.streamingToolCallsByScope.get(scope)
75+
if (!streamingToolCalls) {
76+
streamingToolCalls = new Map()
77+
this.streamingToolCallsByScope.set(scope, streamingToolCalls)
78+
}
79+
return streamingToolCalls
80+
}
81+
82+
private static getRawChunkTracker(scope = this.defaultScope) {
83+
let rawChunkTracker = this.rawChunkTrackersByScope.get(scope)
84+
if (!rawChunkTracker) {
85+
rawChunkTracker = new Map()
86+
this.rawChunkTrackersByScope.set(scope, rawChunkTracker)
87+
}
88+
return rawChunkTracker
89+
}
90+
7691
private static coerceOptionalBoolean(value: unknown): boolean | undefined {
7792
if (typeof value === "boolean") {
7893
return value
@@ -96,16 +111,20 @@ export class NativeToolCallParser {
96111
* This is the entry point for providers that emit tool_call_partial chunks.
97112
* Returns an array of events to be processed by the consumer.
98113
*/
99-
public static processRawChunk(chunk: {
100-
index: number
101-
id?: string
102-
name?: string
103-
arguments?: string
104-
}): ToolCallStreamEvent[] {
114+
public static processRawChunk(
115+
chunk: {
116+
index: number
117+
id?: string
118+
name?: string
119+
arguments?: string
120+
},
121+
scope = this.defaultScope,
122+
): ToolCallStreamEvent[] {
105123
const events: ToolCallStreamEvent[] = []
106124
const { index, id, name, arguments: args } = chunk
125+
const rawChunkTracker = this.getRawChunkTracker(scope)
107126

108-
let tracked = this.rawChunkTracker.get(index)
127+
let tracked = rawChunkTracker.get(index)
109128

110129
// Initialize new tool call tracking when we receive an id
111130
if (id && !tracked) {
@@ -115,7 +134,7 @@ export class NativeToolCallParser {
115134
hasStarted: false,
116135
deltaBuffer: [],
117136
}
118-
this.rawChunkTracker.set(index, tracked)
137+
rawChunkTracker.set(index, tracked)
119138
}
120139

121140
if (!tracked) {
@@ -167,11 +186,15 @@ export class NativeToolCallParser {
167186
* Process stream finish reason.
168187
* Emits end events when finish_reason is 'tool_calls'.
169188
*/
170-
public static processFinishReason(finishReason: string | null | undefined): ToolCallStreamEvent[] {
189+
public static processFinishReason(
190+
finishReason: string | null | undefined,
191+
scope = this.defaultScope,
192+
): ToolCallStreamEvent[] {
171193
const events: ToolCallStreamEvent[] = []
194+
const rawChunkTracker = this.rawChunkTrackersByScope.get(scope)
172195

173-
if (finishReason === "tool_calls" && this.rawChunkTracker.size > 0) {
174-
for (const [, tracked] of this.rawChunkTracker.entries()) {
196+
if (finishReason === "tool_calls" && rawChunkTracker && rawChunkTracker.size > 0) {
197+
for (const [, tracked] of rawChunkTracker.entries()) {
175198
events.push({
176199
type: "tool_call_end",
177200
id: tracked.id,
@@ -186,20 +209,21 @@ export class NativeToolCallParser {
186209
* Finalize any remaining tool calls that weren't explicitly ended.
187210
* Should be called at the end of stream processing.
188211
*/
189-
public static finalizeRawChunks(): ToolCallStreamEvent[] {
212+
public static finalizeRawChunks(scope = this.defaultScope): ToolCallStreamEvent[] {
190213
const events: ToolCallStreamEvent[] = []
214+
const rawChunkTracker = this.rawChunkTrackersByScope.get(scope)
191215

192-
if (this.rawChunkTracker.size > 0) {
193-
for (const [, tracked] of this.rawChunkTracker.entries()) {
216+
if (rawChunkTracker && rawChunkTracker.size > 0) {
217+
for (const [, tracked] of rawChunkTracker.entries()) {
194218
if (tracked.hasStarted) {
195219
events.push({
196220
type: "tool_call_end",
197221
id: tracked.id,
198222
})
199223
}
200224
}
201-
this.rawChunkTracker.clear()
202225
}
226+
this.rawChunkTrackersByScope.delete(scope)
203227

204228
return events
205229
}
@@ -208,17 +232,17 @@ export class NativeToolCallParser {
208232
* Clear all raw chunk tracking state.
209233
* Should be called when a new API request starts.
210234
*/
211-
public static clearRawChunkState(): void {
212-
this.rawChunkTracker.clear()
235+
public static clearRawChunkState(scope = this.defaultScope): void {
236+
this.rawChunkTrackersByScope.delete(scope)
213237
}
214238

215239
/**
216240
* Start streaming a new tool call.
217241
* Initializes tracking for incremental argument parsing.
218242
* Accepts string to support both ToolName and dynamic MCP tools (mcp--serverName--toolName).
219243
*/
220-
public static startStreamingToolCall(id: string, name: string): void {
221-
this.streamingToolCalls.set(id, {
244+
public static startStreamingToolCall(id: string, name: string, scope = this.defaultScope): void {
245+
this.getStreamingToolCalls(scope).set(id, {
222246
id,
223247
name,
224248
argumentsAccumulator: "",
@@ -230,25 +254,25 @@ export class NativeToolCallParser {
230254
* Should be called when a new API request starts to prevent memory leaks
231255
* from interrupted streams.
232256
*/
233-
public static clearAllStreamingToolCalls(): void {
234-
this.streamingToolCalls.clear()
257+
public static clearAllStreamingToolCalls(scope = this.defaultScope): void {
258+
this.streamingToolCallsByScope.delete(scope)
235259
}
236260

237261
/**
238262
* Check if there are any active streaming tool calls.
239263
* Useful for debugging and testing.
240264
*/
241-
public static hasActiveStreamingToolCalls(): boolean {
242-
return this.streamingToolCalls.size > 0
265+
public static hasActiveStreamingToolCalls(scope = this.defaultScope): boolean {
266+
return (this.streamingToolCallsByScope.get(scope)?.size ?? 0) > 0
243267
}
244268

245269
/**
246270
* Process a chunk of JSON arguments for a streaming tool call.
247271
* Uses partial-json-parser to extract values from incomplete JSON immediately.
248272
* Returns a partial ToolUse with currently parsed parameters.
249273
*/
250-
public static processStreamingChunk(id: string, chunk: string): ToolUse | null {
251-
const toolCall = this.streamingToolCalls.get(id)
274+
public static processStreamingChunk(id: string, chunk: string, scope = this.defaultScope): ToolUse | null {
275+
const toolCall = this.streamingToolCallsByScope.get(scope)?.get(id)
252276
if (!toolCall) {
253277
return null
254278
}
@@ -291,8 +315,9 @@ export class NativeToolCallParser {
291315
* Finalize a streaming tool call.
292316
* Parses the complete JSON and returns the final ToolUse or McpToolUse.
293317
*/
294-
public static finalizeStreamingToolCall(id: string): ToolUse | McpToolUse | null {
295-
const toolCall = this.streamingToolCalls.get(id)
318+
public static finalizeStreamingToolCall(id: string, scope = this.defaultScope): ToolUse | McpToolUse | null {
319+
const streamingToolCalls = this.streamingToolCallsByScope.get(scope)
320+
const toolCall = streamingToolCalls?.get(id)
296321
if (!toolCall) {
297322
return null
298323
}
@@ -306,7 +331,10 @@ export class NativeToolCallParser {
306331
})
307332

308333
// Clean up streaming state
309-
this.streamingToolCalls.delete(id)
334+
streamingToolCalls?.delete(id)
335+
if (streamingToolCalls?.size === 0) {
336+
this.streamingToolCallsByScope.delete(scope)
337+
}
310338

311339
return finalToolUse
312340
}

0 commit comments

Comments
 (0)