feature(code-streaming): added code streaming to editor while AI is writing code

app/lib/hooks/useMessageParser.ts

@@ -36,6 +36,10 @@ const messageParser = new StreamingMessageParser({
 
       workbenchStore.runAction(data);
     },
+    onActionStream: (data) => {
+      logger.trace('onActionStream', data.action);
+      workbenchStore.runAction(data, true);
+    },
   },
 });
 

app/lib/runtime/action-runner.ts

@@ -72,7 +72,7 @@ export class ActionRunner {
     });
   }
 
-  async runAction(data: ActionCallbackData) {
+  async runAction(data: ActionCallbackData, isStreaming: boolean = false) {
     const { actionId } = data;
     const action = this.actions.get()[actionId];
 
@@ -83,19 +83,22 @@ export class ActionRunner {
     if (action.executed) {
       return;
     }
+    if (isStreaming && action.type !== 'file') {
+      return;
+    }
 
-    this.#updateAction(actionId, { ...action, ...data.action, executed: true });
+    this.#updateAction(actionId, { ...action, ...data.action, executed: !isStreaming });
 
     this.#currentExecutionPromise = this.#currentExecutionPromise
       .then(() => {
-        return this.#executeAction(actionId);
+        return this.#executeAction(actionId, isStreaming);
       })
       .catch((error) => {
         console.error('Action failed:', error);
       });
   }
 
-  async #executeAction(actionId: string) {
+  async #executeAction(actionId: string, isStreaming: boolean = false) {
     const action = this.actions.get()[actionId];
 
     this.#updateAction(actionId, { status: 'running' });
@@ -112,7 +115,7 @@ export class ActionRunner {
         }
       }
 
-      this.#updateAction(actionId, { status: action.abortSignal.aborted ? 'aborted' : 'complete' });
+      this.#updateAction(actionId, { status: isStreaming ? 'running' : action.abortSignal.aborted ? 'aborted' : 'complete' });
     } catch (error) {
       this.#updateAction(actionId, { status: 'failed', error: 'Action failed' });
 

app/lib/runtime/message-parser.ts

@@ -28,6 +28,7 @@ export interface ParserCallbacks {
   onArtifactOpen?: ArtifactCallback;
   onArtifactClose?: ArtifactCallback;
   onActionOpen?: ActionCallback;
+  onActionStream?: ActionCallback;
   onActionClose?: ActionCallback;
 }
 
@@ -54,7 +55,7 @@ interface MessageState {
 export class StreamingMessageParser {
   #messages = new Map<string, MessageState>();
 
-  constructor(private _options: StreamingMessageParserOptions = {}) {}
+  constructor(private _options: StreamingMessageParserOptions = {}) { }
 
   parse(messageId: string, input: string) {
     let state = this.#messages.get(messageId);
@@ -118,6 +119,21 @@ export class StreamingMessageParser {
 
             i = closeIndex + ARTIFACT_ACTION_TAG_CLOSE.length;
           } else {
+            if ('type' in currentAction && currentAction.type === 'file') {
+              let content = input.slice(i);
+
+              this._options.callbacks?.onActionStream?.({
+                artifactId: currentArtifact.id,
+                messageId,
+                actionId: String(state.actionId - 1),
+                action: {
+                  ...currentAction as FileAction,
+                  content,
+                  filePath: currentAction.filePath,
+                },
+
+              });
+            }
             break;
           }
         } else {

app/lib/stores/workbench.ts

@@ -12,6 +12,7 @@ import { TerminalStore } from './terminal';
 import JSZip from 'jszip';
 import { saveAs } from 'file-saver';
 import { Octokit } from "@octokit/rest";
+import * as nodePath from 'node:path';
 
 export interface ArtifactState {
   id: string;
@@ -258,7 +259,7 @@ export class WorkbenchStore {
     artifact.runner.addAction(data);
   }
 
-  async runAction(data: ActionCallbackData) {
+  async runAction(data: ActionCallbackData, isStreaming: boolean = false) {
     const { messageId } = data;
 
     const artifact = this.#getArtifact(messageId);
@@ -266,8 +267,29 @@ export class WorkbenchStore {
     if (!artifact) {
       unreachable('Artifact not found');
     }
+    if (data.action.type === 'file') {
+      let wc = await webcontainer
+      const fullPath = nodePath.join(wc.workdir, data.action.filePath);
+      if (this.selectedFile.value !== fullPath) {
+        this.setSelectedFile(fullPath);
+      }
+      if (this.currentView.value !== 'code') {
+        this.currentView.set('code');
+      }
+      const doc = this.#editorStore.documents.get()[fullPath];
+      if (!doc) {
+        await artifact.runner.runAction(data, isStreaming);
+      }
 
-    artifact.runner.runAction(data);
+      this.#editorStore.updateFile(fullPath, data.action.content);
+
+      if (!isStreaming) {
+        this.resetCurrentDocument();
+        await artifact.runner.runAction(data);
+      }
+    } else {
+      artifact.runner.runAction(data);
+    }
   }
 
   #getArtifact(id: string) {