bolt.new/packages/bolt/app/lib/runtime/message-parser.ts

173 lines
5.0 KiB
TypeScript
Raw Normal View History

2024-07-10 16:44:39 +00:00
const ARTIFACT_TAG_OPEN = '<boltArtifact';
const ARTIFACT_TAG_CLOSE = '</boltArtifact>';
const ARTIFACT_ACTION_TAG_OPEN = '<boltAction';
const ARTIFACT_ACTION_TAG_CLOSE = '</boltAction>';
interface BoltArtifact {
title: string;
}
type ArtifactOpenCallback = (messageId: string, artifact: BoltArtifact) => void;
type ArtifactCloseCallback = (messageId: string) => void;
type ActionCallback = (messageId: string, action: BoltActionData) => void;
type ActionType = 'file' | 'shell';
export interface BoltActionData {
type?: ActionType;
path?: string;
content: string;
}
interface Callbacks {
onArtifactOpen?: ArtifactOpenCallback;
onArtifactClose?: ArtifactCloseCallback;
onAction?: ActionCallback;
}
type ElementFactory = () => string;
interface StreamingMessageParserOptions {
callbacks?: Callbacks;
artifactElement?: string | ElementFactory;
}
export class StreamingMessageParser {
#lastPositions = new Map<string, number>();
#insideArtifact = false;
#insideAction = false;
#currentAction: BoltActionData = { content: '' };
constructor(private _options: StreamingMessageParserOptions = {}) {}
parse(id: string, input: string) {
let output = '';
let i = this.#lastPositions.get(id) ?? 0;
let earlyBreak = false;
while (i < input.length) {
if (this.#insideArtifact) {
if (this.#insideAction) {
const closeIndex = input.indexOf(ARTIFACT_ACTION_TAG_CLOSE, i);
if (closeIndex !== -1) {
this.#currentAction.content += input.slice(i, closeIndex);
let content = this.#currentAction.content.trim();
if (this.#currentAction.type === 'file') {
content += '\n';
}
this.#currentAction.content = content;
this._options.callbacks?.onAction?.(id, this.#currentAction);
this.#insideAction = false;
this.#currentAction = { content: '' };
i = closeIndex + ARTIFACT_ACTION_TAG_CLOSE.length;
} else {
break;
}
} else {
const actionOpenIndex = input.indexOf(ARTIFACT_ACTION_TAG_OPEN, i);
const artifactCloseIndex = input.indexOf(ARTIFACT_TAG_CLOSE, i);
if (actionOpenIndex !== -1 && (artifactCloseIndex === -1 || actionOpenIndex < artifactCloseIndex)) {
const actionEndIndex = input.indexOf('>', actionOpenIndex);
if (actionEndIndex !== -1) {
const actionTag = input.slice(actionOpenIndex, actionEndIndex + 1);
this.#currentAction.type = this.#extractAttribute(actionTag, 'type') as ActionType;
this.#currentAction.path = this.#extractAttribute(actionTag, 'path');
this.#insideAction = true;
i = actionEndIndex + 1;
} else {
break;
}
} else if (artifactCloseIndex !== -1) {
this.#insideArtifact = false;
this._options.callbacks?.onArtifactClose?.(id);
i = artifactCloseIndex + ARTIFACT_TAG_CLOSE.length;
} else {
break;
}
}
} else if (input[i] === '<' && input[i + 1] !== '/') {
let j = i;
let potentialTag = '';
while (j < input.length && potentialTag.length < ARTIFACT_TAG_OPEN.length) {
potentialTag += input[j];
if (potentialTag === ARTIFACT_TAG_OPEN) {
const nextChar = input[j + 1];
if (nextChar && nextChar !== '>' && nextChar !== ' ') {
output += input.slice(i, j + 1);
i = j + 1;
break;
}
const openTagEnd = input.indexOf('>', j);
if (openTagEnd !== -1) {
const artifactTag = input.slice(i, openTagEnd + 1);
const artifactTitle = this.#extractAttribute(artifactTag, 'title') as string;
this.#insideArtifact = true;
this._options.callbacks?.onArtifactOpen?.(id, { title: artifactTitle });
output += this._options.artifactElement ?? `<div class="__boltArtifact__" data-message-id="${id}"></div>`;
i = openTagEnd + 1;
} else {
earlyBreak = true;
}
break;
} else if (!ARTIFACT_TAG_OPEN.startsWith(potentialTag)) {
output += input.slice(i, j + 1);
i = j + 1;
break;
}
j++;
}
if (j === input.length && ARTIFACT_TAG_OPEN.startsWith(potentialTag)) {
break;
}
} else {
output += input[i];
i++;
}
if (earlyBreak) {
break;
}
}
this.#lastPositions.set(id, i);
return output;
}
reset() {
this.#lastPositions.clear();
this.#insideArtifact = false;
this.#insideAction = false;
this.#currentAction = { content: '' };
}
#extractAttribute(tag: string, attributeName: string): string | undefined {
const match = tag.match(new RegExp(`${attributeName}="([^"]*)"`, 'i'));
return match ? match[1] : undefined;
}
}