refactor(webui): shrink fork implementation

webui/src/components/thread/ThreadMessages.tsx

@@ -8,10 +8,10 @@ import type { CliAppInfo, McpPresetInfo, UIMessage } from "@/lib/types";
 
 interface ThreadMessagesProps {
   messages: UIMessage[];
-  allMessages?: UIMessage[];
   /** When true, agent turn still in flight — keeps activity timeline expanded. */
   isStreaming?: boolean;
   hiddenMessageCount?: number;
+  hiddenUserMessageCount?: number;
   onLoadEarlier?: () => void;
   cliApps?: CliAppInfo[];
   mcpPresets?: McpPresetInfo[];
@@ -65,9 +65,9 @@ export function assistantCopyFlags(units: DisplayUnit[]): boolean[] {
 
 export function ThreadMessages({
   messages,
-  allMessages,
   isStreaming = false,
   hiddenMessageCount = 0,
+  hiddenUserMessageCount = 0,
   onLoadEarlier,
   cliApps = [],
   mcpPresets = [],
@@ -81,15 +81,12 @@ export function ThreadMessages({
     () => unitIndexAfterMessageCount(units, forkBoundaryMessageCount),
     [forkBoundaryMessageCount, units],
   );
-  const assistantForkIndexById = useMemo(
-    () => assistantForkIndexByMessageId(allMessages ?? messages),
-    [allMessages, messages],
-  );
   const copyFlags = useMemo(() => assistantCopyFlags(units), [units]);
   const liveActivityClusterIndices = useMemo(
     () => isStreaming ? currentActivityClusterIndices(units) : new Set<number>(),
     [isStreaming, units],
   );
+  let nextUserIndex = hiddenUserMessageCount;
 
   return (
     <div className="flex w-full flex-col">
@@ -123,6 +120,11 @@ export function ThreadMessages({
           unit.type === "message" && unit.message.role === "user"
             ? unit.message.id
             : undefined;
+        const forkIndex =
+          unit.type === "message" && unit.message.role === "assistant" && copyFlags[index]
+            ? nextUserIndex
+            : undefined;
+        if (unit.type === "message" && unit.message.role === "user") nextUserIndex += 1;
 
         return (
           <Fragment key={unitKey(unit, index)}>
@@ -149,20 +151,15 @@ export function ThreadMessages({
                   mcpPresets={mcpPresets}
                   onOpenFilePreview={onOpenFilePreview}
                   onForkFromHere={
-                    onForkFromMessage
-                      ? forkHandlerForAssistantMessage(
-                          unit.message,
-                          copyFlags[index],
-                          assistantForkIndexById,
-                          onForkFromMessage,
-                        )
+                    onForkFromMessage && forkIndex !== undefined
+                      ? () => onForkFromMessage(forkIndex)
                       : undefined
                   }
                 />
               )}
             </div>
             {index === forkBoundaryAfterUnitIndex ? (
-              <ForkBoundaryDivider label={t("thread.fork.fromHistory")} />
+              <ForkBoundaryDivider label={t("thread.forkedFromHistory")} />
             ) : null}
           </Fragment>
         );
@@ -195,34 +192,6 @@ function ForkBoundaryDivider({ label }: { label: string }) {
   );
 }
 
-function assistantForkIndexByMessageId(messages: UIMessage[]): Map<string, number> {
-  const out = new Map<string, number>();
-  let nextUserIndex = 0;
-  for (const message of messages) {
-    if (message.role === "user") {
-      nextUserIndex += 1;
-    } else if (message.role === "assistant") {
-      out.set(message.id, nextUserIndex);
-    }
-  }
-  return out;
-}
-
-function forkHandlerForAssistantMessage(
-  message: UIMessage,
-  canForkAssistant: boolean,
-  assistantForkIndexById: Map<string, number>,
-  onForkFromMessage: NonNullable<ThreadMessagesProps["onForkFromMessage"]>,
-): (() => void) | undefined {
-  if (message.role === "assistant" && canForkAssistant) {
-    const beforeUserIndex = assistantForkIndexById.get(message.id);
-    return beforeUserIndex === undefined
-      ? undefined
-      : () => onForkFromMessage(beforeUserIndex);
-  }
-  return undefined;
-}
-
 function currentActivityClusterIndices(units: DisplayUnit[]): Set<number> {
   const indices = new Set<number>();
   let markedCurrentActivity = false;