use assistantMessageParser class instead of parseAssistantMessage (#5341)

Co-authored-by: Daniel Riccio <[email protected]>

packages/types/src/experiment.ts

@@ -6,7 +6,7 @@ import type { Keys, Equals, AssertEqual } from "./type-fu.js"
  * ExperimentId
  */
 
-export const experimentIds = ["powerSteering", "multiFileApplyDiff", "preventFocusDisruption"] as const
+export const experimentIds = ["powerSteering", "multiFileApplyDiff", "preventFocusDisruption", "assistantMessageParser"] as const
 
 export const experimentIdsSchema = z.enum(experimentIds)
 
@@ -20,6 +20,7 @@ export const experimentsSchema = z.object({
 	powerSteering: z.boolean().optional(),
 	multiFileApplyDiff: z.boolean().optional(),
 	preventFocusDisruption: z.boolean().optional(),
+	assistantMessageParser: z.boolean().optional(),
 })
 
 export type Experiments = z.infer<typeof experimentsSchema>

src/core/assistant-message/AssistantMessageParser.ts

@@ -0,0 +1,251 @@
+import { type ToolName, toolNames } from "@roo-code/types"
+import { TextContent, ToolUse, ToolParamName, toolParamNames } from "../../shared/tools"
+import { AssistantMessageContent } from "./parseAssistantMessage"
+
+/**
+ * Parser for assistant messages. Maintains state between chunks
+ * to avoid reprocessing the entire message on each update.
+ */
+export class AssistantMessageParser {
+	private contentBlocks: AssistantMessageContent[] = []
+	private currentTextContent: TextContent | undefined = undefined
+	private currentTextContentStartIndex = 0
+	private currentToolUse: ToolUse | undefined = undefined
+	private currentToolUseStartIndex = 0
+	private currentParamName: ToolParamName | undefined = undefined
+	private currentParamValueStartIndex = 0
+	private readonly MAX_ACCUMULATOR_SIZE = 1024 * 1024 // 1MB limit
+	private readonly MAX_PARAM_LENGTH = 1024 * 100 // 100KB per parameter limit
+	private accumulator = ""
+
+	/**
+	 * Initialize a new AssistantMessageParser instance.
+	 */
+	constructor() {
+		this.reset()
+	}
+
+	/**
+	 * Reset the parser state.
+	 */
+	public reset(): void {
+		this.contentBlocks = []
+		this.currentTextContent = undefined
+		this.currentTextContentStartIndex = 0
+		this.currentToolUse = undefined
+		this.currentToolUseStartIndex = 0
+		this.currentParamName = undefined
+		this.currentParamValueStartIndex = 0
+		this.accumulator = ""
+	}
+
+	/**
+	 * Returns the current parsed content blocks
+	 */
+
+	public getContentBlocks(): AssistantMessageContent[] {
+		// Return a shallow copy to prevent external mutation
+		return this.contentBlocks.slice()
+	}
+	/**
+	 * Process a new chunk of text and update the parser state.
+	 * @param chunk The new chunk of text to process.
+	 */
+	public processChunk(chunk: string): AssistantMessageContent[] {
+		if (this.accumulator.length + chunk.length > this.MAX_ACCUMULATOR_SIZE) {
+			throw new Error("Assistant message exceeds maximum allowed size")
+		}
+		// Store the current length of the accumulator before adding the new chunk
+		const accumulatorStartLength = this.accumulator.length
+
+		for (let i = 0; i < chunk.length; i++) {
+			const char = chunk[i]
+			this.accumulator += char
+			const currentPosition = accumulatorStartLength + i
+
+			// There should not be a param without a tool use.
+			if (this.currentToolUse && this.currentParamName) {
+				const currentParamValue = this.accumulator.slice(this.currentParamValueStartIndex)
+				if (currentParamValue.length > this.MAX_PARAM_LENGTH) {
+					// Reset to a safe state
+					this.currentParamName = undefined
+					this.currentParamValueStartIndex = 0
+					continue
+				}
+				const paramClosingTag = `</${this.currentParamName}>`
+				// Streamed param content: always write the currently accumulated value
+				if (currentParamValue.endsWith(paramClosingTag)) {
+					// End of param value.
+					// Do not trim content parameters to preserve newlines, but strip first and last newline only
+					const paramValue = currentParamValue.slice(0, -paramClosingTag.length)
+					this.currentToolUse.params[this.currentParamName] =
+						this.currentParamName === "content"
+							? paramValue.replace(/^\n/, "").replace(/\n$/, "")
+							: paramValue.trim()
+					this.currentParamName = undefined
+					continue
+				} else {
+					// Partial param value is accumulating.
+					// Write the currently accumulated param content in real time
+					this.currentToolUse.params[this.currentParamName] = currentParamValue
+					continue
+				}
+			}
+
+			// No currentParamName.
+
+			if (this.currentToolUse) {
+				const currentToolValue = this.accumulator.slice(this.currentToolUseStartIndex)
+				const toolUseClosingTag = `</${this.currentToolUse.name}>`
+				if (currentToolValue.endsWith(toolUseClosingTag)) {
+					// End of a tool use.
+					this.currentToolUse.partial = false
+
+					this.currentToolUse = undefined
+					continue
+				} else {
+					const possibleParamOpeningTags = toolParamNames.map((name) => `<${name}>`)
+					for (const paramOpeningTag of possibleParamOpeningTags) {
+						if (this.accumulator.endsWith(paramOpeningTag)) {
+							// Start of a new parameter.
+							const paramName = paramOpeningTag.slice(1, -1)
+							if (!toolParamNames.includes(paramName as ToolParamName)) {
+								// Handle invalid parameter name gracefully
+								continue
+							}
+							this.currentParamName = paramName as ToolParamName
+							this.currentParamValueStartIndex = this.accumulator.length
+							break
+						}
+					}
+
+					// There's no current param, and not starting a new param.
+
+					// Special case for write_to_file where file contents could
+					// contain the closing tag, in which case the param would have
+					// closed and we end up with the rest of the file contents here.
+					// To work around this, get the string between the starting
+					// content tag and the LAST content tag.
+					const contentParamName: ToolParamName = "content"
+
+					if (
+						this.currentToolUse.name === "write_to_file" &&
+						this.accumulator.endsWith(`</${contentParamName}>`)
+					) {
+						const toolContent = this.accumulator.slice(this.currentToolUseStartIndex)
+						const contentStartTag = `<${contentParamName}>`
+						const contentEndTag = `</${contentParamName}>`
+						const contentStartIndex = toolContent.indexOf(contentStartTag) + contentStartTag.length
+						const contentEndIndex = toolContent.lastIndexOf(contentEndTag)
+
+						if (contentStartIndex !== -1 && contentEndIndex !== -1 && contentEndIndex > contentStartIndex) {
+							// Don't trim content to preserve newlines, but strip first and last newline only
+							this.currentToolUse.params[contentParamName] = toolContent
+								.slice(contentStartIndex, contentEndIndex)
+								.replace(/^\n/, "")
+								.replace(/\n$/, "")
+						}
+					}
+
+					// Partial tool value is accumulating.
+					continue
+				}
+			}
+
+			// No currentToolUse.
+
+			let didStartToolUse = false
+			const possibleToolUseOpeningTags = toolNames.map((name) => `<${name}>`)
+
+			for (const toolUseOpeningTag of possibleToolUseOpeningTags) {
+				if (this.accumulator.endsWith(toolUseOpeningTag)) {
+					// Extract and validate the tool name
+					const extractedToolName = toolUseOpeningTag.slice(1, -1)
+
+					// Check if the extracted tool name is valid
+					if (!toolNames.includes(extractedToolName as ToolName)) {
+						// Invalid tool name, treat as plain text and continue
+						continue
+					}
+
+					// Start of a new tool use.
+					this.currentToolUse = {
+						type: "tool_use",
+						name: extractedToolName as ToolName,
+						params: {},
+						partial: true,
+					}
+
+					this.currentToolUseStartIndex = this.accumulator.length
+
+					// This also indicates the end of the current text content.
+					if (this.currentTextContent) {
+						this.currentTextContent.partial = false
+
+						// Remove the partially accumulated tool use tag from the
+						// end of text (<tool).
+						this.currentTextContent.content = this.currentTextContent.content
+							.slice(0, -toolUseOpeningTag.slice(0, -1).length)
+							.trim()
+
+						// No need to push, currentTextContent is already in contentBlocks
+						this.currentTextContent = undefined
+					}
+
+					// Immediately push new tool_use block as partial
+					let idx = this.contentBlocks.findIndex((block) => block === this.currentToolUse)
+					if (idx === -1) {
+						this.contentBlocks.push(this.currentToolUse)
+					}
+
+					didStartToolUse = true
+					break
+				}
+			}
+
+			if (!didStartToolUse) {
+				// No tool use, so it must be text either at the beginning or
+				// between tools.
+				if (this.currentTextContent === undefined) {
+					// If this is the first chunk and we're at the beginning of processing,
+					// set the start index to the current position in the accumulator
+					this.currentTextContentStartIndex = currentPosition
+
+					// Create a new text content block and add it to contentBlocks
+					this.currentTextContent = {
+						type: "text",
+						content: this.accumulator.slice(this.currentTextContentStartIndex).trim(),
+						partial: true,
+					}
+
+					// Add the new text content to contentBlocks immediately
+					// Ensures it appears in the UI right away
+					this.contentBlocks.push(this.currentTextContent)
+				} else {
+					// Update the existing text content
+					this.currentTextContent.content = this.accumulator.slice(this.currentTextContentStartIndex).trim()
+				}
+			}
+		}
+		// Do not call finalizeContentBlocks() here.
+		// Instead, update any partial blocks in the array and add new ones as they're completed.
+		// This matches the behavior of the original parseAssistantMessage function.
+		return this.getContentBlocks()
+	}
+
+	/**
+	 * Finalize any partial content blocks.
+	 * Should be called after processing the last chunk.
+	 */
+	public finalizeContentBlocks(): void {
+		// Mark all partial blocks as complete
+		for (const block of this.contentBlocks) {
+			if (block.partial) {
+				block.partial = false
+			}
+			if (block.type === "text" && typeof block.content === "string") {
+				block.content = block.content.trim()
+			}
+		}
+	}
+}

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

@@ -0,0 +1,396 @@
+// npx vitest src/core/assistant-message/__tests__/AssistantMessageParser.spec.ts
+
+import { describe, it, expect, beforeEach } from "vitest"
+import { AssistantMessageParser } from "../AssistantMessageParser"
+import { AssistantMessageContent } from "../parseAssistantMessage"
+import { TextContent, ToolUse } from "../../../shared/tools"
+import { toolNames } from "@roo-code/types"
+
+/**
+ * Helper to filter out empty text content blocks.
+ */
+const isEmptyTextContent = (block: any) => block.type === "text" && (block as TextContent).content === ""
+
+/**
+ * Helper to simulate streaming by feeding the parser deterministic "random"-sized chunks (1-10 chars).
+ * Uses a seeded pseudo-random number generator for deterministic chunking.
+ */
+
+// Simple linear congruential generator (LCG) for deterministic pseudo-random numbers
+function createSeededRandom(seed: number) {
+	let state = seed
+	return {
+		next: () => {
+			// LCG parameters from Numerical Recipes
+			state = (state * 1664525 + 1013904223) % 0x100000000
+			return state / 0x100000000
+		},
+	}
+}
+
+function streamChunks(
+	parser: AssistantMessageParser,
+	message: string,
+): ReturnType<AssistantMessageParser["getContentBlocks"]> {
+	let result: AssistantMessageContent[] = []
+	let i = 0
+	const rng = createSeededRandom(42) // Fixed seed for deterministic tests
+	while (i < message.length) {
+		// Deterministic chunk size between 1 and 10, but not exceeding message length
+		const chunkSize = Math.min(message.length - i, Math.floor(rng.next() * 10) + 1)
+		const chunk = message.slice(i, i + chunkSize)
+		result = parser.processChunk(chunk)
+		i += chunkSize
+	}
+	return result
+}
+
+describe("AssistantMessageParser (streaming)", () => {
+	let parser: AssistantMessageParser
+
+	beforeEach(() => {
+		parser = new AssistantMessageParser()
+	})
+
+	describe("text content streaming", () => {
+		it("should accumulate a simple text message chunk by chunk", () => {
+			const message = "Hello, this is a test."
+			const result = streamChunks(parser, message)
+			expect(result).toHaveLength(1)
+			expect(result[0]).toEqual({
+				type: "text",
+				content: message,
+				partial: true,
+			})
+		})
+
+		it("should accumulate multi-line text message chunk by chunk", () => {
+			const message = "Line 1\nLine 2\nLine 3"
+			const result = streamChunks(parser, message)
+			expect(result).toHaveLength(1)
+			expect(result[0]).toEqual({
+				type: "text",
+				content: message,
+				partial: true,
+			})
+		})
+	})
+
+	describe("tool use streaming", () => {
+		it("should parse a tool use with parameter, streamed char by char", () => {
+			const message = "<read_file><path>src/file.ts</path></read_file>"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("read_file")
+			expect(toolUse.params.path).toBe("src/file.ts")
+			expect(toolUse.partial).toBe(false)
+		})
+
+		it("should mark tool use as partial when not closed", () => {
+			const message = "<read_file><path>src/file.ts</path>"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("read_file")
+			expect(toolUse.params.path).toBe("src/file.ts")
+			expect(toolUse.partial).toBe(true)
+		})
+
+		it("should handle a partial parameter in a tool use", () => {
+			const message = "<read_file><path>src/file"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("read_file")
+			expect(toolUse.params.path).toBe("src/file")
+			expect(toolUse.partial).toBe(true)
+		})
+
+		it("should handle tool use with multiple parameters streamed", () => {
+			const message =
+				"<read_file><path>src/file.ts</path><start_line>10</start_line><end_line>20</end_line></read_file>"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("read_file")
+			expect(toolUse.params.path).toBe("src/file.ts")
+			expect(toolUse.params.start_line).toBe("10")
+			expect(toolUse.params.end_line).toBe("20")
+			expect(toolUse.partial).toBe(false)
+		})
+	})
+
+	describe("mixed content streaming", () => {
+		it("should parse text followed by a tool use, streamed", () => {
+			const message = "Text before tool <read_file><path>src/file.ts</path></read_file>"
+			const result = streamChunks(parser, message)
+			expect(result).toHaveLength(2)
+			const textContent = result[0] as TextContent
+			expect(textContent.type).toBe("text")
+			expect(textContent.content).toBe("Text before tool")
+			expect(textContent.partial).toBe(false)
+			const toolUse = result[1] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("read_file")
+			expect(toolUse.params.path).toBe("src/file.ts")
+			expect(toolUse.partial).toBe(false)
+		})
+
+		it("should parse a tool use followed by text, streamed", () => {
+			const message = "<read_file><path>src/file.ts</path></read_file>Text after tool"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(2)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("read_file")
+			expect(toolUse.params.path).toBe("src/file.ts")
+			expect(toolUse.partial).toBe(false)
+			const textContent = result[1] as TextContent
+			expect(textContent.type).toBe("text")
+			expect(textContent.content).toBe("Text after tool")
+			expect(textContent.partial).toBe(true)
+		})
+
+		it("should parse multiple tool uses separated by text, streamed", () => {
+			const message =
+				"First: <read_file><path>file1.ts</path></read_file>Second: <read_file><path>file2.ts</path></read_file>"
+			const result = streamChunks(parser, message)
+			expect(result).toHaveLength(4)
+			expect(result[0].type).toBe("text")
+			expect((result[0] as TextContent).content).toBe("First:")
+			expect(result[1].type).toBe("tool_use")
+			expect((result[1] as ToolUse).name).toBe("read_file")
+			expect((result[1] as ToolUse).params.path).toBe("file1.ts")
+			expect(result[2].type).toBe("text")
+			expect((result[2] as TextContent).content).toBe("Second:")
+			expect(result[3].type).toBe("tool_use")
+			expect((result[3] as ToolUse).name).toBe("read_file")
+			expect((result[3] as ToolUse).params.path).toBe("file2.ts")
+		})
+	})
+
+	describe("special and edge cases", () => {
+		it("should handle the write_to_file tool with content that contains closing tags", () => {
+			const message = `<write_to_file><path>src/file.ts</path><content>
+	function example() {
+	// This has XML-like content: </content>
+	return true;
+	}
+	</content><line_count>5</line_count></write_to_file>`
+
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("write_to_file")
+			expect(toolUse.params.path).toBe("src/file.ts")
+			expect(toolUse.params.line_count).toBe("5")
+			expect(toolUse.params.content).toContain("function example()")
+			expect(toolUse.params.content).toContain("// This has XML-like content: </content>")
+			expect(toolUse.params.content).toContain("return true;")
+			expect(toolUse.partial).toBe(false)
+		})
+		it("should handle empty messages", () => {
+			const message = ""
+			const result = streamChunks(parser, message)
+			expect(result).toHaveLength(0)
+		})
+
+		it("should handle malformed tool use tags as plain text", () => {
+			const message = "This has a <not_a_tool>malformed tag</not_a_tool>"
+			const result = streamChunks(parser, message)
+			expect(result).toHaveLength(1)
+			expect(result[0].type).toBe("text")
+			expect((result[0] as TextContent).content).toBe(message)
+		})
+
+		it("should handle tool use with no parameters", () => {
+			const message = "<browser_action></browser_action>"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("browser_action")
+			expect(Object.keys(toolUse.params).length).toBe(0)
+			expect(toolUse.partial).toBe(false)
+		})
+
+		it("should handle a tool use with a parameter containing XML-like content", () => {
+			const message = "<search_files><regex><div>.*</div></regex><path>src</path></search_files>"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("search_files")
+			expect(toolUse.params.regex).toBe("<div>.*</div>")
+			expect(toolUse.params.path).toBe("src")
+			expect(toolUse.partial).toBe(false)
+		})
+
+		it("should handle consecutive tool uses without text in between", () => {
+			const message = "<read_file><path>file1.ts</path></read_file><read_file><path>file2.ts</path></read_file>"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(2)
+			const toolUse1 = result[0] as ToolUse
+			expect(toolUse1.type).toBe("tool_use")
+			expect(toolUse1.name).toBe("read_file")
+			expect(toolUse1.params.path).toBe("file1.ts")
+			expect(toolUse1.partial).toBe(false)
+			const toolUse2 = result[1] as ToolUse
+			expect(toolUse2.type).toBe("tool_use")
+			expect(toolUse2.name).toBe("read_file")
+			expect(toolUse2.params.path).toBe("file2.ts")
+			expect(toolUse2.partial).toBe(false)
+		})
+
+		it("should handle whitespace in parameters", () => {
+			const message = "<read_file><path>  src/file.ts  </path></read_file>"
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("read_file")
+			expect(toolUse.params.path).toBe("src/file.ts")
+			expect(toolUse.partial).toBe(false)
+		})
+
+		it("should handle multi-line parameters", () => {
+			const message = `<write_to_file><path>file.ts</path><content>
+	line 1
+	line 2
+	line 3
+	</content><line_count>3</line_count></write_to_file>`
+			const result = streamChunks(parser, message).filter((block) => !isEmptyTextContent(block))
+
+			expect(result).toHaveLength(1)
+			const toolUse = result[0] as ToolUse
+			expect(toolUse.type).toBe("tool_use")
+			expect(toolUse.name).toBe("write_to_file")
+			expect(toolUse.params.path).toBe("file.ts")
+			expect(toolUse.params.content).toContain("line 1")
+			expect(toolUse.params.content).toContain("line 2")
+			expect(toolUse.params.content).toContain("line 3")
+			expect(toolUse.params.line_count).toBe("3")
+			expect(toolUse.partial).toBe(false)
+		})
+		it("should handle a complex message with multiple content types", () => {
+			const message = `I'll help you with that task.
+
+	<read_file><path>src/index.ts</path></read_file>
+
+	Now let's modify the file:
+
+	<write_to_file><path>src/index.ts</path><content>
+	// Updated content
+	console.log("Hello world");
+	</content><line_count>2</line_count></write_to_file>
+
+	Let's run the code:
+
+	<execute_command><command>node src/index.ts</command></execute_command>`
+
+			const result = streamChunks(parser, message)
+
+			expect(result).toHaveLength(6)
+
+			// First text block
+			expect(result[0].type).toBe("text")
+			expect((result[0] as TextContent).content).toBe("I'll help you with that task.")
+
+			// First tool use (read_file)
+			expect(result[1].type).toBe("tool_use")
+			expect((result[1] as ToolUse).name).toBe("read_file")
+
+			// Second text block
+			expect(result[2].type).toBe("text")
+			expect((result[2] as TextContent).content).toContain("Now let's modify the file:")
+
+			// Second tool use (write_to_file)
+			expect(result[3].type).toBe("tool_use")
+			expect((result[3] as ToolUse).name).toBe("write_to_file")
+
+			// Third text block
+			expect(result[4].type).toBe("text")
+			expect((result[4] as TextContent).content).toContain("Let's run the code:")
+
+			// Third tool use (execute_command)
+			expect(result[5].type).toBe("tool_use")
+			expect((result[5] as ToolUse).name).toBe("execute_command")
+		})
+	})
+
+	describe("size limit handling", () => {
+		it("should throw an error when MAX_ACCUMULATOR_SIZE is exceeded", () => {
+			// Create a message that exceeds 1MB (MAX_ACCUMULATOR_SIZE)
+			const largeMessage = "x".repeat(1024 * 1024 + 1) // 1MB + 1 byte
+
+			expect(() => {
+				parser.processChunk(largeMessage)
+			}).toThrow("Assistant message exceeds maximum allowed size")
+		})
+
+		it("should gracefully handle a parameter that exceeds MAX_PARAM_LENGTH", () => {
+			// Create a parameter value that exceeds 100KB (MAX_PARAM_LENGTH)
+			const largeParamValue = "x".repeat(1024 * 100 + 1) // 100KB + 1 byte
+			const message = `<write_to_file><path>test.txt</path><content>${largeParamValue}</content></write_to_file>After tool`
+
+			// Process the message in chunks to simulate streaming
+			let result: AssistantMessageContent[] = []
+			let error: Error | null = null
+
+			try {
+				// Process the opening tags
+				result = parser.processChunk("<write_to_file><path>test.txt</path><content>")
+
+				// Process the large parameter value in chunks
+				const chunkSize = 1000
+				for (let i = 0; i < largeParamValue.length; i += chunkSize) {
+					const chunk = largeParamValue.slice(i, i + chunkSize)
+					result = parser.processChunk(chunk)
+				}
+
+				// Process the closing tags and text after
+				result = parser.processChunk("</content></write_to_file>After tool")
+			} catch (e) {
+				error = e as Error
+			}
+
+			// Should not throw an error
+			expect(error).toBeNull()
+
+			// Should have processed the content
+			expect(result.length).toBeGreaterThan(0)
+
+			// The tool use should exist but the content parameter should be reset/empty
+			const toolUse = result.find((block) => block.type === "tool_use") as ToolUse
+			expect(toolUse).toBeDefined()
+			expect(toolUse.name).toBe("write_to_file")
+			expect(toolUse.params.path).toBe("test.txt")
+
+			// The text after the tool should still be parsed
+			const textAfter = result.find(
+				(block) => block.type === "text" && (block as TextContent).content.includes("After tool"),
+			)
+			expect(textAfter).toBeDefined()
+		})
+	})
+
+	describe("finalizeContentBlocks", () => {
+		it("should mark all partial blocks as complete", () => {
+			const message = "<read_file><path>src/file.ts"
+			streamChunks(parser, message)
+			let blocks = parser.getContentBlocks()
+			// The block may already be partial or not, depending on chunking.
+			// To ensure the test is robust, we only assert after finalizeContentBlocks.
+			parser.finalizeContentBlocks()
+			blocks = parser.getContentBlocks()
+			expect(blocks[0].partial).toBe(false)
+		})
+	})
+})

src/core/task/Task.ts

@@ -77,7 +77,8 @@ import { ToolRepetitionDetector } from "../tools/ToolRepetitionDetector"
 import { FileContextTracker } from "../context-tracking/FileContextTracker"
 import { RooIgnoreController } from "../ignore/RooIgnoreController"
 import { RooProtectedController } from "../protect/RooProtectedController"
-import { type AssistantMessageContent, parseAssistantMessage, presentAssistantMessage } from "../assistant-message"
+import { type AssistantMessageContent, presentAssistantMessage, parseAssistantMessage } from "../assistant-message"
+import { AssistantMessageParser } from "../assistant-message/AssistantMessageParser"
 import { truncateConversationIfNeeded } from "../sliding-window"
 import { ClineProvider } from "../webview/ClineProvider"
 import { MultiSearchReplaceDiffStrategy } from "../diff/strategies/multi-search-replace"
@@ -249,6 +250,8 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	didRejectTool = false
 	didAlreadyUseTool = false
 	didCompleteReadingStream = false
+	assistantMessageParser?: AssistantMessageParser
+	isAssistantMessageParserEnabled = false
 
 	constructor({
 		provider,
@@ -1553,6 +1556,9 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			this.didAlreadyUseTool = false
 			this.presentAssistantMessageLocked = false
 			this.presentAssistantMessageHasPendingUpdates = false
+			if (this.assistantMessageParser) {
+				this.assistantMessageParser.reset()
+			}
 
 			await this.diffViewProvider.reset()
 
@@ -1587,9 +1593,14 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 						case "text": {
 							assistantMessage += chunk.text
 
-							// Parse raw assistant message into content blocks.
+							// Parse raw assistant message chunk into content blocks.
 							const prevLength = this.assistantMessageContent.length
-							this.assistantMessageContent = parseAssistantMessage(assistantMessage)
+							if (this.isAssistantMessageParserEnabled && this.assistantMessageParser) {
+								this.assistantMessageContent = this.assistantMessageParser.processChunk(chunk.text)
+							} else {
+								// Use the old parsing method when experiment is disabled
+								this.assistantMessageContent = parseAssistantMessage(assistantMessage)
+							}
 
 							if (this.assistantMessageContent.length > prevLength) {
 								// New content we need to present, reset to
@@ -1709,6 +1720,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			// Can't just do this b/c a tool could be in the middle of executing.
 			// this.assistantMessageContent.forEach((e) => (e.partial = false))
 
+			// Now that the stream is complete, finalize any remaining partial content blocks
+			if (this.isAssistantMessageParserEnabled && this.assistantMessageParser) {
+				this.assistantMessageParser.finalizeContentBlocks()
+				this.assistantMessageContent = this.assistantMessageParser.getContentBlocks()
+			}
+			// When using old parser, no finalization needed - parsing already happened during streaming
+
 			if (partialBlocks.length > 0) {
 				// If there is content to update then it will complete and
 				// update `this.userMessageContentReady` to true, which we
@@ -1722,6 +1740,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			await this.saveClineMessages()
 			await this.providerRef.deref()?.postStateToWebview()
 
+			// Reset parser after each complete conversation round
+			if (this.assistantMessageParser) {
+				this.assistantMessageParser.reset()
+			}
+
 			// Now add to apiConversationHistory.
 			// Need to save assistant responses to file before proceeding to
 			// tool use since user can exit at any moment and we wouldn't be

src/shared/__tests__/experiments.spec.ts

@@ -29,6 +29,7 @@ describe("experiments", () => {
 				powerSteering: false,
 				multiFileApplyDiff: false,
 				preventFocusDisruption: false,
+				assistantMessageParser: false,
 			}
 			expect(Experiments.isEnabled(experiments, EXPERIMENT_IDS.POWER_STEERING)).toBe(false)
 		})
@@ -38,6 +39,7 @@ describe("experiments", () => {
 				powerSteering: true,
 				multiFileApplyDiff: false,
 				preventFocusDisruption: false,
+				assistantMessageParser: false,
 			}
 			expect(Experiments.isEnabled(experiments, EXPERIMENT_IDS.POWER_STEERING)).toBe(true)
 		})
@@ -47,6 +49,7 @@ describe("experiments", () => {
 				powerSteering: false,
 				multiFileApplyDiff: false,
 				preventFocusDisruption: false,
+				assistantMessageParser: false,
 			}
 			expect(Experiments.isEnabled(experiments, EXPERIMENT_IDS.POWER_STEERING)).toBe(false)
 		})

src/shared/experiments.ts

@@ -4,6 +4,7 @@ export const EXPERIMENT_IDS = {
 	MULTI_FILE_APPLY_DIFF: "multiFileApplyDiff",
 	POWER_STEERING: "powerSteering",
 	PREVENT_FOCUS_DISRUPTION: "preventFocusDisruption",
+	ASSISTANT_MESSAGE_PARSER: "assistantMessageParser",
 } as const satisfies Record<string, ExperimentId>
 
 type _AssertExperimentIds = AssertEqual<Equals<ExperimentId, Values<typeof EXPERIMENT_IDS>>>
@@ -18,6 +19,7 @@ export const experimentConfigsMap: Record<ExperimentKey, ExperimentConfig> = {
 	MULTI_FILE_APPLY_DIFF: { enabled: false },
 	POWER_STEERING: { enabled: false },
 	PREVENT_FOCUS_DISRUPTION: { enabled: false },
+	ASSISTANT_MESSAGE_PARSER: { enabled: false },
 }
 
 export const experimentDefault = Object.fromEntries(

webview-ui/src/context/__tests__/ExtensionStateContext.spec.tsx

@@ -229,6 +229,7 @@ describe("mergeExtensionState", () => {
 				concurrentFileReads: true,
 				multiFileApplyDiff: true,
 				preventFocusDisruption: false,
+				assistantMessageParser: false,
 			} as Record<ExperimentId, boolean>,
 		}
 
@@ -246,6 +247,7 @@ describe("mergeExtensionState", () => {
 			concurrentFileReads: true,
 			multiFileApplyDiff: true,
 			preventFocusDisruption: false,
+			assistantMessageParser: false,
 		})
 	})
 })

webview-ui/src/i18n/locales/ca/settings.json

@@ -686,6 +686,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Edició en segon pla",
 			"description": "Quan s'activa, evita la interrupció del focus de l'editor. Les edicions de fitxers es produeixen en segon pla sense obrir la vista diff o robar el focus. Pots continuar treballant sense interrupcions mentre Roo fa canvis. Els fitxers poden obrir-se sense focus per capturar diagnòstics o romandre completament tancats."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Utilitza el nou analitzador de missatges",
+			"description": "Activa l'analitzador de missatges en streaming experimental que millora el rendiment en respostes llargues processant els missatges de manera més eficient."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/de/settings.json

@@ -686,6 +686,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Hintergrundbearbeitung",
 			"description": "Verhindert Editor-Fokus-Störungen wenn aktiviert. Dateibearbeitungen erfolgen im Hintergrund ohne Öffnung von Diff-Ansichten oder Fokus-Diebstahl. Du kannst ungestört weiterarbeiten, während Roo Änderungen vornimmt. Dateien können ohne Fokus geöffnet werden, um Diagnosen zu erfassen oder vollständig geschlossen bleiben."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Neuen Nachrichtenparser verwenden",
+			"description": "Aktiviere den experimentellen Streaming-Nachrichtenparser, der lange Antworten durch effizientere Verarbeitung spürbar schneller macht."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/en/settings.json

@@ -682,9 +682,13 @@
 			"name": "Enable concurrent file edits",
 			"description": "When enabled, Roo can edit multiple files in a single request. When disabled, Roo must edit files one at a time. Disabling this can help when working with less capable models or when you want more control over file modifications."
 		},
-		"PREVENT_FOCUS_DISRUPTION": {
+"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Background editing",
 			"description": "Prevent editor focus disruption when enabled. File edits happen in the background without opening diff views or stealing focus. You can continue working uninterrupted while Roo makes changes. Files can be opened without focus to capture diagnostics or kept closed entirely."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Use new message parser",
+			"description": "Enable the experimental streaming message parser that provides significant performance improvements for long assistant responses by processing messages more efficiently."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/es/settings.json

@@ -686,6 +686,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Edición en segundo plano",
 			"description": "Previene la interrupción del foco del editor cuando está habilitado. Las ediciones de archivos ocurren en segundo plano sin abrir vistas de diferencias o robar el foco. Puedes continuar trabajando sin interrupciones mientras Roo realiza cambios. Los archivos pueden abrirse sin foco para capturar diagnósticos o mantenerse completamente cerrados."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Usar el nuevo analizador de mensajes",
+			"description": "Activa el analizador de mensajes en streaming experimental que mejora el rendimiento en respuestas largas procesando los mensajes de forma más eficiente."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/fr/settings.json

@@ -686,6 +686,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Édition en arrière-plan",
 			"description": "Empêche la perturbation du focus de l'éditeur lorsqu'activé. Les modifications de fichiers se font en arrière-plan sans ouvrir de vues de différences ou voler le focus. Vous pouvez continuer à travailler sans interruption pendant que Roo effectue des changements. Les fichiers peuvent être ouverts sans focus pour capturer les diagnostics ou rester complètement fermés."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Utiliser le nouveau parseur de messages",
+			"description": "Active le parseur de messages en streaming expérimental qui accélère nettement les longues réponses en traitant les messages plus efficacement."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/hi/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "बैकग्राउंड संपादन",
 			"description": "सक्षम होने पर एडिटर फोकस व्यवधान को रोकता है। फ़ाइल संपादन diff व्यू खोले बिना या फोकस चुराए बिना बैकग्राउंड में होता है। आप Roo के बदलाव करते समय बिना किसी बाधा के काम जारी रख सकते हैं। फ़ाइलें डायग्नोस्टिक्स कैप्चर करने के लिए बिना फोकस के खुल सकती हैं या पूरी तरह बंद रह सकती हैं।"
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "नए मैसेज पार्सर का उपयोग करें",
+			"description": "प्रायोगिक स्ट्रीमिंग मैसेज पार्सर सक्षम करें, जो लंबे उत्तरों के लिए संदेशों को अधिक कुशलता से प्रोसेस करके प्रदर्शन को बेहतर बनाता है।"
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/id/settings.json

@@ -716,6 +716,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Pengeditan Latar Belakang",
 			"description": "Ketika diaktifkan, mencegah gangguan fokus editor. Pengeditan file terjadi di latar belakang tanpa membuka tampilan diff atau mencuri fokus. Anda dapat terus bekerja tanpa gangguan saat Roo melakukan perubahan. File mungkin dibuka tanpa fokus untuk menangkap diagnostik atau tetap tertutup sepenuhnya."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Gunakan parser pesan baru",
+			"description": "Aktifkan parser pesan streaming eksperimental yang meningkatkan kinerja untuk respons panjang dengan memproses pesan lebih efisien."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/it/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Modifica in background",
 			"description": "Previene l'interruzione del focus dell'editor quando abilitato. Le modifiche ai file avvengono in background senza aprire viste di differenze o rubare il focus. Puoi continuare a lavorare senza interruzioni mentre Roo effettua modifiche. I file possono essere aperti senza focus per catturare diagnostiche o rimanere completamente chiusi."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Usa il nuovo parser dei messaggi",
+			"description": "Abilita il parser di messaggi in streaming sperimentale che migliora nettamente le risposte lunghe elaborando i messaggi in modo più efficiente."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/ja/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "バックグラウンド編集",
 			"description": "有効にすると、エディターのフォーカス中断を防ぎます。ファイル編集は差分ビューを開いたりフォーカスを奪ったりすることなく、バックグラウンドで行われます。Rooが変更を行っている間も中断されることなく作業を続けることができます。ファイルは診断をキャプチャするためにフォーカスなしで開くか、完全に閉じたままにできます。"
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "新しいメッセージパーサーを使う",
+			"description": "実験的なストリーミングメッセージパーサーを有効にします。長い回答をより効率的に処理し、遅延を減らします。"
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/ko/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "백그라운드 편집",
 			"description": "활성화하면 편집기 포커스 방해를 방지합니다. 파일 편집이 diff 뷰를 열거나 포커스를 빼앗지 않고 백그라운드에서 수행됩니다. Roo가 변경사항을 적용하는 동안 방해받지 않고 계속 작업할 수 있습니다. 파일은 진단을 캡처하기 위해 포커스 없이 열거나 완전히 닫힌 상태로 유지할 수 있습니다."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "새 메시지 파서 사용",
+			"description": "실험적 스트리밍 메시지 파서를 활성화합니다. 긴 응답을 더 효율적으로 처리해 지연을 줄입니다."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/nl/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Achtergrondbewerking",
 			"description": "Voorkomt editor focus verstoring wanneer ingeschakeld. Bestandsbewerkingen gebeuren op de achtergrond zonder diff-weergaven te openen of focus te stelen. Je kunt ononderbroken doorwerken terwijl Roo wijzigingen aanbrengt. Bestanden kunnen zonder focus worden geopend om diagnostiek vast te leggen of volledig gesloten blijven."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Nieuwe berichtparser gebruiken",
+			"description": "Schakel de experimentele streaming-berichtparser in die lange antwoorden sneller maakt door berichten efficiënter te verwerken."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/pl/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Edycja w tle",
 			"description": "Zapobiega zakłócaniu fokusa edytora gdy włączone. Edycje plików odbywają się w tle bez otwierania widoków różnic lub kradzieży fokusa. Możesz kontynuować pracę bez przeszkód podczas gdy Roo wprowadza zmiany. Pliki mogą być otwierane bez fokusa aby przechwycić diagnostykę lub pozostać całkowicie zamknięte."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Użyj nowego parsera wiadomości",
+			"description": "Włącz eksperymentalny parser wiadomości w strumieniu, który przyspiesza długie odpowiedzi dzięki bardziej wydajnemu przetwarzaniu wiadomości."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/pt-BR/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Edição em segundo plano",
 			"description": "Previne a interrupção do foco do editor quando habilitado. As edições de arquivos acontecem em segundo plano sem abrir visualizações de diferenças ou roubar o foco. Você pode continuar trabalhando sem interrupções enquanto o Roo faz alterações. Os arquivos podem ser abertos sem foco para capturar diagnósticos ou permanecer completamente fechados."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Usar o novo parser de mensagens",
+			"description": "Ativa o parser de mensagens em streaming experimental que acelera respostas longas ao processar as mensagens de forma mais eficiente."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/ru/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Фоновое редактирование",
 			"description": "Предотвращает нарушение фокуса редактора при включении. Редактирование файлов происходит в фоновом режиме без открытия представлений различий или кражи фокуса. Вы можете продолжать работать без перерывов, пока Roo вносит изменения. Файлы могут открываться без фокуса для захвата диагностики или оставаться полностью закрытыми."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Использовать новый парсер сообщений",
+			"description": "Включите экспериментальный потоковый парсер сообщений, который ускоряет длинные ответы благодаря более эффективной обработке сообщений."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/tr/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Arka plan düzenleme",
 			"description": "Etkinleştirildiğinde editör odak kesintisini önler. Dosya düzenlemeleri diff görünümlerini açmadan veya odağı çalmadan arka planda gerçekleşir. Roo değişiklikler yaparken kesintisiz çalışmaya devam edebilirsiniz. Dosyalar tanılamayı yakalamak için odaksız açılabilir veya tamamen kapalı kalabilir."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Yeni mesaj ayrıştırıcıyı kullan",
+			"description": "Uzun yanıtları daha verimli işleyerek hızlandıran deneysel akış mesaj ayrıştırıcısını etkinleştir."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/vi/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "Chỉnh sửa nền",
 			"description": "Khi được bật, ngăn chặn gián đoạn tiêu điểm trình soạn thảo. Việc chỉnh sửa tệp diễn ra ở nền mà không mở chế độ xem diff hoặc chiếm tiêu điểm. Bạn có thể tiếp tục làm việc không bị gián đoạn trong khi Roo thực hiện thay đổi. Các tệp có thể được mở mà không có tiêu điểm để thu thập chẩn đoán hoặc giữ hoàn toàn đóng."
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "Dùng bộ phân tích tin nhắn mới",
+			"description": "Bật bộ phân tích tin nhắn streaming thử nghiệm. Tính năng này tăng tốc phản hồi dài bằng cách xử lý tin nhắn hiệu quả hơn."
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/zh-CN/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "后台编辑",
 			"description": "启用后防止编辑器焦点干扰。文件编辑在后台进行，不会打开差异视图或抢夺焦点。你可以在 Roo 进行更改时继续不受干扰地工作。文件可以在不获取焦点的情况下打开以捕获诊断信息，或保持完全关闭状态。"
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "使用新的消息解析器",
+			"description": "启用实验性的流式消息解析器。通过更高效地处理消息，可显著提升长回复的性能。"
 		}
 	},
 	"promptCaching": {

webview-ui/src/i18n/locales/zh-TW/settings.json

@@ -687,6 +687,10 @@
 		"PREVENT_FOCUS_DISRUPTION": {
 			"name": "背景編輯",
 			"description": "啟用後可防止編輯器焦點中斷。檔案編輯會在背景進行，不會開啟 diff 檢視或搶奪焦點。您可以在 Roo 進行變更時繼續不受干擾地工作。檔案可能會在不獲得焦點的情況下開啟以捕獲診斷，或保持完全關閉。"
+		},
+		"ASSISTANT_MESSAGE_PARSER": {
+			"name": "使用全新訊息解析器",
+			"description": "啟用實驗性的串流訊息解析器。透過更有效率地處理訊息，能顯著提升長回覆的效能。"
 		}
 	},
 	"promptCaching": {