From 6e59b72dcaddc6b14d5482ed2dbad2b4ccf3c30d Mon Sep 17 00:00:00 2001 From: fretuks Date: Tue, 8 Sep 2026 23:09:27 +0200 Subject: [PATCH] fix support for chat branching with source tracking --- public/script.js | 24 +++++++++++++++++------- src/db/repositories.js | 22 +++++++++++++++++----- src/db/schema.js | 8 +++++++- src/services/chatService.js | 1 + test/app.test.js | 3 +++ 5 files changed, 45 insertions(+), 13 deletions(-) diff --git a/public/script.js b/public/script.js index ca7bfe2..347d97c 100644 --- a/public/script.js +++ b/public/script.js @@ -460,6 +460,19 @@ function ensureRetryState(message) { return state; } +async function openCreatedChat(chat) { + if (!chat?.id) throw new Error("The server did not return the newly created chat"); + if (chatSearchInput) chatSearchInput.value = ""; + const chatsResponse = await get("/chats"); + if (chatsResponse.error) throw new Error(chatsResponse.error); + chatSessions = chatsResponse.chats || []; + if (!chatSessions.some((item) => item.id === chat.id)) { + throw new Error("The newly created chat was not persisted"); + } + renderChatList(); + await setActiveChat(chat.id); +} + async function editChatMessage(messageId, fallbackIndex = -1) { const target = await resolveMessageTarget(messageId, fallbackIndex); if (!target) return setNotice("Message is not ready yet. Try again.", "error"); @@ -496,9 +509,7 @@ async function editChatMessage(messageId, fallbackIndex = -1) { confirmed: true }); if (res.error) return setNotice(res.error, "error"); - chatSessions = [res.chat, ...chatSessions.filter((chat) => chat.id !== res.chat.id)]; - renderChatList(); - await setActiveChat(res.chat.id); + await openCreatedChat(res.chat); return setNotice("Created and opened the edited chat branch.", "success"); } finally { setLoadingState(false); @@ -612,9 +623,7 @@ async function branchFromMessage(messageId) { const sourceChatId = activeChatId; const res = await post(`/chats/${sourceChatId}/messages/${messageId}/branch`, {confirmed: true}); if (res.error) return setNotice(res.error, "error"); - chatSessions = [res.chat, ...chatSessions.filter((chat) => chat.id !== res.chat.id)]; - renderChatList(); - await setActiveChat(res.chat.id); + await openCreatedChat(res.chat); msgInput.focus(); setNotice("Created and opened a new chat branch.", "success"); } @@ -1056,9 +1065,10 @@ function renderChatList() { content.className = "chat-item-copy"; title.textContent = chat.title; meta.className = "chat-item-meta"; - meta.textContent = chat.assistant_persona_id + const chatType = chat.assistant_persona_id ? `Roleplay${chat.user_persona_name ? ` · ${chat.user_persona_name}` : ""}` : "Assistant chat"; + meta.textContent = chat.source_chat_id ? `Branch · ${chatType}` : chatType; remove.type = "button"; remove.className = "chat-item-delete"; remove.textContent = "Delete"; diff --git a/src/db/repositories.js b/src/db/repositories.js index 7ad5b59..a86a2e3 100644 --- a/src/db/repositories.js +++ b/src/db/repositories.js @@ -51,7 +51,7 @@ export function createRepositories(db, config) { listChatSessions: db.prepare(` SELECT cs.id, CASE - WHEN cs.assistant_persona_id IS NOT NULL THEN COALESCE(p.name, cs.title) + WHEN cs.assistant_persona_id IS NOT NULL AND cs.source_chat_id IS NULL THEN COALESCE(p.name, cs.title) ELSE cs.title END AS title, cs.title AS stored_title, @@ -67,6 +67,8 @@ export function createRepositories(db, config) { cs.context_summary, cs.context_summary_message_id, cs.active_leaf_message_id, + cs.source_chat_id, + cs.branched_from_message_id, cs.preferred_model, cs.temperature, cs.context_length, @@ -84,7 +86,7 @@ export function createRepositories(db, config) { getChatSession: db.prepare(` SELECT cs.id, CASE - WHEN cs.assistant_persona_id IS NOT NULL THEN COALESCE(p.name, cs.title) + WHEN cs.assistant_persona_id IS NOT NULL AND cs.source_chat_id IS NULL THEN COALESCE(p.name, cs.title) ELSE cs.title END AS title, cs.title AS stored_title, @@ -100,6 +102,8 @@ export function createRepositories(db, config) { cs.context_summary, cs.context_summary_message_id, cs.active_leaf_message_id, + cs.source_chat_id, + cs.branched_from_message_id, cs.preferred_model, cs.temperature, cs.context_length, @@ -116,7 +120,7 @@ export function createRepositories(db, config) { getChatSessionByParticipants: db.prepare(` SELECT cs.id, CASE - WHEN cs.assistant_persona_id IS NOT NULL THEN COALESCE(p.name, cs.title) + WHEN cs.assistant_persona_id IS NOT NULL AND cs.source_chat_id IS NULL THEN COALESCE(p.name, cs.title) ELSE cs.title END AS title, cs.title AS stored_title, @@ -132,6 +136,8 @@ export function createRepositories(db, config) { cs.context_summary, cs.context_summary_message_id, cs.active_leaf_message_id, + cs.source_chat_id, + cs.branched_from_message_id, cs.preferred_model, cs.temperature, cs.context_length, @@ -145,6 +151,7 @@ export function createRepositories(db, config) { WHERE cs.username = ? AND cs.assistant_persona_id = ? AND ((cs.user_persona_id IS NULL AND ? IS NULL) OR cs.user_persona_id = ?) + AND cs.source_chat_id IS NULL ORDER BY cs.is_pinned DESC, datetime(cs.updated_at) DESC LIMIT 1 `), @@ -157,6 +164,7 @@ export function createRepositories(db, config) { WHERE id = ? AND username = ? `), updateChatActiveLeaf: db.prepare("UPDATE chat_sessions SET active_leaf_message_id = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ? AND username = ?"), + updateChatBranchSource: db.prepare("UPDATE chat_sessions SET source_chat_id = ?, branched_from_message_id = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ? AND username = ?"), setChatActiveLeafById: db.prepare("UPDATE chat_sessions SET active_leaf_message_id = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?"), updateChatSessionOrganization: db.prepare(` UPDATE chat_sessions @@ -367,7 +375,7 @@ export function createRepositories(db, config) { searchChats: db.prepare(` SELECT DISTINCT cs.id, CASE - WHEN cs.assistant_persona_id IS NOT NULL THEN COALESCE(p.name, cs.title) + WHEN cs.assistant_persona_id IS NOT NULL AND cs.source_chat_id IS NULL THEN COALESCE(p.name, cs.title) ELSE cs.title END AS title, cs.title AS stored_title, @@ -383,6 +391,8 @@ export function createRepositories(db, config) { cs.context_summary, cs.context_summary_message_id, cs.active_leaf_message_id, + cs.source_chat_id, + cs.branched_from_message_id, cs.preferred_model, cs.temperature, cs.context_length, @@ -398,7 +408,7 @@ export function createRepositories(db, config) { WHERE cs.username = ? AND ( LOWER(CASE - WHEN cs.assistant_persona_id IS NOT NULL THEN COALESCE(p.name, cs.title) + WHEN cs.assistant_persona_id IS NOT NULL AND cs.source_chat_id IS NULL THEN COALESCE(p.name, cs.title) ELSE cs.title END) LIKE ? OR LOWER(COALESCE(cm.content, '')) LIKE ? @@ -1006,6 +1016,8 @@ export function createRepositories(db, config) { failChatTurn: (chatId, userMessageId, errorMessage, parentMessageId = null) => failChatTurnTransaction(chatId, userMessageId, errorMessage, parentMessageId), updateChatActiveLeaf: (chatId, username, messageId) => statements.updateChatActiveLeaf.run(messageId, chatId, username), + updateChatBranchSource: (chatId, username, sourceChatId, messageId) => + statements.updateChatBranchSource.run(sourceChatId, messageId ?? null, chatId, username), updateChatGenerationSettings: (chatId, username, settings) => statements.updateChatGenerationSettings.run( settings.preferredModel ?? null, settings.temperature ?? null, diff --git a/src/db/schema.js b/src/db/schema.js index 7ce2816..1165046 100644 --- a/src/db/schema.js +++ b/src/db/schema.js @@ -27,6 +27,8 @@ export function initializeSchema(db) { context_summary TEXT, context_summary_message_id INTEGER DEFAULT 0, active_leaf_message_id INTEGER, + source_chat_id INTEGER, + branched_from_message_id INTEGER, preferred_model TEXT, temperature REAL, context_length INTEGER, @@ -35,7 +37,9 @@ export function initializeSchema(db) { created_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (assistant_persona_id) REFERENCES personas (id) ON DELETE SET NULL, - FOREIGN KEY (user_persona_id) REFERENCES personas (id) ON DELETE SET NULL + FOREIGN KEY (user_persona_id) REFERENCES personas (id) ON DELETE SET NULL, + FOREIGN KEY (source_chat_id) REFERENCES chat_sessions (id) ON DELETE SET NULL, + FOREIGN KEY (branched_from_message_id) REFERENCES chat_messages (id) ON DELETE SET NULL ) `).run(); @@ -293,6 +297,8 @@ export function runMigrations(db, {dropLegacyChats = true} = {}) { ensureColumn("chat_sessions", chatSessionColumns, "context_summary", "ALTER TABLE chat_sessions ADD COLUMN context_summary TEXT"); ensureColumn("chat_sessions", chatSessionColumns, "context_summary_message_id", "ALTER TABLE chat_sessions ADD COLUMN context_summary_message_id INTEGER DEFAULT 0"); ensureColumn("chat_sessions", chatSessionColumns, "active_leaf_message_id", "ALTER TABLE chat_sessions ADD COLUMN active_leaf_message_id INTEGER"); + ensureColumn("chat_sessions", chatSessionColumns, "source_chat_id", "ALTER TABLE chat_sessions ADD COLUMN source_chat_id INTEGER"); + ensureColumn("chat_sessions", chatSessionColumns, "branched_from_message_id", "ALTER TABLE chat_sessions ADD COLUMN branched_from_message_id INTEGER"); ensureColumn("chat_sessions", chatSessionColumns, "preferred_model", "ALTER TABLE chat_sessions ADD COLUMN preferred_model TEXT"); ensureColumn("chat_sessions", chatSessionColumns, "temperature", "ALTER TABLE chat_sessions ADD COLUMN temperature REAL"); ensureColumn("chat_sessions", chatSessionColumns, "context_length", "ALTER TABLE chat_sessions ADD COLUMN context_length INTEGER"); diff --git a/src/services/chatService.js b/src/services/chatService.js index ad284c1..6e65a02 100644 --- a/src/services/chatService.js +++ b/src/services/chatService.js @@ -396,6 +396,7 @@ export function createChatService(repositories, modelService, config) { ); try { + repositories.updateChatBranchSource(clonedChat.id, user, sourceChat.id, messageId); repositories.updateChatGenerationSettings(clonedChat.id, user, { preferredModel: sourceChat.preferred_model ?? null, temperature: sourceChat.temperature ?? null, diff --git a/test/app.test.js b/test/app.test.js index 8fd6b5d..aae3999 100644 --- a/test/app.test.js +++ b/test/app.test.js @@ -466,6 +466,9 @@ test("conversation branches, edit-resend, and generation settings work", async ( }); assert.equal(branched.status, 200); assert.notEqual(branched.json.chat.id, chatId); + assert.equal(branched.json.chat.source_chat_id, chatId); + assert.equal(branched.json.chat.branched_from_message_id, firstAssistant.id); + assert.match(branched.json.chat.title, /\(branch\)$/); 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);