From 733b2ac89ba8420ce37c4e109f2bc372dfa2e395 Mon Sep 17 00:00:00 2001
From: Anirban Kar <thecodacus@gmail.com>
Date: Sun, 17 Nov 2024 11:33:24 +0530
Subject: [PATCH] fix(project-reload): execution order is fixed, crutial on
 preject reload

---
 app/lib/hooks/useMessageParser.ts | 26 +++++++++++++-------------
 app/lib/runtime/message-parser.ts | 16 ++++++++--------
 2 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/app/lib/hooks/useMessageParser.ts b/app/lib/hooks/useMessageParser.ts
index 97a063da5..e86e70bd3 100644
--- a/app/lib/hooks/useMessageParser.ts
+++ b/app/lib/hooks/useMessageParser.ts
@@ -8,37 +8,37 @@ const logger = createScopedLogger('useMessageParser');
 
 const messageParser = new StreamingMessageParser({
   callbacks: {
-    onArtifactOpen: (data) => {
+    onArtifactOpen: async (data) => {
       logger.trace('onArtifactOpen', data);
 
       workbenchStore.showWorkbench.set(true);
-      workbenchStore.addArtifact(data);
+      await workbenchStore.addArtifact(data);
     },
-    onArtifactClose: (data) => {
+    onArtifactClose: async (data) => {
       logger.trace('onArtifactClose');
 
-      workbenchStore.updateArtifact(data, { closed: true });
+      await workbenchStore.updateArtifact(data, { closed: true });
     },
-    onActionOpen: (data) => {
+    onActionOpen: async (data) => {
       logger.trace('onActionOpen', data.action);
 
       // we only add shell actions when when the close tag got parsed because only then we have the content
       if (data.action.type !== 'shell') {
-        workbenchStore.addAction(data);
+        await workbenchStore.addAction(data);
       }
     },
-    onActionClose: (data) => {
+    onActionClose: async (data) => {
       logger.trace('onActionClose', data.action);
 
       if (data.action.type === 'shell') {
-        workbenchStore.addAction(data);
+        await workbenchStore.addAction(data);
       }
 
-      workbenchStore.runAction(data);
+      await workbenchStore.runAction(data);
     },
-    onActionStream: (data) => {
+    onActionStream: async (data) => {
       logger.trace('onActionStream', data.action);
-      workbenchStore.runAction(data, true);
+      await workbenchStore.runAction(data, true);
     },
   },
 });
@@ -46,7 +46,7 @@ const messageParser = new StreamingMessageParser({
 export function useMessageParser() {
   const [parsedMessages, setParsedMessages] = useState<{ [key: number]: string }>({});
 
-  const parseMessages = useCallback((messages: Message[], isLoading: boolean) => {
+  const parseMessages = useCallback(async (messages: Message[], isLoading: boolean) => {
     let reset = false;
 
     if (import.meta.env.DEV && !isLoading) {
@@ -56,7 +56,7 @@ export function useMessageParser() {
 
     for (const [index, message] of messages.entries()) {
       if (message.role === 'assistant') {
-        const newParsedContent = messageParser.parse(message.id, message.content);
+        const newParsedContent = await messageParser.parse(message.id, message.content);
 
         setParsedMessages((prevParsed) => ({
           ...prevParsed,
diff --git a/app/lib/runtime/message-parser.ts b/app/lib/runtime/message-parser.ts
index 4b564da16..093228775 100644
--- a/app/lib/runtime/message-parser.ts
+++ b/app/lib/runtime/message-parser.ts
@@ -21,8 +21,8 @@ export interface ActionCallbackData {
   action: BoltAction;
 }
 
-export type ArtifactCallback = (data: ArtifactCallbackData) => void;
-export type ActionCallback = (data: ActionCallbackData) => void;
+export type ArtifactCallback = (data: ArtifactCallbackData) => Promise<void>;
+export type ActionCallback = (data: ActionCallbackData) => Promise<void>;
 
 export interface ParserCallbacks {
   onArtifactOpen?: ArtifactCallback;
@@ -57,7 +57,7 @@ export class StreamingMessageParser {
 
   constructor(private _options: StreamingMessageParserOptions = {}) { }
 
-  parse(messageId: string, input: string) {
+  async parse(messageId: string, input: string) {
     let state = this.#messages.get(messageId);
 
     if (!state) {
@@ -100,7 +100,7 @@ export class StreamingMessageParser {
 
             currentAction.content = content;
 
-            this._options.callbacks?.onActionClose?.({
+            await this._options.callbacks?.onActionClose?.({
               artifactId: currentArtifact.id,
               messageId,
 
@@ -122,7 +122,7 @@ export class StreamingMessageParser {
             if ('type' in currentAction && currentAction.type === 'file') {
               let content = input.slice(i);
 
-              this._options.callbacks?.onActionStream?.({
+              await this._options.callbacks?.onActionStream?.({
                 artifactId: currentArtifact.id,
                 messageId,
                 actionId: String(state.actionId - 1),
@@ -148,7 +148,7 @@ export class StreamingMessageParser {
 
               state.currentAction = this.#parseActionTag(input, actionOpenIndex, actionEndIndex);
 
-              this._options.callbacks?.onActionOpen?.({
+              await this._options.callbacks?.onActionOpen?.({
                 artifactId: currentArtifact.id,
                 messageId,
                 actionId: String(state.actionId++),
@@ -160,7 +160,7 @@ export class StreamingMessageParser {
               break;
             }
           } else if (artifactCloseIndex !== -1) {
-            this._options.callbacks?.onArtifactClose?.({ messageId, ...currentArtifact });
+            await this._options.callbacks?.onArtifactClose?.({ messageId, ...currentArtifact });
 
             state.insideArtifact = false;
             state.currentArtifact = undefined;
@@ -211,7 +211,7 @@ export class StreamingMessageParser {
 
               state.currentArtifact = currentArtifact;
 
-              this._options.callbacks?.onArtifactOpen?.({ messageId, ...currentArtifact });
+              await this._options.callbacks?.onArtifactOpen?.({ messageId, ...currentArtifact });
 
               const artifactFactory = this._options.artifactElement ?? createArtifactElement;