diff --git a/src/core/assistant-message/NativeToolCallParser.ts b/src/core/assistant-message/NativeToolCallParser.ts index 10a4446182..122579c2a0 100644 --- a/src/core/assistant-message/NativeToolCallParser.ts +++ b/src/core/assistant-message/NativeToolCallParser.ts @@ -61,7 +61,17 @@ export class NativeToolCallParser { // Raw chunk tracking state (keyed by index from one API stream) private static rawChunkTrackersByScope = new WeakMap< object, - Map + Map< + number, + { + id?: string + name: string + // Track whether the provider sent a name, including an empty name. + nameSeen: boolean + hasStarted: boolean + deltaBuffer: string[] + } + > >() public static createScope(): object { @@ -124,49 +134,54 @@ export class NativeToolCallParser { let tracked = rawChunkTracker.get(index) - // Initialize new tool call tracking when we receive an id - if (id && !tracked) { + // Create the tracker on first sight of this index, independent of whether + // an id has arrived yet. Keying the lifecycle by index (not id) ensures any + // `arguments` that stream before the id is known are buffered rather than dropped. + if (!tracked) { tracked = { id, name: name || "", + nameSeen: name !== undefined, hasStarted: false, deltaBuffer: [], } rawChunkTracker.set(index, tracked) } - if (!tracked) { - return events + // Record id and name as they arrive (they may come in separate chunks). + if (id) { + tracked.id = id } - - // Update name if present in chunk and not yet set - if (name) { + if (name !== undefined) { tracked.name = name + tracked.nameSeen = true } - // Emit start event when we have the name - if (!tracked.hasStarted && tracked.name) { + // Emit start event only once both id and name are known. Using a local + // non-null id keeps emitted events typed as id: string. + if (!tracked.hasStarted && tracked.id && tracked.nameSeen) { + const startedId = tracked.id events.push({ type: "tool_call_start", - id: tracked.id, + id: startedId, name: tracked.name, }) tracked.hasStarted = true - // Flush buffered deltas + // Flush buffered deltas accumulated during the pre-start window. for (const bufferedDelta of tracked.deltaBuffer) { events.push({ type: "tool_call_delta", - id: tracked.id, + id: startedId, delta: bufferedDelta, }) } tracked.deltaBuffer = [] } - // Emit delta event for argument chunks + // Emit delta event for argument chunks, buffering until start is emitted. if (args) { - if (tracked.hasStarted) { + if (tracked.hasStarted && tracked.id) { events.push({ type: "tool_call_delta", id: tracked.id, @@ -190,7 +205,7 @@ export class NativeToolCallParser { if (rawChunkTracker) { for (const [, tracked] of rawChunkTracker.entries()) { - if (tracked.hasStarted) { + if (tracked.hasStarted && tracked.id) { events.push({ type: "tool_call_end", id: tracked.id, diff --git a/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts b/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts index 47ea2a36cc..ee79d7963f 100644 --- a/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts +++ b/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts @@ -1,4 +1,4 @@ -import { NativeToolCallParser } from "../NativeToolCallParser" +import { NativeToolCallParser, type ToolCallStreamEvent } from "../NativeToolCallParser" describe("NativeToolCallParser", () => { describe("parseToolCall", () => { @@ -289,6 +289,69 @@ describe("NativeToolCallParser", () => { }) describe("processStreamingChunk", () => { + it("preserves read_file arguments that arrive before the call id and name", () => { + const scope = NativeToolCallParser.createScope() + const argumentsJson = JSON.stringify({ path: "src/leading.ts", mode: "slice", offset: 1, limit: 2000 }) + const split = 24 + + expect( + NativeToolCallParser.processRawChunk({ index: 0, arguments: argumentsJson.slice(0, split) }, scope), + ).toEqual([]) + + const identifiedEvents = NativeToolCallParser.processRawChunk( + { index: 0, id: "call_late_identity", name: "read_file" }, + scope, + ) + expect(identifiedEvents).toEqual([ + { type: "tool_call_start", id: "call_late_identity", name: "read_file" }, + { + type: "tool_call_delta", + id: "call_late_identity", + delta: argumentsJson.slice(0, split), + }, + ]) + + NativeToolCallParser.startStreamingToolCall("call_late_identity", "read_file", scope) + for (const event of identifiedEvents) { + if (event.type === "tool_call_delta") { + NativeToolCallParser.processStreamingChunk(event.id, event.delta, scope) + } + } + + const trailingEvents = NativeToolCallParser.processRawChunk( + { index: 0, arguments: argumentsJson.slice(split) }, + scope, + ) + expect(trailingEvents).toEqual([ + { + type: "tool_call_delta", + id: "call_late_identity", + delta: argumentsJson.slice(split), + }, + ]) + for (const event of trailingEvents) { + if (event.type === "tool_call_delta") { + NativeToolCallParser.processStreamingChunk(event.id, event.delta, scope) + } + } + + expect(NativeToolCallParser.finalizeRawChunks(scope)).toEqual([ + { type: "tool_call_end", id: "call_late_identity" }, + ]) + const result = NativeToolCallParser.finalizeStreamingToolCall("call_late_identity", scope) + expect(result?.type).toBe("tool_use") + if (result?.type === "tool_use") { + expect(result.nativeArgs).toEqual({ path: "src/leading.ts", mode: "slice", offset: 1, limit: 2000 }) + } + }) + + it("does not emit an end event for argument chunks that never receive an identity", () => { + const scope = NativeToolCallParser.createScope() + NativeToolCallParser.processRawChunk({ index: 0, arguments: '{"path":"orphan.ts"}' }, scope) + + expect(NativeToolCallParser.finalizeRawChunks(scope)).toEqual([]) + }) + it("retains peer calls until each call in a scope is finalized", () => { const scope = NativeToolCallParser.createScope() NativeToolCallParser.startStreamingToolCall("call_first", "read_file", scope) @@ -391,11 +454,14 @@ describe("NativeToolCallParser", () => { expect(NativeToolCallParser.finalizeRawChunks(firstScope)).toEqual([]) expect(NativeToolCallParser.finalizeStreamingToolCall("call_first", firstScope)).toBeNull() expect( - NativeToolCallParser.processRawChunk({ index: 0, arguments: "ignored-after-cleanup" }, firstScope), + NativeToolCallParser.processRawChunk({ index: 0, arguments: "buffered-after-cleanup" }, firstScope), ).toEqual([]) expect( NativeToolCallParser.processRawChunk({ index: 0, id: "call_reprobe", name: "read_file" }, firstScope), - ).toEqual([{ type: "tool_call_start", id: "call_reprobe", name: "read_file" }]) + ).toEqual([ + { type: "tool_call_start", id: "call_reprobe", name: "read_file" }, + { type: "tool_call_delta", id: "call_reprobe", delta: "buffered-after-cleanup" }, + ]) }) describe("read_file tool", () => { @@ -450,4 +516,255 @@ describe("NativeToolCallParser", () => { }) }) }) + + describe("processRawChunk streaming reassembly", () => { + // Mirror the sequencing Task.ts performs: feed each raw chunk through + // processRawChunk, drive startStreamingToolCall on tool_call_start, feed + // tool_call_delta into processStreamingChunk, and emit ends at stream close + // via finalizeRawChunks() (the same call Task.ts makes after the stream ends). + // Returns the ordered event types/ids plus the finalized tool uses by id. + const drive = (rawChunks: Array<{ index: number; id?: string; name?: string; arguments?: string }>) => { + const events: ToolCallStreamEvent[] = [] + const scope = NativeToolCallParser.createScope() + + const handleEvent = (event: ToolCallStreamEvent) => { + events.push(event) + if (event.type === "tool_call_start") { + NativeToolCallParser.startStreamingToolCall(event.id, event.name, scope) + } else if (event.type === "tool_call_delta") { + NativeToolCallParser.processStreamingChunk(event.id, event.delta, scope) + } + } + + for (const chunk of rawChunks) { + for (const event of NativeToolCallParser.processRawChunk(chunk, scope)) { + handleEvent(event) + } + } + + // Task.ts finalizes any tool calls still open at stream end via + // finalizeRawChunks(), which emits the tool_call_end events. + for (const event of NativeToolCallParser.finalizeRawChunks(scope)) { + handleEvent(event) + } + + const finalized = new Map>() + const startIds = events.filter((e) => e.type === "tool_call_start").map((e) => e.id) + for (const id of startIds) { + finalized.set(id, NativeToolCallParser.finalizeStreamingToolCall(id, scope)) + } + + return { events, finalized } + } + + it("preserves leading argument bytes that arrive before the id", () => { + // First chunk carries arguments but NO id; id+name arrive later, then more args. + const fullArgs = JSON.stringify({ path: "src/leading.ts", mode: "slice" }) + const firstHalf = fullArgs.slice(0, 10) + const secondHalf = fullArgs.slice(10) + + const { events, finalized } = drive([ + { index: 0, arguments: firstHalf }, + { index: 0, id: "call_late_id", name: "read_file" }, + { index: 0, arguments: secondHalf }, + ]) + + // Exactly one start, in the right order, with the late id. + const starts = events.filter((e) => e.type === "tool_call_start") + expect(starts).toHaveLength(1) + expect(starts[0].id).toBe("call_late_id") + + // The finalized arguments must contain the complete, uncorrupted payload. + const result = finalized.get("call_late_id") + expect(result).not.toBeNull() + expect(result?.type).toBe("tool_use") + if (result?.type === "tool_use") { + const nativeArgs = result.nativeArgs as { path: string; mode?: string } + expect(nativeArgs.path).toBe("src/leading.ts") + expect(nativeArgs.mode).toBe("slice") + } + }) + + it("handles id and name arriving in separate chunks (issue #218)", () => { + const fullArgs = JSON.stringify({ path: "src/split.ts" }) + + const { events, finalized } = drive([ + { index: 0, id: "call_split" }, + { index: 0, name: "read_file" }, + { index: 0, arguments: fullArgs }, + ]) + + const starts = events.filter((e) => e.type === "tool_call_start") + expect(starts).toHaveLength(1) + expect(starts[0].id).toBe("call_split") + + const result = finalized.get("call_split") + expect(result?.type).toBe("tool_use") + if (result?.type === "tool_use") { + const nativeArgs = result.nativeArgs as { path: string } + expect(nativeArgs.path).toBe("src/split.ts") + } + }) + + it("handles name arriving before id with buffered args in between (reverse ordering)", () => { + const fullArgs = JSON.stringify({ path: "src/reverse.ts" }) + const firstHalf = fullArgs.slice(0, 9) + const secondHalf = fullArgs.slice(9) + + const { events, finalized } = drive([ + { index: 0, name: "read_file" }, + { index: 0, arguments: firstHalf }, + { index: 0, id: "call_reverse" }, + { index: 0, arguments: secondHalf }, + ]) + + // Start must not fire until the id arrives, so exactly one start with the late id. + const starts = events.filter((e) => e.type === "tool_call_start") + expect(starts).toHaveLength(1) + expect(starts[0].id).toBe("call_reverse") + + // The buffered delta must be flushed only after the start event. + const startIndex = events.findIndex((e) => e.type === "tool_call_start") + const firstDeltaIndex = events.findIndex((e) => e.type === "tool_call_delta") + expect(startIndex).toBeLessThan(firstDeltaIndex) + + const result = finalized.get("call_reverse") + expect(result).not.toBeNull() + expect(result?.type).toBe("tool_use") + if (result?.type === "tool_use") { + expect((result.nativeArgs as { path: string }).path).toBe("src/reverse.ts") + } + }) + + it("keeps two parallel tool calls on distinct indices isolated", () => { + const argsA = JSON.stringify({ path: "src/a.ts" }) + const argsB = JSON.stringify({ path: "src/b.ts" }) + + const { events, finalized } = drive([ + { index: 0, arguments: argsA.slice(0, 8) }, + { index: 1, arguments: argsB.slice(0, 8) }, + { index: 0, id: "call_a", name: "read_file" }, + { index: 1, id: "call_b", name: "read_file" }, + { index: 0, arguments: argsA.slice(8) }, + { index: 1, arguments: argsB.slice(8) }, + ]) + + const starts = events.filter((e) => e.type === "tool_call_start") + expect(starts).toHaveLength(2) + + const resultA = finalized.get("call_a") + const resultB = finalized.get("call_b") + expect(resultA).not.toBeNull() + expect(resultB).not.toBeNull() + if (resultA?.type === "tool_use") { + expect((resultA.nativeArgs as { path: string }).path).toBe("src/a.ts") + } + if (resultB?.type === "tool_use") { + expect((resultB.nativeArgs as { path: string }).path).toBe("src/b.ts") + } + }) + + it("emits the same event sequence for the single-chunk-with-id flow (regression guard)", () => { + const fullArgs = JSON.stringify({ path: "src/single.ts" }) + + const { events, finalized } = drive([ + { index: 0, id: "call_single", name: "read_file", arguments: fullArgs }, + ]) + + expect(events.map((e) => e.type)).toEqual(["tool_call_start", "tool_call_delta", "tool_call_end"]) + expect(events.every((e) => e.id === "call_single")).toBe(true) + + const result = finalized.get("call_single") + expect(result).not.toBeNull() + expect(result?.type).toBe("tool_use") + if (result?.type === "tool_use") { + expect((result.nativeArgs as { path: string }).path).toBe("src/single.ts") + } + }) + + it("does not emit a phantom tool_call_end for a tracker that never received an id", () => { + const { events } = drive([{ index: 0, arguments: '{"path":"orphan.ts"}' }]) + + expect(events.filter((e) => e.type === "tool_call_start")).toHaveLength(0) + expect(events.filter((e) => e.type === "tool_call_end")).toHaveLength(0) + }) + + it("finalizeRawChunks() emits end events and guards against missing id", () => { + const scope = NativeToolCallParser.createScope() + // Simulate a started tool call: process chunks to populate state + const chunks = [ + { index: 0, id: "call_finalize", name: "read_file" }, + { index: 0, arguments: '{"path":"file.ts"' }, + { index: 0, arguments: ',"mode":"slice"}' }, + ] + + const events: Array<{ type: string; id?: string }> = [] + for (const chunk of chunks) { + for (const event of NativeToolCallParser.processRawChunk(chunk, scope)) { + events.push(event) + if (event.type === "tool_call_start") { + NativeToolCallParser.startStreamingToolCall(event.id, event.name, scope) + } else if (event.type === "tool_call_delta") { + NativeToolCallParser.processStreamingChunk(event.id, event.delta, scope) + } + } + } + + // Now finalize the raw chunks to emit the end event + const finalizeEvents = NativeToolCallParser.finalizeRawChunks(scope) + for (const event of finalizeEvents) { + events.push(event) + } + + // Verify the end event was produced by finalizeRawChunks + const ends = events.filter((e) => e.type === "tool_call_end") + expect(ends).toHaveLength(1) + expect(ends[0].id).toBe("call_finalize") + + // Finalize the tool call to ensure it contains the complete arguments + const result = NativeToolCallParser.finalizeStreamingToolCall("call_finalize", scope) + expect(result?.type).toBe("tool_use") + if (result?.type === "tool_use") { + expect((result.nativeArgs as { path: string }).path).toBe("file.ts") + } + }) + + it("finalizeRawChunks() does not emit end for tracker without id", () => { + const scope = NativeToolCallParser.createScope() + // Start a tracker with arguments but no id, then finalize + const chunks = [{ index: 0, arguments: '{"incomplete":true}' }] + + for (const chunk of chunks) { + NativeToolCallParser.processRawChunk(chunk, scope) + } + + // Finalize should not emit an end event if id was never set + const finalizeEvents = NativeToolCallParser.finalizeRawChunks(scope) + const ends = finalizeEvents.filter((e) => e.type === "tool_call_end") + expect(ends).toHaveLength(0) + + NativeToolCallParser.clearRawChunkState(scope) + }) + + it("does not double-fire end events across repeated finalizeRawChunks calls", () => { + const scope = NativeToolCallParser.createScope() + // Drive a started tool call through the raw chunk path. + const chunks = [ + { index: 0, id: "call_dup", name: "read_file" }, + { index: 0, arguments: '{"path":"file.ts"}' }, + ] + for (const chunk of chunks) { + NativeToolCallParser.processRawChunk(chunk, scope) + } + + const finishEvents = NativeToolCallParser.finalizeRawChunks(scope) + const finalizeEvents = NativeToolCallParser.finalizeRawChunks(scope) + + const allEnds = [...finishEvents, ...finalizeEvents].filter((e) => e.type === "tool_call_end") + expect(allEnds).toHaveLength(1) + expect(allEnds[0].id).toBe("call_dup") + + NativeToolCallParser.clearRawChunkState(scope) + }) + }) })