diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index ff7c83c..deb8a29 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -128,7 +128,7 @@ function App() {
 
   function doGenerateCode(
     params: CodeGenerationParams,
-    parentVersion?: number
+    parentVersion: number | null
   ) {
     setExecutionConsole([]);
     setAppState(AppState.CODING);
@@ -166,7 +166,6 @@ function App() {
               ...prev,
               {
                 type: "ai_edit",
-                // TODO: It should never be null
                 parent: parentVersion,
                 code,
                 inputs: {
@@ -193,10 +192,13 @@ function App() {
 
     setReferenceImages(referenceImages);
     if (referenceImages.length > 0) {
-      doGenerateCode({
-        generationType: "create",
-        image: referenceImages[0],
-      });
+      doGenerateCode(
+        {
+          generationType: "create",
+          image: referenceImages[0],
+        },
+        currentVersion
+      );
     }
   }
 
diff --git a/frontend/src/components/history/HistoryDisplay.tsx b/frontend/src/components/history/HistoryDisplay.tsx
index aa00920..beb78a2 100644
--- a/frontend/src/components/history/HistoryDisplay.tsx
+++ b/frontend/src/components/history/HistoryDisplay.tsx
@@ -54,9 +54,11 @@ export default function HistoryDisplay({
             >
               <div className="flex gap-x-1">
                 <h2 className="text-sm">{displayHistoryItemType(item.type)}</h2>
-                {item.parent && item.parent !== index - 1 && (
-                  <h2 className="text-sm">(parent: v{item.parent + 1})</h2>
-                )}
+                {item.parent && item.parent !== index - 1 ? (
+                  <h2 className="text-sm">
+                    (parent: v{(item.parent || 0) + 1})
+                  </h2>
+                ) : null}
               </div>
               <h2 className="text-sm">
                 {item.type === "ai_edit"
diff --git a/frontend/src/history_types.ts b/frontend/src/history_types.ts
index a832329..4826595 100644
--- a/frontend/src/history_types.ts
+++ b/frontend/src/history_types.ts
@@ -1,18 +1,19 @@
 export type HistoryItemType = "ai_create" | "ai_edit";
 
+type CommonHistoryItem = {
+  parent: null | number;
+  code: string;
+};
+
 export type HistoryItem =
-  | {
+  | ({
       type: "ai_create";
-      parent: null | number;
-      code: string;
       inputs: AiCreateInputs;
-    }
-  | {
+    } & CommonHistoryItem)
+  | ({
       type: "ai_edit";
-      parent: null | number;
-      code: string;
       inputs: AiEditInputs;
-    };
+    } & CommonHistoryItem);
 
 export type AiCreateInputs = {
   image_url: string;