Refactor DirectiveStreamingParser for better modularity and fix test failures across all related files

src/core/assistant-message/DirectiveHandlerRegistry.ts

@@ -0,0 +1,28 @@
+import { DirectiveHandler } from "./interfaces/DirectiveHandler"
+import { TextDirectiveHandler } from "./handlers/TextDirectiveHandler"
+import { ToolDirectiveHandler } from "./handlers/ToolDirectiveHandler"
+
+export class DirectiveHandlerRegistry {
+	private handlers: Map<string, DirectiveHandler> = new Map()
+	private textHandler = new TextDirectiveHandler()
+
+	register(handler: DirectiveHandler): void {
+		this.handlers.set(handler.tagName, handler)
+	}
+
+	registerTool(toolName: string): void {
+		this.register(new ToolDirectiveHandler(toolName))
+	}
+
+	getHandler(tagName: string): DirectiveHandler | undefined {
+		return this.handlers.get(tagName)
+	}
+
+	getTextHandler(): TextDirectiveHandler {
+		return this.textHandler
+	}
+
+	getAllHandlers(): DirectiveHandler[] {
+		return [this.textHandler, ...Array.from(this.handlers.values())]
+	}
+}

src/core/assistant-message/DirectiveRegistryFactory.ts

@@ -0,0 +1,19 @@
+import { DirectiveHandlerRegistry } from "./DirectiveHandlerRegistry"
+import { LogDirectiveHandler } from "./handlers/LogDirectiveHandler"
+import { toolNames } from "@roo-code/types"
+
+export class DirectiveRegistryFactory {
+	static create(): DirectiveHandlerRegistry {
+		const registry = new DirectiveHandlerRegistry()
+
+		// Register built-in directives
+		registry.register(new LogDirectiveHandler())
+
+		// Register all tool directives
+		toolNames.forEach((toolName) => {
+			registry.registerTool(toolName)
+		})
+
+		return registry
+	}
+}

src/core/assistant-message/DirectiveStreamingParser.ts

@@ -1,296 +1,80 @@
-import { Directive } from "./parsers"
-import { TextDirective, LogDirective } from "./directives"
-import { ToolUse, ToolParamName } from "../../shared/tools"
-import { toolNames } from "@roo-code/types"
 import * as sax from "sax"
+import { Directive } from "./parsers"
+import { ParseContext } from "./interfaces/ParseContext"
+import { DirectiveRegistryFactory } from "./DirectiveRegistryFactory"
+import { FallbackParser } from "./parsers/FallbackParser"
+import { XmlUtils } from "./XmlUtils"
 
 export class DirectiveStreamingParser {
-	static parse(assistantMessage: string): Directive[] {
-		const contentBlocks: Directive[] = []
-		let currentText = ""
-		let currentToolUse: ToolUse | undefined
-		let currentLogMessage: LogDirective | undefined
-		let currentParamName: ToolParamName | undefined
-		let currentParamValue = ""
-		let currentContext: "text" | "logMessage" | "logLevel" | "param" | "none" = "text"
-		let hasXmlTags = false
-		let parseError = false
+	private static registry = DirectiveRegistryFactory.create()
 
-		// Check if the input has incomplete XML (for partial detection)
-		const hasIncompleteXml = this.hasIncompleteXml(assistantMessage)
+	static parse(assistantMessage: string): Directive[] {
+		const context: ParseContext = {
+			currentText: "",
+			contentBlocks: [],
+			hasXmlTags: false,
+			hasIncompleteXml: XmlUtils.hasIncompleteXml(assistantMessage),
+		}
 
 		const parser = sax.parser(false, { lowercase: true })
+		let parseError = false
+		let tagStack: string[] = []
+		let activeHandler: any = null
 
 		parser.onopentag = (node: sax.Tag) => {
-			hasXmlTags = true
-			const tagName = node.name
+			context.hasXmlTags = true
+			tagStack.push(node.name)
+			const handler = this.registry.getHandler(node.name)
 
-			if (tagName === "log_message") {
-				// Push any accumulated text before starting log message
-				if (currentText.trim()) {
-					contentBlocks.push({
-						type: "text",
-						content: currentText.trim(),
-						partial: false,
-					} as TextDirective)
-					currentText = ""
-				}
-
-				currentLogMessage = {
-					type: "log_message",
-					message: "",
-					level: "info",
-					partial: true,
-				}
-				currentContext = "none"
-			} else if (tagName === "message" && currentLogMessage) {
-				currentContext = "logMessage"
-			} else if (tagName === "level" && currentLogMessage) {
-				currentContext = "logLevel"
-			} else if (toolNames.includes(tagName as any)) {
-				// Push any accumulated text before starting tool use
-				if (currentText.trim()) {
-					contentBlocks.push({
-						type: "text",
-						content: currentText.trim(),
-						partial: false,
-					} as TextDirective)
-					currentText = ""
-				}
-
-				currentToolUse = {
-					type: "tool_use",
-					name: tagName as any,
-					params: {},
-					partial: true,
-				}
-				currentContext = "none"
-			} else if (currentToolUse) {
-				currentParamName = tagName as ToolParamName
-				currentParamValue = ""
-				currentContext = "param"
+			if (handler) {
+				activeHandler = handler
+				this.registry.getTextHandler().setState("none")
+			}
+			if (activeHandler) {
+				activeHandler.onOpenTag(node, context)
 			}
 		}
 
 		parser.onclosetag = (tagName: string) => {
-			if (tagName === "log_message" && currentLogMessage) {
-				currentLogMessage.partial = hasIncompleteXml
-				contentBlocks.push(currentLogMessage)
-				currentLogMessage = undefined
-				currentContext = "text"
-			} else if (tagName === "message" && currentLogMessage) {
-				currentContext = "none"
-			} else if (tagName === "level" && currentLogMessage) {
-				currentContext = "none"
-			} else if (currentToolUse && tagName === currentToolUse.name) {
-				currentToolUse.partial = hasIncompleteXml || Object.keys(currentToolUse.params).length === 0
-				contentBlocks.push(currentToolUse)
-				currentToolUse = undefined
-				currentContext = "text"
-			} else if (currentToolUse && currentParamName && tagName === currentParamName) {
-				;(currentToolUse.params as Record<string, string>)[currentParamName] = currentParamValue.trim()
-				currentParamName = undefined
-				currentParamValue = ""
-				currentContext = "none"
+			if (activeHandler) {
+				activeHandler.onCloseTag(tagName, context)
+				if (tagName === activeHandler.tagName) {
+					activeHandler = null
+					this.registry.getTextHandler().setState("text")
+				}
 			}
+			tagStack.pop()
 		}
 
 		parser.ontext = (text: string) => {
-			if (currentContext === "param" && currentParamName && currentToolUse) {
-				currentParamValue += text
-			} else if (currentContext === "logMessage" && currentLogMessage) {
-				currentLogMessage.message += text
-			} else if (currentContext === "logLevel" && currentLogMessage) {
-				const levelText = text.trim()
-				if (["debug", "info", "warn", "error"].includes(levelText)) {
-					currentLogMessage.level = levelText as "debug" | "info" | "warn" | "error"
-				}
-			} else if (currentContext === "text") {
-				currentText += text
+			if (activeHandler) {
+				activeHandler.onText(text, context)
+			} else {
+				this.registry.getTextHandler().onText(text, context)
 			}
 		}
 
 		parser.onend = () => {
-			// Push any remaining text
-			if (currentText.trim()) {
-				contentBlocks.push({
-					type: "text",
-					content: currentText.trim(),
-					partial: true,
-				} as TextDirective)
-			}
-
-			// Handle partial log message at the end
-			if (currentLogMessage) {
-				currentLogMessage.partial = true
-				contentBlocks.push(currentLogMessage)
-			}
-
-			// Handle partial tool use at the end
-			if (currentToolUse) {
-				if (currentParamName && currentParamValue) {
-					;(currentToolUse.params as Record<string, string>)[currentParamName] = currentParamValue.trim()
-				}
-				currentToolUse.partial = true
-				contentBlocks.push(currentToolUse)
+			for (const handler of this.registry.getAllHandlers()) {
+				handler.onEnd(context)
 			}
 		}
 
 		parser.onerror = (error: Error) => {
 			parseError = true
-			// Don't clear content blocks here - let the fallback logic handle it
 		}
 
 		try {
-			// Wrap multiple root elements to make valid XML
 			const wrappedMessage = `<root>${assistantMessage}</root>`
 			parser.write(wrappedMessage).close()
 		} catch (e) {
 			parseError = true
 		}
 
-		// If parsing failed or no XML tags were found, use fallback logic
-		if (parseError || (!hasXmlTags && contentBlocks.length === 0 && assistantMessage.trim())) {
-			// Try to handle partial XML manually for streaming scenarios
-			return this.handlePartialXml(assistantMessage)
+		if (parseError || (!context.hasXmlTags && context.contentBlocks.length === 0 && assistantMessage.trim())) {
+			return FallbackParser.parse(assistantMessage)
 		}
 
-		return contentBlocks
-	}
-
-	private static handlePartialXml(assistantMessage: string): Directive[] {
-		const contentBlocks: Directive[] = []
-
-		// Handle multiple log messages
-		const logMessageRegex = /<log_message>([\s\S]*?)(?:<\/log_message>|$)/g
-		let lastIndex = 0
-		let match
-
-		while ((match = logMessageRegex.exec(assistantMessage)) !== null) {
-			// Add any text before this log message
-			if (match.index > lastIndex) {
-				const textBefore = assistantMessage.substring(lastIndex, match.index).trim()
-				if (textBefore) {
-					contentBlocks.push({
-						type: "text",
-						content: textBefore,
-						partial: false,
-					} as TextDirective)
-				}
-			}
-
-			const logContent = match[1]
-			const isComplete = assistantMessage.includes("</log_message>", match.index)
-
-			// For streaming behavior, preserve raw XML content when incomplete
-			let message = ""
-			let level: "debug" | "info" | "warn" | "error" = "info"
-
-			if (isComplete) {
-				// Complete log message - parse normally
-				const messageMatch = logContent.match(/<message>(.*?)<\/message>/)
-				const levelMatch = logContent.match(/<level>(.*?)<\/level>/)
-
-				message = messageMatch ? messageMatch[1] : ""
-				if (levelMatch && ["debug", "info", "warn", "error"].includes(levelMatch[1])) {
-					level = levelMatch[1] as "debug" | "info" | "warn" | "error"
-				}
-			} else {
-				// Incomplete log message - preserve raw content for streaming behavior
-				message = logContent
-			}
-
-			const logMessage: LogDirective = {
-				type: "log_message",
-				message,
-				level,
-				partial: !isComplete,
-			}
-
-			contentBlocks.push(logMessage)
-			lastIndex = logMessageRegex.lastIndex
-		}
-
-		// If no log messages were found, check for tool use
-		if (contentBlocks.length === 0) {
-			for (const toolName of toolNames) {
-				const toolRegex = new RegExp(`<${toolName}>[\\s\\S]*?(?:<\\/${toolName}>|$)`)
-				const toolMatch = assistantMessage.match(toolRegex)
-				if (toolMatch) {
-					const toolContent = toolMatch[0]
-					const params: Record<string, string> = {}
-
-					// Extract parameters
-					const paramRegex = /<(\w+)>(.*?)(?:<\/\1>|$)/g
-					let paramMatch
-					while ((paramMatch = paramRegex.exec(toolContent)) !== null) {
-						const [, paramName, paramValue] = paramMatch
-						if (paramName !== toolName) {
-							params[paramName] = paramValue
-						}
-					}
-
-					const toolUse: ToolUse = {
-						type: "tool_use",
-						name: toolName as any,
-						params,
-						partial: !assistantMessage.includes(`</${toolName}>`),
-					}
-
-					contentBlocks.push(toolUse)
-					return contentBlocks
-				}
-			}
-		}
-
-		// Add any remaining text after the last log message
-		if (lastIndex < assistantMessage.length) {
-			const remainingText = assistantMessage.substring(lastIndex).trim()
-			if (remainingText) {
-				contentBlocks.push({
-					type: "text",
-					content: remainingText,
-					partial: true,
-				} as TextDirective)
-			}
-		}
-
-		// If no structured content was found, treat as plain text
-		if (contentBlocks.length === 0) {
-			contentBlocks.push({
-				type: "text",
-				content: assistantMessage,
-				partial: true,
-			} as TextDirective)
-		}
-
-		return contentBlocks
-	}
-
-	private static hasIncompleteXml(input: string): boolean {
-		// Check for incomplete XML by looking for opening tags without corresponding closing tags
-		const openTags: string[] = []
-		const tagRegex = /<\/?([a-zA-Z_][a-zA-Z0-9_-]*)[^>]*>/g
-		let match
-
-		while ((match = tagRegex.exec(input)) !== null) {
-			const fullTag = match[0]
-			const tagName = match[1]
-
-			if (fullTag.startsWith("</")) {
-				// Closing tag
-				const lastOpenTag = openTags.pop()
-				if (lastOpenTag !== tagName) {
-					// Mismatched closing tag, consider incomplete
-					return true
-				}
-			} else if (!fullTag.endsWith("/>")) {
-				// Opening tag (not self-closing)
-				openTags.push(tagName)
-			}
-		}
-
-		// If there are unclosed tags, it's incomplete
-		return openTags.length > 0
+		return context.contentBlocks
 	}
 }

src/core/assistant-message/DirectiveStreamingParser_original.ts

@@ -0,0 +1,296 @@
+import { Directive } from "./parsers"
+import { TextDirective, LogDirective } from "./directives"
+import { ToolUse, ToolParamName } from "../../shared/tools"
+import { toolNames } from "@roo-code/types"
+import * as sax from "sax"
+
+export class DirectiveStreamingParser {
+	static parse(assistantMessage: string): Directive[] {
+		const contentBlocks: Directive[] = []
+		let currentText = ""
+		let currentToolUse: ToolUse | undefined
+		let currentLogMessage: LogDirective | undefined
+		let currentParamName: ToolParamName | undefined
+		let currentParamValue = ""
+		let currentContext: "text" | "logMessage" | "logLevel" | "param" | "none" = "text"
+		let hasXmlTags = false
+		let parseError = false
+
+		// Check if the input has incomplete XML (for partial detection)
+		const hasIncompleteXml = this.hasIncompleteXml(assistantMessage)
+
+		const parser = sax.parser(false, { lowercase: true })
+
+		parser.onopentag = (node: sax.Tag) => {
+			hasXmlTags = true
+			const tagName = node.name
+
+			if (tagName === "log_message") {
+				// Push any accumulated text before starting log message
+				if (currentText.trim()) {
+					contentBlocks.push({
+						type: "text",
+						content: currentText.trim(),
+						partial: false,
+					} as TextDirective)
+					currentText = ""
+				}
+
+				currentLogMessage = {
+					type: "log_message",
+					message: "",
+					level: "info",
+					partial: true,
+				}
+				currentContext = "none"
+			} else if (tagName === "message" && currentLogMessage) {
+				currentContext = "logMessage"
+			} else if (tagName === "level" && currentLogMessage) {
+				currentContext = "logLevel"
+			} else if (toolNames.includes(tagName as any)) {
+				// Push any accumulated text before starting tool use
+				if (currentText.trim()) {
+					contentBlocks.push({
+						type: "text",
+						content: currentText.trim(),
+						partial: false,
+					} as TextDirective)
+					currentText = ""
+				}
+
+				currentToolUse = {
+					type: "tool_use",
+					name: tagName as any,
+					params: {},
+					partial: true,
+				}
+				currentContext = "none"
+			} else if (currentToolUse) {
+				currentParamName = tagName as ToolParamName
+				currentParamValue = ""
+				currentContext = "param"
+			}
+		}
+
+		parser.onclosetag = (tagName: string) => {
+			if (tagName === "log_message" && currentLogMessage) {
+				currentLogMessage.partial = hasIncompleteXml
+				contentBlocks.push(currentLogMessage)
+				currentLogMessage = undefined
+				currentContext = "text"
+			} else if (tagName === "message" && currentLogMessage) {
+				currentContext = "none"
+			} else if (tagName === "level" && currentLogMessage) {
+				currentContext = "none"
+			} else if (currentToolUse && tagName === currentToolUse.name) {
+				currentToolUse.partial = hasIncompleteXml || Object.keys(currentToolUse.params).length === 0
+				contentBlocks.push(currentToolUse)
+				currentToolUse = undefined
+				currentContext = "text"
+			} else if (currentToolUse && currentParamName && tagName === currentParamName) {
+				;(currentToolUse.params as Record<string, string>)[currentParamName] = currentParamValue.trim()
+				currentParamName = undefined
+				currentParamValue = ""
+				currentContext = "none"
+			}
+		}
+
+		parser.ontext = (text: string) => {
+			if (currentContext === "param" && currentParamName && currentToolUse) {
+				currentParamValue += text
+			} else if (currentContext === "logMessage" && currentLogMessage) {
+				currentLogMessage.message += text
+			} else if (currentContext === "logLevel" && currentLogMessage) {
+				const levelText = text.trim()
+				if (["debug", "info", "warn", "error"].includes(levelText)) {
+					currentLogMessage.level = levelText as "debug" | "info" | "warn" | "error"
+				}
+			} else if (currentContext === "text") {
+				currentText += text
+			}
+		}
+
+		parser.onend = () => {
+			// Push any remaining text
+			if (currentText.trim()) {
+				contentBlocks.push({
+					type: "text",
+					content: currentText.trim(),
+					partial: true,
+				} as TextDirective)
+			}
+
+			// Handle partial log message at the end
+			if (currentLogMessage) {
+				currentLogMessage.partial = true
+				contentBlocks.push(currentLogMessage)
+			}
+
+			// Handle partial tool use at the end
+			if (currentToolUse) {
+				if (currentParamName && currentParamValue) {
+					;(currentToolUse.params as Record<string, string>)[currentParamName] = currentParamValue.trim()
+				}
+				currentToolUse.partial = true
+				contentBlocks.push(currentToolUse)
+			}
+		}
+
+		parser.onerror = (error: Error) => {
+			parseError = true
+			// Don't clear content blocks here - let the fallback logic handle it
+		}
+
+		try {
+			// Wrap multiple root elements to make valid XML
+			const wrappedMessage = `<root>${assistantMessage}</root>`
+			parser.write(wrappedMessage).close()
+		} catch (e) {
+			parseError = true
+		}
+
+		// If parsing failed or no XML tags were found, use fallback logic
+		if (parseError || (!hasXmlTags && contentBlocks.length === 0 && assistantMessage.trim())) {
+			// Try to handle partial XML manually for streaming scenarios
+			return this.handlePartialXml(assistantMessage)
+		}
+
+		return contentBlocks
+	}
+
+	private static handlePartialXml(assistantMessage: string): Directive[] {
+		const contentBlocks: Directive[] = []
+
+		// Handle multiple log messages
+		const logMessageRegex = /<log_message>([\s\S]*?)(?:<\/log_message>|$)/g
+		let lastIndex = 0
+		let match
+
+		while ((match = logMessageRegex.exec(assistantMessage)) !== null) {
+			// Add any text before this log message
+			if (match.index > lastIndex) {
+				const textBefore = assistantMessage.substring(lastIndex, match.index).trim()
+				if (textBefore) {
+					contentBlocks.push({
+						type: "text",
+						content: textBefore,
+						partial: false,
+					} as TextDirective)
+				}
+			}
+
+			const logContent = match[1]
+			const isComplete = assistantMessage.includes("</log_message>", match.index)
+
+			// For streaming behavior, preserve raw XML content when incomplete
+			let message = ""
+			let level: "debug" | "info" | "warn" | "error" = "info"
+
+			if (isComplete) {
+				// Complete log message - parse normally
+				const messageMatch = logContent.match(/<message>(.*?)<\/message>/)
+				const levelMatch = logContent.match(/<level>(.*?)<\/level>/)
+
+				message = messageMatch ? messageMatch[1] : ""
+				if (levelMatch && ["debug", "info", "warn", "error"].includes(levelMatch[1])) {
+					level = levelMatch[1] as "debug" | "info" | "warn" | "error"
+				}
+			} else {
+				// Incomplete log message - preserve raw content for streaming behavior
+				message = logContent
+			}
+
+			const logMessage: LogDirective = {
+				type: "log_message",
+				message,
+				level,
+				partial: !isComplete,
+			}
+
+			contentBlocks.push(logMessage)
+			lastIndex = logMessageRegex.lastIndex
+		}
+
+		// If no log messages were found, check for tool use
+		if (contentBlocks.length === 0) {
+			for (const toolName of toolNames) {
+				const toolRegex = new RegExp(`<${toolName}>[\\s\\S]*?(?:<\\/${toolName}>|$)`)
+				const toolMatch = assistantMessage.match(toolRegex)
+				if (toolMatch) {
+					const toolContent = toolMatch[0]
+					const params: Record<string, string> = {}
+
+					// Extract parameters
+					const paramRegex = /<(\w+)>(.*?)(?:<\/\1>|$)/g
+					let paramMatch
+					while ((paramMatch = paramRegex.exec(toolContent)) !== null) {
+						const [, paramName, paramValue] = paramMatch
+						if (paramName !== toolName) {
+							params[paramName] = paramValue
+						}
+					}
+
+					const toolUse: ToolUse = {
+						type: "tool_use",
+						name: toolName as any,
+						params,
+						partial: !assistantMessage.includes(`</${toolName}>`),
+					}
+
+					contentBlocks.push(toolUse)
+					return contentBlocks
+				}
+			}
+		}
+
+		// Add any remaining text after the last log message
+		if (lastIndex < assistantMessage.length) {
+			const remainingText = assistantMessage.substring(lastIndex).trim()
+			if (remainingText) {
+				contentBlocks.push({
+					type: "text",
+					content: remainingText,
+					partial: true,
+				} as TextDirective)
+			}
+		}
+
+		// If no structured content was found, treat as plain text
+		if (contentBlocks.length === 0) {
+			contentBlocks.push({
+				type: "text",
+				content: assistantMessage,
+				partial: true,
+			} as TextDirective)
+		}
+
+		return contentBlocks
+	}
+
+	private static hasIncompleteXml(input: string): boolean {
+		// Check for incomplete XML by looking for opening tags without corresponding closing tags
+		const openTags: string[] = []
+		const tagRegex = /<\/?([a-zA-Z_][a-zA-Z0-9_-]*)[^>]*>/g
+		let match
+
+		while ((match = tagRegex.exec(input)) !== null) {
+			const fullTag = match[0]
+			const tagName = match[1]
+
+			if (fullTag.startsWith("</")) {
+				// Closing tag
+				const lastOpenTag = openTags.pop()
+				if (lastOpenTag !== tagName) {
+					// Mismatched closing tag, consider incomplete
+					return true
+				}
+			} else if (!fullTag.endsWith("/>")) {
+				// Opening tag (not self-closing)
+				openTags.push(tagName)
+			}
+		}
+
+		// If there are unclosed tags, it's incomplete
+		return openTags.length > 0
+	}
+}

src/core/assistant-message/XmlUtils.ts

@@ -0,0 +1,23 @@
+export class XmlUtils {
+	static hasIncompleteXml(input: string): boolean {
+		const openTags: string[] = []
+		const tagRegex = /<\/?([a-zA-Z_][a-zA-Z0-9_-]*)[^>]*>/g
+		let match
+
+		while ((match = tagRegex.exec(input)) !== null) {
+			const fullTag = match[0]
+			const tagName = match[1]
+
+			if (fullTag.startsWith("</")) {
+				const lastOpenTag = openTags.pop()
+				if (lastOpenTag !== tagName) {
+					return true
+				}
+			} else if (!fullTag.endsWith("/>")) {
+				openTags.push(tagName)
+			}
+		}
+
+		return openTags.length > 0
+	}
+}

src/core/assistant-message/directives/index.ts

@@ -1,2 +1,2 @@
-export * from "./logDirective"
-export * from "./textDirective"
+export * from "./LogDirective"
+export * from "./TextDirective"

src/core/assistant-message/handlers/BaseDirectiveHandler.ts

@@ -0,0 +1,28 @@
+import * as sax from "sax"
+import { DirectiveHandler } from "../interfaces/DirectiveHandler"
+import { ParseContext } from "../interfaces/ParseContext"
+import { TextDirective } from "../directives"
+
+export abstract class BaseDirectiveHandler implements DirectiveHandler {
+	abstract readonly tagName: string
+
+	canHandle(tagName: string): boolean {
+		return tagName === this.tagName
+	}
+
+	onOpenTag(node: sax.Tag, context: ParseContext): void {}
+	onCloseTag(tagName: string, context: ParseContext): void {}
+	onText(text: string, context: ParseContext): void {}
+	onEnd(context: ParseContext): void {}
+
+	protected flushCurrentText(context: ParseContext): void {
+		if (context.currentText.trim()) {
+			context.contentBlocks.push({
+				type: "text",
+				content: context.currentText.trim(),
+				partial: false,
+			} as TextDirective)
+			context.currentText = ""
+		}
+	}
+}

src/core/assistant-message/handlers/LogDirectiveHandler.ts

@@ -0,0 +1,57 @@
+import * as sax from "sax"
+import { BaseDirectiveHandler } from "./BaseDirectiveHandler"
+import { ParseContext } from "../interfaces/ParseContext"
+import { LogDirective } from "../directives"
+
+export class LogDirectiveHandler extends BaseDirectiveHandler {
+	readonly tagName = "log_message"
+	private currentLogMessage?: LogDirective
+	private currentContext: "message" | "level" | "none" = "none"
+
+	override onOpenTag(node: sax.Tag, context: ParseContext): void {
+		if (node.name === this.tagName) {
+			this.flushCurrentText(context)
+			this.currentLogMessage = {
+				type: "log_message",
+				message: "",
+				level: "info",
+				partial: true,
+			}
+			this.currentContext = "none"
+		} else if (node.name === "message" && this.currentLogMessage) {
+			this.currentContext = "message"
+		} else if (node.name === "level" && this.currentLogMessage) {
+			this.currentContext = "level"
+		}
+	}
+
+	override onCloseTag(tagName: string, context: ParseContext): void {
+		if (tagName === this.tagName && this.currentLogMessage) {
+			this.currentLogMessage.partial = context.hasIncompleteXml
+			context.contentBlocks.push(this.currentLogMessage)
+			this.currentLogMessage = undefined
+		} else if (tagName === "message" || tagName === "level") {
+			this.currentContext = "none"
+		}
+	}
+
+	override onText(text: string, context: ParseContext): void {
+		if (!this.currentLogMessage) return
+
+		if (this.currentContext === "message") {
+			this.currentLogMessage.message += text
+		} else if (this.currentContext === "level") {
+			const levelText = text.trim()
+			if (["debug", "info", "warn", "error"].includes(levelText)) {
+				this.currentLogMessage.level = levelText as "debug" | "info" | "warn" | "error"
+			}
+		}
+	}
+
+	override onEnd(context: ParseContext): void {
+		if (this.currentLogMessage) {
+			this.currentLogMessage.partial = true
+			context.contentBlocks.push(this.currentLogMessage)
+		}
+	}
+}

src/core/assistant-message/handlers/TextDirectiveHandler.ts

@@ -0,0 +1,32 @@
+import { BaseDirectiveHandler } from "./BaseDirectiveHandler"
+import { ParseContext } from "../interfaces/ParseContext"
+import { TextDirective } from "../directives"
+
+export class TextDirectiveHandler extends BaseDirectiveHandler {
+	readonly tagName = "text"
+	private currentState: "text" | "none" = "text"
+
+	override canHandle(tagName: string): boolean {
+		return false // Text handler is fallback
+	}
+
+	override onText(text: string, context: ParseContext): void {
+		if (this.currentState === "text") {
+			context.currentText += text
+		}
+	}
+
+	setState(state: "text" | "none"): void {
+		this.currentState = state
+	}
+
+	override onEnd(context: ParseContext): void {
+		if (context.currentText.trim()) {
+			context.contentBlocks.push({
+				type: "text",
+				content: context.currentText.trim(),
+				partial: true,
+			} as TextDirective)
+		}
+	}
+}

src/core/assistant-message/handlers/ToolDirectiveHandler.ts

@@ -0,0 +1,66 @@
+import * as sax from "sax"
+import { BaseDirectiveHandler } from "./BaseDirectiveHandler"
+import { ParseContext } from "../interfaces/ParseContext"
+import { ToolUse, ToolParamName } from "../../../shared/tools"
+
+export class ToolDirectiveHandler extends BaseDirectiveHandler {
+	readonly tagName: string
+	private currentToolUse?: ToolUse
+	private currentParamName?: ToolParamName
+	private currentParamValue = ""
+	private currentContext: "param" | "none" = "none"
+
+	constructor(toolName: string) {
+		super()
+		this.tagName = toolName
+	}
+
+	override onOpenTag(node: sax.Tag, context: ParseContext): void {
+		if (node.name === this.tagName) {
+			this.flushCurrentText(context)
+			this.currentToolUse = {
+				type: "tool_use",
+				name: this.tagName as any,
+				params: {},
+				partial: true,
+			}
+			this.currentContext = "none"
+		} else if (this.currentToolUse) {
+			this.currentParamName = node.name as ToolParamName
+			this.currentParamValue = ""
+			this.currentContext = "param"
+		}
+	}
+
+	override onCloseTag(tagName: string, context: ParseContext): void {
+		if (tagName === this.tagName && this.currentToolUse) {
+			this.currentToolUse.partial =
+				context.hasIncompleteXml || Object.keys(this.currentToolUse.params).length === 0
+			context.contentBlocks.push(this.currentToolUse)
+			this.currentToolUse = undefined
+		} else if (this.currentToolUse && this.currentParamName && tagName === this.currentParamName) {
+			;(this.currentToolUse.params as Record<string, string>)[this.currentParamName] =
+				this.currentParamValue.trim()
+			this.currentParamName = undefined
+			this.currentParamValue = ""
+			this.currentContext = "none"
+		}
+	}
+
+	override onText(text: string, context: ParseContext): void {
+		if (this.currentContext === "param" && this.currentParamName && this.currentToolUse) {
+			this.currentParamValue += text
+		}
+	}
+
+	override onEnd(context: ParseContext): void {
+		if (this.currentToolUse) {
+			if (this.currentParamName && this.currentParamValue) {
+				;(this.currentToolUse.params as Record<string, string>)[this.currentParamName] =
+					this.currentParamValue.trim()
+			}
+			this.currentToolUse.partial = true
+			context.contentBlocks.push(this.currentToolUse)
+		}
+	}
+}

src/core/assistant-message/index.ts

@@ -1,3 +1,26 @@
 export { type AssistantMessageContent, parseAssistantMessage } from "./parseAssistantMessage"
 export { presentAssistantMessage } from "./presentAssistantMessage"
-export { type LogDirective } from "./directives/logDirective"
+export { type LogDirective } from "./directives/LogDirective"
+
+// Main API
+export { DirectiveStreamingParser } from "./DirectiveStreamingParser"
+
+// Core interfaces and types
+export type { DirectiveHandler } from "./interfaces/DirectiveHandler"
+export type { ParseContext } from "./interfaces/ParseContext"
+
+// Base classes for extension
+export { BaseDirectiveHandler } from "./handlers/BaseDirectiveHandler"
+
+// Registry system
+export { DirectiveHandlerRegistry } from "./DirectiveHandlerRegistry"
+export { DirectiveRegistryFactory } from "./DirectiveRegistryFactory"
+
+// Built-in handlers (for custom registration)
+export { LogDirectiveHandler } from "./handlers/LogDirectiveHandler"
+export { ToolDirectiveHandler } from "./handlers/ToolDirectiveHandler"
+export { TextDirectiveHandler } from "./handlers/TextDirectiveHandler"
+
+// Utilities
+export { XmlUtils } from "./XmlUtils"
+export { FallbackParser } from "./parsers/FallbackParser"

src/core/assistant-message/index_original.ts

@@ -0,0 +1,3 @@
+export { type AssistantMessageContent, parseAssistantMessage } from "./parseAssistantMessage"
+export { presentAssistantMessage } from "./presentAssistantMessage"
+export { type LogDirective } from "./directives/LogDirective"

src/core/assistant-message/interfaces/DirectiveHandler.ts

@@ -0,0 +1,11 @@
+import * as sax from "sax"
+import { ParseContext } from "../interfaces/ParseContext"
+
+export interface DirectiveHandler {
+	readonly tagName: string
+	canHandle(tagName: string): boolean
+	onOpenTag(node: sax.Tag, context: ParseContext): void
+	onCloseTag(tagName: string, context: ParseContext): void
+	onText(text: string, context: ParseContext): void
+	onEnd(context: ParseContext): void
+}

src/core/assistant-message/interfaces/ParseContext.ts

@@ -0,0 +1,8 @@
+import { Directive } from "../parsers"
+
+export interface ParseContext {
+	currentText: string
+	contentBlocks: Directive[]
+	hasXmlTags: boolean
+	hasIncompleteXml: boolean
+}

src/core/assistant-message/parsers/FallbackParser.ts

@@ -0,0 +1,115 @@
+import { Directive } from "./types"
+import { TextDirective, LogDirective } from "../directives"
+import { toolNames } from "@roo-code/types"
+import { ToolUse } from "../../../shared/tools"
+
+export class FallbackParser {
+	static parse(assistantMessage: string): Directive[] {
+		const contentBlocks: Directive[] = []
+
+		// Handle multiple log messages
+		const logMessageRegex = /<log_message>([\s\S]*?)(?:<\/log_message>|$)/g
+		let lastIndex = 0
+		let match
+
+		while ((match = logMessageRegex.exec(assistantMessage)) !== null) {
+			// Add any text before this log message
+			if (match.index > lastIndex) {
+				const textBefore = assistantMessage.substring(lastIndex, match.index).trim()
+				if (textBefore) {
+					contentBlocks.push({
+						type: "text",
+						content: textBefore,
+						partial: false,
+					} as TextDirective)
+				}
+			}
+
+			const logContent = match[1]
+			const isComplete = assistantMessage.includes("</log_message>", match.index)
+
+			// For streaming behavior, preserve raw XML content when incomplete
+			let message = ""
+			let level: "debug" | "info" | "warn" | "error" = "info"
+
+			if (isComplete) {
+				// Complete log message - parse normally
+				const messageMatch = logContent.match(/<message>(.*?)<\/message>/)
+				const levelMatch = logContent.match(/<level>(.*?)<\/level>/)
+
+				message = messageMatch ? messageMatch[1] : ""
+				if (levelMatch && ["debug", "info", "warn", "error"].includes(levelMatch[1])) {
+					level = levelMatch[1] as "debug" | "info" | "warn" | "error"
+				}
+			} else {
+				// Incomplete log message - preserve raw content for streaming behavior
+				message = logContent
+			}
+
+			const logMessage: LogDirective = {
+				type: "log_message",
+				message,
+				level,
+				partial: !isComplete,
+			}
+
+			contentBlocks.push(logMessage)
+			lastIndex = logMessageRegex.lastIndex
+		}
+
+		// If no log messages were found, check for tool use
+		if (contentBlocks.length === 0) {
+			for (const toolName of toolNames) {
+				const toolRegex = new RegExp(`<${toolName}>[\\s\\S]*?(?:<\\/${toolName}>|$)`)
+				const toolMatch = assistantMessage.match(toolRegex)
+				if (toolMatch) {
+					const toolContent = toolMatch[0]
+					const params: Record<string, string> = {}
+
+					// Extract parameters
+					const paramRegex = /<(\w+)>(.*?)(?:<\/\1>|$)/g
+					let paramMatch
+					while ((paramMatch = paramRegex.exec(toolContent)) !== null) {
+						const [, paramName, paramValue] = paramMatch
+						if (paramName !== toolName) {
+							params[paramName] = paramValue
+						}
+					}
+
+					const toolUse: ToolUse = {
+						type: "tool_use",
+						name: toolName as any,
+						params,
+						partial: !assistantMessage.includes(`</${toolName}>`),
+					}
+
+					contentBlocks.push(toolUse)
+					return contentBlocks
+				}
+			}
+		}
+
+		// Add any remaining text after the last log message
+		if (lastIndex < assistantMessage.length) {
+			const remainingText = assistantMessage.substring(lastIndex).trim()
+			if (remainingText) {
+				contentBlocks.push({
+					type: "text",
+					content: remainingText,
+					partial: true,
+				} as TextDirective)
+			}
+		}
+
+		// If no structured content was found, treat as plain text
+		if (contentBlocks.length === 0) {
+			contentBlocks.push({
+				type: "text",
+				content: assistantMessage,
+				partial: true,
+			} as TextDirective)
+		}
+
+		return contentBlocks
+	}
+}

src/core/assistant-message/parsers/LogParser.ts

@@ -1,4 +1,4 @@
-import { LogDirective } from "../directives/logDirective"
+import { LogDirective } from "../directives/LogDirective"
 import { ParsingState } from "./types"
 
 export class LogParser {

src/core/assistant-message/presentAssistantMessage.ts

@@ -4,7 +4,7 @@ import { serializeError } from "serialize-error"
 import type { ToolName, ClineAsk, ToolProgressStatus } from "@roo-code/types"
 import { TelemetryService } from "@roo-code/telemetry"
 
-import type { LogDirective } from "./directives/logDirective"
+import type { LogDirective } from "./directives/LogDirective"
 
 import { defaultModeSlug, getModeBySlug } from "../../shared/modes"
 import type { ToolParamName, ToolResponse } from "../../shared/tools"