diff --git a/src/core/assistant-message/DirectiveHandlerRegistry.ts b/src/core/assistant-message/DirectiveHandlerRegistry.ts new file mode 100644 index 0000000000..835df87f33 --- /dev/null +++ b/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 = 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())] + } +} diff --git a/src/core/assistant-message/DirectiveRegistryFactory.ts b/src/core/assistant-message/DirectiveRegistryFactory.ts new file mode 100644 index 0000000000..8095f482c3 --- /dev/null +++ b/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 + } +} diff --git a/src/core/assistant-message/DirectiveStreamingParser.ts b/src/core/assistant-message/DirectiveStreamingParser.ts index 545a0258b0..49fa3f6815 100644 --- a/src/core/assistant-message/DirectiveStreamingParser.ts +++ b/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)[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)[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 = `${assistantMessage}` 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 = /([\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("", 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>/) - const levelMatch = logContent.match(/(.*?)<\/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 = {} - - // 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(``), - } - - 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("")) { - // Opening tag (not self-closing) - openTags.push(tagName) - } - } - - // If there are unclosed tags, it's incomplete - return openTags.length > 0 + return context.contentBlocks } } diff --git a/src/core/assistant-message/DirectiveStreamingParser_original.ts b/src/core/assistant-message/DirectiveStreamingParser_original.ts new file mode 100644 index 0000000000..545a0258b0 --- /dev/null +++ b/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)[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)[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 = `${assistantMessage}` + 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 = /([\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("", 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>/) + const levelMatch = logContent.match(/(.*?)<\/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 = {} + + // 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(``), + } + + 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("")) { + // Opening tag (not self-closing) + openTags.push(tagName) + } + } + + // If there are unclosed tags, it's incomplete + return openTags.length > 0 + } +} diff --git a/src/core/assistant-message/XmlUtils.ts b/src/core/assistant-message/XmlUtils.ts new file mode 100644 index 0000000000..3c84536f87 --- /dev/null +++ b/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("")) { + openTags.push(tagName) + } + } + + return openTags.length > 0 + } +} diff --git a/src/core/assistant-message/directives/index.ts b/src/core/assistant-message/directives/index.ts index 5a4d4059fa..8d7a468a85 100644 --- a/src/core/assistant-message/directives/index.ts +++ b/src/core/assistant-message/directives/index.ts @@ -1,2 +1,2 @@ -export * from "./logDirective" -export * from "./textDirective" +export * from "./LogDirective" +export * from "./TextDirective" diff --git a/src/core/assistant-message/handlers/BaseDirectiveHandler.ts b/src/core/assistant-message/handlers/BaseDirectiveHandler.ts new file mode 100644 index 0000000000..3f71294867 --- /dev/null +++ b/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 = "" + } + } +} diff --git a/src/core/assistant-message/handlers/LogDirectiveHandler.ts b/src/core/assistant-message/handlers/LogDirectiveHandler.ts new file mode 100644 index 0000000000..e662f2aa62 --- /dev/null +++ b/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) + } + } +} diff --git a/src/core/assistant-message/handlers/TextDirectiveHandler.ts b/src/core/assistant-message/handlers/TextDirectiveHandler.ts new file mode 100644 index 0000000000..976fb57815 --- /dev/null +++ b/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) + } + } +} diff --git a/src/core/assistant-message/handlers/ToolDirectiveHandler.ts b/src/core/assistant-message/handlers/ToolDirectiveHandler.ts new file mode 100644 index 0000000000..655073dfca --- /dev/null +++ b/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)[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)[this.currentParamName] = + this.currentParamValue.trim() + } + this.currentToolUse.partial = true + context.contentBlocks.push(this.currentToolUse) + } + } +} diff --git a/src/core/assistant-message/index.ts b/src/core/assistant-message/index.ts index 52a63ee19d..ce9a5f9477 100644 --- a/src/core/assistant-message/index.ts +++ b/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" diff --git a/src/core/assistant-message/index_original.ts b/src/core/assistant-message/index_original.ts new file mode 100644 index 0000000000..e6fff4fcdb --- /dev/null +++ b/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" diff --git a/src/core/assistant-message/interfaces/DirectiveHandler.ts b/src/core/assistant-message/interfaces/DirectiveHandler.ts new file mode 100644 index 0000000000..9b82cb9d4d --- /dev/null +++ b/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 +} diff --git a/src/core/assistant-message/interfaces/ParseContext.ts b/src/core/assistant-message/interfaces/ParseContext.ts new file mode 100644 index 0000000000..e78c6da9c0 --- /dev/null +++ b/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 +} diff --git a/src/core/assistant-message/parsers/FallbackParser.ts b/src/core/assistant-message/parsers/FallbackParser.ts new file mode 100644 index 0000000000..8757640837 --- /dev/null +++ b/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 = /([\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("", 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>/) + const levelMatch = logContent.match(/(.*?)<\/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 = {} + + // 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(``), + } + + 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 + } +} diff --git a/src/core/assistant-message/parsers/LogParser.ts b/src/core/assistant-message/parsers/LogParser.ts index 3b52100143..327216250c 100644 --- a/src/core/assistant-message/parsers/LogParser.ts +++ b/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 { diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index 80c884a69b..8d64acef66 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/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"