diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index cbefcfe..fc5521a 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -30,7 +30,10 @@ import { USER_CLOSE_WEB_SOCKET_CODE } from "./constants";
 import CodeTab from "./components/CodeTab";
 import OutputSettingsSection from "./components/OutputSettingsSection";
 import { History } from "./history_types";
-import HistoryDisplay from "./components/HistoryDisplay";
+import HistoryDisplay, {
+  extractHistoryTree,
+} from "./components/HistoryDisplay";
+import toast from "react-hot-toast";
 
 const IS_OPENAI_DOWN = false;
 
@@ -41,7 +44,6 @@ function App() {
   const [referenceImages, setReferenceImages] = useState<string[]>([]);
   const [executionConsole, setExecutionConsole] = useState<string[]>([]);
   const [updateInstruction, setUpdateInstruction] = useState("");
-  const [history, setHistory] = useState<string[]>([]);
 
   // Settings
   const [settings, setSettings] = usePersistedState<Settings>(
@@ -61,6 +63,7 @@ function App() {
 
   // App history
   const [appHistory, setAppHistory] = useState<History>([]);
+  // Tracks the currently viewed version from app history
   const [currentVersion, setCurrentVersion] = useState<number | null>(null);
 
   const [shouldIncludeResultImage, setShouldIncludeResultImage] =
@@ -115,7 +118,6 @@ function App() {
     setGeneratedCode("");
     setReferenceImages([]);
     setExecutionConsole([]);
-    setHistory([]);
     setAppHistory([]);
   };
 
@@ -143,7 +145,6 @@ function App() {
             {
               type: "ai_create",
               code,
-              // TODO: Doesn't typecheck correctly
               inputs: { image_url: referenceImages[0] },
             },
           ]);
@@ -154,11 +155,8 @@ function App() {
               {
                 type: "ai_edit",
                 code,
-                // TODO: Doesn't typecheck correctly
                 inputs: {
-                  // TODO: Fix this
-                  previous_commands: [],
-                  new_instruction: updateInstruction,
+                  prompt: updateInstruction,
                 },
               },
               ...prev,
@@ -191,7 +189,18 @@ function App() {
 
   // Subsequent updates
   async function doUpdate() {
-    const updatedHistory = [...history, generatedCode, updateInstruction];
+    if (currentVersion === null) {
+      toast.error(
+        "No current version set. Contact support or open a Github issue."
+      );
+      return;
+    }
+
+    const updatedHistory = [
+      ...extractHistoryTree(appHistory, currentVersion),
+      updateInstruction,
+    ];
+
     if (shouldIncludeResultImage) {
       const resultImage = await takeScreenshot();
       doGenerateCode({
@@ -208,7 +217,6 @@ function App() {
       });
     }
 
-    setHistory(updatedHistory);
     setGeneratedCode("");
     setUpdateInstruction("");
   }
diff --git a/frontend/src/components/HistoryDisplay.tsx b/frontend/src/components/HistoryDisplay.tsx
index 49a295d..69eb025 100644
--- a/frontend/src/components/HistoryDisplay.tsx
+++ b/frontend/src/components/HistoryDisplay.tsx
@@ -10,18 +10,35 @@ interface Props {
   shouldDisableReverts: boolean;
 }
 
+export function extractHistoryTree(
+  history: History,
+  version: number
+): string[] {
+  // History is in reverse chronological order
+
+  // Get all history items up to the current version
+  const extractedHistory = history.slice(version);
+
+  const obj: string[] = [];
+
+  // Reverse the history so that it is in chronological order for the server
+  extractedHistory.reverse().forEach((item) => {
+    // Don't include the image for ai_create since the server gets it passed and will include it directly
+    if (item.type !== "ai_create") {
+      obj.push(item.inputs.prompt);
+    }
+    obj.push(item.code);
+  });
+
+  return obj;
+}
+
 function displayHistoryItemType(itemType: HistoryItemType) {
   switch (itemType) {
     case "ai_create":
       return "Create";
     case "ai_edit":
       return "Edit";
-    case "code_create":
-      return "Create";
-    case "code_edit":
-      return "Code Edit";
-    case "revert":
-      return "Revert";
     default:
       // TODO: Error out since this is exhaustive
       return "Unknown";
diff --git a/frontend/src/history_types.ts b/frontend/src/history_types.ts
index 6fbaafe..ab19825 100644
--- a/frontend/src/history_types.ts
+++ b/frontend/src/history_types.ts
@@ -1,40 +1,23 @@
-export type HistoryItemType =
-  | "ai_create"
-  | "code_create"
-  | "ai_edit"
-  | "revert"
-  | "code_edit";
+export type HistoryItemType = "ai_create" | "ai_edit";
 
-export type HistoryItem = {
-  type: HistoryItemType;
-  code: string;
-  inputs:
-    | AiCreateInputs
-    | CodeCreateInputs
-    | AiEditInputs
-    | RevertInputs
-    | CodeEditInputs;
-};
+export type HistoryItem =
+  | {
+      type: "ai_create";
+      code: string;
+      inputs: AiCreateInputs;
+    }
+  | {
+      type: "ai_edit";
+      code: string;
+      inputs: AiEditInputs;
+    };
 
 export type AiCreateInputs = {
-  image_url?: string;
-};
-
-export type CodeCreateInputs = {
-  // Define specific properties relevant for code creation
+  image_url: string;
 };
 
 export type AiEditInputs = {
-  previous_commands: string[];
-  new_instruction: string;
-};
-
-export type RevertInputs = {
-  parent: number;
-};
-
-export type CodeEditInputs = {
-  // TODO: Fill in
+  prompt: string;
 };
 
 export type History = HistoryItem[];