Skip to content

Commit eb77d0d

Browse files
Gerome-ElassaadGerome-Elassaad
committed
Update message-parser.ts
1 parent 6d42821 commit eb77d0d

1 file changed

Lines changed: 117 additions & 43 deletions

File tree

‎app/lib/runtime/message-parser.ts‎

Lines changed: 117 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,16 @@ import { unreachable } from '~/utils/unreachable';
1212

1313
const ARTIFACT_TAG_OPEN = '<codinitArtifact';
1414
const ARTIFACT_TAG_CLOSE = '</codinitArtifact>';
15-
const ARTIFACT_ACTION_TAG_OPEN = '<CodinitAction';
16-
const ARTIFACT_ACTION_TAG_CLOSE = '</CodinitAction>';
15+
const ARTIFACT_ACTION_TAG_OPEN = '<codinitAction';
16+
const ARTIFACT_ACTION_TAG_CLOSE = '</codinitAction>';
17+
const CODINIT_QUICK_ACTIONS_OPEN = '<codinit-quick-actions>';
18+
const CODINIT_QUICK_ACTIONS_CLOSE = '</codinit-quick-actions>';
1719

1820
const logger = createScopedLogger('MessageParser');
1921

2022
export interface ArtifactCallbackData extends codinitArticactData {
2123
messageId: string;
24+
artifactId?: string;
2225
}
2326

2427
export interface ActionCallbackData {
@@ -41,6 +44,7 @@ export interface ParserCallbacks {
4144

4245
interface ElementFactoryProps {
4346
messageId: string;
47+
artifactId?: string;
4448
}
4549

4650
type ElementFactory = (props: ElementFactoryProps) => string;
@@ -54,6 +58,7 @@ interface MessageState {
5458
position: number;
5559
insideArtifact: boolean;
5660
insideAction: boolean;
61+
artifactCounter: number;
5762
currentArtifact?: codinitArticactData;
5863
currentAction: CodinitActionData;
5964
actionId: number;
@@ -77,6 +82,8 @@ function cleanEscapedTags(content: string) {
7782
}
7883
export class StreamingMessageParser {
7984
#messages = new Map<string, MessageState>();
85+
#artifactCounter = 0;
86+
#buffer = ''; // Add a buffer to accumulate incoming chunks
8087

8188
constructor(private _options: StreamingMessageParserOptions = {}) {}
8289

@@ -88,18 +95,58 @@ export class StreamingMessageParser {
8895
position: 0,
8996
insideAction: false,
9097
insideArtifact: false,
98+
artifactCounter: 0,
9199
currentAction: { content: '' },
92100
actionId: 0,
93101
};
94102

95103
this.#messages.set(messageId, state);
96104
}
97105

98-
let output = '';
99-
let i = state.position;
100-
let earlyBreak = false;
106+
this.#buffer += input; // Append new input to the buffer
107+
108+
let parsedOutput = ''; // Use a new variable for the output of this parse call
109+
let i = 0; // Current position in the buffer
110+
let lastOutputIndex = 0; // Tracks the last index from which content was appended to parsedOutput
111+
112+
while (i < this.#buffer.length) {
113+
// Handle CODINIT_QUICK_ACTIONS_OPEN
114+
if (this.#buffer.startsWith(CODINIT_QUICK_ACTIONS_OPEN, i)) {
115+
const actionsBlockEnd = this.#buffer.indexOf(CODINIT_QUICK_ACTIONS_CLOSE, i);
116+
117+
if (actionsBlockEnd !== -1) {
118+
parsedOutput += this.#buffer.slice(lastOutputIndex, i); // Add content before the quick actions block
119+
120+
const actionsBlockContent = this.#buffer.slice(i + CODINIT_QUICK_ACTIONS_OPEN.length, actionsBlockEnd);
121+
const quickActionRegex = /<codinit-quick-action([^>]*)>([\s\S]*?)<\/codinit-quick-action>/g;
122+
let match;
123+
const buttons = [];
124+
125+
while ((match = quickActionRegex.exec(actionsBlockContent)) !== null) {
126+
const tagAttrs = match[1];
127+
const label = match[2];
128+
const type = this.#extractAttribute(tagAttrs, 'type');
129+
const message = this.#extractAttribute(tagAttrs, 'message');
130+
const path = this.#extractAttribute(tagAttrs, 'path');
131+
const href = this.#extractAttribute(tagAttrs, 'href');
132+
buttons.push(
133+
createQuickActionElement(
134+
{ type: type || '', message: message || '', path: path || '', href: href || '' },
135+
label,
136+
),
137+
);
138+
}
139+
parsedOutput += createQuickActionGroup(buttons);
140+
i = actionsBlockEnd + CODINIT_QUICK_ACTIONS_CLOSE.length;
141+
lastOutputIndex = i; // Update lastOutputIndex after processing quick actions
142+
continue;
143+
} else {
144+
// Incomplete quick actions block, wait for more data
145+
break;
146+
}
147+
}
101148

102-
while (i < input.length) {
149+
// Handle insideArtifact state
103150
if (state.insideArtifact) {
104151
const currentArtifact = state.currentArtifact;
105152

@@ -108,12 +155,11 @@ export class StreamingMessageParser {
108155
}
109156

110157
if (state.insideAction) {
111-
const closeIndex = input.indexOf(ARTIFACT_ACTION_TAG_CLOSE, i);
112-
158+
const closeIndex = this.#buffer.indexOf(ARTIFACT_ACTION_TAG_CLOSE, i);
113159
const currentAction = state.currentAction;
114160

115161
if (closeIndex !== -1) {
116-
currentAction.content += input.slice(i, closeIndex);
162+
currentAction.content += this.#buffer.slice(i, closeIndex);
117163

118164
let content = currentAction.content.trim();
119165

@@ -132,24 +178,19 @@ export class StreamingMessageParser {
132178
this._options.callbacks?.onActionClose?.({
133179
artifactId: currentArtifact.id,
134180
messageId,
135-
136-
/**
137-
* We decrement the id because it's been incremented already
138-
* when `onActionOpen` was emitted to make sure the ids are
139-
* the same.
140-
*/
141181
actionId: String(state.actionId - 1),
142-
143182
action: currentAction as CodinitAction,
144183
});
145184

146185
state.insideAction = false;
147186
state.currentAction = { content: '' };
148187

149188
i = closeIndex + ARTIFACT_ACTION_TAG_CLOSE.length;
189+
lastOutputIndex = i; // Update lastOutputIndex after processing action close
150190
} else {
191+
// Incomplete action, stream content and wait for more data
151192
if ('type' in currentAction && currentAction.type === 'file') {
152-
let content = input.slice(i);
193+
let content = this.#buffer.slice(i);
153194

154195
if (!currentAction.filePath.endsWith('.md')) {
155196
content = cleanoutMarkdownSyntax(content);
@@ -171,16 +212,16 @@ export class StreamingMessageParser {
171212
break;
172213
}
173214
} else {
174-
const actionOpenIndex = input.indexOf(ARTIFACT_ACTION_TAG_OPEN, i);
175-
const artifactCloseIndex = input.indexOf(ARTIFACT_TAG_CLOSE, i);
215+
const actionOpenIndex = this.#buffer.indexOf(ARTIFACT_ACTION_TAG_OPEN, i);
216+
const artifactCloseIndex = this.#buffer.indexOf(ARTIFACT_TAG_CLOSE, i);
176217

177218
if (actionOpenIndex !== -1 && (artifactCloseIndex === -1 || actionOpenIndex < artifactCloseIndex)) {
178-
const actionEndIndex = input.indexOf('>', actionOpenIndex);
219+
const actionEndIndex = this.#buffer.indexOf('>', actionOpenIndex);
179220

180221
if (actionEndIndex !== -1) {
222+
parsedOutput += this.#buffer.slice(lastOutputIndex, actionOpenIndex); // Add content before action open
181223
state.insideAction = true;
182-
183-
state.currentAction = this.#parseActionTag(input, actionOpenIndex, actionEndIndex);
224+
state.currentAction = this.#parseActionTag(this.#buffer, actionOpenIndex, actionEndIndex);
184225

185226
this._options.callbacks?.onActionOpen?.({
186227
artifactId: currentArtifact.id,
@@ -190,47 +231,60 @@ export class StreamingMessageParser {
190231
});
191232

192233
i = actionEndIndex + 1;
234+
lastOutputIndex = i; // Update lastOutputIndex after processing action open
193235
} else {
236+
// Incomplete action open tag, wait for more data
194237
break;
195238
}
196239
} else if (artifactCloseIndex !== -1) {
240+
parsedOutput += this.#buffer.slice(lastOutputIndex, artifactCloseIndex); // Add content before artifact close
197241
this._options.callbacks?.onArtifactClose?.({
198242
messageId,
243+
artifactId: currentArtifact.id,
199244
...currentArtifact,
200245
});
201246

202247
state.insideArtifact = false;
203248
state.currentArtifact = undefined;
204249

205250
i = artifactCloseIndex + ARTIFACT_TAG_CLOSE.length;
251+
lastOutputIndex = i; // Update lastOutputIndex after processing artifact close
206252
} else {
253+
// Incomplete artifact, wait for more data
207254
break;
208255
}
209256
}
210-
} else if (input[i] === '<' && input[i + 1] !== '/') {
257+
} else if (this.#buffer[i] === '<' && this.#buffer[i + 1] !== '/') {
211258
let j = i;
212259
let potentialTag = '';
260+
let tagFound = false;
213261

214-
while (j < input.length && potentialTag.length < ARTIFACT_TAG_OPEN.length) {
215-
potentialTag += input[j];
262+
while (j < this.#buffer.length && potentialTag.length < ARTIFACT_TAG_OPEN.length) {
263+
potentialTag += this.#buffer[j];
216264

217265
if (potentialTag === ARTIFACT_TAG_OPEN) {
218-
const nextChar = input[j + 1];
266+
const nextChar = this.#buffer[j + 1];
219267

220268
if (nextChar && nextChar !== '>' && nextChar !== ' ') {
221-
output += input.slice(i, j + 1);
269+
// This is not a codinitArtifact tag, treat as normal text
270+
parsedOutput += this.#buffer.slice(lastOutputIndex, j + 1);
222271
i = j + 1;
272+
lastOutputIndex = i;
273+
tagFound = true;
223274
break;
224275
}
225276

226-
const openTagEnd = input.indexOf('>', j);
277+
const openTagEnd = this.#buffer.indexOf('>', j);
227278

228279
if (openTagEnd !== -1) {
229-
const artifactTag = input.slice(i, openTagEnd + 1);
280+
parsedOutput += this.#buffer.slice(lastOutputIndex, i); // Add content before artifact open
281+
282+
const artifactTag = this.#buffer.slice(i, openTagEnd + 1);
230283

231284
const artifactTitle = this.#extractAttribute(artifactTag, 'title') as string;
232285
const type = this.#extractAttribute(artifactTag, 'type') as string;
233-
const artifactId = this.#extractAttribute(artifactTag, 'id') as string;
286+
287+
const artifactId = `${messageId}-${state.artifactCounter++}`;
234288

235289
if (!artifactTitle) {
236290
logger.warn('Artifact title missing');
@@ -252,44 +306,50 @@ export class StreamingMessageParser {
252306

253307
this._options.callbacks?.onArtifactOpen?.({
254308
messageId,
309+
artifactId: currentArtifact.id,
255310
...currentArtifact,
256311
});
257312

258313
const artifactFactory = this._options.artifactElement ?? createArtifactElement;
259314

260-
output += artifactFactory({ messageId });
315+
parsedOutput += artifactFactory({ messageId, artifactId });
261316

262317
i = openTagEnd + 1;
318+
lastOutputIndex = i; // Update lastOutputIndex after processing artifact open
319+
tagFound = true;
263320
} else {
264-
earlyBreak = true;
321+
// Incomplete artifact open tag, wait for more data
322+
break;
265323
}
266324

267325
break;
268326
} else if (!ARTIFACT_TAG_OPEN.startsWith(potentialTag)) {
269-
output += input.slice(i, j + 1);
327+
// Not a codinitArtifact tag, treat as normal text
328+
parsedOutput += this.#buffer.slice(lastOutputIndex, j + 1);
270329
i = j + 1;
330+
lastOutputIndex = i;
331+
tagFound = true;
271332
break;
272333
}
273334

274335
j++;
275336
}
276337

277-
if (j === input.length && ARTIFACT_TAG_OPEN.startsWith(potentialTag)) {
338+
if (!tagFound) {
339+
// If no tag was found or it's an incomplete potential tag, break and wait for more data
278340
break;
279341
}
280342
} else {
281-
output += input[i];
343+
// Normal text character
282344
i++;
283345
}
284-
285-
if (earlyBreak) {
286-
break;
287-
}
288346
}
289347

290-
state.position = i;
348+
// Append any remaining non-processed content to the output
349+
parsedOutput += this.#buffer.slice(lastOutputIndex, i);
350+
this.#buffer = this.#buffer.slice(i); // Remove processed content from the buffer
291351

292-
return output;
352+
return parsedOutput;
293353
}
294354

295355
reset() {
@@ -334,7 +394,7 @@ export class StreamingMessageParser {
334394
}
335395

336396
(actionAttributes as FileAction).filePath = filePath;
337-
} else if (!['shell', 'start'].includes(actionType)) {
397+
} else if (!['shell', 'start', 'build'].includes(actionType)) {
338398
logger.warn(`Unknown action type '${actionType}'`);
339399
}
340400

@@ -349,7 +409,7 @@ export class StreamingMessageParser {
349409

350410
const createArtifactElement: ElementFactory = (props) => {
351411
const elementProps = [
352-
'class="__codinitArticact__"',
412+
'class="__codinitArtifact__"',
353413
...Object.entries(props).map(([key, value]) => {
354414
return `data-${camelToDashCase(key)}=${JSON.stringify(value)}`;
355415
}),
@@ -361,3 +421,17 @@ const createArtifactElement: ElementFactory = (props) => {
361421
function camelToDashCase(input: string) {
362422
return input.replace(/([a-z])([A-Z])/g, '$1-$2').toLowerCase();
363423
}
424+
425+
function createQuickActionElement(props: Record<string, string>, label: string) {
426+
const elementProps = [
427+
'class="__codinitQuickAction__"',
428+
'data-codinit-quick-action="true"',
429+
...Object.entries(props).map(([key, value]) => `data-${camelToDashCase(key)}=${JSON.stringify(value)}`),
430+
];
431+
432+
return `<button ${elementProps.join(' ')}>${label}</button>`;
433+
}
434+
435+
function createQuickActionGroup(buttons: string[]) {
436+
return `<div class=\"__codinitQuickAction__\" data-codinit-quick-action=\"true\">${buttons.join('')}</div>`;
437+
}

0 commit comments

Comments
 (0)