Make undo more robust by using sourceCommitHash (#2015)

e2e-tests/fixtures/no-code-response.md

+I understand your request. This is a response without any code changes.
+

e2e-tests/undo.spec.ts

@@ -40,3 +40,29 @@ testSkipIfWindows("undo", async ({ po }) => {
 testSkipIfWindows("undo with native git", async ({ po }) => {
   await runUndoTest(po, true);
 });
+
+testSkipIfWindows("undo after assistant with no code", async ({ po }) => {
+  await po.setUp({ autoApprove: true, nativeGit: false });
+
+  // First prompt - no code generated
+  await po.sendPrompt("tc=no-code-response");
+
+  // Second prompt - generates code
+  await po.sendPrompt("tc=write-index");
+
+  const iframe = po.getPreviewIframeElement();
+  await expect(
+    iframe.contentFrame().getByText("Testing:write-index!"),
+  ).toBeVisible({
+    timeout: Timeout.LONG,
+  });
+
+  // Undo should work even though first assistant had no commit
+  await po.clickUndo();
+
+  await expect(
+    iframe.contentFrame().getByText("Welcome to Your Blank App"),
+  ).toBeVisible({
+    timeout: Timeout.LONG,
+  });
+});

src/components/chat/MessagesList.tsx

@@ -88,8 +88,7 @@ function FooterComponent({ context }: { context?: FooterContext }) {
       {!isStreaming && (
         <div className="flex max-w-3xl mx-auto gap-2">
           {!!messages.length &&
-            messages[messages.length - 1].role === "assistant" &&
-            messages[messages.length - 1].commitHash && (
+            messages[messages.length - 1].role === "assistant" && (
               <Button
                 variant="outline"
                 size="sm"
@@ -102,18 +101,22 @@ function FooterComponent({ context }: { context?: FooterContext }) {
 
                   setIsUndoLoading(true);
                   try {
-                    if (messages.length >= 3) {
-                      const previousAssistantMessage =
-                        messages[messages.length - 3];
-                      if (
-                        previousAssistantMessage?.role === "assistant" &&
-                        previousAssistantMessage?.commitHash
-                      ) {
+                    const currentMessage = messages[messages.length - 1];
+                    // The user message that triggered this assistant response
+                    const userMessage = messages[messages.length - 2];
+                    if (currentMessage?.sourceCommitHash) {
                       console.debug(
-                          "Reverting to previous assistant version",
+                        "Reverting to source commit hash",
+                        currentMessage.sourceCommitHash,
                       );
                       await revertVersion({
-                          versionId: previousAssistantMessage.commitHash,
+                        versionId: currentMessage.sourceCommitHash,
+                        currentChatMessageId: userMessage
+                          ? {
+                              chatId: selectedChatId,
+                              messageId: userMessage.id,
+                            }
+                          : undefined,
                       });
                       const chat =
                         await IpcClient.getInstance().getChat(selectedChatId);
@@ -122,32 +125,11 @@ function FooterComponent({ context }: { context?: FooterContext }) {
                         next.set(selectedChatId, chat.messages);
                         return next;
                       });
-                      }
-                    } else {
-                      const chat =
-                        await IpcClient.getInstance().getChat(selectedChatId);
-                      if (chat.initialCommitHash) {
-                        await revertVersion({
-                          versionId: chat.initialCommitHash,
-                        });
-                        try {
-                          await IpcClient.getInstance().deleteMessages(
-                            selectedChatId,
-                          );
-                          setMessagesById((prev) => {
-                            const next = new Map(prev);
-                            next.set(selectedChatId, []);
-                            return next;
-                          });
-                        } catch (err) {
-                          showError(err);
-                        }
                     } else {
                       showWarning(
-                          "No initial commit hash found for chat. Need to manually undo code changes",
+                        "No source commit hash found for message. Need to manually undo code changes",
                       );
                     }
-                    }
                   } catch (error) {
                     console.error("Error during undo operation:", error);
                     showError("Failed to undo changes");

src/hooks/useVersions.ts

@@ -42,9 +42,18 @@ export function useVersions(appId: number | null) {
   const revertVersionMutation = useMutation<
     RevertVersionResponse,
     Error,
-    { versionId: string }
+    {
+      versionId: string;
+      currentChatMessageId?: { chatId: number; messageId: number };
+    }
   >({
-    mutationFn: async ({ versionId }: { versionId: string }) => {
+    mutationFn: async ({
+      versionId,
+      currentChatMessageId,
+    }: {
+      versionId: string;
+      currentChatMessageId?: { chatId: number; messageId: number };
+    }) => {
       const currentAppId = appId;
       if (currentAppId === null) {
         throw new Error("App ID is null");
@@ -53,6 +62,7 @@ export function useVersions(appId: number | null) {
       return ipcClient.revertVersion({
         appId: currentAppId,
         previousVersionId: versionId,
+        currentChatMessageId,
       });
     },
     onSuccess: async (result) => {

src/ipc/handlers/version_handlers.ts

 import { db } from "../../db";
 import { apps, messages, versions } from "../../db/schema";
-import { desc, eq, and, gt } from "drizzle-orm";
+import { desc, eq, and, gt, gte } from "drizzle-orm";
 import type {
   Version,
   BranchResult,
@@ -23,6 +23,7 @@ import {
   getCurrentCommitHash,
   gitCurrentBranch,
   gitLog,
+  isGitStatusClean,
 } from "../utils/git_utils";
 
 import {
@@ -156,7 +157,7 @@ export function registerVersionHandlers() {
     "revert-version",
     async (
       _,
-      { appId, previousVersionId }: RevertVersionParams,
+      { appId, previousVersionId, currentChatMessageId }: RevertVersionParams,
     ): Promise<RevertVersionResponse> => {
       return withLock(appId, async () => {
         let successMessage = "Restored version";
@@ -193,12 +194,39 @@ export function registerVersionHandlers() {
           path: appPath,
           targetOid: previousVersionId,
         });
-
+        const isClean = await isGitStatusClean({ path: appPath });
+        if (!isClean) {
           await gitCommit({
             path: appPath,
             message: `Reverted all changes back to version ${previousVersionId}`,
           });
+        }
+
+        // Delete messages based on currentChatMessageId if provided, otherwise use commit hash lookup
+        if (currentChatMessageId) {
+          // Delete all messages including and after the specified message
+          const { chatId, messageId } = currentChatMessageId;
 
+          const messagesToDelete = await db.query.messages.findMany({
+            where: and(
+              eq(messages.chatId, chatId),
+              gte(messages.id, messageId),
+            ),
+            orderBy: desc(messages.id),
+          });
+
+          logger.log(
+            `Deleting ${messagesToDelete.length} messages (id >= ${messageId}) from chat ${chatId}`,
+          );
+
+          if (messagesToDelete.length > 0) {
+            await db
+              .delete(messages)
+              .where(
+                and(eq(messages.chatId, chatId), gte(messages.id, messageId)),
+              );
+          }
+        } else {
           // Find the chat and message associated with the commit hash
           const messageWithCommit = await db.query.messages.findFirst({
             where: eq(messages.commitHash, previousVersionId),
@@ -236,6 +264,7 @@ export function registerVersionHandlers() {
                 );
             }
           }
+        }
 
         if (app.neonProjectId && app.neonDevelopmentBranchId) {
           const version = await db.query.versions.findFirst({

src/ipc/ipc_types.ts

@@ -80,6 +80,7 @@ export interface Message {
   content: string;
   approvalState?: "approved" | "rejected" | null;
   commitHash?: string | null;
+  sourceCommitHash?: string | null;
   dbTimestamp?: string | null;
   createdAt?: Date | string;
   requestId?: string | null;
@@ -442,6 +443,10 @@ export interface GetNeonProjectResponse {
 export interface RevertVersionParams {
   appId: number;
   previousVersionId: string;
+  currentChatMessageId?: {
+    chatId: number;
+    messageId: number;
+  };
 }
 
 export type RevertVersionResponse =