From c203b1d00924dc23183c3d39ec47fc1870389775 Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Thu, 8 May 2025 23:23:24 -0700
Subject: [PATCH] make checkout version and revert version fit pattern (#118)

---
 src/components/ChatPanel.tsx         |  1 +
 src/components/chat/ChatHeader.tsx   | 29 ++++-----
 src/components/chat/VersionPane.tsx  | 95 ++++++++++++++++++++--------
 src/hooks/useCheckoutVersion.ts      | 38 +++++++++++
 src/hooks/useVersions.ts             | 31 ++++-----
 src/ipc/handlers/version_handlers.ts | 11 ++--
 src/ipc/ipc_client.ts                | 32 +++-------
 src/lib/assert.ts                    |  8 +++
 src/renderer.tsx                     |  8 +++
 9 files changed, 161 insertions(+), 92 deletions(-)
 create mode 100644 src/hooks/useCheckoutVersion.ts
 create mode 100644 src/lib/assert.ts

diff --git a/src/components/ChatPanel.tsx b/src/components/ChatPanel.tsx
index 8f65c24..50ee448 100644
--- a/src/components/ChatPanel.tsx
+++ b/src/components/ChatPanel.tsx
@@ -117,6 +117,7 @@ export function ChatPanel({
   return (
     <div className="flex flex-col h-full">
       <ChatHeader
+        isVersionPaneOpen={isVersionPaneOpen}
         isPreviewOpen={isPreviewOpen}
         onTogglePreview={onTogglePreview}
         onVersionClick={() => setIsVersionPaneOpen(!isVersionPaneOpen)}
diff --git a/src/components/chat/ChatHeader.tsx b/src/components/chat/ChatHeader.tsx
index 11578c6..c2dc6cb 100644
--- a/src/components/chat/ChatHeader.tsx
+++ b/src/components/chat/ChatHeader.tsx
@@ -21,17 +21,20 @@ import { useRouter } from "@tanstack/react-router";
 import { selectedChatIdAtom } from "@/atoms/chatAtoms";
 import { useChats } from "@/hooks/useChats";
 import { showError } from "@/lib/toast";
-import { useEffect, useState } from "react";
+import { useEffect } from "react";
 import { useStreamChat } from "@/hooks/useStreamChat";
 import { useCurrentBranch } from "@/hooks/useCurrentBranch";
+import { useCheckoutVersion } from "@/hooks/useCheckoutVersion";
 
 interface ChatHeaderProps {
+  isVersionPaneOpen: boolean;
   isPreviewOpen: boolean;
   onTogglePreview: () => void;
   onVersionClick: () => void;
 }
 
 export function ChatHeader({
+  isVersionPaneOpen,
   isPreviewOpen,
   onTogglePreview,
   onVersionClick,
@@ -41,7 +44,6 @@ export function ChatHeader({
   const { navigate } = useRouter();
   const [selectedChatId, setSelectedChatId] = useAtom(selectedChatIdAtom);
   const { refreshChats } = useChats(appId);
-  const [checkingOutMain, setCheckingOutMain] = useState(false);
   const { isStreaming } = useStreamChat();
 
   const {
@@ -50,6 +52,8 @@ export function ChatHeader({
     refetchBranchInfo,
   } = useCurrentBranch(appId);
 
+  const { checkoutVersion, isCheckingOutVersion } = useCheckoutVersion();
+
   useEffect(() => {
     if (appId) {
       refetchBranchInfo();
@@ -58,19 +62,7 @@ export function ChatHeader({
 
   const handleCheckoutMainBranch = async () => {
     if (!appId) return;
-
-    try {
-      setCheckingOutMain(true);
-      await IpcClient.getInstance().checkoutVersion({
-        appId,
-        versionId: "main",
-      });
-      await refetchBranchInfo();
-    } catch (error) {
-      showError(`Failed to checkout main branch: ${error}`);
-    } finally {
-      setCheckingOutMain(false);
-    }
+    await checkoutVersion({ appId, versionId: "main" });
   };
 
   const handleNewChat = async () => {
@@ -100,7 +92,8 @@ export function ChatHeader({
 
   return (
     <div className="flex flex-col w-full @container">
-      {isNotMainBranch && (
+      {/* If the version pane is open, it's expected to not always be on the main branch. */}
+      {isNotMainBranch && !isVersionPaneOpen && (
         <div className="flex flex-col @sm:flex-row items-center justify-between px-4 py-2 bg-amber-100 dark:bg-amber-900 text-amber-800 dark:text-amber-200">
           <div className="flex items-center gap-2 text-sm">
             <GitBranch size={16} />
@@ -138,9 +131,9 @@ export function ChatHeader({
             variant="outline"
             size="sm"
             onClick={handleCheckoutMainBranch}
-            disabled={checkingOutMain || branchInfoLoading}
+            disabled={isCheckingOutVersion || branchInfoLoading}
           >
-            {checkingOutMain ? "Checking out..." : "Switch to main branch"}
+            {isCheckingOutVersion ? "Checking out..." : "Switch to main branch"}
           </Button>
         </div>
       )}
diff --git a/src/components/chat/VersionPane.tsx b/src/components/chat/VersionPane.tsx
index 81adaa7..6a521d4 100644
--- a/src/components/chat/VersionPane.tsx
+++ b/src/components/chat/VersionPane.tsx
@@ -4,9 +4,9 @@ import { useVersions } from "@/hooks/useVersions";
 import { formatDistanceToNow } from "date-fns";
 import { RotateCcw, X } from "lucide-react";
 import type { Version } from "@/ipc/ipc_types";
-import { IpcClient } from "@/ipc/ipc_client";
 import { cn } from "@/lib/utils";
-import { useEffect } from "react";
+import { useEffect, useRef, useState } from "react";
+import { useCheckoutVersion } from "@/hooks/useCheckoutVersion";
 
 interface VersionPaneProps {
   isVisible: boolean;
@@ -15,31 +15,69 @@ interface VersionPaneProps {
 
 export function VersionPane({ isVisible, onClose }: VersionPaneProps) {
   const appId = useAtomValue(selectedAppIdAtom);
-  const { versions, loading, refreshVersions, revertVersion } =
-    useVersions(appId);
+  const {
+    versions: liveVersions,
+    refreshVersions,
+    revertVersion,
+  } = useVersions(appId);
   const [selectedVersionId, setSelectedVersionId] = useAtom(
     selectedVersionIdAtom,
   );
+  const { checkoutVersion, isCheckingOutVersion } = useCheckoutVersion();
+  const wasVisibleRef = useRef(false);
+  const [cachedVersions, setCachedVersions] = useState<Version[]>([]);
+
   useEffect(() => {
-    async function updateVersions() {
-      // Refresh versions in case the user updated versions outside of the app
-      // (e.g. manually using git).
-      // Avoid loading state which causes brief flash of loading state.
+    async function updatePaneState() {
+      // When pane becomes visible after being closed
+      if (isVisible && !wasVisibleRef.current) {
+        if (appId) {
+          await refreshVersions();
+          setCachedVersions(liveVersions);
+        }
+      }
+
+      // Reset when closing
       if (!isVisible && selectedVersionId) {
         setSelectedVersionId(null);
-        await IpcClient.getInstance().checkoutVersion({
-          appId: appId!,
-          versionId: "main",
-        });
+        if (appId) {
+          await checkoutVersion({ appId, versionId: "main" });
+        }
       }
-      refreshVersions();
+
+      wasVisibleRef.current = isVisible;
     }
-    updateVersions();
-  }, [isVisible, refreshVersions]);
+    updatePaneState();
+  }, [
+    isVisible,
+    selectedVersionId,
+    setSelectedVersionId,
+    appId,
+    checkoutVersion,
+    refreshVersions,
+    liveVersions,
+  ]);
+
+  // Initial load of cached versions when live versions become available
+  useEffect(() => {
+    if (isVisible && liveVersions.length > 0 && cachedVersions.length === 0) {
+      setCachedVersions(liveVersions);
+    }
+  }, [isVisible, liveVersions, cachedVersions.length]);
+
   if (!isVisible) {
     return null;
   }
 
+  const handleVersionClick = async (versionOid: string) => {
+    if (appId) {
+      setSelectedVersionId(versionOid);
+      await checkoutVersion({ appId, versionId: versionOid });
+    }
+  };
+
+  const versions = cachedVersions.length > 0 ? cachedVersions : liveVersions;
+
   return (
     <div className="h-full border-t border-2 border-border w-full">
       <div className="p-2 border-b border-border flex items-center justify-between">
@@ -53,26 +91,25 @@ export function VersionPane({ isVisible, onClose }: VersionPaneProps) {
         </button>
       </div>
       <div className="overflow-y-auto h-[calc(100%-60px)]">
-        {loading ? (
-          <div className="p-4 ">Loading versions...</div>
-        ) : versions.length === 0 ? (
+        {versions.length === 0 ? (
           <div className="p-4 ">No versions available</div>
         ) : (
           <div className="divide-y divide-border">
             {versions.map((version: Version, index) => (
               <div
                 key={version.oid}
-                className={`px-4 py-2 hover:bg-(--background-lightest) cursor-pointer ${
-                  selectedVersionId === version.oid
-                    ? "bg-(--background-lightest)"
-                    : ""
-                }`}
+                className={cn(
+                  "px-4 py-2 hover:bg-(--background-lightest) cursor-pointer",
+                  selectedVersionId === version.oid &&
+                    "bg-(--background-lightest)",
+                  isCheckingOutVersion &&
+                    selectedVersionId === version.oid &&
+                    "opacity-50 cursor-not-allowed",
+                )}
                 onClick={() => {
-                  IpcClient.getInstance().checkoutVersion({
-                    appId: appId!,
-                    versionId: version.oid,
-                  });
-                  setSelectedVersionId(version.oid);
+                  if (!isCheckingOutVersion) {
+                    handleVersionClick(version.oid);
+                  }
                 }}
               >
                 <div className="flex items-center justify-between">
@@ -115,6 +152,8 @@ export function VersionPane({ isVisible, onClose }: VersionPaneProps) {
                       await revertVersion({
                         versionId: version.oid,
                       });
+                      // Close the pane after revert to force a refresh on next open
+                      onClose();
                     }}
                     className={cn(
                       "invisible mt-1 flex items-center gap-1 px-2 py-0.5 text-sm font-medium bg-(--primary) text-(--primary-foreground) hover:bg-background-lightest rounded-md transition-colors",
diff --git a/src/hooks/useCheckoutVersion.ts b/src/hooks/useCheckoutVersion.ts
new file mode 100644
index 0000000..b470843
--- /dev/null
+++ b/src/hooks/useCheckoutVersion.ts
@@ -0,0 +1,38 @@
+import { useMutation, useQueryClient } from "@tanstack/react-query";
+import { IpcClient } from "@/ipc/ipc_client";
+
+interface CheckoutVersionVariables {
+  appId: number;
+  versionId: string;
+}
+
+export function useCheckoutVersion() {
+  const queryClient = useQueryClient();
+
+  const { isPending: isCheckingOutVersion, mutateAsync: checkoutVersion } =
+    useMutation<void, Error, CheckoutVersionVariables>({
+      mutationFn: async ({ appId, versionId }) => {
+        if (appId === null) {
+          // Should be caught by UI logic before calling, but as a safeguard.
+          throw new Error("App ID is null, cannot checkout version.");
+        }
+        const ipcClient = IpcClient.getInstance();
+        await ipcClient.checkoutVersion({ appId, versionId });
+      },
+      onSuccess: (_, variables) => {
+        // Invalidate queries that depend on the current version/branch
+        queryClient.invalidateQueries({
+          queryKey: ["currentBranch", variables.appId],
+        });
+        queryClient.invalidateQueries({
+          queryKey: ["versions", variables.appId],
+        });
+      },
+      meta: { showErrorToast: true },
+    });
+
+  return {
+    checkoutVersion,
+    isCheckingOutVersion,
+  };
+}
diff --git a/src/hooks/useVersions.ts b/src/hooks/useVersions.ts
index 42f6d7c..6526dc7 100644
--- a/src/hooks/useVersions.ts
+++ b/src/hooks/useVersions.ts
@@ -1,8 +1,8 @@
-import { useCallback, useEffect } from "react";
+import { useEffect } from "react";
 import { useAtom, useAtomValue } from "jotai";
 import { versionsListAtom } from "@/atoms/appAtoms";
 import { IpcClient } from "@/ipc/ipc_client";
-import { showError } from "@/lib/toast";
+
 import { chatMessagesAtom, selectedChatIdAtom } from "@/atoms/chatAtoms";
 import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
 import type { Version } from "@/ipc/ipc_types";
@@ -41,40 +41,35 @@ export function useVersions(appId: number | null) {
   const revertVersionMutation = useMutation<void, Error, { versionId: string }>(
     {
       mutationFn: async ({ versionId }: { versionId: string }) => {
-        if (appId === null) {
+        const currentAppId = appId;
+        if (currentAppId === null) {
           throw new Error("App ID is null");
         }
         const ipcClient = IpcClient.getInstance();
-        await ipcClient.revertVersion({ appId, previousVersionId: versionId });
+        await ipcClient.revertVersion({
+          appId: currentAppId,
+          previousVersionId: versionId,
+        });
       },
       onSuccess: async () => {
         await queryClient.invalidateQueries({ queryKey: ["versions", appId] });
+        await queryClient.invalidateQueries({
+          queryKey: ["currentBranch", appId],
+        });
         if (selectedChatId) {
           const chat = await IpcClient.getInstance().getChat(selectedChatId);
           setMessages(chat.messages);
         }
       },
-      onError: (e: Error) => {
-        showError(e);
-      },
+      meta: { showErrorToast: true },
     },
   );
 
-  const revertVersion = useCallback(
-    async ({ versionId }: { versionId: string }) => {
-      if (appId === null) {
-        return;
-      }
-      await revertVersionMutation.mutateAsync({ versionId });
-    },
-    [appId, revertVersionMutation],
-  );
-
   return {
     versions: versions || [],
     loading,
     error,
     refreshVersions,
-    revertVersion,
+    revertVersion: revertVersionMutation.mutateAsync,
   };
 }
diff --git a/src/ipc/handlers/version_handlers.ts b/src/ipc/handlers/version_handlers.ts
index d07821b..8e5cd96 100644
--- a/src/ipc/handlers/version_handlers.ts
+++ b/src/ipc/handlers/version_handlers.ts
@@ -94,7 +94,7 @@ export function registerVersionHandlers() {
         appId,
         previousVersionId,
       }: { appId: number; previousVersionId: string },
-    ) => {
+    ): Promise<void> => {
       return withLock(appId, async () => {
         const app = await db.query.apps.findFirst({
           where: eq(apps.id, appId),
@@ -205,8 +205,6 @@ export function registerVersionHandlers() {
                 );
             }
           }
-
-          return { success: true };
         } catch (error: any) {
           logger.error(
             `Error reverting to version ${previousVersionId} for app ${appId}:`,
@@ -220,7 +218,10 @@ export function registerVersionHandlers() {
 
   ipcMain.handle(
     "checkout-version",
-    async (_, { appId, versionId }: { appId: number; versionId: string }) => {
+    async (
+      _,
+      { appId, versionId }: { appId: number; versionId: string },
+    ): Promise<void> => {
       return withLock(appId, async () => {
         const app = await db.query.apps.findFirst({
           where: eq(apps.id, appId),
@@ -240,8 +241,6 @@ export function registerVersionHandlers() {
             ref: versionId,
             force: true,
           });
-
-          return { success: true };
         } catch (error: any) {
           logger.error(
             `Error checking out version ${versionId} for app ${appId}:`,
diff --git a/src/ipc/ipc_client.ts b/src/ipc/ipc_client.ts
index c51899b..5ef7a50 100644
--- a/src/ipc/ipc_client.ts
+++ b/src/ipc/ipc_client.ts
@@ -458,17 +458,11 @@ export class IpcClient {
   }: {
     appId: number;
     previousVersionId: string;
-  }): Promise<{ success: boolean }> {
-    try {
-      const result = await this.ipcRenderer.invoke("revert-version", {
-        appId,
-        previousVersionId,
-      });
-      return result;
-    } catch (error) {
-      showError(error);
-      throw error;
-    }
+  }): Promise<void> {
+    await this.ipcRenderer.invoke("revert-version", {
+      appId,
+      previousVersionId,
+    });
   }
 
   // Checkout a specific version without creating a revert commit
@@ -478,17 +472,11 @@ export class IpcClient {
   }: {
     appId: number;
     versionId: string;
-  }): Promise<{ success: boolean }> {
-    try {
-      const result = await this.ipcRenderer.invoke("checkout-version", {
-        appId,
-        versionId,
-      });
-      return result;
-    } catch (error) {
-      showError(error);
-      throw error;
-    }
+  }): Promise<void> {
+    await this.ipcRenderer.invoke("checkout-version", {
+      appId,
+      versionId,
+    });
   }
 
   // Get the current branch of an app
diff --git a/src/lib/assert.ts b/src/lib/assert.ts
new file mode 100644
index 0000000..eaa0bcd
--- /dev/null
+++ b/src/lib/assert.ts
@@ -0,0 +1,8 @@
+export function assertExists<T>(
+  value: T,
+  message: string,
+): asserts value is NonNullable<T> {
+  if (value === undefined || value === null) {
+    throw new Error(message);
+  }
+}
diff --git a/src/renderer.tsx b/src/renderer.tsx
index 08e2883..e504416 100644
--- a/src/renderer.tsx
+++ b/src/renderer.tsx
@@ -9,6 +9,7 @@ import {
   QueryCache,
   QueryClient,
   QueryClientProvider,
+  MutationCache,
 } from "@tanstack/react-query";
 import { showError } from "./lib/toast";
 
@@ -39,6 +40,13 @@ const queryClient = new QueryClient({
       }
     },
   }),
+  mutationCache: new MutationCache({
+    onError: (error, _variables, _context, mutation) => {
+      if (mutation.meta?.showErrorToast) {
+        showError(error);
+      }
+    },
+  }),
 });
 
 const posthogClient = posthog.init(