import { describe, expect, test } from "bun:test"; import { buildFocusedMessageTree, buildMessageTree, getActiveMessage, getCurrentChildIndex, getMessageCount, getNextActiveChildId, } from "../utils/message"; describe("message focus and versioning", () => { test("reproduces and fixes the 3/2 focus bug after deleting root message and regenerating", () => { // Step 1: Generate initial assistant message (A) const msgA = { id: "msg_a", userId: "user_1", topicId: "topic_1", role: "assistant" as const, content: "Response 1", parentMessageId: null, generationId: null, activeChildId: null, deleted: false, createdAt: new Date(), updatedAt: new Date(), parts: [], generation: null, attachments: [], }; let tree = buildMessageTree([msgA]); expect(tree.length).toBe(1); expect(getMessageCount(tree[0]!)).toBe(1); expect(getCurrentChildIndex(tree[0]!)).toBe(0); expect(getActiveMessage(tree[0]!)?.id).toBe("msg_a"); // Step 2: Regenerate the message -> creates child B, activeChildId points to B const msgB = { ...msgA, id: "msg_b", content: "Response 2", parentMessageId: "msg_a", }; msgA.activeChildId = "msg_b"; tree = buildMessageTree([msgA, msgB]); expect(tree.length).toBe(1); expect(tree[0]!.children.length).toBe(1); expect(getMessageCount(tree[0]!)).toBe(2); // Variant 2 of 2: index should be 1 (displayed as 1 + 1 = 2) expect(getCurrentChildIndex(tree[0]!)).toBe(1); expect(getActiveMessage(tree[0]!)?.id).toBe("msg_b"); // Step 3: Switch focus to first, root message (A) const prevId = getNextActiveChildId(tree[0]!, -1); expect(prevId).toBeNull(); msgA.activeChildId = null; tree = buildMessageTree([msgA, msgB]); expect(getMessageCount(tree[0]!)).toBe(2); // Variant 1 of 2: index should be 0 (displayed as 0 + 1 = 1) expect(getCurrentChildIndex(tree[0]!)).toBe(0); expect(getActiveMessage(tree[0]!)?.id).toBe("msg_a"); // Step 4: Delete the root message (A) // Root is soft-deleted, active child switches to B msgA.deleted = true; msgA.activeChildId = "msg_b"; tree = buildMessageTree([msgA, msgB]); expect(tree.length).toBe(1); expect(tree[0]!.deleted).toBe(true); expect(getMessageCount(tree[0]!)).toBe(1); expect(getCurrentChildIndex(tree[0]!)).toBe(0); expect(getActiveMessage(tree[0]!)?.id).toBe("msg_b"); // Step 5: Regenerate the message -> creates child C with parentMessageId: msg_a const msgC = { ...msgA, id: "msg_c", content: "Response 3", parentMessageId: "msg_a", deleted: false, }; msgA.activeChildId = "msg_c"; tree = buildMessageTree([msgA, msgB, msgC]); expect(tree.length).toBe(1); expect(tree[0]!.children.length).toBe(2); expect(getMessageCount(tree[0]!)).toBe(2); // Crucial test: When C (second child) is active, index must be 1 (displayed as 2/2, NOT 3/2!) expect(getCurrentChildIndex(tree[0]!)).toBe(1); expect(getActiveMessage(tree[0]!)?.id).toBe("msg_c"); // Verify switching focus back to B works seamlessly const backToB = getNextActiveChildId(tree[0]!, -1); expect(backToB).toBe("msg_b"); msgA.activeChildId = backToB; tree = buildMessageTree([msgA, msgB, msgC]); expect(getCurrentChildIndex(tree[0]!)).toBe(0); expect(getActiveMessage(tree[0]!)?.id).toBe("msg_b"); // Verify switching forward to C works seamlessly const forwardToC = getNextActiveChildId(tree[0]!, 1); expect(forwardToC).toBe("msg_c"); msgA.activeChildId = forwardToC; tree = buildMessageTree([msgA, msgB, msgC]); expect(getCurrentChildIndex(tree[0]!)).toBe(1); expect(getActiveMessage(tree[0]!)?.id).toBe("msg_c"); }); test("buildFocusedMessageTree excludes soft-deleted root messages", () => { const root = { id: "root", userId: "user_1", topicId: "topic_1", role: "assistant" as const, content: "Root", parentMessageId: null, generationId: null, activeChildId: "child_1", deleted: true, createdAt: new Date(), updatedAt: new Date(), parts: [], generation: null, attachments: [], }; const child1 = { ...root, id: "child_1", content: "Child 1", parentMessageId: "root", activeChildId: null, deleted: false, }; const tree = buildMessageTree([root, child1]); const focused = buildFocusedMessageTree(tree); expect(focused.length).toBe(1); expect(focused[0]!.id).toBe("child_1"); expect(focused[0]!.deleted).toBe(false); }); test("buildMessageTree filters out deleted roots with no children", () => { const deletedEmptyRoot = { id: "deleted_root", userId: "user_1", topicId: "topic_1", role: "assistant" as const, content: "Empty", parentMessageId: null, generationId: null, activeChildId: null, deleted: true, createdAt: new Date(), updatedAt: new Date(), parts: [], generation: null, attachments: [], }; const tree = buildMessageTree([deletedEmptyRoot]); expect(tree.length).toBe(0); }); test("getNextActiveChildId clamps to bounds and handles null fallback", () => { const root = { id: "root", deleted: true, activeChildId: null as string | null, children: [ { id: "child_a" }, { id: "child_b" }, ], }; // If activeChildId is null and root is deleted, going left stays at child_a expect(getNextActiveChildId(root, -1)).toBe("child_a"); // Going right moves to child_b expect(getNextActiveChildId(root, 1)).toBe("child_b"); // When at child_b, going right clamps to child_b root.activeChildId = "child_b"; expect(getNextActiveChildId(root, 1)).toBe("child_b"); }); });