Add warning banner if user is not on main branch (#91)

src/components/chat/ChatHeader.tsx

@@ -1,14 +1,30 @@
-import { PanelRightOpen, History, PlusCircle } from "lucide-react";
+import {
+  PanelRightOpen,
+  History,
+  PlusCircle,
+  GitBranch,
+  AlertCircle,
+  Info,
+} from "lucide-react";
 import { PanelRightClose } from "lucide-react";
-import { useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue, useSetAtom } from "jotai";
 import { selectedAppIdAtom } from "@/atoms/appAtoms";
 import { useVersions } from "@/hooks/useVersions";
 import { Button } from "../ui/button";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger,
+} from "../ui/tooltip";
 import { IpcClient } from "@/ipc/ipc_client";
 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 { BranchResult } from "@/ipc/ipc_types";
+import { useStreamChat } from "@/hooks/useStreamChat";
 
 interface ChatHeaderProps {
   isPreviewOpen: boolean;
@@ -24,8 +40,59 @@ export function ChatHeader({
   const appId = useAtomValue(selectedAppIdAtom);
   const { versions, loading } = useVersions(appId);
   const { navigate } = useRouter();
-  const setSelectedChatId = useSetAtom(selectedChatIdAtom);
+  const [selectedChatId, setSelectedChatId] = useAtom(selectedChatIdAtom);
   const { refreshChats } = useChats(appId);
+  const [branchInfo, setBranchInfo] = useState<BranchResult | null>(null);
+  const [checkingOutMain, setCheckingOutMain] = useState(false);
+  const { isStreaming } = useStreamChat();
+
+  // Fetch the current branch when appId changes
+  useEffect(() => {
+    if (!appId) return;
+
+    const fetchBranch = async () => {
+      try {
+        const result = await IpcClient.getInstance().getCurrentBranch(appId);
+        if (result.success) {
+          setBranchInfo(result);
+        } else {
+          showError("Failed to get current branch: " + result.errorMessage);
+        }
+      } catch (error) {
+        showError(`Failed to get current branch: ${error}`);
+      }
+    };
+
+    fetchBranch();
+    // The use of selectedChatId and isStreaming is a hack to ensure that
+    // the branch info is relatively up to date.
+  }, [appId, selectedChatId, isStreaming]);
+
+  const handleCheckoutMainBranch = async () => {
+    if (!appId) return;
+
+    try {
+      setCheckingOutMain(true);
+      // Find the latest commit on main branch
+      // For simplicity, we'll just checkout to "main" directly
+      await IpcClient.getInstance().checkoutVersion({
+        appId,
+        versionId: "main",
+      });
+
+      // Refresh branch info
+      const result = await IpcClient.getInstance().getCurrentBranch(appId);
+      if (result.success) {
+        setBranchInfo(result);
+      } else {
+        showError(result.errorMessage);
+      }
+    } catch (error) {
+      showError(`Failed to checkout main branch: ${error}`);
+    } finally {
+      setCheckingOutMain(false);
+    }
+  };
 
   const handleNewChat = async () => {
     // Only create a new chat if an app is selected
@@ -54,37 +121,83 @@ export function ChatHeader({
   };
   // TODO: KEEP UP TO DATE WITH app_handlers.ts
   const versionPostfix = versions.length === 10_000 ? `+` : "";
-  return (
-    <div className="@container flex items-center justify-between py-1.5">
-      <div className="flex items-center space-x-2">
-        <Button
-          onClick={handleNewChat}
-          variant="ghost"
-          className="hidden @2xs:flex items-center justify-start gap-2 mx-2 py-3"
-        >
-          <PlusCircle size={16} />
-          <span>New Chat</span>
-        </Button>
-        <Button
-          onClick={onVersionClick}
-          variant="ghost"
-          className="hidden @6xs:flex cursor-pointer items-center gap-1 text-sm px-2 py-1 rounded-md"
-        >
-          <History size={16} />
-          {loading ? "..." : `Version ${versions.length}${versionPostfix}`}
-        </Button>
-      </div>
 
-      <button
-        onClick={onTogglePreview}
-        className="cursor-pointer p-2 hover:bg-(--background-lightest) rounded-md"
-      >
-        {isPreviewOpen ? (
-          <PanelRightClose size={20} />
-        ) : (
-          <PanelRightOpen size={20} />
-        )}
-      </button>
+  // Check if we're not on the main branch
+  const isNotMainBranch =
+    branchInfo?.success && branchInfo.data.branch !== "main";
+
+  return (
+    <div className="flex flex-col w-full @container">
+      {isNotMainBranch && (
+        <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} />
+            <span>
+              {branchInfo?.data.branch === "<no-branch>" && (
+                <>
+                  <TooltipProvider>
+                    <Tooltip>
+                      <TooltipTrigger asChild>
+                        <span className="flex items-center  gap-1">
+                          <strong>Warning:</strong>
+                          <span>You are not on a branch</span>
+                          <Info size={14} />
+                        </span>
+                      </TooltipTrigger>
+                      <TooltipContent>
+                        <p>
+                          Checkout main branch, otherwise changes will not be
+                          saved properly
+                        </p>
+                      </TooltipContent>
+                    </Tooltip>
+                  </TooltipProvider>
+                </>
+              )}
+            </span>
+          </div>
+          <Button
+            variant="outline"
+            size="sm"
+            onClick={handleCheckoutMainBranch}
+            disabled={checkingOutMain}
+          >
+            {checkingOutMain ? "Checking out..." : "Switch to main branch"}
+          </Button>
+        </div>
+      )}
+
+      <div className="@container flex items-center justify-between py-1.5">
+        <div className="flex items-center space-x-2">
+          <Button
+            onClick={handleNewChat}
+            variant="ghost"
+            className="hidden @2xs:flex items-center justify-start gap-2 mx-2 py-3"
+          >
+            <PlusCircle size={16} />
+            <span>New Chat</span>
+          </Button>
+          <Button
+            onClick={onVersionClick}
+            variant="ghost"
+            className="hidden @6xs:flex cursor-pointer items-center gap-1 text-sm px-2 py-1 rounded-md"
+          >
+            <History size={16} />
+            {loading ? "..." : `Version ${versions.length}${versionPostfix}`}
+          </Button>
+        </div>
+
+        <button
+          onClick={onTogglePreview}
+          className="cursor-pointer p-2 hover:bg-(--background-lightest) rounded-md"
+        >
+          {isPreviewOpen ? (
+            <PanelRightClose size={20} />
+          ) : (
+            <PanelRightOpen size={20} />
+          )}
+        </button>
+      </div>
     </div>
   );
 }

src/ipc/handlers/version_handlers.ts

@@ -2,7 +2,7 @@ import { ipcMain } from "electron";
 import { db } from "../../db";
 import { apps, messages } from "../../db/schema";
 import { desc, eq, and, gt } from "drizzle-orm";
-import type { Version } from "../ipc_types";
+import type { Version, BranchResult } from "../ipc_types";
 import fs from "node:fs";
 import path from "node:path";
 import { getDyadAppPath } from "../../paths/paths";
@@ -51,6 +51,53 @@ export function registerVersionHandlers() {
     }
   });
 
+  ipcMain.handle(
+    "get-current-branch",
+    async (_, { appId }: { appId: number }): Promise<BranchResult> => {
+      const app = await db.query.apps.findFirst({
+        where: eq(apps.id, appId),
+      });
+
+      if (!app) {
+        return {
+          success: false,
+          errorMessage: "App not found",
+        };
+      }
+
+      const appPath = getDyadAppPath(app.path);
+
+      // Return appropriate result if the app is not a git repo
+      if (!fs.existsSync(path.join(appPath, ".git"))) {
+        return {
+          success: false,
+          errorMessage: "Not a git repository",
+        };
+      }
+
+      try {
+        const currentBranch = await git.currentBranch({
+          fs,
+          dir: appPath,
+          fullname: false,
+        });
+
+        return {
+          success: true,
+          data: {
+            branch: currentBranch || "<no-branch>",
+          },
+        };
+      } catch (error: any) {
+        logger.error(`Error getting current branch for app ${appId}:`, error);
+        return {
+          success: false,
+          errorMessage: `Failed to get current branch: ${error.message}`,
+        };
+      }
+    }
+  );
+
   ipcMain.handle(
     "revert-version",
     async (

src/ipc/ipc_client.ts

@@ -22,6 +22,7 @@ import type {
   TokenCountParams,
   TokenCountResult,
   ChatLogsData,
+  BranchResult,
 } from "./ipc_types";
 import type { CodeProposal, ProposalResult } from "@/lib/schemas";
 import { showError } from "@/lib/toast";
@@ -478,6 +479,14 @@ export class IpcClient {
     }
   }
 
+  // Get the current branch of an app
+  public async getCurrentBranch(appId: number): Promise<BranchResult> {
+    const result = await this.ipcRenderer.invoke("get-current-branch", {
+      appId,
+    });
+    return result;
+  }
+
   // Get user settings
   public async getUserSettings(): Promise<UserSettings> {
     try {

src/ipc/ipc_types.ts

@@ -74,6 +74,18 @@ export interface Version {
   timestamp: number;
 }
 
+export type Result<T> =
+  | {
+      success: true;
+      data: T;
+    }
+  | {
+      success: false;
+      errorMessage: string;
+    };
+
+export type BranchResult = Result<{ branch: string }>;
+
 export interface SandboxConfig {
   files: Record<string, string>;
   dependencies: Record<string, string>;

src/preload.ts

@@ -25,6 +25,7 @@ const validInvokeChannels = [
   "list-versions",
   "revert-version",
   "checkout-version",
+  "get-current-branch",
   "delete-app",
   "rename-app",
   "get-user-settings",