feat: improve prompt, add ability to abort streaming, improve message parser

packages/bolt/app/lib/stores/workbench.ts

@@ -4,25 +4,30 @@ import { unreachable } from '../../utils/unreachable';
 import { ActionRunner } from '../runtime/action-runner';
 import type { ActionCallbackData, ArtifactCallbackData } from '../runtime/message-parser';
 import { webcontainer } from '../webcontainer';
+import { chatStore } from './chat';
 import { PreviewsStore } from './previews';
 
-export type RunningState = BoltAction & {
+const MIN_SPINNER_TIME = 200;
+
+export type BaseActionState = BoltAction & {
   status: 'running' | 'complete' | 'pending' | 'aborted';
+  executing: boolean;
   abort?: () => void;
 };
 
-export type FailedState = BoltAction & {
-  status: 'failed';
-  error: string;
-  abort?: () => void;
-};
+export type FailedActionState = BoltAction &
+  Omit<BaseActionState, 'status'> & {
+    status: 'failed';
+    error: string;
+  };
 
-export type ActionState = RunningState | FailedState;
+export type ActionState = BaseActionState | FailedActionState;
+
+type BaseActionUpdate = Partial<Pick<BaseActionState, 'status' | 'executing' | 'abort'>>;
 
 export type ActionStateUpdate =
-  | { status: 'running' | 'complete' | 'pending' | 'aborted'; abort?: () => void }
-  | { status: 'failed'; error: string; abort?: () => void }
-  | { abort?: () => void };
+  | BaseActionUpdate
+  | (Omit<BaseActionUpdate, 'status'> & { status: 'failed'; error: string });
 
 export interface ArtifactState {
   title: string;
@@ -49,6 +54,16 @@ export class WorkbenchStore {
     this.showWorkbench.set(show);
   }
 
+  abortAllActions() {
+    for (const [, artifact] of Object.entries(this.artifacts.get())) {
+      for (const [, action] of Object.entries(artifact.actions.get())) {
+        if (action.status === 'running') {
+          action.abort?.();
+        }
+      }
+    }
+  }
+
   addArtifact({ id, messageId, title }: ArtifactCallbackData) {
     const artifacts = this.artifacts.get();
     const artifactKey = getArtifactKey(id, messageId);
@@ -78,7 +93,7 @@ export class WorkbenchStore {
     this.artifacts.setKey(key, { ...artifact, ...state });
   }
 
-  async runAction(data: ActionCallbackData) {
+  async addAction(data: ActionCallbackData) {
     const { artifactId, messageId, actionId } = data;
 
     const artifacts = this.artifacts.get();
@@ -96,33 +111,70 @@ export class WorkbenchStore {
       return;
     }
 
-    artifact.actions.setKey(actionId, { ...data.action, status: 'pending' });
+    artifact.actions.setKey(actionId, { ...data.action, status: 'pending', executing: false });
+
+    artifact.currentActionPromise.then(() => {
+      if (chatStore.get().aborted) {
+        return;
+      }
+
+      this.#updateAction(key, actionId, { status: 'running' });
+    });
+  }
+
+  async runAction(data: ActionCallbackData) {
+    const { artifactId, messageId, actionId } = data;
+
+    const artifacts = this.artifacts.get();
+    const key = getArtifactKey(artifactId, messageId);
+    const artifact = artifacts[key];
+
+    if (!artifact) {
+      unreachable('Artifact not found');
+    }
+
+    const actions = artifact.actions.get();
+    const action = actions[actionId];
+
+    if (!action) {
+      unreachable('Expected action to exist');
+    }
+
+    if (action.executing || action.status === 'complete' || action.status === 'failed' || action.status === 'aborted') {
+      return;
+    }
 
     artifact.currentActionPromise = artifact.currentActionPromise.then(async () => {
+      if (chatStore.get().aborted) {
+        return;
+      }
+
+      const abortController = new AbortController();
+
+      this.#updateAction(key, actionId, {
+        status: 'running',
+        executing: true,
+        abort: () => {
+          abortController.abort();
+          this.#updateAction(key, actionId, { status: 'aborted' });
+        },
+      });
+
       try {
-        let abortController: AbortController | undefined;
+        await Promise.all([
+          this.#actionRunner.runAction(data, abortController.signal),
+          new Promise((resolve) => setTimeout(resolve, MIN_SPINNER_TIME)),
+        ]);
 
-        if (data.action.type === 'shell') {
-          abortController = new AbortController();
+        if (!abortController.signal.aborted) {
+          this.#updateAction(key, actionId, { status: 'complete' });
         }
-
-        let aborted = false;
-
-        this.#updateAction(key, actionId, {
-          status: 'running',
-          abort: () => {
-            aborted = true;
-            abortController?.abort();
-          },
-        });
-
-        await this.#actionRunner.runAction(data, abortController?.signal);
-
-        this.#updateAction(key, actionId, { status: aborted ? 'aborted' : 'complete' });
       } catch (error) {
         this.#updateAction(key, actionId, { status: 'failed', error: 'Action failed' });
 
         throw error;
+      } finally {
+        this.#updateAction(key, actionId, { executing: false });
       }
     });
   }