From 46cb7ca577f8856ff44189259b2f9fca85f47c69 Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Fri, 2 May 2025 14:55:46 -0700
Subject: [PATCH] redo reverts code (when possible) (#75)

---
 src/components/chat/ChatInput.tsx    |  4 +++
 src/components/chat/MessagesList.tsx | 45 ++++++++++++++++++++++++++--
 src/ipc/handlers/app_handlers.ts     |  1 +
 src/ipc/ipc_types.ts                 |  1 +
 4 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/src/components/chat/ChatInput.tsx b/src/components/chat/ChatInput.tsx
index d77e892..8b349d4 100644
--- a/src/components/chat/ChatInput.tsx
+++ b/src/components/chat/ChatInput.tsx
@@ -53,6 +53,7 @@ import {
   TooltipTrigger,
 } from "../ui/tooltip";
 import { useNavigate } from "@tanstack/react-router";
+import { useVersions } from "@/hooks/useVersions";
 
 const showTokenBarAtom = atom(false);
 
@@ -61,6 +62,8 @@ export function ChatInput({ chatId }: { chatId?: number }) {
   const [inputValue, setInputValue] = useAtom(chatInputValueAtom);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const { settings, updateSettings, isAnyProviderSetup } = useSettings();
+  const appId = useAtomValue(selectedAppIdAtom);
+  const { refreshVersions } = useVersions(appId);
   const { streamMessage, isStreaming, setIsStreaming, error, setError } =
     useStreamChat();
   const [showError, setShowError] = useState(true);
@@ -162,6 +165,7 @@ export function ChatInput({ chatId }: { chatId?: number }) {
     } finally {
       setIsApproving(false);
       setIsPreviewOpen(true);
+      refreshVersions();
 
       // Keep same as handleReject
       refreshProposal();
diff --git a/src/components/chat/MessagesList.tsx b/src/components/chat/MessagesList.tsx
index 85758c5..7836a34 100644
--- a/src/components/chat/MessagesList.tsx
+++ b/src/components/chat/MessagesList.tsx
@@ -9,6 +9,9 @@ import { selectedChatIdAtom } from "@/atoms/chatAtoms";
 import { useAtom, useAtomValue } from "jotai";
 import { RefreshCw } from "lucide-react";
 import { Button } from "@/components/ui/button";
+import { useVersions } from "@/hooks/useVersions";
+import { selectedAppIdAtom } from "@/atoms/appAtoms";
+import { showError } from "@/lib/toast";
 
 interface MessagesListProps {
   messages: Message[];
@@ -17,6 +20,8 @@ interface MessagesListProps {
 
 export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
   function MessagesList({ messages, messagesEndRef }, ref) {
+    const appId = useAtomValue(selectedAppIdAtom);
+    const { versions, revertVersion } = useVersions(appId);
     const { streamMessage, isStreaming, error, setError } = useStreamChat();
     const { isAnyProviderSetup } = useSettings();
     const selectedChatId = useAtomValue(selectedChatIdAtom);
@@ -39,11 +44,42 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
             <Button
               variant="ghost"
               size="sm"
-              onClick={() => {
+              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");
+                    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()
@@ -52,10 +88,15 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
                   console.error("No user message found");
                   return;
                 }
+                // 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: true,
+                  redo,
                 });
               }}
             >
diff --git a/src/ipc/handlers/app_handlers.ts b/src/ipc/handlers/app_handlers.ts
index 5249795..8d22326 100644
--- a/src/ipc/handlers/app_handlers.ts
+++ b/src/ipc/handlers/app_handlers.ts
@@ -494,6 +494,7 @@ 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_types.ts b/src/ipc/ipc_types.ts
index 7fe7c90..972d74a 100644
--- a/src/ipc/ipc_types.ts
+++ b/src/ipc/ipc_types.ts
@@ -41,6 +41,7 @@ export interface Message {
   role: "user" | "assistant";
   content: string;
   approvalState?: "approved" | "rejected" | null;
+  commitHash?: string | null;
 }
 
 export interface Chat {