Add better xml parsing to handle params and values on the same line

src/core/ClaudeDev.ts

@@ -1489,115 +1489,106 @@ export class ClaudeDev {
 	}
 
 	parseAssistantMessage(assistantMessage: string) {
-		// let text = ""
 		let textContent: TextContent = {
 			type: "text",
 			content: "",
 			partial: true,
 		}
 		let toolUses: ToolUse[] = []
-
 		let currentToolUse: ToolUse | undefined = undefined
+		let currentToolUseStartIndex = 0
 		let currentParamName: ToolParamName | undefined = undefined
-		let currentParamValueLines: string[] = []
+		let currentParamValueStartIndex = 0
-		let textContentLines: string[] = []
+		let accumulator = ""
 
-		const rawLines = assistantMessage.split("\n")
+		for (let i = 0; i < assistantMessage.length; i++) {
+			const char = assistantMessage[i]
+			accumulator += char
 
-		if (rawLines.length === 1) {
+			// there should not be a param without a tool use
-			const firstLine = rawLines[0].trim()
+			if (currentToolUse && currentParamName) {
-			if (!firstLine.startsWith("<t") && firstLine.startsWith("<")) {
+				const currentParamValue = accumulator.slice(currentParamValueStartIndex)
-				// (we ignore tags that start with <t since it's most like a <thinking> tag (and none of our tags start with t)
+				const paramClosingTag = `</${currentParamName}>`
-				// content is just starting, if it starts with < we can assume it's a tool call, so we'll wait for the next line
+				if (currentParamValue.endsWith(paramClosingTag)) {
-				return
+					// end of param value
-			}
+					currentToolUse.params[currentParamName] = currentParamValue.slice(0, -paramClosingTag.length).trim()
-		}
+					currentParamName = undefined
-
-		if (
-			this.assistantMessageContent.length === 1 &&
-			this.assistantMessageContent[0].partial // first element is always TextContent
-		) {
-			// we're updating text content, so if we have a partial xml tag on the last line we can ignore it until we get the full line.
-			const lastLine = rawLines.at(-1)?.trim()
-			if (lastLine && !lastLine.startsWith("<t") && lastLine.startsWith("<") && !lastLine.endsWith(">")) {
-				return
-			}
-		}
-
-		for (const line of rawLines) {
-			const trimmed = line.trim()
-			// if currenttoolcall or currentparamname look for closing tag, more efficient and safe
-			if (currentToolUse && currentParamName && trimmed === `</${currentParamName}>`) {
-				// End of a tool parameter
-				currentToolUse.params[currentParamName] = currentParamValueLines.join("\n")
-				currentParamName = undefined
-				currentParamValueLines = []
-				// currentParamValue = undefined
-				continue
-			} else if (currentToolUse && !currentParamName && trimmed === `</${currentToolUse.name}>`) {
-				// End of a tool call
-				currentToolUse.partial = false
-				toolUses.push(currentToolUse)
-				currentToolUse = undefined
-				continue
-			}
-			if (!currentParamName && trimmed.startsWith("<") && trimmed.endsWith(">")) {
-				const tag = trimmed.slice(1, -1)
-				if (toolUseNames.includes(tag as ToolUseName)) {
-					// Start of a new tool call
-					currentToolUse = {
-						type: "tool_use",
-						name: tag as ToolUseName,
-						params: {},
-						partial: true,
-					} satisfies ToolUse
-					// This also indicates the end of the text content
-					textContent.partial = false
 					continue
-				} else if (currentToolUse && toolParamNames.includes(tag as ToolParamName)) {
+				} else {
-					// Start of a parameter
+					// partial param value is accumulating
-					currentParamName = tag as ToolParamName
-					// currentToolUse.params[currentParamName] = ""
 					continue
 				}
 			}
 
-			if (currentToolUse && !currentParamName) {
+			// no currentParamName
-				// Even though system prompt instructs to put tags on separate lines, sometimes model outputs small non-file params on single lines (have not seen this happen with a tool use tag though)
-				// E.g. <path>file</path>
-				// We're making some assumptions here, like if we do match then the entire param will be on this line.
 
-				// Try to match a parameter tag with content, even if the closing tag is missing or partial
+			if (currentToolUse) {
-				// matches <paramName> and rest of line as paramContent
+				const currentToolValue = accumulator.slice(currentToolUseStartIndex)
-				// const paramMatch = trimmed.match(/^<(\w+)>(.*)$/)
+				const toolUseClosingTag = `</${currentToolUse.name}>`
-				// if (paramMatch) {
+				if (currentToolValue.endsWith(toolUseClosingTag)) {
-				// 	const paramName = paramMatch[1] as ToolParamName
+					// end of a tool use
-				// 	let paramContent = paramMatch[2]
+					currentToolUse.partial = false
+					toolUses.push(currentToolUse)
+					currentToolUse = undefined
+					continue
+				} else {
+					const possibleParamOpeningTags = toolParamNames.map((name) => `<${name}>`)
+					for (const paramOpeningTag of possibleParamOpeningTags) {
+						if (accumulator.endsWith(paramOpeningTag)) {
+							// start of a new parameter
+							currentParamName = paramOpeningTag.slice(1, -1) as ToolParamName
+							currentParamValueStartIndex = accumulator.length
+							break
+						}
+					}
 
-				// 	// Remove any closing tag or partial closing tag from paramContent
+					// there's no current param, and not starting a new param
-				// 	// replaces any sequence that starts with </ (a closing tag) to the end of the line with an empty string
-				// 	paramContent = paramContent?.replace(/<\/.*$/, "").trim()
 
-				// 	if (paramName && paramContent && toolParamNames.includes(paramName)) {
+					// 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, we get the string between the starting content tag and the LAST content tag.
-				// 		currentToolUse.params[paramName] = paramContent
+					const contentParamName: ToolParamName = "content"
-				// 	}
+					if (currentToolUse.name === "write_to_file" && accumulator.endsWith(`</${contentParamName}>`)) {
+						const toolContent = accumulator.slice(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) {
+							currentToolUse.params[contentParamName] = toolContent
+								.slice(contentStartIndex, contentEndIndex)
+								.trim()
+						}
+					}
 
-				// 	// Assuming the entire parameter is on this line, we don't need to set currentParamName
+					// partial tool value is accumulating
-				// }
+					continue
-
+				}
-				// If no param name, assume it's a partial and wait for more output
-				continue
 			}
 
-			if (currentToolUse && currentParamName) {
+			// no currentToolUse
-				// add line to current param value
+
-				currentParamValueLines.push(line)
+			const possibleToolUseOpeningTags = toolUseNames.map((name) => `<${name}>`)
-				continue
+			for (const toolUseOpeningTag of possibleToolUseOpeningTags) {
+				if (accumulator.endsWith(toolUseOpeningTag)) {
+					// start of a new tool use
+					currentToolUse = {
+						type: "tool_use",
+						name: toolUseOpeningTag.slice(1, -1) as ToolUseName,
+						params: {},
+						partial: true,
+					}
+					currentToolUseStartIndex = accumulator.length
+					// this also indicates the end of the text content
+					textContent.partial = false
+					// remove the partially accumulated tool use tag from the end of text (<tool)
+					console.log("removing from text", toolUseOpeningTag.slice(0, -1))
+					textContent.content = textContent.content.slice(0, -toolUseOpeningTag.slice(0, -1).length).trim()
+					break
+				}
 			}
 
 			// only add text content if we haven't started a tool yet
 			if (textContent.partial) {
-				textContentLines.push(line)
+				textContent.content = accumulator.trim()
 			}
 		}
 
@@ -1605,12 +1596,13 @@ export class ClaudeDev {
 			// stream did not complete tool call, add it as partial
 			if (currentParamName) {
 				// tool call has a parameter that was not completed
-				currentToolUse.params[currentParamName] = currentParamValueLines.join("\n")
+				currentToolUse.params[currentParamName] = accumulator.slice(currentParamValueStartIndex)
 			}
 			toolUses.push(currentToolUse)
 		}
 
-		textContent.content = textContentLines.join("\n")
+		console.log(assistantMessage, textContent, toolUses)
+
 		const prevLength = this.assistantMessageContent.length
 		this.assistantMessageContent = [textContent, ...toolUses]
 		if (this.assistantMessageContent.length > prevLength) {