fix: regenerate from target turn

src/routes/chat.ts

@@ -129,6 +129,26 @@ const createInitialStreamingMessages = (existingMessages: unknown[], userContent
   ];
 };
 
+const countFrontendUserMessages = (messages: unknown[]) =>
+  messages.filter(
+    (message) => isObjectRecord(message) && message.role === "user",
+  ).length;
+
+const pruneBranchGroupsForMessageIndex = (
+  branchGroups: unknown[],
+  messageIndex: number | undefined,
+) => {
+  if (messageIndex === undefined) {
+    return branchGroups;
+  }
+  return branchGroups.filter(
+    (group) =>
+      !isObjectRecord(group) ||
+      typeof group.parentCount !== "number" ||
+      group.parentCount < messageIndex,
+  );
+};
+
 const upsertBackendProgress = (
   progress: unknown,
   payload: Record<string, unknown>,
@@ -952,15 +972,24 @@ export const buildChatRouter = (
         projectKey: requestContext.projectKey,
         sessionId: requestContext.clientSessionId,
       };
-      const recentTurns = await sessionTranscriptStore.getRecentTurns(historyContext, 8);
       const initialSessionState = await sessionUiStateStore.read(
         toSessionUiStateContext(activeSessionRecord),
       );
       const persistedMessages = initialSessionState?.messages ?? [];
+      const isRegenerate =
+        parsed.data.regenerate_from_message_index !== undefined;
       const baseMessages =
-        parsed.data.regenerate_from_message_index !== undefined
+        isRegenerate
           ? persistedMessages.slice(0, parsed.data.regenerate_from_message_index)
           : persistedMessages;
+      const targetUserOrdinal = isRegenerate
+        ? countFrontendUserMessages(
+            persistedMessages.slice(
+              0,
+              (parsed.data.regenerate_from_message_index ?? 0) + 1,
+            ),
+          )
+        : undefined;
       if (activeRuns.get(activeSessionRecord.sessionId)?.status === "running") {
         res.status(409).json({
           message: "session is already streaming",
@@ -968,6 +997,15 @@ export const buildChatRouter = (
         });
         return;
       }
+      if (isRegenerate) {
+        await sessionTranscriptStore.truncateThread(
+          historyContext,
+          parsed.data.regenerate_from_message_index ?? 0,
+        );
+      }
+      const recentTurns = isRegenerate
+        ? []
+        : await sessionTranscriptStore.getRecentTurns(historyContext, 8);
 
       logger.info(
         {
@@ -976,6 +1014,7 @@ export const buildChatRouter = (
           created: created || sessionCreated,
           model: parsed.data.model,
           approvalMode: parsed.data.approval_mode,
+          regenerateFromMessageIndex: parsed.data.regenerate_from_message_index,
           traceId: requestContext.traceId,
           projectId: requestContext.projectId,
         },
@@ -997,7 +1036,10 @@ export const buildChatRouter = (
         baseMessages,
         parsed.data.message,
       );
-      const branchGroups = initialSessionState?.branchGroups ?? [];
+      const branchGroups = pruneBranchGroupsForMessageIndex(
+        initialSessionState?.branchGroups ?? [],
+        parsed.data.regenerate_from_message_index,
+      );
       const activeRun: ActiveRun = {
         clientSessionId,
         controller: abortController,
@@ -1128,8 +1170,13 @@ export const buildChatRouter = (
       };
 
       try {
-        if (parsed.data.regenerate_from_message_index !== undefined) {
-          await runtime.revertLastUserMessage(binding.sessionId);
+        if (isRegenerate) {
+          if (!targetUserOrdinal || targetUserOrdinal < 1) {
+            throw new Error("target user message not found for regeneration");
+          }
+          await runtime.revertToUserMessage(binding.sessionId, {
+            userOrdinal: targetUserOrdinal,
+          });
         }
 
         const preparedMessage = await buildPromptWithLearningContext(