fix support for chat branching with source tracking

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)];
+            await openCreatedChat(res.chat);
-            renderChatList();
-            await setActiveChat(res.chat.id);
             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)];
+    await openCreatedChat(res.chat);
-    renderChatList();
-    await setActiveChat(res.chat.id);
     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";

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,

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");

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,

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);