veridian/tests/message-focus.test.ts

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