Skip to content

Commit 9f515e1

Browse files
committed
feat(core/assistant-message): fallback normalizer for VSCode-LM function_calls/invoke → native <tool> XML; integrate in streaming/non-stream parsers; add minimal telemetry and tests
1 parent ed45d1c commit 9f515e1

File tree

7 files changed

+450
-21
lines changed

7 files changed

+450
-21
lines changed

src/core/assistant-message/AssistantMessageParser.ts

Lines changed: 33 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { type ToolName, toolNames } from "@roo-code/types"
22
import { TextContent, ToolUse, ToolParamName, toolParamNames } from "../../shared/tools"
33
import { AssistantMessageContent } from "./parseAssistantMessage"
4+
import { FunctionCallsStreamingNormalizer } from "./functionCallsNormalizer"
45

56
/**
67
* Parser for assistant messages. Maintains state between chunks
@@ -17,6 +18,11 @@ export class AssistantMessageParser {
1718
private readonly MAX_ACCUMULATOR_SIZE = 1024 * 1024 // 1MB limit
1819
private readonly MAX_PARAM_LENGTH = 1024 * 100 // 100KB per parameter limit
1920
private accumulator = ""
21+
// VSCode-LM function_calls/invoke streaming normalizer
22+
private normalizer = new FunctionCallsStreamingNormalizer()
23+
// Minimal telemetry flags (readable by caller if needed)
24+
public functionCallsNormalized = false
25+
public functionCallsToolNamesEncountered = new Set<string>()
2026

2127
/**
2228
* Initialize a new AssistantMessageParser instance.
@@ -37,6 +43,10 @@ export class AssistantMessageParser {
3743
this.currentParamName = undefined
3844
this.currentParamValueStartIndex = 0
3945
this.accumulator = ""
46+
// Reset normalizer and telemetry
47+
this.normalizer.reset()
48+
this.functionCallsNormalized = false
49+
this.functionCallsToolNamesEncountered.clear()
4050
}
4151

4252
/**
@@ -52,14 +62,24 @@ export class AssistantMessageParser {
5262
* @param chunk The new chunk of text to process.
5363
*/
5464
public processChunk(chunk: string): AssistantMessageContent[] {
55-
if (this.accumulator.length + chunk.length > this.MAX_ACCUMULATOR_SIZE) {
65+
// Pre-normalize VSCode-LM function_calls/invoke XML to native tool XML
66+
const normalizedChunk = this.normalizer.process(chunk)
67+
// Collect minimal telemetry
68+
if (this.normalizer.normalizedInLastChunk) {
69+
this.functionCallsNormalized = true
70+
}
71+
for (const name of this.normalizer.toolNamesEncountered) {
72+
this.functionCallsToolNamesEncountered.add(name)
73+
}
74+
75+
if (this.accumulator.length + normalizedChunk.length > this.MAX_ACCUMULATOR_SIZE) {
5676
throw new Error("Assistant message exceeds maximum allowed size")
5777
}
5878
// Store the current length of the accumulator before adding the new chunk
5979
const accumulatorStartLength = this.accumulator.length
6080

61-
for (let i = 0; i < chunk.length; i++) {
62-
const char = chunk[i]
81+
for (let i = 0; i < normalizedChunk.length; i++) {
82+
const char = normalizedChunk[i]
6383
this.accumulator += char
6484
const currentPosition = accumulatorStartLength + i
6585

@@ -78,10 +98,16 @@ export class AssistantMessageParser {
7898
// End of param value.
7999
// Do not trim content parameters to preserve newlines, but strip first and last newline only
80100
const paramValue = currentParamValue.slice(0, -paramClosingTag.length)
81-
this.currentToolUse.params[this.currentParamName] =
82-
this.currentParamName === "content"
83-
? paramValue.replace(/^\n/, "").replace(/\n$/, "")
84-
: paramValue.trim()
101+
if (this.currentParamName === "content") {
102+
this.currentToolUse.params[this.currentParamName] = paramValue
103+
.replace(/^\n/, "")
104+
.replace(/\n$/, "")
105+
} else if (this.currentParamName === "args") {
106+
// Preserve args exactly, including whitespace/newlines
107+
this.currentToolUse.params[this.currentParamName] = paramValue
108+
} else {
109+
this.currentToolUse.params[this.currentParamName] = paramValue.trim()
110+
}
85111
this.currentParamName = undefined
86112
continue
87113
} else {

src/core/assistant-message/__tests__/AssistantMessageParser.spec.ts

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,3 +392,69 @@ describe("AssistantMessageParser (streaming)", () => {
392392
})
393393
})
394394
})
395+
396+
// VSCode-LM function_calls normalizer tests (streaming)
397+
describe("VSCode-LM function_calls normalizer (streaming)", () => {
398+
it("should normalize single invoke with args preserved", () => {
399+
const parser = new AssistantMessageParser()
400+
const argsXml = "<file><path>src/a.ts</path></file>"
401+
const message = `<function_calls><invoke name="read_file"><args>${argsXml}</args></invoke></function_calls>`
402+
const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
403+
expect(result).toHaveLength(1)
404+
const toolUse = result[0] as ToolUse
405+
expect(toolUse.type).toBe("tool_use")
406+
expect(toolUse.name).toBe("read_file")
407+
expect(toolUse.params.args).toBe(argsXml)
408+
expect(toolUse.partial).toBe(false)
409+
})
410+
411+
it("should handle multiple invokes with surrounding text", () => {
412+
const parser = new AssistantMessageParser()
413+
const args1 = "<file><path>file1.ts</path></file>"
414+
const args2 = "<file><path>file2.ts</path></file>"
415+
const message = `Before <function_calls><invoke name="read_file"><args>${args1}</args></invoke></function_calls> Middle <function_calls><invoke name="read_file"><args>${args2}</args></invoke></function_calls> After`
416+
const result = streamChunks(parser, message)
417+
expect(result).toHaveLength(5)
418+
419+
expect(result[0].type).toBe("text")
420+
expect((result[0] as TextContent).content).toBe("Before")
421+
422+
const toolUse1 = result[1] as ToolUse
423+
expect(toolUse1.type).toBe("tool_use")
424+
expect(toolUse1.name).toBe("read_file")
425+
expect(toolUse1.params.args).toBe(args1)
426+
427+
expect(result[2].type).toBe("text")
428+
expect((result[2] as TextContent).content).toBe("Middle")
429+
430+
const toolUse2 = result[3] as ToolUse
431+
expect(toolUse2.type).toBe("tool_use")
432+
expect(toolUse2.name).toBe("read_file")
433+
expect(toolUse2.params.args).toBe(args2)
434+
435+
expect(result[4].type).toBe("text")
436+
expect((result[4] as TextContent).content).toBe("After")
437+
})
438+
439+
it("should pass through unknown invoke as text and not create tool_use", () => {
440+
const parser = new AssistantMessageParser()
441+
const message = `<function_calls><invoke name="unknown_tool"><args><x>y</x></args></invoke></function_calls>`
442+
const result = streamChunks(parser, message)
443+
expect(result).toHaveLength(1)
444+
const text = result[0] as TextContent
445+
expect(text.type).toBe("text")
446+
expect(text.content).toContain('<invoke name="unknown_tool">')
447+
})
448+
449+
it("should preserve multi-file args xml exactly", () => {
450+
const parser = new AssistantMessageParser()
451+
const argsXml = "<file><path>a.ts</path></file><file><path>b.ts</path></file>"
452+
const message = `<function_calls><invoke name="read_file"><args>${argsXml}</args></invoke></function_calls>`
453+
const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
454+
expect(result).toHaveLength(1)
455+
const toolUse = result[0] as ToolUse
456+
expect(toolUse.type).toBe("tool_use")
457+
expect(toolUse.name).toBe("read_file")
458+
expect(toolUse.params.args).toBe(argsXml)
459+
})
460+
})

src/core/assistant-message/__tests__/parseAssistantMessage.spec.ts

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,3 +338,78 @@ const isEmptyTextContent = (block: AssistantMessageContent) =>
338338
})
339339
})
340340
})
341+
342+
// VSCode-LM function_calls normalizer tests (non-stream)
343+
;[parseAssistantMessageV1, parseAssistantMessageV2].forEach((parser, index) => {
344+
describe(`VSCode-LM function_calls normalizer (non-stream) V${index + 1}`, () => {
345+
it("should normalize single invoke with args preserved", () => {
346+
const argsXml = "<file><path>src/a.ts</path></file>"
347+
const message = `<function_calls><invoke name="read_file"><args>${argsXml}</args></invoke></function_calls>`
348+
const result = parser(message).filter((block) => !isEmptyTextContent(block))
349+
expect(result).toHaveLength(1)
350+
const toolUse = result[0] as ToolUse
351+
expect(toolUse.type).toBe("tool_use")
352+
expect(toolUse.name).toBe("read_file")
353+
expect(toolUse.params.args).toBe(argsXml)
354+
expect(toolUse.partial).toBe(false)
355+
})
356+
357+
it("should handle multiple invokes with surrounding text", () => {
358+
const args1 = "<file><path>file1.ts</path></file>"
359+
const args2 = "<file><path>file2.ts</path></file>"
360+
const message = `Before <function_calls><invoke name="read_file"><args>${args1}</args></invoke></function_calls> Middle <function_calls><invoke name="read_file"><args>${args2}</args></invoke></function_calls> After`
361+
const result = parser(message)
362+
363+
expect(result).toHaveLength(5)
364+
365+
expect(result[0].type).toBe("text")
366+
expect((result[0] as TextContent).content).toBe("Before")
367+
368+
const toolUse1 = result[1] as ToolUse
369+
expect(toolUse1.type).toBe("tool_use")
370+
expect(toolUse1.name).toBe("read_file")
371+
expect(toolUse1.params.args).toBe(args1)
372+
373+
expect(result[2].type).toBe("text")
374+
expect((result[2] as TextContent).content).toBe("Middle")
375+
376+
const toolUse2 = result[3] as ToolUse
377+
expect(toolUse2.type).toBe("tool_use")
378+
expect(toolUse2.name).toBe("read_file")
379+
expect(toolUse2.params.args).toBe(args2)
380+
381+
expect(result[4].type).toBe("text")
382+
expect((result[4] as TextContent).content).toBe("After")
383+
})
384+
385+
it("should pass through unknown invoke as text and not create tool_use", () => {
386+
const message = `<function_calls><invoke name="unknown_tool"><args><x>y</x></args></invoke></function_calls>`
387+
const result = parser(message)
388+
expect(result).toHaveLength(1)
389+
const text = result[0] as TextContent
390+
expect(text.type).toBe("text")
391+
expect(text.content).toContain('<invoke name="unknown_tool">')
392+
})
393+
394+
it("should preserve multi-file args xml exactly", () => {
395+
const argsXml = "<file><path>a.ts</path></file><file><path>b.ts</path></file>"
396+
const message = `<function_calls><invoke name="read_file"><args>${argsXml}</args></invoke></function_calls>`
397+
const result = parser(message).filter((block) => !isEmptyTextContent(block))
398+
expect(result).toHaveLength(1)
399+
const toolUse = result[0] as ToolUse
400+
expect(toolUse.type).toBe("tool_use")
401+
expect(toolUse.name).toBe("read_file")
402+
expect(toolUse.params.args).toBe(argsXml)
403+
})
404+
405+
it("should be idempotent for native tool XML (no changes)", () => {
406+
const native = "<read_file><path>src/x.ts</path></read_file>"
407+
const result = parser(native).filter((b) => !isEmptyTextContent(b))
408+
expect(result).toHaveLength(1)
409+
const toolUse = result[0] as ToolUse
410+
expect(toolUse.name).toBe("read_file")
411+
expect(toolUse.params.path).toBe("src/x.ts")
412+
expect(toolUse.partial).toBe(false)
413+
})
414+
})
415+
})

0 commit comments

Comments
 (0)