From 993ff88e2be6b88d2343439a5e880c052ae57507 Mon Sep 17 00:00:00 2001
From: Abi Raja <abimanyuraja@gmail.com>
Date: Thu, 25 Jul 2024 14:35:04 -0400
Subject: [PATCH] abstract into more components

---
 frontend/src/App.tsx                          | 273 ++----------------
 frontend/src/components/preview/Preview.tsx   | 114 ++++++++
 .../PreviewComponent.tsx}                     |   8 +-
 frontend/src/components/sidebar/Sidebar.tsx   | 174 +++++++++++
 frontend/src/store/app-store.ts               |  21 ++
 frontend/src/store/project-store.ts           |  13 +
 6 files changed, 346 insertions(+), 257 deletions(-)
 create mode 100644 frontend/src/components/preview/Preview.tsx
 rename frontend/src/components/{Preview.tsx => preview/PreviewComponent.tsx} (87%)
 create mode 100644 frontend/src/components/sidebar/Sidebar.tsx

diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index ae58b89..bbd55f5 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -1,21 +1,6 @@
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useRef } from "react";
 import ImageUpload from "./components/ImageUpload";
-import CodePreview from "./components/CodePreview";
-import Preview from "./components/Preview";
 import { generateCode } from "./generateCode";
-import Spinner from "./components/Spinner";
-import classNames from "classnames";
-import {
-  FaCode,
-  FaDesktop,
-  FaDownload,
-  FaMobile,
-  FaUndo,
-} from "react-icons/fa";
-import { Switch } from "./components/ui/switch";
-import { Button } from "@/components/ui/button";
-import { Textarea } from "@/components/ui/textarea";
-import { Tabs, TabsContent, TabsList, TabsTrigger } from "./components/ui/tabs";
 import SettingsDialog from "./components/SettingsDialog";
 import { AppState, CodeGenerationParams, EditorTheme, Settings } from "./types";
 import { IS_RUNNING_ON_CLOUD } from "./config";
@@ -26,7 +11,6 @@ import { UrlInputSection } from "./components/UrlInputSection";
 import TermsOfServiceDialog from "./components/TermsOfServiceDialog";
 import html2canvas from "html2canvas";
 import { USER_CLOSE_WEB_SOCKET_CODE } from "./constants";
-import CodeTab from "./components/CodeTab";
 import OutputSettingsSection from "./components/OutputSettingsSection";
 import { History } from "./components/history/history_types";
 import HistoryDisplay from "./components/history/HistoryDisplay";
@@ -36,17 +20,14 @@ import ImportCodeSection from "./components/ImportCodeSection";
 import { Stack } from "./lib/stacks";
 import { CodeGenerationModel } from "./lib/models";
 import ModelSettingsSection from "./components/ModelSettingsSection";
-import { extractHtml } from "./components/preview/extractHtml";
 import useBrowserTabIndicator from "./hooks/useBrowserTabIndicator";
 import TipLink from "./components/core/TipLink";
-import SelectAndEditModeToggleButton from "./components/select-and-edit/SelectAndEditModeToggleButton";
 import { useAppStore } from "./store/app-store";
-import KeyboardShortcutBadge from "./components/core/KeyboardShortcutBadge";
 import { useProjectStore } from "./store/project-store";
+import Sidebar from "./components/sidebar/Sidebar";
+import Preview from "./components/preview/Preview";
 
 function App() {
-  const [appState, setAppState] = useState<AppState>(AppState.INITIAL);
-
   const {
     // Inputs
     inputMode,
@@ -57,20 +38,22 @@ function App() {
     setReferenceImages,
 
     // Outputs
-    generatedCode,
     setGeneratedCode,
+    setExecutionConsole,
     currentVersion,
     setCurrentVersion,
     appHistory,
     setAppHistory,
   } = useProjectStore();
 
-  const [executionConsole, setExecutionConsole] = useState<string[]>([]);
-  const [updateInstruction, setUpdateInstruction] = useState("");
-
-  const textareaRef = useRef<HTMLTextAreaElement>(null);
-
-  const { disableInSelectAndEditMode } = useAppStore();
+  const {
+    disableInSelectAndEditMode,
+    setUpdateInstruction,
+    appState,
+    setAppState,
+    shouldIncludeResultImage,
+    setShouldIncludeResultImage,
+  } = useAppStore();
 
   // Settings
   const [settings, setSettings] = usePersistedState<Settings>(
@@ -93,9 +76,6 @@ function App() {
   const selectedCodeGenerationModel =
     settings.codeGenerationModel || CodeGenerationModel.GPT_4_VISION;
 
-  const [shouldIncludeResultImage, setShouldIncludeResultImage] =
-    useState<boolean>(false);
-
   const wsRef = useRef<WebSocket>(null);
 
   const showBetterModelMessage =
@@ -139,23 +119,6 @@ function App() {
     return png;
   };
 
-  const downloadCode = () => {
-    // Create a blob from the generated code
-    const blob = new Blob([generatedCode], { type: "text/html" });
-    const url = URL.createObjectURL(blob);
-
-    // Create an anchor element and set properties for download
-    const a = document.createElement("a");
-    a.href = url;
-    a.download = "index.html"; // Set the file name for download
-    document.body.appendChild(a); // Append to the document
-    a.click(); // Programmatically click the anchor to trigger download
-
-    // Clean up by removing the anchor and revoking the Blob URL
-    document.body.removeChild(a);
-    URL.revokeObjectURL(url);
-  };
-
   const reset = () => {
     setAppState(AppState.INITIAL);
     setGeneratedCode("");
@@ -194,11 +157,6 @@ function App() {
     cancelCodeGenerationAndReset();
   };
 
-  const previewCode =
-    inputMode === "video" && appState === AppState.CODING
-      ? extractHtml(generatedCode)
-      : generatedCode;
-
   const cancelCodeGenerationAndReset = () => {
     // When this is the first version, reset the entire app state
     if (currentVersion === null) {
@@ -257,7 +215,7 @@ function App() {
                 inputs: {
                   prompt: params.history
                     ? params.history[params.history.length - 1]
-                    : updateInstruction,
+                    : "",
                 },
               },
             ];
@@ -407,13 +365,6 @@ function App() {
     setAppState(AppState.CODE_READY);
   }
 
-  // When coding is complete, focus on the update instruction textarea
-  useEffect(() => {
-    if (appState === AppState.CODE_READY && textareaRef.current) {
-      textareaRef.current.focus();
-    }
-  }, [appState]);
-
   return (
     <div className="mt-2 dark:bg-black dark:text-white">
       {IS_RUNNING_ON_CLOUD && <PicoBadge />}
@@ -461,135 +412,12 @@ function App() {
 
           {(appState === AppState.CODING ||
             appState === AppState.CODE_READY) && (
-            <>
-              {/* Show code preview only when coding */}
-              {appState === AppState.CODING && (
-                <div className="flex flex-col">
-                  {/* Speed disclaimer for video mode */}
-                  {inputMode === "video" && (
-                    <div
-                      className="bg-yellow-100 border-l-4 border-yellow-500 text-yellow-700
-                    p-2 text-xs mb-4 mt-1"
-                    >
-                      Code generation from videos can take 3-4 minutes. We do
-                      multiple passes to get the best result. Please be patient.
-                    </div>
-                  )}
-
-                  <div className="flex items-center gap-x-1">
-                    <Spinner />
-                    {executionConsole.slice(-1)[0]}
-                  </div>
-
-                  <CodePreview code={generatedCode} />
-
-                  <div className="flex w-full">
-                    <Button
-                      onClick={cancelCodeGeneration}
-                      className="w-full dark:text-white dark:bg-gray-700"
-                    >
-                      Cancel
-                    </Button>
-                  </div>
-                </div>
-              )}
-
-              {appState === AppState.CODE_READY && (
-                <div>
-                  <div className="grid w-full gap-2">
-                    <Textarea
-                      ref={textareaRef}
-                      placeholder="Tell the AI what to change..."
-                      onChange={(e) => setUpdateInstruction(e.target.value)}
-                      onKeyDown={(e) => {
-                        if (e.key === "Enter" && !e.shiftKey) {
-                          doUpdate(updateInstruction);
-                        }
-                      }}
-                      value={updateInstruction}
-                    />
-                    <div className="flex justify-between items-center gap-x-2">
-                      <div className="font-500 text-xs text-slate-700 dark:text-white">
-                        Include screenshot of current version?
-                      </div>
-                      <Switch
-                        checked={shouldIncludeResultImage}
-                        onCheckedChange={setShouldIncludeResultImage}
-                        className="dark:bg-gray-700"
-                      />
-                    </div>
-                    <Button
-                      onClick={() => doUpdate(updateInstruction)}
-                      className="dark:text-white dark:bg-gray-700 update-btn"
-                    >
-                      Update <KeyboardShortcutBadge letter="enter" />
-                    </Button>
-                  </div>
-                  <div className="flex items-center justify-end gap-x-2 mt-2">
-                    <Button
-                      onClick={regenerate}
-                      className="flex items-center gap-x-2 dark:text-white dark:bg-gray-700 regenerate-btn"
-                    >
-                      🔄 Regenerate
-                    </Button>
-                    {showSelectAndEditFeature && (
-                      <SelectAndEditModeToggleButton />
-                    )}
-                  </div>
-                  <div className="flex justify-end items-center mt-2">
-                    <TipLink />
-                  </div>
-                </div>
-              )}
-
-              {/* Reference image display */}
-              <div className="flex gap-x-2 mt-2">
-                {referenceImages.length > 0 && (
-                  <div className="flex flex-col">
-                    <div
-                      className={classNames({
-                        "scanning relative": appState === AppState.CODING,
-                      })}
-                    >
-                      {inputMode === "image" && (
-                        <img
-                          className="w-[340px] border border-gray-200 rounded-md"
-                          src={referenceImages[0]}
-                          alt="Reference"
-                        />
-                      )}
-                      {inputMode === "video" && (
-                        <video
-                          muted
-                          autoPlay
-                          loop
-                          className="w-[340px] border border-gray-200 rounded-md"
-                          src={referenceImages[0]}
-                        />
-                      )}
-                    </div>
-                    <div className="text-gray-400 uppercase text-sm text-center mt-1">
-                      {inputMode === "video"
-                        ? "Original Video"
-                        : "Original Screenshot"}
-                    </div>
-                  </div>
-                )}
-                <div className="bg-gray-400 px-4 py-2 rounded text-sm">
-                  <h2 className="text-lg mb-4 border-b border-gray-800">
-                    Console
-                  </h2>
-                  {executionConsole.map((line, index) => (
-                    <div
-                      key={index}
-                      className="border-b border-gray-400 mb-2 text-gray-600 font-mono"
-                    >
-                      {line}
-                    </div>
-                  ))}
-                </div>
-              </div>
-            </>
+            <Sidebar
+              showSelectAndEditFeature={showSelectAndEditFeature}
+              doUpdate={doUpdate}
+              regenerate={regenerate}
+              cancelCodeGeneration={cancelCodeGeneration}
+            />
           )}
           {
             <HistoryDisplay
@@ -612,68 +440,7 @@ function App() {
         )}
 
         {(appState === AppState.CODING || appState === AppState.CODE_READY) && (
-          // Right side preview and code
-          <div className="ml-4">
-            <Tabs defaultValue="desktop">
-              <div className="flex justify-between mr-8 mb-4">
-                <div className="flex items-center gap-x-2">
-                  {appState === AppState.CODE_READY && (
-                    <>
-                      <Button
-                        onClick={reset}
-                        className="flex items-center ml-4 gap-x-2 dark:text-white dark:bg-gray-700"
-                      >
-                        <FaUndo />
-                        Reset
-                      </Button>
-                      <Button
-                        onClick={downloadCode}
-                        variant="secondary"
-                        className="flex items-center gap-x-2 mr-4 dark:text-white dark:bg-gray-700 download-btn"
-                      >
-                        <FaDownload /> Download
-                      </Button>
-                    </>
-                  )}
-                </div>
-                <div className="flex items-center">
-                  <TabsList>
-                    <TabsTrigger value="desktop" className="flex gap-x-2">
-                      <FaDesktop /> Desktop
-                    </TabsTrigger>
-                    <TabsTrigger value="mobile" className="flex gap-x-2">
-                      <FaMobile /> Mobile
-                    </TabsTrigger>
-                    <TabsTrigger value="code" className="flex gap-x-2">
-                      <FaCode />
-                      Code
-                    </TabsTrigger>
-                  </TabsList>
-                </div>
-              </div>
-              <TabsContent value="desktop">
-                <Preview
-                  code={previewCode}
-                  device="desktop"
-                  doUpdate={doUpdate}
-                />
-              </TabsContent>
-              <TabsContent value="mobile">
-                <Preview
-                  code={previewCode}
-                  device="mobile"
-                  doUpdate={doUpdate}
-                />
-              </TabsContent>
-              <TabsContent value="code">
-                <CodeTab
-                  code={previewCode}
-                  setCode={setGeneratedCode}
-                  settings={settings}
-                />
-              </TabsContent>
-            </Tabs>
-          </div>
+          <Preview doUpdate={doUpdate} reset={reset} settings={settings} />
         )}
       </main>
     </div>
diff --git a/frontend/src/components/preview/Preview.tsx b/frontend/src/components/preview/Preview.tsx
new file mode 100644
index 0000000..8f145f2
--- /dev/null
+++ b/frontend/src/components/preview/Preview.tsx
@@ -0,0 +1,114 @@
+import { Tabs, TabsList, TabsTrigger, TabsContent } from "../ui/tabs";
+import {
+  FaUndo,
+  FaDownload,
+  FaDesktop,
+  FaMobile,
+  FaCode,
+} from "react-icons/fa";
+import { AppState, Settings } from "../../types";
+import CodeTab from "../CodeTab";
+import { Button } from "../ui/button";
+import { useAppStore } from "../../store/app-store";
+import { useProjectStore } from "../../store/project-store";
+import { extractHtml } from "./extractHtml";
+import PreviewComponent from "./PreviewComponent";
+
+interface Props {
+  doUpdate: (instruction: string) => void;
+  reset: () => void;
+  settings: Settings;
+}
+
+function Preview({ doUpdate, reset, settings }: Props) {
+  const downloadCode = (code: string) => {
+    // Create a blob from the generated code
+    const blob = new Blob([code], { type: "text/html" });
+    const url = URL.createObjectURL(blob);
+
+    // Create an anchor element and set properties for download
+    const a = document.createElement("a");
+    a.href = url;
+    a.download = "index.html"; // Set the file name for download
+    document.body.appendChild(a); // Append to the document
+    a.click(); // Programmatically click the anchor to trigger download
+
+    // Clean up by removing the anchor and revoking the Blob URL
+    document.body.removeChild(a);
+    URL.revokeObjectURL(url);
+  };
+
+  const { appState } = useAppStore();
+  const { inputMode, generatedCode, setGeneratedCode } = useProjectStore();
+
+  const previewCode =
+    inputMode === "video" && appState === AppState.CODING
+      ? extractHtml(generatedCode)
+      : generatedCode;
+
+  return (
+    <div className="ml-4">
+      <Tabs defaultValue="desktop">
+        <div className="flex justify-between mr-8 mb-4">
+          <div className="flex items-center gap-x-2">
+            {appState === AppState.CODE_READY && (
+              <>
+                <Button
+                  onClick={reset}
+                  className="flex items-center ml-4 gap-x-2 dark:text-white dark:bg-gray-700"
+                >
+                  <FaUndo />
+                  Reset
+                </Button>
+                <Button
+                  onClick={() => downloadCode(generatedCode)}
+                  variant="secondary"
+                  className="flex items-center gap-x-2 mr-4 dark:text-white dark:bg-gray-700 download-btn"
+                >
+                  <FaDownload /> Download
+                </Button>
+              </>
+            )}
+          </div>
+          <div className="flex items-center">
+            <TabsList>
+              <TabsTrigger value="desktop" className="flex gap-x-2">
+                <FaDesktop /> Desktop
+              </TabsTrigger>
+              <TabsTrigger value="mobile" className="flex gap-x-2">
+                <FaMobile /> Mobile
+              </TabsTrigger>
+              <TabsTrigger value="code" className="flex gap-x-2">
+                <FaCode />
+                Code
+              </TabsTrigger>
+            </TabsList>
+          </div>
+        </div>
+        <TabsContent value="desktop">
+          <PreviewComponent
+            code={previewCode}
+            device="desktop"
+            doUpdate={doUpdate}
+          />
+        </TabsContent>
+        <TabsContent value="mobile">
+          <PreviewComponent
+            code={previewCode}
+            device="mobile"
+            doUpdate={doUpdate}
+          />
+        </TabsContent>
+        <TabsContent value="code">
+          <CodeTab
+            code={previewCode}
+            setCode={setGeneratedCode}
+            settings={settings}
+          />
+        </TabsContent>
+      </Tabs>
+    </div>
+  );
+}
+
+export default Preview;
diff --git a/frontend/src/components/Preview.tsx b/frontend/src/components/preview/PreviewComponent.tsx
similarity index 87%
rename from frontend/src/components/Preview.tsx
rename to frontend/src/components/preview/PreviewComponent.tsx
index d601f78..1c9e5ba 100644
--- a/frontend/src/components/Preview.tsx
+++ b/frontend/src/components/preview/PreviewComponent.tsx
@@ -1,7 +1,7 @@
 import { useEffect, useRef, useState } from "react";
 import classNames from "classnames";
-import useThrottle from "../hooks/useThrottle";
-import EditPopup from "./select-and-edit/EditPopup";
+import useThrottle from "../../hooks/useThrottle";
+import EditPopup from "../select-and-edit/EditPopup";
 
 interface Props {
   code: string;
@@ -9,7 +9,7 @@ interface Props {
   doUpdate: (updateInstruction: string, selectedElement?: HTMLElement) => void;
 }
 
-function Preview({ code, device, doUpdate }: Props) {
+function PreviewComponent({ code, device, doUpdate }: Props) {
   const iframeRef = useRef<HTMLIFrameElement | null>(null);
 
   // Don't update code more often than every 200ms.
@@ -53,4 +53,4 @@ function Preview({ code, device, doUpdate }: Props) {
   );
 }
 
-export default Preview;
+export default PreviewComponent;
diff --git a/frontend/src/components/sidebar/Sidebar.tsx b/frontend/src/components/sidebar/Sidebar.tsx
new file mode 100644
index 0000000..f8ba3c4
--- /dev/null
+++ b/frontend/src/components/sidebar/Sidebar.tsx
@@ -0,0 +1,174 @@
+import { Switch } from "@radix-ui/react-switch";
+import classNames from "classnames";
+import { useAppStore } from "../../store/app-store";
+import { useProjectStore } from "../../store/project-store";
+import { AppState } from "../../types";
+import CodePreview from "../CodePreview";
+import Spinner from "../Spinner";
+import KeyboardShortcutBadge from "../core/KeyboardShortcutBadge";
+import TipLink from "../core/TipLink";
+import SelectAndEditModeToggleButton from "../select-and-edit/SelectAndEditModeToggleButton";
+import { Button } from "../ui/button";
+import { Textarea } from "../ui/textarea";
+import { useEffect, useRef } from "react";
+
+interface SidebarProps {
+  showSelectAndEditFeature: boolean;
+  doUpdate: (instruction: string) => void;
+  regenerate: () => void;
+  cancelCodeGeneration: () => void;
+}
+
+function Sidebar({
+  showSelectAndEditFeature,
+  doUpdate,
+  regenerate,
+  cancelCodeGeneration,
+}: SidebarProps) {
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+
+  const {
+    appState,
+    updateInstruction,
+    setUpdateInstruction,
+    shouldIncludeResultImage,
+    setShouldIncludeResultImage,
+  } = useAppStore();
+  const { inputMode, generatedCode, referenceImages, executionConsole } =
+    useProjectStore();
+
+  // When coding is complete, focus on the update instruction textarea
+  useEffect(() => {
+    if (appState === AppState.CODE_READY && textareaRef.current) {
+      textareaRef.current.focus();
+    }
+  }, [appState]);
+
+  return (
+    <>
+      {/* Show code preview only when coding */}
+      {appState === AppState.CODING && (
+        <div className="flex flex-col">
+          {/* Speed disclaimer for video mode */}
+          {inputMode === "video" && (
+            <div
+              className="bg-yellow-100 border-l-4 border-yellow-500 text-yellow-700
+            p-2 text-xs mb-4 mt-1"
+            >
+              Code generation from videos can take 3-4 minutes. We do multiple
+              passes to get the best result. Please be patient.
+            </div>
+          )}
+
+          <div className="flex items-center gap-x-1">
+            <Spinner />
+            {executionConsole.slice(-1)[0]}
+          </div>
+
+          <CodePreview code={generatedCode} />
+
+          <div className="flex w-full">
+            <Button
+              onClick={cancelCodeGeneration}
+              className="w-full dark:text-white dark:bg-gray-700"
+            >
+              Cancel
+            </Button>
+          </div>
+        </div>
+      )}
+
+      {appState === AppState.CODE_READY && (
+        <div>
+          <div className="grid w-full gap-2">
+            <Textarea
+              ref={textareaRef}
+              placeholder="Tell the AI what to change..."
+              onChange={(e) => setUpdateInstruction(e.target.value)}
+              onKeyDown={(e) => {
+                if (e.key === "Enter" && !e.shiftKey) {
+                  doUpdate(updateInstruction);
+                }
+              }}
+              value={updateInstruction}
+            />
+            <div className="flex justify-between items-center gap-x-2">
+              <div className="font-500 text-xs text-slate-700 dark:text-white">
+                Include screenshot of current version?
+              </div>
+              <Switch
+                checked={shouldIncludeResultImage}
+                onCheckedChange={setShouldIncludeResultImage}
+                className="dark:bg-gray-700"
+              />
+            </div>
+            <Button
+              onClick={() => doUpdate(updateInstruction)}
+              className="dark:text-white dark:bg-gray-700 update-btn"
+            >
+              Update <KeyboardShortcutBadge letter="enter" />
+            </Button>
+          </div>
+          <div className="flex items-center justify-end gap-x-2 mt-2">
+            <Button
+              onClick={regenerate}
+              className="flex items-center gap-x-2 dark:text-white dark:bg-gray-700 regenerate-btn"
+            >
+              🔄 Regenerate
+            </Button>
+            {showSelectAndEditFeature && <SelectAndEditModeToggleButton />}
+          </div>
+          <div className="flex justify-end items-center mt-2">
+            <TipLink />
+          </div>
+        </div>
+      )}
+
+      {/* Reference image display */}
+      <div className="flex gap-x-2 mt-2">
+        {referenceImages.length > 0 && (
+          <div className="flex flex-col">
+            <div
+              className={classNames({
+                "scanning relative": appState === AppState.CODING,
+              })}
+            >
+              {inputMode === "image" && (
+                <img
+                  className="w-[340px] border border-gray-200 rounded-md"
+                  src={referenceImages[0]}
+                  alt="Reference"
+                />
+              )}
+              {inputMode === "video" && (
+                <video
+                  muted
+                  autoPlay
+                  loop
+                  className="w-[340px] border border-gray-200 rounded-md"
+                  src={referenceImages[0]}
+                />
+              )}
+            </div>
+            <div className="text-gray-400 uppercase text-sm text-center mt-1">
+              {inputMode === "video" ? "Original Video" : "Original Screenshot"}
+            </div>
+          </div>
+        )}
+        <div className="bg-gray-400 px-4 py-2 rounded text-sm">
+          <h2 className="text-lg mb-4 border-b border-gray-800">Console</h2>
+          {executionConsole.map((line, index) => (
+            <div
+              key={index}
+              className="border-b border-gray-400 mb-2 text-gray-600 font-mono"
+            >
+              {line}
+            </div>
+          ))}
+        </div>
+      </div>
+    </>
+  );
+}
+
+export default Sidebar;
diff --git a/frontend/src/store/app-store.ts b/frontend/src/store/app-store.ts
index 834c47a..cf2048a 100644
--- a/frontend/src/store/app-store.ts
+++ b/frontend/src/store/app-store.ts
@@ -1,13 +1,34 @@
 import { create } from "zustand";
+import { AppState } from "../types";
 
 // Store for app-wide state
 interface AppStore {
+  appState: AppState;
+  setAppState: (state: AppState) => void;
+
+  // UI state
+  updateInstruction: string;
+  setUpdateInstruction: (instruction: string) => void;
+  shouldIncludeResultImage: boolean;
+  setShouldIncludeResultImage: (shouldInclude: boolean) => void;
+
   inSelectAndEditMode: boolean;
   toggleInSelectAndEditMode: () => void;
   disableInSelectAndEditMode: () => void;
 }
 
 export const useAppStore = create<AppStore>((set) => ({
+  appState: AppState.INITIAL,
+  setAppState: (state: AppState) => set({ appState: state }),
+
+  // UI state
+  updateInstruction: "",
+  setUpdateInstruction: (instruction: string) =>
+    set({ updateInstruction: instruction }),
+  shouldIncludeResultImage: true,
+  setShouldIncludeResultImage: (shouldInclude: boolean) =>
+    set({ shouldIncludeResultImage: shouldInclude }),
+
   inSelectAndEditMode: false,
   toggleInSelectAndEditMode: () =>
     set((state) => ({ inSelectAndEditMode: !state.inSelectAndEditMode })),
diff --git a/frontend/src/store/project-store.ts b/frontend/src/store/project-store.ts
index c583cb2..8a416f0 100644
--- a/frontend/src/store/project-store.ts
+++ b/frontend/src/store/project-store.ts
@@ -16,6 +16,10 @@ interface ProjectStore {
   setGeneratedCode: (
     updater: string | ((currentCode: string) => string)
   ) => void;
+  executionConsole: string[];
+  setExecutionConsole: (
+    updater: string[] | ((currentConsole: string[]) => string[])
+  ) => void;
 
   // Tracks the currently shown version from app history
   // TODO: might want to move to appStore
@@ -44,6 +48,15 @@ export const useProjectStore = create<ProjectStore>((set) => ({
       generatedCode:
         typeof updater === "function" ? updater(state.generatedCode) : updater,
     })),
+  executionConsole: [],
+  setExecutionConsole: (updater) =>
+    set((state) => ({
+      executionConsole:
+        typeof updater === "function"
+          ? updater(state.executionConsole)
+          : updater,
+    })),
+
   currentVersion: null,
   setCurrentVersion: (version) => set({ currentVersion: version }),
   appHistory: [],