Reorganize assistant-message

This commit is contained in:
Steven T. Cramer 2025-06-15 13:31:54 +07:00
parent e3f673a064
commit 6d2b212a88
20 changed files with 27 additions and 586 deletions

View file

@ -1,5 +1,5 @@
import * as sax from "sax"
import { ParseContext } from "../interfaces/ParseContext"
import { ParseContext } from "./ParseContext"
export interface DirectiveHandler {
readonly tagName: string

View file

@ -1,4 +1,4 @@
import { DirectiveHandler } from "./interfaces/DirectiveHandler"
import { DirectiveHandler } from "./DirectiveHandler"
import { TextDirectiveHandler } from "./handlers/TextDirectiveHandler"
import { ToolDirectiveHandler } from "./handlers/ToolDirectiveHandler"

View file

@ -1,8 +1,8 @@
import * as sax from "sax"
import { Directive } from "./parsers"
import { ParseContext } from "./interfaces/ParseContext"
import { Directive } from "./directives"
import { ParseContext } from "./ParseContext"
import { DirectiveRegistryFactory } from "./DirectiveRegistryFactory"
import { FallbackParser } from "./parsers/FallbackParser"
import { FallbackParser } from "./FallbackParser"
import { XmlUtils } from "./XmlUtils"
export class DirectiveStreamingParser {

View file

@ -1,296 +0,0 @@
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
}
}

View file

@ -1,7 +1,7 @@
import { Directive } from "./types"
import { TextDirective, LogDirective } from "../directives"
import { Directive } from "./directives"
import { TextDirective, LogDirective } from "./directives"
import { toolNames } from "@roo-code/types"
import { ToolUse } from "../../../shared/tools"
import { ToolUse } from "../../shared/tools"
export class FallbackParser {
static parse(assistantMessage: string): Directive[] {

View file

@ -1,4 +1,4 @@
import { Directive } from "../parsers"
import { Directive } from "./directives"
export interface ParseContext {
currentText: string

View file

@ -1,10 +1,5 @@
import { TextDirective, LogDirective } from "../directives"
import { ToolUse, ToolParamName } from "../../../shared/tools"
// Type aliases for directive parsing
export type ToolDirective = ToolUse
export type Directive = TextDirective | ToolDirective | LogDirective
import { TextDirective, Directive, LogDirective, ToolDirective } from "./directives"
import { ToolParamName } from "../../shared/tools"
export interface ParsingState {
contentBlocks: Directive[]

View file

@ -1,2 +1,8 @@
import { TextDirective, LogDirective } from "../directives"
import { ToolUse } from "../../../shared/tools"
export * from "./LogDirective"
export * from "./TextDirective"
export type ToolDirective = ToolUse
export type Directive = TextDirective | ToolDirective | LogDirective

View file

@ -1,6 +1,6 @@
import * as sax from "sax"
import { DirectiveHandler } from "../interfaces/DirectiveHandler"
import { ParseContext } from "../interfaces/ParseContext"
import { DirectiveHandler } from "../DirectiveHandler"
import { ParseContext } from "../ParseContext"
import { TextDirective } from "../directives"
export abstract class BaseDirectiveHandler implements DirectiveHandler {

View file

@ -1,6 +1,6 @@
import * as sax from "sax"
import { BaseDirectiveHandler } from "./BaseDirectiveHandler"
import { ParseContext } from "../interfaces/ParseContext"
import { ParseContext } from "../ParseContext"
import { LogDirective } from "../directives"
export class LogDirectiveHandler extends BaseDirectiveHandler {

View file

@ -1,5 +1,5 @@
import { BaseDirectiveHandler } from "./BaseDirectiveHandler"
import { ParseContext } from "../interfaces/ParseContext"
import { ParseContext } from "../ParseContext"
import { TextDirective } from "../directives"
export class TextDirectiveHandler extends BaseDirectiveHandler {

View file

@ -1,6 +1,6 @@
import * as sax from "sax"
import { BaseDirectiveHandler } from "./BaseDirectiveHandler"
import { ParseContext } from "../interfaces/ParseContext"
import { ParseContext } from "../ParseContext"
import { ToolUse, ToolParamName } from "../../../shared/tools"
export class ToolDirectiveHandler extends BaseDirectiveHandler {

View file

@ -6,8 +6,8 @@ export { type LogDirective } from "./directives/LogDirective"
export { DirectiveStreamingParser } from "./DirectiveStreamingParser"
// Core interfaces and types
export type { DirectiveHandler } from "./interfaces/DirectiveHandler"
export type { ParseContext } from "./interfaces/ParseContext"
export type { DirectiveHandler } from "./DirectiveHandler"
export type { ParseContext } from "./ParseContext"
// Base classes for extension
export { BaseDirectiveHandler } from "./handlers/BaseDirectiveHandler"
@ -23,4 +23,4 @@ export { TextDirectiveHandler } from "./handlers/TextDirectiveHandler"
// Utilities
export { XmlUtils } from "./XmlUtils"
export { FallbackParser } from "./parsers/FallbackParser"
export { FallbackParser } from "./FallbackParser"

View file

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

View file

@ -1,9 +1,8 @@
import { DirectiveStreamingParser } from "./DirectiveStreamingParser"
import type { Directive } from "./parsers"
import type { Directive } from "./directives"
export type { TextDirective } from "./directives"
// Re-export types for backward compatibility
export type { ToolDirective, Directive } from "./parsers"
export type { TextDirective, ToolDirective, Directive } from "./directives"
// Backward compatibility alias
export type AssistantMessageContent = Directive

View file

@ -1,113 +0,0 @@
import { LogDirective } from "../directives/LogDirective"
import { ParsingState } from "./types"
export class LogParser {
static parse(state: ParsingState): boolean {
if (!state.currentToolUse && !state.currentTextContent && !state.currentLogMessage) {
if (state.accumulator.includes("<log_message")) {
const startIndex = state.accumulator.indexOf("<log_message")
// Use a separate property for log directive to avoid type issues with TextDirective
state.currentTextContent = undefined
state.currentToolUse = undefined
// Create a new log directive
const logDirective: LogDirective = {
type: "log_message",
message: "",
level: "info",
partial: true,
}
state.currentLogMessage = logDirective
state.currentLogMessageStartIndex = startIndex
// Only add to contentBlocks if not already added by DirectiveStreamingParser
if (!state.contentBlocks.includes(logDirective)) {
state.contentBlocks.push(logDirective)
}
return true
}
}
// Check if there is a current log message being parsed
if (state.currentLogMessage) {
const logMessage = state.currentLogMessage
const currentContent = state.accumulator.slice(state.currentLogMessageStartIndex)
const messageStartIndex = currentContent.indexOf("<message>")
if (messageStartIndex !== -1) {
const messageContentStart = messageStartIndex + "<message>".length
const messageEndIndex = currentContent.indexOf("</message>", messageContentStart)
const logEndMatchDeclared = currentContent.match(/<\/log_message>/)
if (messageEndIndex !== -1 && logEndMatchDeclared) {
// Complete message tag found and log message is complete
logMessage.message = currentContent.slice(messageContentStart, messageEndIndex).trim()
} else if (messageEndIndex !== -1) {
// Message tag is complete but log message is not
const afterMessageTag = currentContent.slice(messageEndIndex + "</message>".length)
if (afterMessageTag.trim().length > 0) {
// There's additional content after </message> (like <level> tags) - include everything for streaming
logMessage.message = currentContent.slice(messageContentStart).trim()
} else if (afterMessageTag.length > 0) {
// There's whitespace/newline after </message> - this is streaming behavior, include the closing tag
logMessage.message = currentContent
.slice(messageContentStart, messageEndIndex + "</message>".length)
.trim()
} else {
// No content after </message> - exclude the closing tag for partial entries
logMessage.message = currentContent.slice(messageContentStart, messageEndIndex).trim()
}
} else {
// Partial message, include content without closing tag
logMessage.message = currentContent.slice(messageContentStart).trim()
}
}
// Check for log message completion before updating level
const logEndMatchDeclared = currentContent.match(/<\/log_message>/)
// Update level only if log message is complete
if (logEndMatchDeclared) {
const levelMatch = currentContent.match(/<level>(.*?)(?:<\/level>|$)/s)
if (levelMatch && levelMatch[1]) {
const levelValue = levelMatch[1].trim()
if (["debug", "info", "warn", "error"].includes(levelValue)) {
logMessage.level = levelValue as "debug" | "info" | "warn" | "error"
}
}
}
const logEndMatch = currentContent.match(/<\/log_message>/)
if (logEndMatch) {
logMessage.partial = false
state.currentLogMessage = undefined
// Reset accumulator to after the closing tag to handle multiple log entries
// Find the exact position of the closing tag in the current content
const logEndIndex = currentContent.indexOf("</log_message>") + "</log_message>".length
const absoluteEndIndex = state.currentLogMessageStartIndex + logEndIndex
// Keep any remaining content after this log message
const remainingContent = state.accumulator.slice(absoluteEndIndex)
state.accumulator = remainingContent
state.currentLogMessageStartIndex = 0 // Reset start index for next log message
// Ensure state is fully reset to detect new log messages
state.currentTextContent = undefined
state.currentToolUse = undefined
}
return true
}
return false
}
static checkForLogStart(state: ParsingState): boolean {
// Check if there's a new log_message tag that hasn't been processed yet
const logStartIndex = state.accumulator.indexOf("<log_message>")
if (logStartIndex === -1) {
return false
}
// If we already have a current log message, don't start a new one
if (state.currentLogMessage) {
return false
}
return true
}
}

View file

@ -1,22 +0,0 @@
import { ParsingState } from "./types"
export class ParameterParser {
static parse(state: ParsingState): boolean {
if (!state.currentToolUse || !state.currentParamName) return false
const currentParamValue = state.accumulator.slice(state.currentParamValueStartIndex)
const paramClosingTag = `</${state.currentParamName}>`
if (currentParamValue.endsWith(paramClosingTag)) {
// End of param value.
state.currentToolUse.params[state.currentParamName] = currentParamValue
.slice(0, -paramClosingTag.length)
.trim()
state.currentParamName = undefined
return true
} else {
// Partial param value is accumulating.
return true
}
}
}

View file

@ -1,32 +0,0 @@
import { ParsingState } from "./types"
export class TextContentParser {
static parse(state: ParsingState, currentIndex: number, didStartToolUse: boolean): void {
if (!didStartToolUse) {
// No tool use, so it must be text either at the beginning or between tools.
if (state.currentTextContent === undefined) {
state.currentTextContentStartIndex = currentIndex
}
state.currentTextContent = {
type: "text",
content: state.accumulator.slice(state.currentTextContentStartIndex).trim(),
partial: true,
}
}
}
static finalize(state: ParsingState, toolUseOpeningTag: string): void {
if (state.currentTextContent) {
state.currentTextContent.partial = false
// Remove the partially accumulated tool use tag from the end of text (<tool).
state.currentTextContent.content = state.currentTextContent.content
.slice(0, -toolUseOpeningTag.slice(0, -1).length)
.trim()
state.contentBlocks.push(state.currentTextContent)
state.currentTextContent = undefined
}
}
}

View file

@ -1,89 +0,0 @@
import { type ToolName, toolNames } from "@roo-code/types"
import { ToolParamName, toolParamNames } from "../../../shared/tools"
import { ParsingState } from "./types"
import { TextContentParser } from "./TextContentParser"
export class ToolUseParser {
static checkForToolStart(state: ParsingState): boolean {
let didStartToolUse = false
const possibleToolUseOpeningTags = toolNames.map((name) => `<${name}>`)
for (const toolUseOpeningTag of possibleToolUseOpeningTags) {
if (state.accumulator.endsWith(toolUseOpeningTag)) {
// Start of a new tool use.
state.currentToolUse = {
type: "tool_use",
name: toolUseOpeningTag.slice(1, -1) as ToolName,
params: {},
partial: true,
}
state.currentToolUseStartIndex = state.accumulator.length
// This also indicates the end of the current text content.
TextContentParser.finalize(state, toolUseOpeningTag)
didStartToolUse = true
break
}
}
return didStartToolUse
}
static parse(state: ParsingState): boolean {
if (!state.currentToolUse) return false
const currentToolValue = state.accumulator.slice(state.currentToolUseStartIndex)
const toolUseClosingTag = `</${state.currentToolUse.name}>`
if (currentToolValue.endsWith(toolUseClosingTag)) {
// End of a tool use.
state.currentToolUse.partial = false
state.contentBlocks.push(state.currentToolUse)
state.currentToolUse = undefined
return true
} else {
this.parseParameter(state)
this.parseSpecialCases(state)
return true // Continue processing
}
}
private static parseParameter(state: ParsingState): void {
const possibleParamOpeningTags = toolParamNames.map((name) => `<${name}>`)
for (const paramOpeningTag of possibleParamOpeningTags) {
if (state.accumulator.endsWith(paramOpeningTag)) {
// Start of a new parameter.
state.currentParamName = paramOpeningTag.slice(1, -1) as ToolParamName
state.currentParamValueStartIndex = state.accumulator.length
break
}
}
}
private static parseSpecialCases(state: ParsingState): void {
if (!state.currentToolUse) return
// 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.
const contentParamName: ToolParamName = "content"
if (state.currentToolUse.name === "write_to_file" && state.accumulator.endsWith(`</${contentParamName}>`)) {
const toolContent = state.accumulator.slice(state.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) {
state.currentToolUse.params[contentParamName] = toolContent
.slice(contentStartIndex, contentEndIndex)
.trim()
}
}
}
}

View file

@ -1,4 +0,0 @@
export { TextContentParser } from "./TextContentParser"
export { ToolUseParser } from "./ToolUseParser"
export { ParameterParser } from "./ParameterParser"
export type { ToolDirective, Directive, ParsingState } from "./types"