From b9dc2cc0f9a88dae1494cb308f40d59c8cd8007c Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Fri, 2 May 2025 15:35:39 -0700
Subject: [PATCH] Undo button (#76)

---
 src/components/chat/MessagesList.tsx | 177 +++++++++++++++++++--------
 src/ipc/handlers/app_handlers.ts     |   1 -
 src/ipc/ipc_client.ts                |  14 ++-
 3 files changed, 136 insertions(+), 56 deletions(-)

diff --git a/src/components/chat/MessagesList.tsx b/src/components/chat/MessagesList.tsx
index 7836a34..e44edda 100644
--- a/src/components/chat/MessagesList.tsx
+++ b/src/components/chat/MessagesList.tsx
@@ -1,17 +1,19 @@
 import type React from "react";
 import type { Message } from "@/ipc/ipc_types";
-import { forwardRef } from "react";
+import { forwardRef, useState } from "react";
 import ChatMessage from "./ChatMessage";
 import { SetupBanner } from "../SetupBanner";
 import { useSettings } from "@/hooks/useSettings";
 import { useStreamChat } from "@/hooks/useStreamChat";
 import { selectedChatIdAtom } from "@/atoms/chatAtoms";
-import { useAtom, useAtomValue } from "jotai";
-import { RefreshCw } from "lucide-react";
+import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import { Loader2, RefreshCw, Undo } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { useVersions } from "@/hooks/useVersions";
 import { selectedAppIdAtom } from "@/atoms/appAtoms";
 import { showError } from "@/lib/toast";
+import { IpcClient } from "@/ipc/ipc_client";
+import { chatMessagesAtom } from "@/atoms/chatAtoms";
 
 interface MessagesListProps {
   messages: Message[];
@@ -25,6 +27,10 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
     const { streamMessage, isStreaming, error, setError } = useStreamChat();
     const { isAnyProviderSetup } = useSettings();
     const selectedChatId = useAtomValue(selectedChatIdAtom);
+    const setMessages = useSetAtom(chatMessagesAtom);
+    const [isUndoLoading, setIsUndoLoading] = useState(false);
+    const [isRetryLoading, setIsRetryLoading] = useState(false);
+
     return (
       <div className="flex-1 overflow-y-auto p-4" ref={ref}>
         {messages.length > 0 ? (
@@ -40,67 +46,136 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
           </div>
         )}
         {messages.length > 0 && !isStreaming && (
-          <div className="flex max-w-3xl mx-auto">
+          <div className="flex max-w-3xl mx-auto gap-2">
+            {messages[messages.length - 1].role === "assistant" && (
+              <Button
+                variant="outline"
+                size="sm"
+                disabled={isUndoLoading}
+                onClick={async () => {
+                  if (!selectedChatId || !appId) {
+                    console.error("No chat selected or app ID not available");
+                    return;
+                  }
+                  if (versions.length < 2) {
+                    showError("Cannot undo; no previous version");
+                    return;
+                  }
+
+                  setIsUndoLoading(true);
+                  try {
+                    const previousAssistantMessage =
+                      messages[messages.length - 3];
+                    if (
+                      previousAssistantMessage?.role === "assistant" &&
+                      previousAssistantMessage?.commitHash
+                    ) {
+                      console.debug("Reverting to previous assistant version");
+                      await revertVersion({
+                        versionId: previousAssistantMessage.commitHash,
+                      });
+                    } else {
+                      // Revert to the previous version
+                      await revertVersion({
+                        versionId: versions[1].oid,
+                      });
+                    }
+                    if (selectedChatId) {
+                      const chat = await IpcClient.getInstance().getChat(
+                        selectedChatId
+                      );
+                      setMessages(chat.messages);
+                    }
+                  } catch (error) {
+                    console.error("Error during undo operation:", error);
+                    showError("Failed to undo changes");
+                  } finally {
+                    setIsUndoLoading(false);
+                  }
+                }}
+              >
+                {isUndoLoading ? (
+                  <Loader2 size={16} className="mr-1 animate-spin" />
+                ) : (
+                  <Undo size={16} />
+                )}
+                Undo
+              </Button>
+            )}
             <Button
-              variant="ghost"
+              variant="outline"
               size="sm"
+              disabled={isRetryLoading}
               onClick={async () => {
                 if (!selectedChatId) {
                   console.error("No chat selected");
                   return;
                 }
-                // The last message is usually an assistant, but it might not be.
-                const lastVersion = versions[0];
-                const lastMessage = messages[messages.length - 1];
-                let reverted = false;
-                if (
-                  lastVersion.oid === lastMessage.commitHash &&
-                  lastMessage.role === "assistant"
-                ) {
-                  if (versions.length < 2) {
-                    showError("Cannot retry message; no previous version");
+
+                setIsRetryLoading(true);
+                try {
+                  // The last message is usually an assistant, but it might not be.
+                  const lastVersion = versions[0];
+                  const lastMessage = messages[messages.length - 1];
+                  let reverted = false;
+                  if (
+                    lastVersion.oid === lastMessage.commitHash &&
+                    lastMessage.role === "assistant"
+                  ) {
+                    if (versions.length < 2) {
+                      showError("Cannot retry message; no previous version");
+                      return;
+                    }
+                    const previousAssistantMessage =
+                      messages[messages.length - 3];
+                    if (
+                      previousAssistantMessage?.role === "assistant" &&
+                      previousAssistantMessage?.commitHash
+                    ) {
+                      console.debug("Reverting to previous assistant version");
+                      await revertVersion({
+                        versionId: previousAssistantMessage.commitHash,
+                      });
+                      reverted = true;
+                    } else {
+                      console.debug("Reverting to previous version");
+                      await revertVersion({
+                        versionId: versions[1].oid,
+                      });
+                    }
+                  }
+
+                  // Find the last user message
+                  const lastUserMessage = [...messages]
+                    .reverse()
+                    .find((message) => message.role === "user");
+                  if (!lastUserMessage) {
+                    console.error("No user message found");
                     return;
                   }
-                  const previousAssistantMessage =
-                    messages[messages.length - 3];
-                  if (
-                    previousAssistantMessage?.role === "assistant" &&
-                    previousAssistantMessage?.commitHash
-                  ) {
-                    console.debug("Reverting to previous assistant version");
-                    await revertVersion({
-                      versionId: previousAssistantMessage.commitHash,
-                    });
-                    reverted = true;
-                  } else {
-                    console.debug("Reverting to previous version");
-                    await revertVersion({
-                      versionId: versions[1].oid,
-                    });
-                  }
-                }
+                  // If we reverted, we don't need to mark "redo" because
+                  // the old message has already been deleted.
+                  const redo = !reverted;
+                  console.debug("Streaming message with redo", redo);
 
-                // Find the last user message
-                const lastUserMessage = [...messages]
-                  .reverse()
-                  .find((message) => message.role === "user");
-                if (!lastUserMessage) {
-                  console.error("No user message found");
-                  return;
+                  streamMessage({
+                    prompt: lastUserMessage.content,
+                    chatId: selectedChatId,
+                    redo,
+                  });
+                } catch (error) {
+                  console.error("Error during retry operation:", error);
+                  showError("Failed to retry message");
+                } finally {
+                  setIsRetryLoading(false);
                 }
-                // If we reverted, we don't need to mark "redo" because
-                // the old message has already been deleted.
-                const redo = !reverted;
-                console.debug("Streaming message with redo", redo);
-
-                streamMessage({
-                  prompt: lastUserMessage.content,
-                  chatId: selectedChatId,
-                  redo,
-                });
               }}
             >
-              <RefreshCw size={16} />
+              {isRetryLoading ? (
+                <Loader2 size={16} className="mr-1 animate-spin" />
+              ) : (
+                <RefreshCw size={16} />
+              )}
               Retry
             </Button>
           </div>
diff --git a/src/ipc/handlers/app_handlers.ts b/src/ipc/handlers/app_handlers.ts
index 8d22326..5249795 100644
--- a/src/ipc/handlers/app_handlers.ts
+++ b/src/ipc/handlers/app_handlers.ts
@@ -494,7 +494,6 @@ export function registerAppHandlers() {
       _,
       { appId, previousVersionId }: { appId: number; previousVersionId: string }
     ) => {
-      logger.log(`Reverting to version ${previousVersionId} for app ${appId}`);
       return withLock(appId, async () => {
         const app = await db.query.apps.findFirst({
           where: eq(apps.id, appId),
diff --git a/src/ipc/ipc_client.ts b/src/ipc/ipc_client.ts
index ab55a47..a11932a 100644
--- a/src/ipc/ipc_client.ts
+++ b/src/ipc/ipc_client.ts
@@ -787,25 +787,31 @@ export class IpcClient {
 
   public async listLocalOllamaModels(): Promise<LocalModel[]> {
     try {
-      const response = await this.ipcRenderer.invoke("local-models:list-ollama");
+      const response = await this.ipcRenderer.invoke(
+        "local-models:list-ollama"
+      );
       return response?.models || [];
     } catch (error) {
       if (error instanceof Error) {
         throw new Error(`Failed to fetch Ollama models: ${error.message}`);
       }
-      throw new Error('Failed to fetch Ollama models: Unknown error occurred');
+      throw new Error("Failed to fetch Ollama models: Unknown error occurred");
     }
   }
 
   public async listLocalLMStudioModels(): Promise<LocalModel[]> {
     try {
-      const response = await this.ipcRenderer.invoke("local-models:list-lmstudio");
+      const response = await this.ipcRenderer.invoke(
+        "local-models:list-lmstudio"
+      );
       return response?.models || [];
     } catch (error) {
       if (error instanceof Error) {
         throw new Error(`Failed to fetch LM Studio models: ${error.message}`);
       }
-      throw new Error('Failed to fetch LM Studio models: Unknown error occurred');
+      throw new Error(
+        "Failed to fetch LM Studio models: Unknown error occurred"
+      );
     }
   }