better UX for different app states

backend/main.py

@@ -46,3 +46,5 @@ async def stream_code_test(websocket: WebSocket):
 
     # Write the messages dict into a log so that we can debug later
     write_logs(prompt_messages, completion)
+
+    await websocket.close()

frontend/src/App.tsx

@@ -5,31 +5,44 @@ import Preview from "./components/Preview";
 import { generateCode } from "./generateCode";
 
 function App() {
+  const [appState, setAppState] = useState<"INITIAL" | "CODING" | "CODE_READY">(
+    "INITIAL"
+  );
   const [generatedCode, setGeneratedCode] = useState<string>("");
   const [referenceImages, setReferenceImages] = useState<string[]>([]);
 
   function startCodeGeneration(referenceImages: string[]) {
+    setAppState("CODING");
     setReferenceImages(referenceImages);
-    generateCode(referenceImages[0], function (token) {
-      setGeneratedCode((prev) => prev + token);
-    });
+    generateCode(
+      referenceImages[0],
+      function (token) {
+        setGeneratedCode((prev) => prev + token);
+      },
+      function () {
+        setAppState("CODE_READY");
+      }
+    );
   }
 
   return (
     <div className="mx-auto mt-10 max-w-[1000px]">
       <h1 className="text-2xl mb-4">Drag & Drop a Screenshot</h1>
-      {referenceImages.length > 0 && (
-        <img className="w-[300px]" src={referenceImages[0]} alt="Reference" />
-      )}
 
-      {referenceImages.length === 0 && (
+      {appState === "INITIAL" && (
         <>
           <ImageUpload setReferenceImages={startCodeGeneration} />
         </>
       )}
 
-      <CodePreview code={generatedCode} />
-      <Preview code={generatedCode} />
+      {(appState === "CODING" || appState === "CODE_READY") && (
+        <>
+          <img className="w-[300px]" src={referenceImages[0]} alt="Reference" />
+          {/* Show code preview only when coding */}
+          {appState === "CODING" && <CodePreview code={generatedCode} />}
+          <Preview code={generatedCode} />
+        </>
+      )}
     </div>
   );
 }

frontend/src/generateCode.ts

@@ -4,7 +4,8 @@ const WS_BACKEND_URL = import.meta.env.VITE_WS_BACKEND_URL;
 
 export function generateCode(
   imageUrl: string,
-  onChange: (chunk: string) => void
+  onChange: (chunk: string) => void,
+  onComplete: () => void
 ) {
   const wsUrl = `${WS_BACKEND_URL}/generate-code`;
   console.log("Connecting to backend @ ", wsUrl);
@@ -33,6 +34,8 @@ export function generateCode(
       toast.error(
         "We ran into an error. Try again or contact support at support@picoapps.xyz"
       );
+    } else {
+      onComplete();
     }
   });