From f4987a4ecdf143ce51b6c29e93d4c529aa4af1e2 Mon Sep 17 00:00:00 2001
From: Dominic Elm <d3lm@users.noreply.github.com>
Date: Mon, 22 Jul 2024 17:40:28 +0200
Subject: [PATCH] refactor: workbench store and move logic into action runner
 (#4)

---
 .../bolt/app/components/chat/Artifact.tsx     | 114 +++++-----
 .../bolt/app/components/chat/Markdown.tsx     |   9 +-
 .../bolt/app/components/workbench/Preview.tsx |  17 +-
 .../app/lib/.server/llm/switchable-stream.ts  |   2 +
 .../bolt/app/lib/runtime/action-runner.ts     | 195 ++++++++++++++----
 .../app/lib/runtime/message-parser.spec.ts    |   2 +-
 .../bolt/app/lib/runtime/message-parser.ts    |  29 ++-
 packages/bolt/app/lib/stores/previews.ts      |   7 +
 packages/bolt/app/lib/stores/workbench.ts     | 147 ++-----------
 packages/bolt/app/routes/api.chat.ts          |  10 +-
 10 files changed, 295 insertions(+), 237 deletions(-)

diff --git a/packages/bolt/app/components/chat/Artifact.tsx b/packages/bolt/app/components/chat/Artifact.tsx
index 130b184..a321372 100644
--- a/packages/bolt/app/components/chat/Artifact.tsx
+++ b/packages/bolt/app/components/chat/Artifact.tsx
@@ -3,11 +3,11 @@ import { AnimatePresence, motion } from 'framer-motion';
 import { computed } from 'nanostores';
 import { memo, useEffect, useRef, useState } from 'react';
 import { createHighlighter, type BundledLanguage, type BundledTheme, type HighlighterGeneric } from 'shiki';
+import type { ActionState } from '../../lib/runtime/action-runner';
 import { chatStore } from '../../lib/stores/chat';
-import { getArtifactKey, workbenchStore, type ActionState } from '../../lib/stores/workbench';
+import { workbenchStore } from '../../lib/stores/workbench';
 import { classNames } from '../../utils/classNames';
 import { cubicEasingFn } from '../../utils/easings';
-import { IconButton } from '../ui/IconButton';
 
 const highlighterOptions = {
   langs: ['shell'],
@@ -22,20 +22,19 @@ if (import.meta.hot) {
 }
 
 interface ArtifactProps {
-  artifactId: string;
   messageId: string;
 }
 
-export const Artifact = memo(({ artifactId, messageId }: ArtifactProps) => {
+export const Artifact = memo(({ messageId }: ArtifactProps) => {
   const userToggledActions = useRef(false);
   const [showActions, setShowActions] = useState(false);
 
   const chat = useStore(chatStore);
   const artifacts = useStore(workbenchStore.artifacts);
-  const artifact = artifacts[getArtifactKey(artifactId, messageId)];
+  const artifact = artifacts[messageId];
 
   const actions = useStore(
-    computed(artifact.actions, (actions) => {
+    computed(artifact.runner.actions, (actions) => {
       return Object.values(actions);
     }),
   );
@@ -100,50 +99,7 @@ export const Artifact = memo(({ artifactId, messageId }: ArtifactProps) => {
             transition={{ duration: 0.15 }}
           >
             <div className="p-4 text-left border-t">
-              <motion.div
-                initial={{ opacity: 0 }}
-                animate={{ opacity: 1 }}
-                exit={{ opacity: 0 }}
-                transition={{ duration: 0.15 }}
-              >
-                <h4 className="font-semibold mb-2">Actions</h4>
-                <ul className="list-none space-y-2.5">
-                  {actions.map((action, index) => {
-                    const { status, type, content, abort } = action;
-
-                    return (
-                      <li key={index}>
-                        <div className="flex items-center gap-1.5">
-                          <div className={classNames('text-lg', getTextColor(action.status))}>
-                            {status === 'running' ? (
-                              <div className="i-svg-spinners:90-ring-with-bg"></div>
-                            ) : status === 'pending' ? (
-                              <div className="i-ph:circle-duotone"></div>
-                            ) : status === 'complete' ? (
-                              <div className="i-ph:check-circle-duotone"></div>
-                            ) : status === 'failed' || status === 'aborted' ? (
-                              <div className="i-ph:x-circle-duotone"></div>
-                            ) : null}
-                          </div>
-                          {type === 'file' ? (
-                            <div>
-                              Create <code className="bg-gray-100 text-gray-700">{action.filePath}</code>
-                            </div>
-                          ) : type === 'shell' ? (
-                            <div className="flex items-center w-full min-h-[28px]">
-                              <span className="flex-1">Run command</span>
-                              {abort !== undefined && status === 'running' && (
-                                <IconButton icon="i-ph:x-circle" size="xl" onClick={() => abort()} />
-                              )}
-                            </div>
-                          ) : null}
-                        </div>
-                        {type === 'shell' && <ShellCodeBlock classsName="mt-1" code={content} />}
-                      </li>
-                    );
-                  })}
-                </ul>
-              </motion.div>
+              <ActionList actions={actions} />
             </div>
           </motion.div>
         )}
@@ -188,3 +144,61 @@ function ShellCodeBlock({ classsName, code }: ShellCodeBlockProps) {
     ></div>
   );
 }
+
+interface ActionListProps {
+  actions: ActionState[];
+}
+
+const actionVariants = {
+  hidden: { opacity: 0, y: 20 },
+  visible: { opacity: 1, y: 0 },
+};
+
+const ActionList = memo(({ actions }: ActionListProps) => {
+  return (
+    <motion.div initial={{ opacity: 0 }} animate={{ opacity: 1 }} exit={{ opacity: 0 }} transition={{ duration: 0.15 }}>
+      <ul className="list-none space-y-2.5">
+        {actions.map((action, index) => {
+          const { status, type, content } = action;
+
+          return (
+            <motion.li
+              key={index}
+              variants={actionVariants}
+              initial="hidden"
+              animate="visible"
+              transition={{
+                duration: 0.2,
+                ease: cubicEasingFn,
+              }}
+            >
+              <div className="flex items-center gap-1.5">
+                <div className={classNames('text-lg', getTextColor(action.status))}>
+                  {status === 'running' ? (
+                    <div className="i-svg-spinners:90-ring-with-bg"></div>
+                  ) : status === 'pending' ? (
+                    <div className="i-ph:circle-duotone"></div>
+                  ) : status === 'complete' ? (
+                    <div className="i-ph:check-circle-duotone"></div>
+                  ) : status === 'failed' || status === 'aborted' ? (
+                    <div className="i-ph:x-circle-duotone"></div>
+                  ) : null}
+                </div>
+                {type === 'file' ? (
+                  <div>
+                    Create <code className="bg-gray-100 text-gray-700">{action.filePath}</code>
+                  </div>
+                ) : type === 'shell' ? (
+                  <div className="flex items-center w-full min-h-[28px]">
+                    <span className="flex-1">Run command</span>
+                  </div>
+                ) : null}
+              </div>
+              {type === 'shell' && <ShellCodeBlock classsName="mt-1" code={content} />}
+            </motion.li>
+          );
+        })}
+      </ul>
+    </motion.div>
+  );
+});
diff --git a/packages/bolt/app/components/chat/Markdown.tsx b/packages/bolt/app/components/chat/Markdown.tsx
index a6e9c42..609e070 100644
--- a/packages/bolt/app/components/chat/Markdown.tsx
+++ b/packages/bolt/app/components/chat/Markdown.tsx
@@ -21,18 +21,13 @@ export const Markdown = memo(({ children }: MarkdownProps) => {
     return {
       div: ({ className, children, node, ...props }) => {
         if (className?.includes('__boltArtifact__')) {
-          const artifactId = node?.properties.dataArtifactId as string;
           const messageId = node?.properties.dataMessageId as string;
 
-          if (!artifactId) {
-            logger.debug(`Invalid artifact id ${messageId}`);
-          }
-
           if (!messageId) {
-            logger.debug(`Invalid message id ${messageId}`);
+            logger.error(`Invalid message id ${messageId}`);
           }
 
-          return <Artifact artifactId={artifactId} messageId={messageId} />;
+          return <Artifact messageId={messageId} />;
         }
 
         return (
diff --git a/packages/bolt/app/components/workbench/Preview.tsx b/packages/bolt/app/components/workbench/Preview.tsx
index 0b523c0..3a211f0 100644
--- a/packages/bolt/app/components/workbench/Preview.tsx
+++ b/packages/bolt/app/components/workbench/Preview.tsx
@@ -13,7 +13,14 @@ export const Preview = memo(() => {
   const [iframeUrl, setIframeUrl] = useState<string | undefined>();
 
   useEffect(() => {
-    if (activePreview && !iframeUrl) {
+    if (!activePreview) {
+      setUrl('');
+      setIframeUrl(undefined);
+
+      return;
+    }
+
+    if (!iframeUrl) {
       const { baseUrl } = activePreview;
 
       setUrl(baseUrl);
@@ -31,16 +38,16 @@ export const Preview = memo(() => {
     <div className="w-full h-full flex flex-col">
       <div className="bg-gray-100 rounded-t-lg p-2 flex items-center space-x-1.5">
         <div className="flex items-center gap-2 text-gray-800">
-          <div className="i-ph:app-window-duotone scale-130 ml-1.5"></div>
+          <div className="i-ph:app-window-duotone scale-130 ml-1.5" />
           <span className="text-sm">Preview</span>
         </div>
-        <div className="flex-grow"></div>
+        <div className="flex-grow" />
       </div>
       <div className="bg-white p-2 flex items-center gap-1.5">
         <IconButton icon="i-ph:arrow-clockwise" onClick={reloadPreview} />
         <div className="flex items-center gap-1 flex-grow bg-gray-100 rounded-full px-3 py-1 text-sm text-gray-600 hover:bg-gray-200 hover:focus-within:bg-white focus-within:bg-white focus-within:ring-2 focus-within:ring-accent">
           <div className="bg-white rounded-full p-[2px] -ml-1">
-            <div className="i-ph:info-bold text-lg"></div>
+            <div className="i-ph:info-bold text-lg" />
           </div>
           <input
             className="w-full bg-transparent outline-none"
@@ -54,7 +61,7 @@ export const Preview = memo(() => {
       </div>
       <div className="flex-1 bg-white border-t">
         {activePreview ? (
-          <iframe ref={iframeRef} className="border-none w-full h-full" src={iframeUrl}></iframe>
+          <iframe ref={iframeRef} className="border-none w-full h-full" src={iframeUrl} />
         ) : (
           <div className="flex w-full h-full justify-center items-center">No preview available</div>
         )}
diff --git a/packages/bolt/app/lib/.server/llm/switchable-stream.ts b/packages/bolt/app/lib/.server/llm/switchable-stream.ts
index bedc883..94b5ee6 100644
--- a/packages/bolt/app/lib/.server/llm/switchable-stream.ts
+++ b/packages/bolt/app/lib/.server/llm/switchable-stream.ts
@@ -24,6 +24,8 @@ export default class SwitchableStream extends TransformStream {
       await this._currentReader.cancel();
     }
 
+    console.log('Switching stream');
+
     this._currentReader = newStream.getReader();
 
     this._pumpStream();
diff --git a/packages/bolt/app/lib/runtime/action-runner.ts b/packages/bolt/app/lib/runtime/action-runner.ts
index 0f6735e..2b92b62 100644
--- a/packages/bolt/app/lib/runtime/action-runner.ts
+++ b/packages/bolt/app/lib/runtime/action-runner.ts
@@ -1,68 +1,187 @@
 import { WebContainer } from '@webcontainer/api';
+import { map, type MapStore } from 'nanostores';
 import * as nodePath from 'node:path';
+import type { BoltAction } from '../../types/actions';
 import { createScopedLogger } from '../../utils/logger';
+import { unreachable } from '../../utils/unreachable';
 import type { ActionCallbackData } from './message-parser';
 
 const logger = createScopedLogger('ActionRunner');
 
+export type ActionStatus = 'pending' | 'running' | 'complete' | 'aborted' | 'failed';
+
+export type BaseActionState = BoltAction & {
+  status: Exclude<ActionStatus, 'failed'>;
+  abort: () => void;
+  executed: boolean;
+  abortSignal: AbortSignal;
+};
+
+export type FailedActionState = BoltAction &
+  Omit<BaseActionState, 'status'> & {
+    status: Extract<ActionStatus, 'failed'>;
+    error: string;
+  };
+
+export type ActionState = BaseActionState | FailedActionState;
+
+type BaseActionUpdate = Partial<Pick<BaseActionState, 'status' | 'abort' | 'executed'>>;
+
+export type ActionStateUpdate =
+  | BaseActionUpdate
+  | (Omit<BaseActionUpdate, 'status'> & { status: 'failed'; error: string });
+
+type ActionsMap = MapStore<Record<string, ActionState>>;
+
 export class ActionRunner {
   #webcontainer: Promise<WebContainer>;
+  #currentExecutionPromise: Promise<void> = Promise.resolve();
+
+  actions: ActionsMap = import.meta.hot?.data.actions ?? map({});
 
   constructor(webcontainerPromise: Promise<WebContainer>) {
     this.#webcontainer = webcontainerPromise;
+
+    if (import.meta.hot) {
+      import.meta.hot.data.actions = this.actions;
+    }
   }
 
-  async runAction({ action }: ActionCallbackData, abortSignal?: AbortSignal) {
-    logger.trace('Running action', action);
+  addAction(data: ActionCallbackData) {
+    const { actionId } = data;
 
-    const { content } = action;
+    const action = this.actions.get()[actionId];
+
+    if (action) {
+      // action already added
+      return;
+    }
+
+    const abortController = new AbortController();
+
+    this.actions.setKey(actionId, {
+      ...data.action,
+      status: 'pending',
+      executed: false,
+      abort: () => {
+        abortController.abort();
+        this.#updateAction(actionId, { status: 'aborted' });
+      },
+      abortSignal: abortController.signal,
+    });
+
+    this.#currentExecutionPromise.then(() => {
+      this.#updateAction(actionId, { status: 'running' });
+    });
+  }
+
+  async runAction(data: ActionCallbackData) {
+    const { actionId } = data;
+    const action = this.actions.get()[actionId];
+
+    if (!action) {
+      unreachable(`Action ${actionId} not found`);
+    }
+
+    if (action.executed) {
+      return;
+    }
+
+    this.#updateAction(actionId, { ...action, ...data.action, executed: true });
+
+    this.#currentExecutionPromise = this.#currentExecutionPromise
+      .then(() => {
+        return this.#executeAction(actionId);
+      })
+      .catch((error) => {
+        console.error('Action execution failed:', error);
+      });
+  }
+
+  async #executeAction(actionId: string) {
+    const action = this.actions.get()[actionId];
+
+    this.#updateAction(actionId, { status: 'running' });
+
+    try {
+      switch (action.type) {
+        case 'shell': {
+          await this.#runShellAction(action);
+          break;
+        }
+        case 'file': {
+          await this.#runFileAction(action);
+          break;
+        }
+      }
+
+      this.#updateAction(actionId, { status: action.abortSignal.aborted ? 'aborted' : 'complete' });
+    } catch (error) {
+      this.#updateAction(actionId, { status: 'failed', error: 'Action failed' });
+
+      // re-throw the error to be caught in the promise chain
+      throw error;
+    }
+  }
+
+  async #runShellAction(action: ActionState) {
+    if (action.type !== 'shell') {
+      unreachable('Expected shell action');
+    }
 
     const webcontainer = await this.#webcontainer;
 
-    switch (action.type) {
-      case 'file': {
-        let folder = nodePath.dirname(action.filePath);
+    const process = await webcontainer.spawn('jsh', ['-c', action.content]);
 
-        // remove trailing slashes
-        folder = folder.replace(/\/$/g, '');
+    action.abortSignal.addEventListener('abort', () => {
+      process.kill();
+    });
 
-        if (folder !== '.') {
-          try {
-            await webcontainer.fs.mkdir(folder, { recursive: true });
-            logger.debug('Created folder', folder);
-          } catch (error) {
-            logger.error('Failed to create folder\n', error);
-          }
-        }
+    process.output.pipeTo(
+      new WritableStream({
+        write(data) {
+          console.log(data);
+        },
+      }),
+    );
 
-        try {
-          await webcontainer.fs.writeFile(action.filePath, content);
-          logger.debug(`File written ${action.filePath}`);
-        } catch (error) {
-          logger.error('Failed to write file\n', error);
-        }
+    const exitCode = await process.exit;
 
-        break;
-      }
-      case 'shell': {
-        const process = await webcontainer.spawn('jsh', ['-c', content]);
+    logger.debug(`Process terminated with code ${exitCode}`);
+  }
 
-        abortSignal?.addEventListener('abort', () => {
-          process.kill();
-        });
+  async #runFileAction(action: ActionState) {
+    if (action.type !== 'file') {
+      unreachable('Expected file action');
+    }
 
-        process.output.pipeTo(
-          new WritableStream({
-            write(data) {
-              console.log(data);
-            },
-          }),
-        );
+    const webcontainer = await this.#webcontainer;
 
-        const exitCode = await process.exit;
+    let folder = nodePath.dirname(action.filePath);
 
-        logger.debug(`Process terminated with code ${exitCode}`);
+    // remove trailing slashes
+    folder = folder.replace(/\/+$/g, '');
+
+    if (folder !== '.') {
+      try {
+        await webcontainer.fs.mkdir(folder, { recursive: true });
+        logger.debug('Created folder', folder);
+      } catch (error) {
+        logger.error('Failed to create folder\n', error);
       }
     }
+
+    try {
+      await webcontainer.fs.writeFile(action.filePath, action.content);
+      logger.debug(`File written ${action.filePath}`);
+    } catch (error) {
+      logger.error('Failed to write file\n', error);
+    }
+  }
+
+  #updateAction(id: string, newState: ActionStateUpdate) {
+    const actions = this.actions.get();
+
+    this.actions.setKey(id, { ...actions[id], ...newState });
   }
 }
diff --git a/packages/bolt/app/lib/runtime/message-parser.spec.ts b/packages/bolt/app/lib/runtime/message-parser.spec.ts
index bc8080f..739604b 100644
--- a/packages/bolt/app/lib/runtime/message-parser.spec.ts
+++ b/packages/bolt/app/lib/runtime/message-parser.spec.ts
@@ -179,7 +179,7 @@ function runTest(input: string | string[], outputOrExpectedResult: string | Expe
   };
 
   const parser = new StreamingMessageParser({
-    artifactElement: '',
+    artifactElement: () => '',
     callbacks,
   });
 
diff --git a/packages/bolt/app/lib/runtime/message-parser.ts b/packages/bolt/app/lib/runtime/message-parser.ts
index a607759..fd954cb 100644
--- a/packages/bolt/app/lib/runtime/message-parser.ts
+++ b/packages/bolt/app/lib/runtime/message-parser.ts
@@ -31,11 +31,15 @@ export interface ParserCallbacks {
   onActionClose?: ActionCallback;
 }
 
-type ElementFactory = () => string;
+interface ElementFactoryProps {
+  messageId: string;
+}
+
+type ElementFactory = (props: ElementFactoryProps) => string;
 
 export interface StreamingMessageParserOptions {
   callbacks?: ParserCallbacks;
-  artifactElement?: string | ElementFactory;
+  artifactElement?: ElementFactory;
 }
 
 interface MessageState {
@@ -193,9 +197,9 @@ export class StreamingMessageParser {
 
               this._options.callbacks?.onArtifactOpen?.({ messageId, ...currentArtifact });
 
-              output +=
-                this._options.artifactElement ??
-                `<div class="__boltArtifact__" data-artifact-id="${artifactId}" data-message-id="${messageId}"></div>`;
+              const artifactFactory = this._options.artifactElement ?? createArtifactElement;
+
+              output += artifactFactory({ messageId });
 
               i = openTagEnd + 1;
             } else {
@@ -264,3 +268,18 @@ export class StreamingMessageParser {
     return match ? match[1] : undefined;
   }
 }
+
+const createArtifactElement: ElementFactory = (props) => {
+  const elementProps = [
+    'class="__boltArtifact__"',
+    Object.entries(props).map(([key, value]) => {
+      return `data-${camelToDashCase(key)}=${JSON.stringify(value)}`;
+    }),
+  ];
+
+  return `<div ${elementProps.join(' ')}></div>`;
+};
+
+function camelToDashCase(input: string) {
+  return input.replace(/([a-z])([A-Z])/g, '$1-$2').toLowerCase();
+}
diff --git a/packages/bolt/app/lib/stores/previews.ts b/packages/bolt/app/lib/stores/previews.ts
index f64d1da..79480d1 100644
--- a/packages/bolt/app/lib/stores/previews.ts
+++ b/packages/bolt/app/lib/stores/previews.ts
@@ -25,6 +25,13 @@ export class PreviewsStore {
     webcontainer.on('port', (port, type, url) => {
       let previewInfo = this.#availablePreviews.get(port);
 
+      if (type === 'close' && previewInfo) {
+        this.#availablePreviews.delete(port);
+        this.previews.set(this.previews.get().filter((preview) => preview.port !== port));
+
+        return;
+      }
+
       const previews = this.previews.get();
 
       if (!previewInfo) {
diff --git a/packages/bolt/app/lib/stores/workbench.ts b/packages/bolt/app/lib/stores/workbench.ts
index 4a86539..5cebd73 100644
--- a/packages/bolt/app/lib/stores/workbench.ts
+++ b/packages/bolt/app/lib/stores/workbench.ts
@@ -1,48 +1,24 @@
 import { atom, map, type MapStore, type ReadableAtom, type WritableAtom } from 'nanostores';
 import type { EditorDocument, ScrollPosition } from '../../components/editor/codemirror/CodeMirrorEditor';
-import type { BoltAction } from '../../types/actions';
 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 { EditorStore } from './editor';
 import { FilesStore, type FileMap } from './files';
 import { PreviewsStore } from './previews';
 
-const MIN_SPINNER_TIME = 200;
-
-export type BaseActionState = BoltAction & {
-  status: 'running' | 'complete' | 'pending' | 'aborted';
-  executing: boolean;
-  abort?: () => void;
-};
-
-export type FailedActionState = BoltAction &
-  Omit<BaseActionState, 'status'> & {
-    status: 'failed';
-    error: string;
-  };
-
-export type ActionState = BaseActionState | FailedActionState;
-
-type BaseActionUpdate = Partial<Pick<BaseActionState, 'status' | 'executing' | 'abort'>>;
-
-export type ActionStateUpdate =
-  | BaseActionUpdate
-  | (Omit<BaseActionUpdate, 'status'> & { status: 'failed'; error: string });
-
 export interface ArtifactState {
   title: string;
   closed: boolean;
-  currentActionPromise: Promise<void>;
-  actions: MapStore<Record<string, ActionState>>;
+  runner: ActionRunner;
 }
 
+export type ArtifactUpdateState = Pick<ArtifactState, 'title' | 'closed'>;
+
 type Artifacts = MapStore<Record<string, ArtifactState>>;
 
 export class WorkbenchStore {
-  #actionRunner = new ActionRunner(webcontainer);
   #previewsStore = new PreviewsStore(webcontainer);
   #filesStore = new FilesStore(webcontainer);
   #editorStore = new EditorStore(webcontainer);
@@ -102,148 +78,63 @@ export class WorkbenchStore {
   }
 
   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?.();
-        }
-      }
-    }
+    // TODO: what do we wanna do and how do we wanna recover from this?
   }
 
-  addArtifact({ id, messageId, title }: ArtifactCallbackData) {
-    const artifacts = this.artifacts.get();
-    const artifactKey = getArtifactKey(id, messageId);
-    const artifact = artifacts[artifactKey];
+  addArtifact({ messageId, title }: ArtifactCallbackData) {
+    const artifact = this.#getArtifact(messageId);
 
     if (artifact) {
       return;
     }
 
-    this.artifacts.setKey(artifactKey, {
+    this.artifacts.setKey(messageId, {
       title,
       closed: false,
-      actions: map({}),
-      currentActionPromise: Promise.resolve(),
+      runner: new ActionRunner(webcontainer),
     });
   }
 
-  updateArtifact({ id, messageId }: ArtifactCallbackData, state: Partial<ArtifactState>) {
-    const artifacts = this.artifacts.get();
-    const key = getArtifactKey(id, messageId);
-    const artifact = artifacts[key];
+  updateArtifact({ messageId }: ArtifactCallbackData, state: Partial<ArtifactUpdateState>) {
+    const artifact = this.#getArtifact(messageId);
 
     if (!artifact) {
       return;
     }
 
-    this.artifacts.setKey(key, { ...artifact, ...state });
+    this.artifacts.setKey(messageId, { ...artifact, ...state });
   }
 
   async addAction(data: ActionCallbackData) {
-    const { artifactId, messageId, actionId } = data;
+    const { messageId } = data;
 
-    const artifacts = this.artifacts.get();
-    const key = getArtifactKey(artifactId, messageId);
-    const artifact = artifacts[key];
+    const artifact = this.#getArtifact(messageId);
 
     if (!artifact) {
       unreachable('Artifact not found');
     }
 
-    const actions = artifact.actions.get();
-    const action = actions[actionId];
-
-    if (action) {
-      return;
-    }
-
-    artifact.actions.setKey(actionId, { ...data.action, status: 'pending', executing: false });
-
-    artifact.currentActionPromise.then(() => {
-      if (chatStore.get().aborted) {
-        return;
-      }
-
-      this.#updateAction(key, actionId, { status: 'running' });
-    });
+    artifact.runner.addAction(data);
   }
 
   async runAction(data: ActionCallbackData) {
-    const { artifactId, messageId, actionId } = data;
+    const { messageId } = data;
 
-    const artifacts = this.artifacts.get();
-    const key = getArtifactKey(artifactId, messageId);
-    const artifact = artifacts[key];
+    const artifact = this.#getArtifact(messageId);
 
     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 {
-        await Promise.all([
-          this.#actionRunner.runAction(data, abortController.signal),
-          new Promise((resolve) => setTimeout(resolve, MIN_SPINNER_TIME)),
-        ]);
-
-        if (!abortController.signal.aborted) {
-          this.#updateAction(key, actionId, { status: 'complete' });
-        }
-      } catch (error) {
-        this.#updateAction(key, actionId, { status: 'failed', error: 'Action failed' });
-
-        throw error;
-      } finally {
-        this.#updateAction(key, actionId, { executing: false });
-      }
-    });
+    artifact.runner.runAction(data);
   }
 
-  #updateAction(artifactId: string, actionId: string, newState: ActionStateUpdate) {
+  #getArtifact(id: string) {
     const artifacts = this.artifacts.get();
-    const artifact = artifacts[artifactId];
-
-    if (!artifact) {
-      return;
-    }
-
-    const actions = artifact.actions.get();
-
-    artifact.actions.setKey(actionId, { ...actions[actionId], ...newState });
+    return artifacts[id];
   }
 }
 
-export function getArtifactKey(artifactId: string, messageId: string) {
-  return `${artifactId}_${messageId}`;
-}
-
 export const workbenchStore = new WorkbenchStore();
 
 if (import.meta.hot) {
diff --git a/packages/bolt/app/routes/api.chat.ts b/packages/bolt/app/routes/api.chat.ts
index b75d4e5..e554eaf 100644
--- a/packages/bolt/app/routes/api.chat.ts
+++ b/packages/bolt/app/routes/api.chat.ts
@@ -1,9 +1,9 @@
 import { type ActionFunctionArgs } from '@remix-run/cloudflare';
-import { MAX_RESPONSE_SEGMENTS } from '../lib/.server/llm/constants';
+import { StreamingTextResponse } from 'ai';
+import { MAX_RESPONSE_SEGMENTS, MAX_TOKENS } from '../lib/.server/llm/constants';
 import { CONTINUE_PROMPT } from '../lib/.server/llm/prompts';
 import { streamText, type Messages, type StreamingOptions } from '../lib/.server/llm/stream-text';
 import SwitchableStream from '../lib/.server/llm/switchable-stream';
-import { StreamingTextResponse } from 'ai';
 
 export async function action({ context, request }: ActionFunctionArgs) {
   const { messages } = await request.json<{ messages: Messages }>();
@@ -18,9 +18,13 @@ export async function action({ context, request }: ActionFunctionArgs) {
         }
 
         if (stream.switches >= MAX_RESPONSE_SEGMENTS) {
-          throw Error('Cannot continue message: maximum segments reached');
+          throw Error('Cannot continue message: Maximum segments reached');
         }
 
+        const switchesLeft = MAX_RESPONSE_SEGMENTS - stream.switches;
+
+        console.log(`Reached max token limit (${MAX_TOKENS}): Continuing message (${switchesLeft} switches left)`);
+
         messages.push({ role: 'assistant', content });
         messages.push({ role: 'user', content: CONTINUE_PROMPT });