fix: improve AssistantMessageParser with validation, consistent trimming, and test coverage

- Add tool name validation to prevent invalid casts
- Ensure consistent text content trimming for both new and updated blocks
- Add test coverage for MAX_ACCUMULATOR_SIZE and MAX_PARAM_LENGTH limits
- Fix auto-approve regression by keeping blocks partial until finalization

Author: daniel-lxs

src/core/assistant-message/AssistantMessageParser.ts

@@ -159,10 +159,19 @@ export class AssistantMessageParser {
 
 			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: toolUseOpeningTag.slice(1, -1) as ToolName,
+						name: extractedToolName as ToolName,
 						params: {},
 						partial: true,
 					}
@@ -179,8 +188,6 @@ export class AssistantMessageParser {
 							.slice(0, -toolUseOpeningTag.slice(0, -1).length)
 							.trim()
 
-						this.currentTextContent.content = this.currentTextContent.content.trim()
-
 						// No need to push, currentTextContent is already in contentBlocks
 						this.currentTextContent = undefined
 					}
@@ -207,7 +214,7 @@ export class AssistantMessageParser {
 					// Create a new text content block and add it to contentBlocks
 					this.currentTextContent = {
 						type: "text",
-						content: this.accumulator.slice(this.currentTextContentStartIndex),
+						content: this.accumulator.slice(this.currentTextContentStartIndex).trim(),
 						partial: true,
 					}
 
@@ -216,7 +223,7 @@ export class AssistantMessageParser {
 					this.contentBlocks.push(this.currentTextContent)
 				} else {
 					// Update the existing text content
-					this.currentTextContent.content = this.accumulator.slice(this.currentTextContentStartIndex)
+					this.currentTextContent.content = this.accumulator.slice(this.currentTextContentStartIndex).trim()
 				}
 			}
 		}

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

@@ -325,6 +325,62 @@ describe("AssistantMessageParser (streaming)", () => {
 		})
 	})
 
+	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"