Add chat thread forking

server/src/routes.ts

@@ -1,3 +1,4 @@
+import { randomUUID } from "node:crypto";
 import { performance } from "node:perf_hooks";
 import { z } from "zod";
 import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
@@ -18,6 +19,7 @@ import type { ChatAttachment } from "./llm/types.js";
 
 const ProviderSchema = z.enum(["openai", "anthropic", "xai", "gemini", "hermes-agent"]);
 const MAX_ADDITIONAL_SYSTEM_PROMPT_CHARS = 12_000;
+const MAX_FORK_MESSAGE_SNIPPET_CHARS = 48;
 const EnabledToolsSchema = z.array(z.string().trim().min(1).max(80)).max(20).transform((value) => normalizeEnabledChatTools(value));
 
 type IncomingChatMessage = {
@@ -94,7 +96,7 @@ async function storeNonAssistantMessages(chatId: string, messages: IncomingChatM
 
   const existing = await prisma.message.findMany({
     where: { chatId },
-    orderBy: { createdAt: "asc" },
+    orderBy: [{ createdAt: "asc" }, { id: "asc" }],
     select: { role: true, content: true, name: true, metadata: true },
   });
   const existingNonAssistant = existing.filter((m) => m.role !== "assistant" && !isToolCallLogMessage(m));
@@ -317,6 +319,21 @@ function normalizeSuggestedTitle(raw: string, fallback: string) {
   return words.slice(0, 4).join(" ").slice(0, 64).trim() || fallback;
 }
 
+export function buildForkTitle(originalTitle: string | null, messageContent?: string) {
+  if (messageContent !== undefined) {
+    const snippet = truncateContextPart(messageContent.replace(/\s+/g, " "), MAX_FORK_MESSAGE_SNIPPET_CHARS) ?? "message";
+    return `Fork of '${snippet}'`;
+  }
+  return `Fork of ${originalTitle?.trim() || "Untitled chat"}`;
+}
+
+export function copyForkMessageMetadata(metadata: unknown) {
+  if (metadata === null || metadata === undefined) return undefined;
+  if (typeof metadata !== "object" || Array.isArray(metadata)) return metadata;
+  const { clientRequestId: _clientRequestId, ...copied } = metadata as Record<string, unknown>;
+  return Object.keys(copied).length ? copied : undefined;
+}
+
 async function generateChatTitle(content: string) {
   const systemPrompt =
     "You create short chat titles. Return exactly one line, maximum 4 words, no quotes, no trailing punctuation.";
@@ -415,6 +432,7 @@ type SearchRunRequest = z.infer<typeof SearchRunBody>;
 
 const activeChatStreams = new Map<string, ActiveSseStream>();
 const activeChatStreamRequestIds = new Map<string, string>();
+const chatDeletionRoots = new Set<string>();
 const activeSearchStreams = new Map<string, ActiveSseStream>();
 const STARRED_PROJECT_ID = "starred";
 
@@ -427,6 +445,8 @@ const starredProjectItemsSelect = {
 const chatSummarySelect = {
   id: true,
   title: true,
+  titleGenerationPending: true,
+  parentChatId: true,
   createdAt: true,
   updatedAt: true,
   initiatedProvider: true,
@@ -507,6 +527,29 @@ async function getSearchSummary(searchId: string) {
   return search ? serializeSearchLike(search) : null;
 }
 
+async function listRecentChatsWithRoots() {
+  const chats = await prisma.chat.findMany({
+    orderBy: { updatedAt: "desc" },
+    take: 100,
+    select: chatSummarySelect,
+  });
+  const includedIds = new Set(chats.map((chat) => chat.id));
+  const missingRootIds = [
+    ...new Set(
+      chats
+        .map((chat) => chat.parentChatId)
+        .filter((id): id is string => id !== null && !includedIds.has(id))
+    ),
+  ];
+  if (!missingRootIds.length) return chats;
+
+  const roots = await prisma.chat.findMany({
+    where: { id: { in: missingRootIds } },
+    select: chatSummarySelect,
+  });
+  return [...chats, ...roots].sort(compareUpdatedAtDesc);
+}
+
 async function setChatStarred(chatId: string, starred: boolean) {
   const exists = await prisma.chat.findUnique({ where: { id: chatId }, select: { id: true } });
   if (!exists) return null;
@@ -545,11 +588,7 @@ async function setSearchStarred(searchId: string, starred: boolean) {
 
 async function listWorkspaceItems() {
   const [chats, searches] = await Promise.all([
-    prisma.chat.findMany({
-      orderBy: { updatedAt: "desc" },
-      take: 100,
-      select: chatSummarySelect,
-    }),
+    listRecentChatsWithRoots(),
     prisma.search.findMany({
       orderBy: { updatedAt: "desc" },
       take: 100,
@@ -605,7 +644,7 @@ function mapChatStreamEvent(ev: StreamEvent): SseStreamEvent {
   return { event: ev.type, data: ev };
 }
 
-function registerActiveChatStream(chatId: string, clientRequestId?: string) {
+export function registerActiveChatStream(chatId: string, clientRequestId?: string) {
   const stream = new ActiveSseStream();
   activeChatStreams.set(chatId, stream);
   if (clientRequestId) {
@@ -616,7 +655,7 @@ function registerActiveChatStream(chatId: string, clientRequestId?: string) {
   return stream;
 }
 
-function clearActiveChatStream(chatId: string, stream: ActiveSseStream) {
+export function clearActiveChatStream(chatId: string, stream: ActiveSseStream) {
   if (activeChatStreams.get(chatId) !== stream) return;
   activeChatStreams.delete(chatId);
   activeChatStreamRequestIds.delete(chatId);
@@ -647,12 +686,6 @@ function executeActiveChatStream(chatId: string, body: z.infer<typeof Completion
   })();
 }
 
-function startActiveChatStream(chatId: string, body: z.infer<typeof CompletionStreamBody>) {
-  const stream = registerActiveChatStream(chatId, body.clientRequestId);
-  executeActiveChatStream(chatId, body, stream);
-  return stream;
-}
-
 function getMetadataClientRequestId(metadata: unknown) {
   if (!metadata || typeof metadata !== "object" || Array.isArray(metadata)) return null;
   const clientRequestId = (metadata as Record<string, unknown>).clientRequestId;
@@ -888,11 +921,7 @@ export async function registerRoutes(app: FastifyInstance) {
 
   app.get("/v1/chats", async (req) => {
     requireAdmin(req);
-    const chats = await prisma.chat.findMany({
-      orderBy: { updatedAt: "desc" },
-      take: 100,
-      select: chatSummarySelect,
-    });
+    const chats = await listRecentChatsWithRoots();
     return { chats: chats.map((chat) => serializeChatLike(chat)) };
   });
 
@@ -951,6 +980,80 @@ export async function registerRoutes(app: FastifyInstance) {
     return { chat: serializeChatLike(chat) };
   });
 
+  app.post("/v1/chats/:chatId/fork", async (req) => {
+    requireAdmin(req);
+    const Params = z.object({ chatId: z.string() });
+    const Body = z.object({ messageId: z.string().trim().min(1).optional() });
+    const { chatId } = Params.parse(req.params);
+    const parsed = Body.safeParse(req.body ?? {});
+    if (!parsed.success) return app.httpErrors.badRequest(parsed.error.message);
+
+    const result = await prisma.$transaction(async (tx) => {
+      const source = await tx.chat.findUnique({
+        where: { id: chatId },
+        select: {
+          id: true,
+          title: true,
+          parentChatId: true,
+          initiatedProvider: true,
+          initiatedModel: true,
+          lastUsedProvider: true,
+          lastUsedModel: true,
+          additionalSystemPrompt: true,
+          enabledTools: true,
+          userId: true,
+          messages: { orderBy: [{ createdAt: "asc" }, { id: "asc" }] },
+        },
+      });
+      if (!source) return { status: "chat-not-found" as const };
+
+      let messages = source.messages;
+      let selectedMessage: (typeof source.messages)[number] | undefined;
+      if (parsed.data.messageId) {
+        const selectedIndex = messages.findIndex((message) => message.id === parsed.data.messageId);
+        if (selectedIndex < 0) return { status: "message-not-found" as const };
+        selectedMessage = messages[selectedIndex];
+        if (selectedMessage.role !== "assistant") return { status: "message-not-assistant" as const };
+        messages = messages.slice(0, selectedIndex + 1);
+      }
+
+      const forkMessageIdPrefix = `fork-${randomUUID()}`;
+      const chat = await tx.chat.create({
+        data: {
+          title: buildForkTitle(source.title, selectedMessage?.content),
+          titleGenerationPending: true,
+          parentChatId: source.parentChatId ?? source.id,
+          initiatedProvider: source.initiatedProvider,
+          initiatedModel: source.initiatedModel,
+          lastUsedProvider: source.lastUsedProvider,
+          lastUsedModel: source.lastUsedModel,
+          additionalSystemPrompt: source.additionalSystemPrompt,
+          enabledTools: (source.enabledTools ?? undefined) as any,
+          userId: source.userId,
+          messages: messages.length
+            ? {
+                create: messages.map((message, index) => ({
+                  id: `${forkMessageIdPrefix}-${String(index).padStart(8, "0")}`,
+                  createdAt: message.createdAt,
+                  role: message.role,
+                  content: message.content,
+                  name: message.name,
+                  metadata: copyForkMessageMetadata(message.metadata) as any,
+                })),
+              }
+            : undefined,
+        },
+        select: chatSummarySelect,
+      });
+      return { status: "created" as const, chat };
+    });
+
+    if (result.status === "chat-not-found") return app.httpErrors.notFound("chat not found");
+    if (result.status === "message-not-found") return app.httpErrors.notFound("message not found in chat");
+    if (result.status === "message-not-assistant") return app.httpErrors.badRequest("fork message must be an assistant response");
+    return { chat: serializeChatLike(result.chat) };
+  });
+
   app.patch("/v1/chats/:chatId", async (req) => {
     requireAdmin(req);
     const Params = z.object({ chatId: z.string() });
@@ -963,7 +1066,10 @@ export async function registerRoutes(app: FastifyInstance) {
     const body = Body.parse(req.body ?? {});
 
     const data: Record<string, unknown> = {};
-    if (body.title !== undefined) data.title = body.title;
+    if (body.title !== undefined) {
+      data.title = body.title;
+      data.titleGenerationPending = false;
+    }
     if (body.additionalSystemPrompt !== undefined) data.additionalSystemPrompt = normalizeAdditionalSystemPrompt(body.additionalSystemPrompt);
     if (body.enabledTools !== undefined) data.enabledTools = body.enabledTools;
 
@@ -1004,7 +1110,7 @@ export async function registerRoutes(app: FastifyInstance) {
       select: chatSummarySelect,
     });
     if (!existing) return app.httpErrors.notFound("chat not found");
-    if (existing.title?.trim()) return { chat: serializeChatLike(existing) };
+    if (existing.title?.trim() && !existing.titleGenerationPending) return { chat: serializeChatLike(existing) };
 
     const fallback = body.content.split(/\r?\n/)[0]?.trim().slice(0, 48) || "New chat";
     let suggestedRaw = "";
@@ -1022,8 +1128,12 @@ export async function registerRoutes(app: FastifyInstance) {
     const title = normalizeSuggestedTitle(suggestedRaw, fallback);
 
     await prisma.chat.updateMany({
-      where: { id: body.chatId, title: existing.title },
-      data: { title },
+      where: {
+        id: body.chatId,
+        title: existing.title,
+        titleGenerationPending: existing.titleGenerationPending,
+      },
+      data: { title, titleGenerationPending: false },
     });
 
     const chat = await getChatSummary(body.chatId);
@@ -1039,14 +1149,46 @@ export async function registerRoutes(app: FastifyInstance) {
 
     req.log.info({ chatId }, "delete chat requested");
 
-    const result = await prisma.chat.deleteMany({ where: { id: chatId } });
-    if (result.count === 0) {
+    const target = await prisma.chat.findUnique({
+      where: { id: chatId },
+      select: { parentChatId: true },
+    });
+    if (!target) {
       req.log.warn({ chatId }, "delete chat target not found");
       return app.httpErrors.notFound("chat not found");
     }
 
-    req.log.info({ chatId }, "chat deleted");
-    return { deleted: true };
+    const familyRootId = target.parentChatId ?? chatId;
+    if (chatDeletionRoots.has(familyRootId)) {
+      return app.httpErrors.conflict("chat family deletion already in progress");
+    }
+
+    chatDeletionRoots.add(familyRootId);
+    try {
+      const familyIds = target.parentChatId
+        ? [chatId]
+        : (
+            await prisma.chat.findMany({
+              where: { OR: [{ id: chatId }, { parentChatId: chatId }] },
+              select: { id: true },
+            })
+          ).map((chat) => chat.id);
+      if (familyIds.some((id) => activeChatStreams.has(id))) {
+        req.log.warn({ chatId }, "delete chat rejected while chat family is active");
+        return app.httpErrors.conflict("chat or fork has an active stream");
+      }
+
+      const result = await prisma.chat.deleteMany({ where: { id: chatId } });
+      if (result.count === 0) {
+        req.log.warn({ chatId }, "delete chat target no longer exists");
+        return app.httpErrors.notFound("chat not found");
+      }
+
+      req.log.info({ chatId }, "chat deleted");
+      return { deleted: true };
+    } finally {
+      chatDeletionRoots.delete(familyRootId);
+    }
   });
 
   app.get("/v1/searches", async (req) => {
@@ -1336,7 +1478,7 @@ export async function registerRoutes(app: FastifyInstance) {
     const chat = await prisma.chat.findUnique({
       where: { id: chatId },
       include: {
-        messages: { orderBy: { createdAt: "asc" } },
+        messages: { orderBy: [{ createdAt: "asc" }, { id: "asc" }] },
         calls: { orderBy: { createdAt: "desc" } },
         projectItems: starredProjectItemsSelect,
       },
@@ -1430,10 +1572,15 @@ export async function registerRoutes(app: FastifyInstance) {
     if (!parsed.success) return app.httpErrors.badRequest(parsed.error.message);
     const body = withRequestUserLocation(parsed.data, req);
 
-    // ensure chat exists if provided
+    // Ensure the chat exists and identify its family before reserving a stream.
+    let chatFamilyRootId: string | null = null;
     if (body.chatId) {
-      const exists = await prisma.chat.findUnique({ where: { id: body.chatId }, select: { id: true } });
+      const exists = await prisma.chat.findUnique({
+        where: { id: body.chatId },
+        select: { id: true, parentChatId: true },
+      });
       if (!exists) return app.httpErrors.notFound("chat not found");
+      chatFamilyRootId = exists.parentChatId ?? exists.id;
     }
 
     if (body.persist !== false && body.chatId) {
@@ -1445,32 +1592,31 @@ export async function registerRoutes(app: FastifyInstance) {
         return app.httpErrors.conflict("chat completion already running");
       }
 
-      if (body.clientRequestId) {
-        const reservedStream = registerActiveChatStream(body.chatId, body.clientRequestId);
-        try {
+      if (chatFamilyRootId && chatDeletionRoots.has(chatFamilyRootId)) {
+        return app.httpErrors.conflict("chat family deletion already in progress");
+      }
+
+      const reservedStream = registerActiveChatStream(body.chatId, body.clientRequestId);
+      try {
+        if (body.clientRequestId) {
           const completedSubmission = await findCompletedChatSubmission(body.chatId, body.clientRequestId);
           if (completedSubmission) {
             completeChatSubmissionStream(reservedStream, body.chatId, body, completedSubmission.content);
             clearActiveChatStream(body.chatId, reservedStream);
             return streamActiveRun(req, reply, reservedStream);
           }
-
-          // Store only new non-assistant messages to avoid duplicate history entries.
-          await storeNonAssistantMessages(body.chatId, body.messages, body.clientRequestId);
-          const configuredBody = await applyStoredChatSettings(body);
-          executeActiveChatStream(body.chatId, configuredBody, reservedStream);
-          return streamActiveRun(req, reply, reservedStream);
-        } catch (err) {
-          reservedStream.complete({ event: "error", data: { message: getErrorMessage(err) } });
-          clearActiveChatStream(body.chatId, reservedStream);
-          throw err;
         }
-      }
 
-      // Legacy requests without an idempotency key retain the original behavior.
-      await storeNonAssistantMessages(body.chatId, body.messages);
-      const stream = startActiveChatStream(body.chatId, await applyStoredChatSettings(body));
-      return streamActiveRun(req, reply, stream);
+        // Reserve the stream before persistence so deletion cannot interleave with setup.
+        await storeNonAssistantMessages(body.chatId, body.messages, body.clientRequestId);
+        const configuredBody = await applyStoredChatSettings(body);
+        executeActiveChatStream(body.chatId, configuredBody, reservedStream);
+        return streamActiveRun(req, reply, reservedStream);
+      } catch (err) {
+        reservedStream.complete({ event: "error", data: { message: getErrorMessage(err) } });
+        clearActiveChatStream(body.chatId, reservedStream);
+        throw err;
+      }
     }
 
     reply.raw.writeHead(200, buildSseHeaders(typeof req.headers.origin === "string" ? req.headers.origin : undefined));