From c0e3a6b693f8970e8bbb1bbe0c7b15ef1ef13daa Mon Sep 17 00:00:00 2001 From: fretuks Date: Sat, 5 Sep 2026 21:02:03 +0200 Subject: [PATCH] fix branching --- public/script.js | 40 ++++++++++++---- src/db/repositories.js | 14 ++++++ src/routes/chats.js | 12 +++++ src/services/chatService.js | 94 +++++++++++++++++++++++++++++++++++-- test/app.test.js | 23 +++++++-- 5 files changed, 166 insertions(+), 17 deletions(-) diff --git a/public/script.js b/public/script.js index c4dd685..ca7bfe2 100644 --- a/public/script.js +++ b/public/script.js @@ -475,20 +475,31 @@ async function editChatMessage(messageId, fallbackIndex = -1) { const content = (nextContent || "").trim(); if (!content) return; if (message.role === "user") { + const confirmed = await confirmPopup({ + eyebrow: "Create branch", + title: "Create an edited chat branch?", + description: "This creates a new chat from the messages before this one and sends your edited message there. The original chat stays unchanged.", + confirmLabel: "Create branch" + }); + if (!confirmed) return; + const sourceChatId = activeChatId; setLoadingState(true, { eyebrow: "Branching", title: "Editing and resending", - detail: "Creating a new branch from this point." + detail: "Creating and opening a new chat branch." }); try { - const res = await post(`/chats/${activeChatId}/messages/${message.id}/edit-resend`, { - chatId: activeChatId, + const res = await post(`/chats/${sourceChatId}/messages/${message.id}/edit-resend`, { + chatId: sourceChatId, message: content, - model: modelSelect.value + model: modelSelect.value, + confirmed: true }); if (res.error) return setNotice(res.error, "error"); - await setActiveChat(activeChatId); - return setNotice("Created a new branch from the edited message.", "success"); + chatSessions = [res.chat, ...chatSessions.filter((chat) => chat.id !== res.chat.id)]; + renderChatList(); + await setActiveChat(res.chat.id); + return setNotice("Created and opened the edited chat branch.", "success"); } finally { setLoadingState(false); } @@ -591,12 +602,21 @@ async function switchRetryVariant(messageId, direction, fallbackIndex = -1) { async function branchFromMessage(messageId) { if (!activeChatId || !messageId) return; - const res = await post(`/chats/${activeChatId}/branches/${messageId}/activate`, {}); + const confirmed = await confirmPopup({ + eyebrow: "Create branch", + title: "Create a new chat branch?", + description: "The new chat will contain the conversation through this message. The original chat stays unchanged.", + confirmLabel: "Create branch" + }); + if (!confirmed) return; + const sourceChatId = activeChatId; + const res = await post(`/chats/${sourceChatId}/messages/${messageId}/branch`, {confirmed: true}); if (res.error) return setNotice(res.error, "error"); - currentMessages = res.messages || []; - renderMessages(); + chatSessions = [res.chat, ...chatSessions.filter((chat) => chat.id !== res.chat.id)]; + renderChatList(); + await setActiveChat(res.chat.id); msgInput.focus(); - setNotice("Branch point selected. Your next message continues from here.", "success"); + setNotice("Created and opened a new chat branch.", "success"); } function addMessage(contentOrMessage, isUser = false, isLoading = false, options = {}) { diff --git a/src/db/repositories.js b/src/db/repositories.js index 26ff2d0..7ad5b59 100644 --- a/src/db/repositories.js +++ b/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), diff --git a/src/routes/chats.js b/src/routes/chats.js index 8e4cb0a..14230fc 100644 --- a/src/routes/chats.js +++ b/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"}); diff --git a/src/services/chatService.js b/src/services/chatService.js index 000b2ac..ad284c1 100644 --- a/src/services/chatService.js +++ b/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); diff --git a/test/app.test.js b/test/app.test.js index 649cceb..8fd6b5d 100644 --- a/test/app.test.js +++ b/test/app.test.js @@ -456,13 +456,30 @@ test("conversation branches, edit-resend, and generation settings work", async ( active = await request(server.baseUrl, `/chats/${chatId}/messages`, {jar}); assert.equal(active.json.messages.length, 4); + const unconfirmedBranch = await request(server.baseUrl, `/chats/${chatId}/messages/${firstAssistant.id}/branch`, { + method: "POST", body: {}, jar + }); + assert.equal(unconfirmedBranch.status, 400); + + const branched = await request(server.baseUrl, `/chats/${chatId}/messages/${firstAssistant.id}/branch`, { + method: "POST", body: {confirmed: true}, jar + }); + assert.equal(branched.status, 200); + assert.notEqual(branched.json.chat.id, chatId); + assert.deepEqual(branched.json.messages.map((message) => message.content), ["first", "reply-1"]); + assert.equal(branched.json.chat.preferred_model, "preferred:model"); + assert.equal(branched.json.chat.temperature, 0.4); + const edited = await request(server.baseUrl, `/chats/${chatId}/messages/${active.json.messages[0].id}/edit-resend`, { - method: "POST", body: {chatId, message: "edited first", model: "ignored:model"}, jar + method: "POST", body: {chatId, message: "edited first", model: "ignored:model", confirmed: true}, jar }); assert.equal(edited.status, 200); + assert.notEqual(edited.json.chat.id, chatId); + assert.notEqual(edited.json.chat.id, branched.json.chat.id); + assert.deepEqual(edited.json.messages.map((message) => message.content), ["edited first", "reply-4"]); + active = await request(server.baseUrl, `/chats/${chatId}/messages`, {jar}); - assert.deepEqual(active.json.messages.map((message) => message.content), ["edited first", "reply-4"]); - assert.equal(active.json.messages[0].siblingCount, 2); + assert.deepEqual(active.json.messages.map((message) => message.content), ["first", "reply-1", "second", "reply-2"]); await server.close(); });