Skip to content

Commit 773075e

Browse files
StevenTCramerStevenTCramer
committed
refactor: extract directive parsing into focused classes
- Extract parsing logic from monolithic parseAssistantMessage function - Create StreamingParser with specialized handler classes: - TextContentHandler: handles text content parsing - ToolUseHandler: manages tool directive detection and parsing - ParameterHandler: processes tool parameters - Maintain original streaming behavior and backward compatibility - All tests pass, no breaking changes - Improve code organization and maintainability
1 parent 2928dd3 commit 773075e

File tree

3 files changed

+224
-157
lines changed

3 files changed

+224
-157
lines changed

‎src/core/assistant-message/directives/StreamingParser.ts‎

Lines changed: 216 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,216 @@
1+
import { type ToolName, toolNames } from "@roo-code/types"
2+
import { TextContent, ToolUse, ToolParamName, toolParamNames } from "../../../shared/tools"
3+
4+
// Type aliases for directive parsing
5+
export type TextDirective = TextContent
6+
export type ToolDirective = ToolUse
7+
export type Directive = TextDirective | ToolDirective
8+
9+
export interface ParsingState {
10+
contentBlocks: Directive[]
11+
currentTextContent?: TextDirective
12+
currentTextContentStartIndex: number
13+
currentToolUse?: ToolDirective
14+
currentToolUseStartIndex: number
15+
currentParamName?: ToolParamName
16+
currentParamValueStartIndex: number
17+
accumulator: string
18+
}
19+
20+
export class TextContentHandler {
21+
static handleTextContent(state: ParsingState, currentIndex: number, didStartToolUse: boolean): void {
22+
if (!didStartToolUse) {
23+
// No tool use, so it must be text either at the beginning or between tools.
24+
if (state.currentTextContent === undefined) {
25+
state.currentTextContentStartIndex = currentIndex
26+
}
27+
28+
state.currentTextContent = {
29+
type: "text",
30+
content: state.accumulator.slice(state.currentTextContentStartIndex).trim(),
31+
partial: true,
32+
}
33+
}
34+
}
35+
36+
static finalizeTextContent(state: ParsingState, toolUseOpeningTag: string): void {
37+
if (state.currentTextContent) {
38+
state.currentTextContent.partial = false
39+
40+
// Remove the partially accumulated tool use tag from the end of text (<tool).
41+
state.currentTextContent.content = state.currentTextContent.content
42+
.slice(0, -toolUseOpeningTag.slice(0, -1).length)
43+
.trim()
44+
45+
state.contentBlocks.push(state.currentTextContent)
46+
state.currentTextContent = undefined
47+
}
48+
}
49+
}
50+
51+
export class ToolUseHandler {
52+
static checkForToolStart(state: ParsingState): boolean {
53+
let didStartToolUse = false
54+
const possibleToolUseOpeningTags = toolNames.map((name) => `<${name}>`)
55+
56+
for (const toolUseOpeningTag of possibleToolUseOpeningTags) {
57+
if (state.accumulator.endsWith(toolUseOpeningTag)) {
58+
// Start of a new tool use.
59+
state.currentToolUse = {
60+
type: "tool_use",
61+
name: toolUseOpeningTag.slice(1, -1) as ToolName,
62+
params: {},
63+
partial: true,
64+
}
65+
66+
state.currentToolUseStartIndex = state.accumulator.length
67+
68+
// This also indicates the end of the current text content.
69+
TextContentHandler.finalizeTextContent(state, toolUseOpeningTag)
70+
71+
didStartToolUse = true
72+
break
73+
}
74+
}
75+
76+
return didStartToolUse
77+
}
78+
79+
static handleToolUse(state: ParsingState): boolean {
80+
if (!state.currentToolUse) return false
81+
82+
const currentToolValue = state.accumulator.slice(state.currentToolUseStartIndex)
83+
const toolUseClosingTag = `</${state.currentToolUse.name}>`
84+
85+
if (currentToolValue.endsWith(toolUseClosingTag)) {
86+
// End of a tool use.
87+
state.currentToolUse.partial = false
88+
state.contentBlocks.push(state.currentToolUse)
89+
state.currentToolUse = undefined
90+
return true
91+
} else {
92+
this.handleParameterParsing(state)
93+
this.handleSpecialCases(state)
94+
return true // Continue processing
95+
}
96+
}
97+
98+
private static handleParameterParsing(state: ParsingState): void {
99+
const possibleParamOpeningTags = toolParamNames.map((name) => `<${name}>`)
100+
for (const paramOpeningTag of possibleParamOpeningTags) {
101+
if (state.accumulator.endsWith(paramOpeningTag)) {
102+
// Start of a new parameter.
103+
state.currentParamName = paramOpeningTag.slice(1, -1) as ToolParamName
104+
state.currentParamValueStartIndex = state.accumulator.length
105+
break
106+
}
107+
}
108+
}
109+
110+
private static handleSpecialCases(state: ParsingState): void {
111+
if (!state.currentToolUse) return
112+
113+
// Special case for write_to_file where file contents could
114+
// contain the closing tag, in which case the param would have
115+
// closed and we end up with the rest of the file contents here.
116+
// To work around this, we get the string between the starting
117+
// content tag and the LAST content tag.
118+
const contentParamName: ToolParamName = "content"
119+
120+
if (state.currentToolUse.name === "write_to_file" && state.accumulator.endsWith(`</${contentParamName}>`)) {
121+
const toolContent = state.accumulator.slice(state.currentToolUseStartIndex)
122+
const contentStartTag = `<${contentParamName}>`
123+
const contentEndTag = `</${contentParamName}>`
124+
const contentStartIndex = toolContent.indexOf(contentStartTag) + contentStartTag.length
125+
const contentEndIndex = toolContent.lastIndexOf(contentEndTag)
126+
127+
if (contentStartIndex !== -1 && contentEndIndex !== -1 && contentEndIndex > contentStartIndex) {
128+
state.currentToolUse.params[contentParamName] = toolContent
129+
.slice(contentStartIndex, contentEndIndex)
130+
.trim()
131+
}
132+
}
133+
}
134+
}
135+
136+
export class ParameterHandler {
137+
static handleParameter(state: ParsingState): boolean {
138+
if (!state.currentToolUse || !state.currentParamName) return false
139+
140+
const currentParamValue = state.accumulator.slice(state.currentParamValueStartIndex)
141+
const paramClosingTag = `</${state.currentParamName}>`
142+
143+
if (currentParamValue.endsWith(paramClosingTag)) {
144+
// End of param value.
145+
state.currentToolUse.params[state.currentParamName] = currentParamValue
146+
.slice(0, -paramClosingTag.length)
147+
.trim()
148+
state.currentParamName = undefined
149+
return true
150+
} else {
151+
// Partial param value is accumulating.
152+
return true
153+
}
154+
}
155+
}
156+
157+
export class StreamingParser {
158+
static parse(assistantMessage: string): Directive[] {
159+
const state: ParsingState = {
160+
contentBlocks: [],
161+
currentTextContent: undefined,
162+
currentTextContentStartIndex: 0,
163+
currentToolUse: undefined,
164+
currentToolUseStartIndex: 0,
165+
currentParamName: undefined,
166+
currentParamValueStartIndex: 0,
167+
accumulator: "",
168+
}
169+
170+
for (let i = 0; i < assistantMessage.length; i++) {
171+
const char = assistantMessage[i]
172+
state.accumulator += char
173+
174+
// There should not be a param without a tool use.
175+
if (ParameterHandler.handleParameter(state)) {
176+
continue
177+
}
178+
179+
// No currentParamName.
180+
if (ToolUseHandler.handleToolUse(state)) {
181+
continue
182+
}
183+
184+
// No currentToolUse.
185+
const didStartToolUse = ToolUseHandler.checkForToolStart(state)
186+
TextContentHandler.handleTextContent(state, i, didStartToolUse)
187+
}
188+
189+
// Handle remaining partial content
190+
this.handlePartialContent(state)
191+
192+
return state.contentBlocks
193+
}
194+
195+
private static handlePartialContent(state: ParsingState): void {
196+
if (state.currentToolUse) {
197+
// Stream did not complete tool call, add it as partial.
198+
if (state.currentParamName) {
199+
// Tool call has a parameter that was not completed.
200+
state.currentToolUse.params[state.currentParamName] = state.accumulator
201+
.slice(state.currentParamValueStartIndex)
202+
.trim()
203+
}
204+
205+
state.contentBlocks.push(state.currentToolUse)
206+
}
207+
208+
// NOTE: It doesn't matter if check for currentToolUse or
209+
// currentTextContent, only one of them will be defined since only one can
210+
// be partial at a time.
211+
if (state.currentTextContent) {
212+
// Stream did not complete text content, add it as partial.
213+
state.contentBlocks.push(state.currentTextContent)
214+
}
215+
}
216+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export { StreamingParser, TextContentHandler, ToolUseHandler, ParameterHandler } from "./StreamingParser"
2+
export type { TextDirective, ToolDirective, Directive, ParsingState } from "./StreamingParser"

‎src/core/assistant-message/parseAssistantMessage.ts‎

Lines changed: 6 additions & 157 deletions
Original file line numberDiff line numberDiff line change
@@ -1,162 +1,11 @@
1-
import { type ToolName, toolNames } from "@roo-code/types"
1+
import { StreamingParser } from "./directives/StreamingParser"
22

3-
import { TextContent, ToolUse, ToolParamName, toolParamNames } from "../../shared/tools"
4-
5-
// Type aliases for this file
6-
export type TextDirective = TextContent
7-
export type ToolDirective = ToolUse
8-
export type Directive = TextDirective | ToolDirective
3+
// Re-export types for backward compatibility
4+
export type { TextDirective, ToolDirective, Directive } from "./directives"
95

106
// Backward compatibility alias
11-
export type AssistantMessageContent = Directive
12-
13-
export function parseAssistantMessage(assistantMessage: string): Directive[] {
14-
let contentBlocks: Directive[] = []
15-
let currentTextContent: TextDirective | undefined = undefined
16-
let currentTextContentStartIndex = 0
17-
let currentToolUse: ToolDirective | undefined = undefined
18-
let currentToolUseStartIndex = 0
19-
let currentParamName: ToolParamName | undefined = undefined
20-
let currentParamValueStartIndex = 0
21-
let accumulator = ""
22-
23-
for (let i = 0; i < assistantMessage.length; i++) {
24-
const char = assistantMessage[i]
25-
accumulator += char
26-
27-
// There should not be a param without a tool use.
28-
if (currentToolUse && currentParamName) {
29-
const currentParamValue = accumulator.slice(currentParamValueStartIndex)
30-
const paramClosingTag = `</${currentParamName}>`
31-
if (currentParamValue.endsWith(paramClosingTag)) {
32-
// End of param value.
33-
currentToolUse.params[currentParamName] = currentParamValue.slice(0, -paramClosingTag.length).trim()
34-
currentParamName = undefined
35-
continue
36-
} else {
37-
// Partial param value is accumulating.
38-
continue
39-
}
40-
}
41-
42-
// No currentParamName.
43-
44-
if (currentToolUse) {
45-
const currentToolValue = accumulator.slice(currentToolUseStartIndex)
46-
const toolUseClosingTag = `</${currentToolUse.name}>`
47-
if (currentToolValue.endsWith(toolUseClosingTag)) {
48-
// End of a tool use.
49-
currentToolUse.partial = false
50-
contentBlocks.push(currentToolUse)
51-
currentToolUse = undefined
52-
continue
53-
} else {
54-
const possibleParamOpeningTags = toolParamNames.map((name) => `<${name}>`)
55-
for (const paramOpeningTag of possibleParamOpeningTags) {
56-
if (accumulator.endsWith(paramOpeningTag)) {
57-
// Start of a new parameter.
58-
currentParamName = paramOpeningTag.slice(1, -1) as ToolParamName
59-
currentParamValueStartIndex = accumulator.length
60-
break
61-
}
62-
}
63-
64-
// There's no current param, and not starting a new param.
65-
66-
// Special case for write_to_file where file contents could
67-
// contain the closing tag, in which case the param would have
68-
// closed and we end up with the rest of the file contents here.
69-
// To work around this, we get the string between the starting
70-
// content tag and the LAST content tag.
71-
const contentParamName: ToolParamName = "content"
72-
73-
if (currentToolUse.name === "write_to_file" && accumulator.endsWith(`</${contentParamName}>`)) {
74-
const toolContent = accumulator.slice(currentToolUseStartIndex)
75-
const contentStartTag = `<${contentParamName}>`
76-
const contentEndTag = `</${contentParamName}>`
77-
const contentStartIndex = toolContent.indexOf(contentStartTag) + contentStartTag.length
78-
const contentEndIndex = toolContent.lastIndexOf(contentEndTag)
79-
80-
if (contentStartIndex !== -1 && contentEndIndex !== -1 && contentEndIndex > contentStartIndex) {
81-
currentToolUse.params[contentParamName] = toolContent
82-
.slice(contentStartIndex, contentEndIndex)
83-
.trim()
84-
}
85-
}
86-
87-
// Partial tool value is accumulating.
88-
continue
89-
}
90-
}
91-
92-
// No currentToolUse.
93-
94-
let didStartToolUse = false
95-
const possibleToolUseOpeningTags = toolNames.map((name) => `<${name}>`)
96-
97-
for (const toolUseOpeningTag of possibleToolUseOpeningTags) {
98-
if (accumulator.endsWith(toolUseOpeningTag)) {
99-
// Start of a new tool use.
100-
currentToolUse = {
101-
type: "tool_use",
102-
name: toolUseOpeningTag.slice(1, -1) as ToolName,
103-
params: {},
104-
partial: true,
105-
}
106-
107-
currentToolUseStartIndex = accumulator.length
108-
109-
// This also indicates the end of the current text content.
110-
if (currentTextContent) {
111-
currentTextContent.partial = false
112-
113-
// Remove the partially accumulated tool use tag from the
114-
// end of text (<tool).
115-
currentTextContent.content = currentTextContent.content
116-
.slice(0, -toolUseOpeningTag.slice(0, -1).length)
117-
.trim()
118-
119-
contentBlocks.push(currentTextContent)
120-
currentTextContent = undefined
121-
}
122-
123-
didStartToolUse = true
124-
break
125-
}
126-
}
127-
128-
if (!didStartToolUse) {
129-
// No tool use, so it must be text either at the beginning or
130-
// between tools.
131-
if (currentTextContent === undefined) {
132-
currentTextContentStartIndex = i
133-
}
134-
135-
currentTextContent = {
136-
type: "text",
137-
content: accumulator.slice(currentTextContentStartIndex).trim(),
138-
partial: true,
139-
}
140-
}
141-
}
142-
143-
if (currentToolUse) {
144-
// Stream did not complete tool call, add it as partial.
145-
if (currentParamName) {
146-
// Tool call has a parameter that was not completed.
147-
currentToolUse.params[currentParamName] = accumulator.slice(currentParamValueStartIndex).trim()
148-
}
149-
150-
contentBlocks.push(currentToolUse)
151-
}
152-
153-
// NOTE: It doesn't matter if check for currentToolUse or
154-
// currentTextContent, only one of them will be defined since only one can
155-
// be partial at a time.
156-
if (currentTextContent) {
157-
// Stream did not complete text content, add it as partial.
158-
contentBlocks.push(currentTextContent)
159-
}
7+
export type AssistantMessageContent = import("./directives").Directive
1608

161-
return contentBlocks
9+
export function parseAssistantMessage(assistantMessage: string): import("./directives").Directive[] {
10+
return StreamingParser.parse(assistantMessage)
16211
}

0 commit comments

Comments
 (0)