Merge pull request #213 from thecodacus/code-streaming

feat(code-streaming): added code streaming to editor while AI is writing files
stackblitz-labs · Nov 12, 2024 · a081f8b · a081f8b
2 parents f1c5fbf + 8c6e420
commit a081f8b
Show file tree

Hide file tree

Showing 6 changed files with 76 additions and 26 deletions.
diff --git a/app/components/chat/BaseChat.tsx b/app/components/chat/BaseChat.tsx
@@ -24,16 +24,16 @@ const EXAMPLE_PROMPTS = [
   { text: 'How do I center a div?' },
 ];
 
-const providerList = [...new Set(MODEL_LIST.map((model) => model.provider))]
+const providerList = [...new Set(MODEL_LIST.map((model) => model.provider))];
 
 const ModelSelector = ({ model, setModel, provider, setProvider, modelList, providerList }) => {
   return (
     <div className="mb-2 flex gap-2">
-      <select 
+      <select
         value={provider}
         onChange={(e) => {
           setProvider(e.target.value);
-          const firstModel = [...modelList].find(m => m.provider == e.target.value);
+          const firstModel = [...modelList].find((m) => m.provider == e.target.value);
           setModel(firstModel ? firstModel.name : '');
         }}
         className="flex-1 p-2 rounded-lg border border-bolt-elements-borderColor bg-bolt-elements-prompt-background text-bolt-elements-textPrimary focus:outline-none focus:ring-2 focus:ring-bolt-elements-focus transition-all"
@@ -58,11 +58,13 @@ const ModelSelector = ({ model, setModel, provider, setProvider, modelList, prov
         onChange={(e) => setModel(e.target.value)}
         className="flex-1 p-2 rounded-lg border border-bolt-elements-borderColor bg-bolt-elements-prompt-background text-bolt-elements-textPrimary focus:outline-none focus:ring-2 focus:ring-bolt-elements-focus transition-all"
       >
-        {[...modelList].filter(e => e.provider == provider && e.name).map((modelOption) => (
-          <option key={modelOption.name} value={modelOption.name}>
-            {modelOption.label}
-          </option>
-        ))}
+        {[...modelList]
+          .filter((e) => e.provider == provider && e.name)
+          .map((modelOption) => (
+            <option key={modelOption.name} value={modelOption.name}>
+              {modelOption.label}
+            </option>
+          ))}
       </select>
     </div>
   );
@@ -81,10 +83,10 @@ interface BaseChatProps {
   enhancingPrompt?: boolean;
   promptEnhanced?: boolean;
   input?: string;
-  model: string;
-  setModel: (model: string) => void;
-  provider: string;
-  setProvider: (provider: string) => void;
+  model?: string;
+  setModel?: (model: string) => void;
+  provider?: string;
+  setProvider?: (provider: string) => void;
   handleStop?: () => void;
   sendMessage?: (event: React.UIEvent, messageInput?: string) => void;
   handleInputChange?: (event: React.ChangeEvent<HTMLTextAreaElement>) => void;
@@ -144,7 +146,7 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
           expires: 30, // 30 days
           secure: true, // Only send over HTTPS
           sameSite: 'strict', // Protect against CSRF
-          path: '/' // Accessible across the site
+          path: '/', // Accessible across the site
         });
       } catch (error) {
         console.error('Error saving API keys to cookies:', error);
@@ -281,7 +283,9 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                     </div>
                     {input.length > 3 ? (
                       <div className="text-xs text-bolt-elements-textTertiary">
-                        Use <kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Shift</kbd> + <kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Return</kbd> for a new line
+                        Use <kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Shift</kbd> +{' '}
+                        <kbd className="kdb px-1.5 py-0.5 rounded bg-bolt-elements-background-depth-2">Return</kbd> for
+                        a new line
                       </div>
                     ) : null}
                   </div>
@@ -315,4 +319,4 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
       </div>
     );
   },
-);
+);
diff --git a/app/lib/hooks/useMessageParser.ts b/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);
+    },
   },
 });
 

diff --git a/app/lib/runtime/action-runner.ts b/app/lib/runtime/action-runner.ts
@@ -77,7 +77,7 @@ export class ActionRunner {
     });
   }
 
-  async runAction(data: ActionCallbackData) {
+  async runAction(data: ActionCallbackData, isStreaming: boolean = false) {
     const { actionId } = data;
     const action = this.actions.get()[actionId];
 
@@ -88,19 +88,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' });
@@ -121,7 +124,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' });
       logger.error(`[${action.type}]:Action failed\n\n`, error);

diff --git a/app/lib/runtime/message-parser.ts b/app/lib/runtime/message-parser.ts
@@ -28,6 +28,7 @@ export interface ParserCallbacks {
   onArtifactOpen?: ArtifactCallback;
   onArtifactClose?: ArtifactCallback;
   onActionOpen?: ActionCallback;
+  onActionStream?: ActionCallback;
   onActionClose?: ActionCallback;
 }
 
@@ -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 {

diff --git a/app/lib/stores/workbench.ts b/app/lib/stores/workbench.ts
@@ -11,7 +11,8 @@ import { PreviewsStore } from './previews';
 import { TerminalStore } from './terminal';
 import JSZip from 'jszip';
 import { saveAs } from 'file-saver';
-import { Octokit } from "@octokit/rest";
+import { Octokit, type RestEndpointMethodTypes } from "@octokit/rest";
+import * as nodePath from 'node:path';
 import type { WebContainerProcess } from '@webcontainer/api';
 
 export interface ArtifactState {
@@ -267,16 +268,37 @@ 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);
 
     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) {
@@ -360,9 +382,10 @@ export class WorkbenchStore {
       const octokit = new Octokit({ auth: githubToken });
 
       // Check if the repository already exists before creating it
-      let repo
+      let repo: RestEndpointMethodTypes["repos"]["get"]["response"]['data']
       try {
-        repo = await octokit.repos.get({ owner: owner, repo: repoName });
+        let resp = await octokit.repos.get({ owner: owner, repo: repoName });
+        repo = resp.data
       } catch (error) {
         if (error instanceof Error && 'status' in error && error.status === 404) {
           // Repository doesn't exist, so create a new one

diff --git a/package.json b/package.json
@@ -117,5 +117,5 @@
   "resolutions": {
     "@typescript-eslint/utils": "^8.0.0-alpha.30"
   },
-  "packageManager": "pnpm@9.12.2+sha512.22721b3a11f81661ae1ec68ce1a7b879425a1ca5b991c975b074ac220b187ce56c708fe5db69f4c962c989452eee76c82877f4ee80f474cebd61ee13461b6228"
+  "packageManager": "pnpm@9.4.0"
 }