veridian/server/api/topic/[topicId]/fork.post.ts

import * as z from 'zod';
import { messages, messageParts, attachments, topics } from '~~/drizzle/schema';
import { db } from '~~/server/lib/db';
import { nanoid } from 'nanoid';
import { eq, and, inArray } from 'drizzle-orm';

export default defineEventHandler(async (event) => {
    await protectRoute(event);

    const userId = event.context.user!.id;
    const topicId = getRouterParam(event, 'topicId')!;

    const result = await readValidatedBody(event, (body) =>
        z
            .object({
                messageId: z.string(),
            })
            .safeParse(body),
    );

    if (!result.success) {
        throw createError({
            statusCode: 400,
            message: result.error.issues[0]!.message,
        });
    }

    const { messageId } = result.data;

    const topic = await db.select().from(topics).where(and(
        eq(topics.id, topicId),
        eq(topics.userId, userId),
    )).limit(1).then(rows => rows[0]);

    if (!topic) {
        throw createError({
            statusCode: 404,
            statusMessage: 'Not Found',
            message: 'Topic not found',
        });
    }

    const topicMessages = await db.select().from(messages)
        .where(eq(messages.topicId, topicId))
        .orderBy(messages.createdAt);

    const forkMessage = topicMessages.find(m => m.id === messageId);
    if (!forkMessage) {
        throw createError({
            statusCode: 404,
            statusMessage: 'Not Found',
            message: 'Message not found in topic',
        });
    }

    const childrenMap = new Map<string, typeof topicMessages>();
    for (const msg of topicMessages) {
        if (msg.parentMessageId) {
            const siblings = childrenMap.get(msg.parentMessageId) || [];
            siblings.push(msg);
            childrenMap.set(msg.parentMessageId, siblings);
        }
    }

    const messageIdsToCopy = new Set<string>();
    const messageMap = new Map(topicMessages.map(m => [m.id, m]));

    const addMessageAndSiblings = (msgId: string) => {
        const msg = messageMap.get(msgId);
        if (!msg) return;

        if (msg.parentMessageId) {
            const siblings = childrenMap.get(msg.parentMessageId) || [];
            for (const sibling of siblings) {
                messageIdsToCopy.add(sibling.id);
            }
        } else {
            messageIdsToCopy.add(msgId);
        }
    };

    let current: typeof forkMessage | undefined = forkMessage;
    while (current) {
        addMessageAndSiblings(current.id);
        current = current.parentMessageId ? messageMap.get(current.parentMessageId) : undefined;
    }

    const forkIndex = topicMessages.findIndex(m => m.id === messageId);
    for (let i = 0; i <= forkIndex; i++) {
        const msg = topicMessages[i]!;
        if (!msg.parentMessageId) {
            messageIdsToCopy.add(msg.id);
            const children = childrenMap.get(msg.id) || [];
            for (const child of children) {
                messageIdsToCopy.add(child.id);
            }
        }
    }

    const addDescendants = (msgId: string) => {
        const children = childrenMap.get(msgId) || [];
        for (const child of children) {
            if (!messageIdsToCopy.has(child.id)) {
                messageIdsToCopy.add(child.id);
                addDescendants(child.id);
            }
        }
    };

    const initialIds = Array.from(messageIdsToCopy);
    for (const id of initialIds) {
        addDescendants(id);
    }

    const messagesToCopy = topicMessages
        .filter(m => messageIdsToCopy.has(m.id))
        .sort((a, b) => {
            if (!a.parentMessageId && b.parentMessageId) return -1;
            if (a.parentMessageId && !b.parentMessageId) return 1;
            return new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime();
        });

    const newTopicId = nanoid();
    const newTopicName = `${topic.name} (fork)`;

    await db.insert(topics).values({
        id: newTopicId,
        userId,
        name: newTopicName,
        agentId: topic.agentId,
        createdAt: new Date(),
    });

    const messageIdMap = new Map<string, string>();

    for (const msg of messagesToCopy) {
        const newMessageId = nanoid();
        messageIdMap.set(msg.id, newMessageId);

        await db.insert(messages).values({
            id: newMessageId,
            userId,
            topicId: newTopicId,
            parentMessageId: msg.parentMessageId ? (messageIdMap.get(msg.parentMessageId) ?? null) : null,
            generationId: null,
            role: msg.role,
            content: msg.content,
            activeChildId: null,
            deleted: msg.deleted,
            createdAt: msg.createdAt,
            updatedAt: new Date(),
        });
    }

    for (const msg of messagesToCopy) {
        if (msg.activeChildId) {
            const newMessageId = messageIdMap.get(msg.id);
            const newActiveChildId = messageIdMap.get(msg.activeChildId);
            if (newMessageId && newActiveChildId) {
                await db.update(messages)
                    .set({ activeChildId: newActiveChildId })
                    .where(eq(messages.id, newMessageId));
            }
        }
    }

    const originalMessageIds = Array.from(messageIdMap.keys());
    if (originalMessageIds.length > 0) {
        const parts = await db.select().from(messageParts)
            .where(inArray(messageParts.messageId, originalMessageIds));

        for (const part of parts) {
            const newMessageId = messageIdMap.get(part.messageId);
            if (!newMessageId) continue;

            await db.insert(messageParts).values({
                userId,
                topicId: newTopicId,
                messageId: newMessageId,
                type: part.type,
                content: part.content,
                providerOptions: part.providerOptions,
                finished: part.finished,
                createdAt: part.createdAt,
                lastUpdatedAt: new Date(),
            });
        }

        const attachmentsResult = await db.select().from(attachments)
            .where(inArray(attachments.messageId, originalMessageIds));

        for (const attachment of attachmentsResult) {
            const newMessageId = messageIdMap.get(attachment.messageId);
            if (!newMessageId) continue;

            await db.insert(attachments).values({
                userId,
                topicId: newTopicId,
                messageId: newMessageId,
                fileId: attachment.fileId,
                createdAt: new Date(),
            });
        }
    }

    return {
        ok: true,
        topicId: newTopicId,
        name: newTopicName,
    };
});