fix: improve message editing UX

components/chat/ChatMessages.tsx

@@ -297,78 +297,148 @@ export default function ChatMessages({
     return groups;
   }, [messages, systemGroupsMap]);
 
+  // Find the index of the last user message slot (used for hiding stale responses during loading)
+  const lastUserMessageSlotIndex = useMemo(() => {
+    for (let i = messageVersions.size - 1; i >= 0; i--) {
+      const versionsAtI = messageVersions.get(i);
+      if (versionsAtI && versionsAtI.length > 0) {
+        const msgAtI = versionsAtI[versionsAtI.length - 1];
+        if (msgAtI.role === "user") {
+          return i;
+        }
+      }
+    }
+    return -1;
+  }, [messageVersions]);
+
+  // Get the currently selected/displayed message's eventId at a given depth
+  const getSelectedEventIdAtDepth = useCallback(
+    (depth: number): string | undefined => {
+      const versions = messageVersions.get(depth);
+      if (!versions || versions.length === 0) return undefined;
+
+      const selectedId = selectedVersions.get(depth);
+      if (selectedId) {
+        const found = versions.find((v) => v._eventId === selectedId);
+        if (found) return found._eventId;
+      }
+
+      // Default to the last version (most recent)
+      return versions[versions.length - 1]._eventId;
+    },
+    [messageVersions, selectedVersions]
+  );
+
+  // Filter versions at a given depth to only those that descend from the selected parent
+  const getFilteredVersions = useCallback(
+    (index: number): Message[] => {
+      const versions = messageVersions.get(index);
+      if (!versions || versions.length === 0) return [];
+
+      // For root messages (depth 0), return all versions
+      if (index === 0) return versions;
+
+      // Get the selected parent's eventId at the previous depth
+      const parentEventId = getSelectedEventIdAtDepth(index - 1);
+      if (!parentEventId) return versions;
+
+      // Filter to only messages whose _prevId matches the parent's eventId
+      const filtered = versions.filter((v) => v._prevId === parentEventId);
+
+      // If no matches (shouldn't happen in normal cases), return all versions
+      return filtered.length > 0 ? filtered : versions;
+    },
+    [messageVersions, getSelectedEventIdAtDepth]
+  );
+
   const getMessageToDisplay = (message: Message, index: number) => {
-    const versions = messageVersions.get(index);
+    const filteredVersions = getFilteredVersions(index);
 
-    if (!versions || versions.length <= 1) {
-      return { msg: message, currentVersion: 1, totalVersions: 1 };
+    if (!filteredVersions || filteredVersions.length <= 1) {
+      // If only one version in the filtered list, show it without version navigation
+      const msgToShow =
+        filteredVersions.length === 1 ? filteredVersions[0] : message;
+      return { msg: msgToShow, currentVersion: 1, totalVersions: 1 };
     }
 
     // Check if a specific version is selected for this "slot" (identified by index)
     const selectedId = selectedVersions.get(index);
 
     if (selectedId) {
-      const selectedMsg = versions.find((v) => v._eventId === selectedId);
+      const selectedMsg = filteredVersions.find(
+        (v) => v._eventId === selectedId
+      );
       if (selectedMsg) {
-        const versionIndex = versions.findIndex(
+        const versionIndex = filteredVersions.findIndex(
           (v) => v._eventId === selectedId
         );
         return {
           msg: selectedMsg,
           currentVersion: versionIndex + 1,
-          totalVersions: versions.length,
+          totalVersions: filteredVersions.length,
         };
       }
     }
 
     // Default to the message passed in (which comes from the main thread)
     // We need to find its index in the sorted versions array
-    const currentIndex = versions.findIndex(
+    const currentIndex = filteredVersions.findIndex(
       (v) => v._eventId === message._eventId
     );
 
-    // If for some reason the message isn't in the group (shouldn't happen), default to last
+    // If for some reason the message isn't in the filtered group, default to last
     if (currentIndex === -1) {
       return {
-        msg: versions[versions.length - 1],
-        currentVersion: versions.length,
-        totalVersions: versions.length,
+        msg: filteredVersions[filteredVersions.length - 1],
+        currentVersion: filteredVersions.length,
+        totalVersions: filteredVersions.length,
       };
     }
 
     return {
       msg: message,
       currentVersion: currentIndex + 1,
-      totalVersions: versions.length,
+      totalVersions: filteredVersions.length,
     };
   };
 
   const handleVersionChange = useCallback(
     (index: number, direction: "prev" | "next", currentMessageId: string) => {
-      const versions = messageVersions.get(index);
-      console.log("ed", versions, index, messageVersions);
-      if (!versions) return;
+      const filteredVersions = getFilteredVersions(index);
+      if (!filteredVersions || filteredVersions.length === 0) return;
 
       const currentSelectedId = selectedVersions.get(index) || currentMessageId;
-      const currentIndex = versions.findIndex(
+      const currentIndex = filteredVersions.findIndex(
         (v) => v._eventId === currentSelectedId
       );
-      console.log(currentIndex, currentSelectedId, selectedVersions);
 
       if (currentIndex === -1) return;
 
       let newIndex = direction === "prev" ? currentIndex - 1 : currentIndex + 1;
 
       // Clamp index
       if (newIndex < 0) newIndex = 0;
-      if (newIndex >= versions.length) newIndex = versions.length - 1;
+      if (newIndex >= filteredVersions.length)
+        newIndex = filteredVersions.length - 1;
 
-      const newVersionId = versions[newIndex]._eventId;
+      const newVersionId = filteredVersions[newIndex]._eventId;
       if (newVersionId) {
-        setSelectedVersions((prev) => new Map(prev).set(index, newVersionId));
+        // When changing a parent's version, clear all child selections
+        // so they automatically follow the new branch
+        setSelectedVersions((prev) => {
+          const newMap = new Map(prev);
+          newMap.set(index, newVersionId);
+          // Clear selections for all deeper depths
+          Array.from(newMap.keys()).forEach((key) => {
+            if (key > index) {
+              newMap.delete(key);
+            }
+          });
+          return newMap;
+        });
       }
     },
-    [messageVersions, selectedVersions]
+    [getFilteredVersions, selectedVersions]
   );
 
   // Toggle a specific system message group
@@ -618,6 +688,101 @@ export default function ChatMessages({
                 totalVersions,
               } = getMessageToDisplay(originalMessage, index);
 
+              // Determine if we are currently generating a response at this specific index
+              // This happens when we are at the slot immediately following the last user message
+              const isGeneratingAtThisIndex =
+                isLoading &&
+                (thinkingContent || streamingContent || isPaymentProcessing) &&
+                lastUserMessageSlotIndex >= 0 &&
+                index === lastUserMessageSlotIndex + 1;
+
+              // When loading a new response, hide all messages after the last user message
+              // UNLESS it is the slot where we are currently generating (which we will handle specially)
+              if (
+                isLoading &&
+                (thinkingContent || streamingContent || isPaymentProcessing) &&
+                lastUserMessageSlotIndex >= 0 &&
+                index > lastUserMessageSlotIndex + 1
+              ) {
+                return null;
+              }
+
+              // Special handling for the slot where generation is happening
+              // We want to show the streaming content as a "new version"
+              if (isGeneratingAtThisIndex) {
+                // Get existing versions at this slot
+                const versions = messageVersions.get(index) || [];
+                const effectiveTotalVersions = versions.length + 1;
+
+                // Check if user has explicitly selected a previous version
+                const selectedId = selectedVersions.get(index);
+                const showingHistory =
+                  selectedId &&
+                  versions.some((v) => v._eventId === selectedId);
+
+                // If we are NOT showing history (i.e. showing the stream), render the streaming UI
+                if (!showingHistory) {
+                  return (
+                    <div
+                      key={`streaming-${index}`}
+                      className="mb-8 last:mb-0"
+                      ref={messagesEndRef} // Ensure we scroll to this
+                    >
+                      <div className="flex justify-start mb-2">
+                        <VersionNavigator
+                          currentVersion={effectiveTotalVersions}
+                          totalVersions={effectiveTotalVersions}
+                          onNavigate={(direction) => {
+                            if (direction === "prev") {
+                              // Switch to the last existing version
+                              const lastExisting =
+                                versions[versions.length - 1];
+                              if (lastExisting && lastExisting._eventId) {
+                                setSelectedVersions((prev) => {
+                                  const newMap = new Map(prev);
+                                  newMap.set(index, lastExisting._eventId!);
+                                  return newMap;
+                                });
+                              }
+                            }
+                          }}
+                          className="ml-2"
+                        />
+                      </div>
+
+                      <div className="flex flex-col items-start mb-6 group">
+                        {isPaymentProcessing &&
+                          !streamingContent &&
+                          !thinkingContent && (
+                            <div className="flex flex-col items-start mb-6">
+                              <div className="flex items-center gap-2 text-sm text-muted-foreground animate-pulse">
+                                <div className="h-4 w-4 rounded-full border-2 border-primary border-t-transparent animate-spin" />
+                                Processing payment...
+                              </div>
+                            </div>
+                          )}
+
+                        {thinkingContent && (
+                          <ThinkingSection
+                            thinkingContent={thinkingContent}
+                            isStreaming={streamingContent == ""}
+                          />
+                        )}
+
+                        {streamingContent && (
+                          <div className="w-full text-foreground py-2 px-0 text-[18px]">
+                            <MarkdownRenderer content={streamingContent} />
+                          </div>
+                        )}
+                      </div>
+                    </div>
+                  );
+                }
+
+                // If we ARE showing history, we fall through to the standard rendering below
+                // BUT we need to override the totalVersions and onNavigate to allow returning to the stream
+              }
+
               // Check if this message represents a system message group
               // We need to match by the message itself, not by index
               const messageIndex = messages.findIndex(
@@ -630,6 +795,32 @@ export default function ChatMessages({
                 message.role === "system" &&
                 !shouldAlwaysShowSystemMessage(message.content);
 
+              // If we are showing history at the generating index, we need to adjust props
+              const displayTotalVersions = isGeneratingAtThisIndex
+                ? (messageVersions.get(index)?.length || 0) + 1
+                : totalVersions;
+
+              const handleCustomNavigate = (direction: "prev" | "next") => {
+                if (
+                  isGeneratingAtThisIndex &&
+                  direction === "next" &&
+                  currentVersion === displayTotalVersions - 1
+                ) {
+                  // Switch to streaming (clear selection)
+                  setSelectedVersions((prev) => {
+                    const newMap = new Map(prev);
+                    newMap.delete(index);
+                    return newMap;
+                  });
+                } else {
+                  handleVersionChange(
+                    index,
+                    direction,
+                    originalMessage._eventId!
+                  );
+                }
+              };
+
               return (
                 <div
                   key={`msg-${index}-${originalMessage._eventId}`}
@@ -986,14 +1177,8 @@ export default function ChatMessages({
                         <div className="flex justify-start mb-2">
                           <VersionNavigator
                             currentVersion={currentVersion}
-                            totalVersions={totalVersions}
-                            onNavigate={(direction) =>
-                              handleVersionChange(
-                                index,
-                                direction,
-                                originalMessage._eventId!
-                              )
-                            }
+                            totalVersions={displayTotalVersions}
+                            onNavigate={handleCustomNavigate}
                             className="ml-2"
                           />
                         </div>
@@ -1108,33 +1293,34 @@ export default function ChatMessages({
             })
           )}
 
-          {isPaymentProcessing &&
-            !streamingContent &&
-            !thinkingContent &&
-            messages.length > 0 && (
+          {(streamingContent ||
+            thinkingContent ||
+            (isPaymentProcessing && !thinkingContent)) &&
+            (!messageVersions.size ||
+              messageVersions.size <= lastUserMessageSlotIndex + 1) && (
               <div className="flex flex-col items-start mb-6">
-                <div className="flex items-center gap-2 text-sm text-muted-foreground animate-pulse">
-                  <div className="h-4 w-4 rounded-full border-2 border-primary border-t-transparent animate-spin" />
-                  Processing payment...
-                </div>
+                {isPaymentProcessing && !streamingContent && !thinkingContent && (
+                  <div className="flex items-center gap-2 text-sm text-muted-foreground animate-pulse mb-2">
+                    <div className="h-4 w-4 rounded-full border-2 border-primary border-t-transparent animate-spin" />
+                    Processing payment...
+                  </div>
+                )}
+
+                {thinkingContent && (
+                  <ThinkingSection
+                    thinkingContent={thinkingContent}
+                    isStreaming={streamingContent == ""}
+                  />
+                )}
+
+                {streamingContent && (
+                  <div className="w-full text-foreground py-2 px-0 text-[18px]">
+                    <MarkdownRenderer content={streamingContent} />
+                  </div>
+                )}
               </div>
             )}
 
-          {thinkingContent && (
-            <ThinkingSection
-              thinkingContent={thinkingContent}
-              isStreaming={streamingContent == ""}
-            />
-          )}
-
-          {streamingContent && (
-            <div className="flex flex-col items-start mb-6">
-              <div className="w-full text-foreground py-2 px-0 text-[18px]">
-                <MarkdownRenderer content={streamingContent} />
-              </div>
-            </div>
-          )}
-
           <div ref={messagesEndRef} />
         </div>
         {/* Flexible spacer - grows to fill space when content is short, has min-height for input area */}