feat: implement branching functionality for chat threads

.prettierrc

@@ -0,0 +1,7 @@
+{
+    "singleQuote": false,
+    "tabWidth": 2,
+    "trailingComma": "all",
+    "semi": true,
+    "printWidth": 80
+}

src/app/(chat)/project/[id]/page.tsx

@@ -16,6 +16,7 @@ import {
   FileUp,
   Pencil,
   MessagesSquare,
+  GitBranch,
 } from "lucide-react";
 import { useTranslations } from "next-intl";
 import Link from "next/link";
@@ -225,8 +226,10 @@ export default function ProjectPage() {
                   >
                     <MessagesSquare size={16} className="text-primary" />
                     <div className="flex-1 truncate">
-                      <div className="font-medium truncate">
-                        {thread.title}
+                      <div className="font-medium truncate flex items-center">
+                        {thread.parentThreadId && (
+                          <GitBranch className="mr-2 h-3 w-3 text-muted-foreground flex-shrink-0" />
+                        )}
                         {thread.title}
                       </div>
                     </div>

src/app/api/chat/actions.ts

@@ -13,7 +13,12 @@ import {
   generateExampleToolSchemaPrompt,
 } from "lib/ai/prompts";
 
-import type { ChatModel, ChatThread, Project } from "app-types/chat";
+import type {
+  ChatMessage,
+  ChatModel,
+  ChatThread,
+  Project,
+} from "app-types/chat";
 
 import {
   chatRepository,
@@ -31,6 +36,7 @@ import logger from "logger";
 import { JSONSchema7 } from "json-schema";
 import { ObjectJsonSchema7 } from "app-types/util";
 import { jsonSchemaToZod } from "lib/json-schema-to-zod";
+import { randomUUID } from "crypto";
 
 export async function getUserId() {
   const session = await getSession();
@@ -44,7 +50,10 @@ export async function getUserId() {
 export async function generateTitleFromUserMessageAction({
   message,
   model,
-}: { message: Message; model: LanguageModel }) {
+}: {
+  message: Message;
+  model: LanguageModel;
+}) {
   await getSession();
   const prompt = toAny(message.parts?.at(-1))?.text || "unknown";
 
@@ -87,6 +96,64 @@ export async function deleteMessagesByChatIdAfterTimestampAction(
   await chatRepository.deleteMessagesByChatIdAfterTimestamp(messageId);
 }
 
+export async function branchOutAction(
+  threadId: string,
+  messageId: string,
+): Promise<{
+  id: string;
+}> {
+  const userId = await getUserId();
+  console.log("userId", userId);
+
+  if (!userId) {
+    throw new Error("User not found");
+  }
+
+  const threadDetails = await chatRepository.selectThreadDetails(threadId);
+  if (!threadDetails) {
+    throw new Error("Thread not found");
+  }
+
+  const isMessageInThread = threadDetails.messages.some(
+    (message) => message.id === messageId,
+  );
+  if (!isMessageInThread) {
+    throw new Error("Message not found in thread");
+  }
+
+  const messagesForNewThread: ChatMessage[] = [];
+
+  for (const message of threadDetails.messages) {
+    if (message.id === messageId) {
+      messagesForNewThread.push(message);
+      break;
+    }
+    messagesForNewThread.push(message);
+  }
+
+  const newThread = await chatRepository.insertThread({
+    title: `Branch - ${threadDetails.title}`,
+    userId: threadDetails.userId,
+    projectId: threadDetails.projectId,
+    id: randomUUID(),
+    parentThreadId: threadDetails.id,
+  });
+
+  await chatRepository.insertMessages(
+    messagesForNewThread.map((message) => ({
+      role: message.role,
+      parts: message.parts,
+      model: message.model,
+      attachments: message.attachments,
+      annotations: message.annotations,
+      threadId: newThread.id,
+      id: randomUUID(),
+    })),
+  );
+
+  return { id: newThread.id };
+}
+
 export async function selectThreadListByUserIdAction() {
   const userId = await getUserId();
   const threads = await chatRepository.selectThreadsByUserId(userId);

src/app/api/chat/branch/route.ts

@@ -0,0 +1,14 @@
+import { NextRequest } from "next/server";
+
+export async function POST(_request: NextRequest) {
+  try {
+  } catch (error: any) {
+    console.error("Error:", error);
+    return new Response(
+      JSON.stringify({ error: error.message || "Internal Server Error" }),
+      {
+        status: 500,
+      },
+    );
+  }
+}

src/components/layouts/app-header.tsx

@@ -9,6 +9,7 @@ import {
   ChevronRight,
   MessageCircleDashed,
   PanelLeft,
+  GitBranch,
 } from "lucide-react";
 import { Button } from "ui/button";
 import { Separator } from "ui/separator";
@@ -187,6 +188,9 @@ function ThreadDropdownComponent() {
           variant="ghost"
           className="data-[state=open]:bg-input! hover:text-foreground cursor-pointer flex gap-1 items-center px-2 py-1 rounded-md hover:bg-accent"
         >
+          {currentThread.parentThreadId && (
+            <GitBranch className="mr-1 h-3 w-3 text-muted-foreground flex-shrink-0" />
+          )}
           <p className="truncate max-w-60 min-w-0">{currentThread.title}</p>
 
           <ChevronDown size={14} />

src/components/layouts/app-sidebar-threads.tsx

@@ -11,7 +11,13 @@ import {
 import { SidebarGroupContent, SidebarMenu, SidebarMenuItem } from "ui/sidebar";
 import { SidebarGroup } from "ui/sidebar";
 import { ThreadDropdown } from "../thread-dropdown";
-import { ChevronDown, ChevronUp, MoreHorizontal, Trash } from "lucide-react";
+import {
+  ChevronDown,
+  ChevronUp,
+  MoreHorizontal,
+  Trash,
+  GitBranch,
+} from "lucide-react";
 import { useMounted } from "@/hooks/use-mounted";
 import { appStore } from "@/app/store";
 import { Button } from "ui/button";
@@ -234,6 +240,9 @@ export function AppSidebarThreads() {
                                 href={`/chat/${thread.id}`}
                                 className="flex items-center"
                               >
+                                {thread.parentThreadId && (
+                                  <GitBranch className="mr-2 h-3 w-3 text-muted-foreground flex-shrink-0" />
+                                )}
                                 <p className="truncate ">{thread.title}</p>
                               </Link>
                             </SidebarMenuButton>

src/components/message-parts.tsx

@@ -16,6 +16,7 @@ import {
   Loader2,
   AlertTriangleIcon,
   Percent,
+  GitBranch,
   HammerIcon,
 } from "lucide-react";
 import { Tooltip, TooltipContent, TooltipTrigger } from "ui/tooltip";
@@ -39,6 +40,7 @@ import { useCopy } from "@/hooks/use-copy";
 import { AnimatePresence, motion } from "framer-motion";
 import { SelectModel } from "./select-model";
 import {
+  branchOutAction,
   deleteMessageAction,
   deleteMessagesByChatIdAfterTimestampAction,
 } from "@/app/api/chat/actions";
@@ -75,6 +77,7 @@ import { TavilyResponse } from "lib/ai/tools/web/web-search";
 
 import { CodeBlock } from "ui/CodeBlock";
 import { SafeJsExecutionResult, safeJsRun } from "lib/safe-js-run";
+import { useRouter } from "next/navigation";
 
 type MessagePart = UIMessage["parts"][number];
 
@@ -247,6 +250,9 @@ export const AssistMessagePart = memo(function AssistMessagePart({
   const { copied, copy } = useCopy();
   const [isLoading, setIsLoading] = useState(false);
   const [isDeleting, setIsDeleting] = useState(false);
+  const [isBranching, setIsBranching] = useState(false);
+
+  const router = useRouter();
 
   const deleteMessage = useCallback(() => {
     safe(() => setIsDeleting(true))
@@ -295,6 +301,20 @@ export const AssistMessagePart = memo(function AssistMessagePart({
       .unwrap();
   };
 
+  const handleBranchOut = useCallback(async () => {
+    safe(() => setIsBranching(true))
+      .ifOk(async () => {
+        if (!threadId) {
+          throw new Error("Thread ID is required");
+        }
+        const newThread = await branchOutAction(threadId, message.id);
+        router.push(`/chat/${newThread.id}`);
+      })
+      .ifFail((error) => toast.error(error.message))
+      .watch(() => setIsBranching(false))
+      .unwrap();
+  }, [message.id]);
+
   return (
     <div
       className={cn(isLoading && "animate-pulse", "flex flex-col gap-2 group")}
@@ -344,6 +364,24 @@ export const AssistMessagePart = memo(function AssistMessagePart({
             </TooltipTrigger>
             <TooltipContent>Change Model</TooltipContent>
           </Tooltip>
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <Button
+                variant="ghost"
+                size="icon"
+                disabled={isBranching}
+                onClick={handleBranchOut}
+                className="size-3! p-4! opacity-0 group-hover/message:opacity-100"
+              >
+                {isBranching ? (
+                  <Loader className="animate-spin" />
+                ) : (
+                  <GitBranch />
+                )}
+              </Button>
+            </TooltipTrigger>
+            <TooltipContent>Branch Out</TooltipContent>
+          </Tooltip>
           <Tooltip>
             <TooltipTrigger asChild>
               <Button

src/lib/db/migrations/pg/0008_amazing_chat.sql

@@ -0,0 +1 @@
+ALTER TABLE "chat_thread" ADD COLUMN "parent_thread_id" uuid;