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,51 +101,34 @@ function FooterComponent({ context }: { context?: FooterContext }) {
 
                   setIsUndoLoading(true);
                   try {
-                    if (messages.length >= 3) {
-                      const previousAssistantMessage =
-                        messages[messages.length - 3];
-                      if (
-                        previousAssistantMessage?.role === "assistant" &&
-                        previousAssistantMessage?.commitHash
-                      ) {
-                        console.debug(
-                          "Reverting to previous assistant version",
-                        );
-                        await revertVersion({
-                          versionId: previousAssistantMessage.commitHash,
-                        });
-                        const chat =
-                          await IpcClient.getInstance().getChat(selectedChatId);
-                        setMessagesById((prev) => {
-                          const next = new Map(prev);
-                          next.set(selectedChatId, chat.messages);
-                          return next;
-                        });
-                      }
-                    } else {
+                    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 source commit hash",
+                        currentMessage.sourceCommitHash,
+                      );
+                      await revertVersion({
+                        versionId: currentMessage.sourceCommitHash,
+                        currentChatMessageId: userMessage
+                          ? {
+                              chatId: selectedChatId,
+                              messageId: userMessage.id,
+                            }
+                          : undefined,
+                      });
                       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",
-                        );
-                      }
+                      setMessagesById((prev) => {
+                        const next = new Map(prev);
+                        next.set(selectedChatId, chat.messages);
+                        return next;
+                      });
+                    } else {
+                      showWarning(
+                        "No source commit hash found for message. Need to manually undo code changes",
+                      );
                     }
                   } catch (error) {
                     console.error("Error during undo operation:", error);

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,48 +194,76 @@ 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}`,
+          });
+        }
 
-        await gitCommit({
-          path: appPath,
-          message: `Reverted all changes back to version ${previousVersionId}`,
-        });
-
-        // Find the chat and message associated with the commit hash
-        const messageWithCommit = await db.query.messages.findFirst({
-          where: eq(messages.commitHash, previousVersionId),
-          with: {
-            chat: true,
-          },
-        });
-
-        // If we found a message with this commit hash, delete all subsequent messages (but keep this message)
-        if (messageWithCommit) {
-          const chatId = messageWithCommit.chatId;
+        // 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;
 
-          // Find all messages in this chat with IDs > the one with our commit hash
           const messagesToDelete = await db.query.messages.findMany({
             where: and(
               eq(messages.chatId, chatId),
-              gt(messages.id, messageWithCommit.id),
+              gte(messages.id, messageId),
             ),
             orderBy: desc(messages.id),
           });
 
           logger.log(
-            `Deleting ${messagesToDelete.length} messages after commit ${previousVersionId} from chat ${chatId}`,
+            `Deleting ${messagesToDelete.length} messages (id >= ${messageId}) from chat ${chatId}`,
           );
 
-          // Delete the messages
           if (messagesToDelete.length > 0) {
             await db
               .delete(messages)
               .where(
-                and(
-                  eq(messages.chatId, chatId),
-                  gt(messages.id, messageWithCommit.id),
-                ),
+                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),
+            with: {
+              chat: true,
+            },
+          });
+
+          // If we found a message with this commit hash, delete all subsequent messages (but keep this message)
+          if (messageWithCommit) {
+            const chatId = messageWithCommit.chatId;
+
+            // Find all messages in this chat with IDs > the one with our commit hash
+            const messagesToDelete = await db.query.messages.findMany({
+              where: and(
+                eq(messages.chatId, chatId),
+                gt(messages.id, messageWithCommit.id),
+              ),
+              orderBy: desc(messages.id),
+            });
+
+            logger.log(
+              `Deleting ${messagesToDelete.length} messages after commit ${previousVersionId} from chat ${chatId}`,
+            );
+
+            // Delete the messages
+            if (messagesToDelete.length > 0) {
+              await db
+                .delete(messages)
+                .where(
+                  and(
+                    eq(messages.chatId, chatId),
+                    gt(messages.id, messageWithCommit.id),
+                  ),
+                );
+            }
+          }
         }
 
         if (app.neonProjectId && app.neonDevelopmentBranchId) {

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 =