1+ import { TextContent , ToolParamName , ToolUse , toolParamNames } from "../../shared/tools"
2+ import { toolNames , ToolName } from "../../schemas"
3+
4+ // We keep parsing state between chunks so we don't repeatedly
5+ // re-scan the entire assistant message each time a new token arrives.
6+ export interface ParserState {
7+ currentToolUse ?: ToolUse
8+ currentParamName ?: ToolParamName
9+ currentTextContent ?: TextContent
10+ }
11+
12+ const toolOpeningTags = toolNames . map ( ( name ) => `<${ name } )
13+ const paramOpeningTags = toolParamNames . map ( ( name ) => `<${ name } )
14+
15+ /**
16+ * Parses a chunk of streaming assistant text, mutating `state` and
17+ * returning *completed* content blocks in the order they were closed.
18+ */
19+ export function parseAssistantMessageChunk (
20+ state : ParserState ,
21+ chunk : string ,
22+ ) : ( ToolUse | TextContent ) [ ] {
23+ let accumulator = ""
24+ const completedBlocks : ( ToolUse | TextContent ) [ ] = [ ]
25+
26+ // Helper – finalise and push current text block.
27+ const flushText = ( ) => {
28+ if ( state . currentTextContent ) {
29+ state . currentTextContent . content = state . currentTextContent . content + accumulator . trimEnd ( )
30+ accumulator = ""
31+ // Only push when a tool starts or chunk ends (done below)
32+ }
33+ }
34+
35+ for ( let i = 0 ; i < chunk . length ; i ++ ) {
36+ const char = chunk [ i ]
37+ accumulator += char
38+
39+ if ( state . currentToolUse && state . currentParamName ) {
40+ // Inside a param – look for param closing tag.
41+ const paramClosing = `</${ state . currentParamName }
42+ if ( accumulator . endsWith ( paramClosing ) ) {
43+ // Strip closing tag from param value.
44+ const value = accumulator . slice ( 0 , - paramClosing . length )
45+ state . currentToolUse . params [ state . currentParamName ] =
46+ ( state . currentToolUse . params [ state . currentParamName ] ?? "" ) + value
47+
48+ state . currentParamName = undefined
49+ accumulator = "" // reset accumulator for next parsing portion
50+ }
51+ continue // Still inside param; don't treat anything else.
52+ }
53+
54+ // Inside a tool but not inside param.
55+ if ( state . currentToolUse ) {
56+ const toolClosing = `</${ state . currentToolUse . name }
57+ if ( accumulator . endsWith ( toolClosing ) ) {
58+ // Tool block complete
59+ state . currentToolUse . partial = false
60+ state . currentToolUse = undefined
61+ accumulator = ""
62+ continue
63+ }
64+
65+ // Look for param opening tags
66+ for ( const open of paramOpeningTags ) {
67+ if ( accumulator . endsWith ( open ) ) {
68+ state . currentParamName = open . slice ( 1 , - 1 ) as ToolParamName
69+ accumulator = "" // reset; param value starts next char
70+ // Initialize empty param value
71+ state . currentToolUse . params [ state . currentParamName ] = ""
72+ break
73+ }
74+ }
75+ continue
76+ }
77+
78+ // Not in a tool – look for opening of a tool.
79+ for ( const open of toolOpeningTags ) {
80+ if ( accumulator . endsWith ( open ) ) {
81+ // Flush any text accumulated *before* the tool tag.
82+ if ( state . currentTextContent ) {
83+ state . currentTextContent . partial = false
84+ completedBlocks . push ( state . currentTextContent )
85+ state . currentTextContent = undefined
86+ }
87+
88+ state . currentToolUse = {
89+ type : "tool_use" ,
90+ name : open . slice ( 1 , - 1 ) as ToolName ,
91+ params : { } ,
92+ partial : true ,
93+ }
94+ // Expose the new (partial) tool block immediately so the UI can
95+ // respond (e.g. opening diff view for write_to_file).
96+ completedBlocks . push ( state . currentToolUse )
97+ accumulator = "" // reset; tool content begins next char
98+ break
99+ }
100+ }
101+
102+ // If we just started a tool, skip further text processing.
103+ if ( state . currentToolUse ) {
104+ continue
105+ }
106+
107+ // Regular text content.
108+ if ( ! state . currentTextContent ) {
109+ state . currentTextContent = {
110+ type : "text" ,
111+ content : "" ,
112+ partial : true ,
113+ }
114+ // Immediately push the partial text block so downstream consumers
115+ // can stream it progressively.
116+ completedBlocks . push ( state . currentTextContent )
117+ }
118+ // Text will be flushed at the end or before next block.
119+ }
120+
121+ // End of chunk – append remaining accumulator to current context.
122+ if ( state . currentParamName && state . currentToolUse ) {
123+ // Still inside param value.
124+ state . currentToolUse . params [ state . currentParamName ] =
125+ ( state . currentToolUse . params [ state . currentParamName ] ?? "" ) + accumulator
126+ accumulator = ""
127+ } else if ( state . currentToolUse ) {
128+ // Inside tool but outside param – append to a special placeholder param? (ignored)
129+ // Nothing to do.
130+ } else {
131+ // Plain text
132+ if ( ! state . currentTextContent ) {
133+ state . currentTextContent = {
134+ type : "text" ,
135+ content : "" ,
136+ partial : true ,
137+ }
138+ }
139+ state . currentTextContent . content += accumulator
140+ accumulator = ""
141+ }
142+
143+ // If end of chunk finalises a text block fully (no open tool starting later)
144+ flushText ( )
145+
146+ return completedBlocks
147+ }
148+
149+ export function createInitialParserState ( ) : ParserState {
150+ return {
151+ currentToolUse : undefined ,
152+ currentParamName : undefined ,
153+ currentTextContent : undefined ,
154+ }
155+ }
0 commit comments