diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index 3e4701c..248e1ff 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -37,10 +37,13 @@ const IS_OPENAI_DOWN = false;
 function App() {
   const [appState, setAppState] = useState<AppState>(AppState.INITIAL);
   const [generatedCode, setGeneratedCode] = useState<string>("");
+
   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>(
     {
       openAiApiKey: null,
@@ -132,15 +135,13 @@ function App() {
       wsRef,
       updatedParams,
       (token) => setGeneratedCode((prev) => prev + token),
-      (code) => setGeneratedCode(code),
-      (line) => setExecutionConsole((prev) => [...prev, line]),
-      () => {
-        setAppState(AppState.CODE_READY);
+      (code) => {
+        setGeneratedCode(code);
         if (params.generationType === "create") {
           setAppHistory([
             {
               type: "ai_create",
-              code: generatedCode,
+              code,
               // TODO: Doesn't typecheck correctly
               inputs: { image_url: referenceImages[0] },
             },
@@ -149,7 +150,7 @@ function App() {
           setAppHistory((prev) => [
             {
               type: "ai_edit",
-              code: generatedCode,
+              code,
               // TODO: Doesn't typecheck correctly
               inputs: {
                 // TODO: Fix this
@@ -160,6 +161,10 @@ function App() {
             ...prev,
           ]);
         }
+      },
+      (line) => setExecutionConsole((prev) => [...prev, line]),
+      () => {
+        setAppState(AppState.CODE_READY);
       }
     );
   }
@@ -349,7 +354,20 @@ function App() {
               </div>
             </>
           )}
-          {<HistoryDisplay history={appHistory} />}
+          {
+            <HistoryDisplay
+              history={appHistory}
+              revertToVersion={(index) => {
+                if (
+                  index < 0 ||
+                  index >= appHistory.length ||
+                  !appHistory[index]
+                )
+                  return;
+                setGeneratedCode(appHistory[index].code);
+              }}
+            />
+          }
         </div>
       </div>
 
diff --git a/frontend/src/components/HistoryDisplay.tsx b/frontend/src/components/HistoryDisplay.tsx
index 26a428f..2d76783 100644
--- a/frontend/src/components/HistoryDisplay.tsx
+++ b/frontend/src/components/HistoryDisplay.tsx
@@ -3,6 +3,7 @@ import { History, HistoryItemType } from "../history_types";
 
 interface Props {
   history: History;
+  revertToVersion: (version: number) => void;
 }
 
 function displayHistoryItemType(itemType: HistoryItemType) {
@@ -23,14 +24,18 @@ function displayHistoryItemType(itemType: HistoryItemType) {
   }
 }
 
-export default function HistoryDisplay({ history }: Props) {
+export default function HistoryDisplay({ history, revertToVersion }: Props) {
   return (
     <div className="flex flex-col h-screen p-4">
       <h1 className="font-bold mb-4">History</h1>
       <ScrollArea className="flex-1 overflow-y-auto">
         <ul className="space-y-4">
           {history.map((item, index) => (
-            <li className="flex items-center space-x-4 justify-between border-b pb-1">
+            <li
+              key={index}
+              className="flex items-center space-x-4 justify-between border-b pb-1 cursor-pointer"
+              onClick={() => revertToVersion(index)}
+            >
               <h2 className="text-sm">{displayHistoryItemType(item.type)}</h2>
               <h2 className="text-sm">v{history.length - index}</h2>
             </li>
diff --git a/frontend/src/history_types.ts b/frontend/src/history_types.ts
index 1486a99..6fbaafe 100644
--- a/frontend/src/history_types.ts
+++ b/frontend/src/history_types.ts
@@ -7,7 +7,7 @@ export type HistoryItemType =
 
 export type HistoryItem = {
   type: HistoryItemType;
-  code?: string;
+  code: string;
   inputs:
     | AiCreateInputs
     | CodeCreateInputs