fix branching

src/db/repositories.js

@@ -218,6 +218,19 @@ export function createRepositories(db, config) {
             ORDER BY id DESC
             LIMIT ?
         `),
+        getChatBranchThroughMessage: db.prepare(`
+            WITH RECURSIVE branch AS (
+                SELECT * FROM chat_messages WHERE chat_id = ? AND id = ?
+                UNION ALL
+                SELECT parent.* FROM chat_messages parent
+                JOIN branch child ON child.parent_message_id = parent.id
+            )
+            SELECT id, role, content, model_name, retry_variants, retry_active_index, retry_retries_used,
+                   retry_prompt_message_id, delivery_status, error_message, parent_message_id, created_at
+            FROM branch
+            WHERE delivery_status = 'complete'
+            ORDER BY id
+        `),
         getActiveBranchMessages: db.prepare(`
             WITH RECURSIVE branch AS (
                 SELECT cm.* FROM chat_messages cm
@@ -946,6 +959,7 @@ export function createRepositories(db, config) {
         deleteChat: (chatId, username) => statements.deleteChatSession.run(chatId, username),
         listChatMessages: (chatId) => statements.getChatMessages.all(chatId),
         listActiveBranchMessages: (chatId) => statements.getActiveBranchMessages.all(chatId),
+        listChatBranchThroughMessage: (chatId, messageId) => statements.getChatBranchThroughMessage.all(chatId, messageId),
         listSiblingChatMessages: (chatId, role, parentMessageId) => statements.getSiblingChatMessages.all(chatId, role, parentMessageId ?? null),
         getLatestDescendantMessage: (chatId, messageId) => statements.getLatestDescendantMessage.get(chatId, messageId, chatId),
         getRecentChatMessages: (chatId) => statements.getRecentChatMessages.all(chatId, config.chatHistoryLimit),

src/routes/chats.js

@@ -191,6 +191,17 @@ export function createChatsRouter({repositories, chatService, modelService, conf
         return res.json(result);
     });
 
+    router.post("/chats/:id/messages/:messageId/branch", (req, res) => {
+        if (req.body?.confirmed !== true) return res.status(400).json({error: "Confirmation is required before creating a branch"});
+        const result = chatService.createBranchedChat({
+            user: req.session.user,
+            chatId: Number(req.params.id),
+            messageId: Number(req.params.messageId)
+        });
+        if (result.error) return res.status(result.status).json({error: result.error});
+        return res.json(result);
+    });
+
     router.post("/chats/:id/stop", (req, res) => {
         const result = chatService.stopGeneration({user: req.session.user, chatId: Number(req.params.id)});
         if (result.error) return res.status(result.status).json({error: result.error});
@@ -198,6 +209,7 @@ export function createChatsRouter({repositories, chatService, modelService, conf
     });
 
     router.post("/chats/:id/messages/:messageId/edit-resend", validateBody(sendMessageValidator), async (req, res) => {
+        if (req.body?.confirmed !== true) return res.status(400).json({error: "Confirmation is required before creating an edited branch"});
         const chatId = Number(req.params.id);
         const targetMessage = repositories.getChatMessage(chatId, Number(req.params.messageId));
         if (!repositories.getChat(chatId, req.session.user)) return res.status(404).json({error: "Chat not found"});

src/services/chatService.js

@@ -372,6 +372,70 @@ export function createChatService(repositories, modelService, config) {
         };
     });
 
+    const cloneChatThroughMessage = ({user, sourceChat, messageId = null}) => {
+        if (repositories.countChats(user) >= config.chatLimit) {
+            return {error: `Maximum of ${config.chatLimit} chats reached. Please delete an old chat to create a new one.`, status: 400};
+        }
+
+        const sourceMessages = messageId == null
+            ? []
+            : repositories.listChatBranchThroughMessage(sourceChat.id, messageId);
+        if (messageId != null && !sourceMessages.some((message) => message.id === messageId)) {
+            return {error: "Message not found", status: 404};
+        }
+
+        const sourceTitle = String(sourceChat.stored_title || sourceChat.title || "Chat").trim();
+        const clonedChat = repositories.createChat(
+            user,
+            `${sourceTitle} (branch)`,
+            sourceChat.assistant_persona_id ?? null,
+            sourceChat.user_persona_id ?? null,
+            sourceChat.scenario_prompt ?? null,
+            sourceChat.scenario_summary ?? null,
+            {folderName: sourceChat.folder_name ?? null}
+        );
+
+        try {
+            repositories.updateChatGenerationSettings(clonedChat.id, user, {
+                preferredModel: sourceChat.preferred_model ?? null,
+                temperature: sourceChat.temperature ?? null,
+                contextLength: sourceChat.context_length ?? null,
+                responseLength: sourceChat.response_length ?? null,
+                systemInstruction: sourceChat.system_instruction ?? null
+            });
+
+            let parentMessageId = null;
+            for (const sourceMessage of sourceMessages) {
+                const inserted = repositories.insertBranchChatMessage(
+                    clonedChat.id,
+                    sourceMessage.role,
+                    sourceMessage.content,
+                    {
+                        retryVariants: [sourceMessage.content],
+                        retryActiveIndex: 0,
+                        retryRetriesUsed: 0,
+                        retryPromptMessageId: null
+                    },
+                    sourceMessage.model_name ?? null,
+                    "complete",
+                    null,
+                    parentMessageId
+                );
+                parentMessageId = Number(inserted.lastInsertRowid);
+            }
+        } catch (error) {
+            repositories.deleteChat(clonedChat.id, user);
+            throw error;
+        }
+
+        return {
+            chat: repositories.getChat(clonedChat.id, user),
+            messages: listActiveBranchMessages(clonedChat.id),
+            sourceChatId: sourceChat.id,
+            branchedFromMessageId: messageId
+        };
+    };
+
     return {
         formatChatMessage,
         getChatMessageByIndex,
@@ -427,15 +491,37 @@ export function createChatService(repositories, modelService, config) {
         },
         async editAndResendMessage({user, chatId, targetMessage, content, model, onChunk = null, signal = null}) {
             if (!targetMessage || targetMessage.role !== "user") return {error: "Only user messages can be edited and resent", status: 400};
-            return generateChatTurn({
+            const sourceChat = repositories.getChat(chatId, user);
+            if (!sourceChat) return {error: "Chat not found", status: 404};
+            const cloned = cloneChatThroughMessage({
                 user,
-                chatId,
+                sourceChat,
+                messageId: targetMessage.parent_message_id ?? null
+            });
+            if (cloned.error) return cloned;
+
+            const generated = await generateChatTurn({
+                user,
+                chatId: cloned.chat.id,
                 message: content,
                 model,
                 onChunk,
-                signal,
-                parentMessageId: targetMessage.parent_message_id ?? null
+                signal
             });
+            if (generated.error) return {...generated, chat: cloned.chat};
+            return {
+                ...generated,
+                chat: repositories.getChat(cloned.chat.id, user),
+                messages: listActiveBranchMessages(cloned.chat.id),
+                sourceChatId: chatId,
+                branchedFromMessageId: targetMessage.id
+            };
+        },
+        createBranchedChat({user, chatId, messageId}) {
+            const sourceChat = repositories.getChat(chatId, user);
+            if (!sourceChat) return {error: "Chat not found", status: 404};
+            if (!repositories.getChatMessage(chatId, messageId)) return {error: "Message not found", status: 404};
+            return cloneChatThroughMessage({user, sourceChat, messageId});
         },
         activateBranch({user, chatId, messageId, includeDescendants = false}) {
             const session = repositories.getChat(chatId, user);