fix(frontend): resolve session rename bugs and prevent message loss

Author: keval718

src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/components/chat-header.tsx

@@ -119,6 +119,7 @@ export function ChatHeader({
         onMessageLogs={onMessageLogs}
         onClearChat={handleClearChat}
         onDelete={handleDeleteSessionInternal}
+        showRename={!isDefaultSession}
         showClearChat={isDefaultSession}
         showDelete={!isDefaultSession}
         side="bottom"

src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/components/session-more-menu.tsx

@@ -15,6 +15,7 @@ export interface SessionMoreMenuProps {
   onDelete: () => void;
   onClearChat?: () => void;
   showMessageLogs?: boolean;
+  showRename?: boolean;
   showDelete?: boolean;
   showClearChat?: boolean;
   // Positioning props
@@ -36,6 +37,7 @@ export function SessionMoreMenu({
   onDelete,
   onClearChat,
   showMessageLogs = true,
+  showRename = true,
   showDelete = true,
   showClearChat = false,
   side = "bottom",
@@ -102,15 +104,17 @@ export function SessionMoreMenu({
           sideOffset={sideOffset}
           className={cn("p-0", contentClassName)}
         >
-          <SelectItem value="rename" className="session-more-menu-item">
-            <div className="flex items-center">
-              <ForwardedIconComponent
-                name="SquarePen"
-                className="mr-2 h-4 w-4"
-              />
-              Rename
-            </div>
-          </SelectItem>
+          {showRename && (
+            <SelectItem value="rename" className="session-more-menu-item">
+              <div className="flex items-center">
+                <ForwardedIconComponent
+                  name="SquarePen"
+                  className="mr-2 h-4 w-4"
+                />
+                Rename
+              </div>
+            </SelectItem>
+          )}
           {showMessageLogs && (
             <SelectItem value="messageLogs" className="session-more-menu-item">
               <div className="flex items-center">

src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/components/session-selector.tsx

@@ -111,6 +111,7 @@ export function SessionSelector({
           onRename={handleEditClick}
           onMessageLogs={() => inspectSession?.(session)}
           onDelete={() => deleteSession(session)}
+          showRename={session !== currentFlowId}
           showDelete={session !== currentFlowId}
           side="bottom"
           align="end"

src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-edit-session-info.ts

@@ -3,6 +3,9 @@ import { useGetSessionsFromFlowQuery } from "@/controllers/API/queries/messages/
 import { useUpdateSessionName } from "@/controllers/API/queries/messages/use-rename-session";
 import { usePlaygroundStore } from "@/stores/playgroundStore";
 
+const LOCAL_SESSIONS_STORAGE_KEY = (flowId: string) =>
+  `langflow_local_sessions_${flowId}`;
+
 export const useEditSessionInfo = ({
   flowId,
   dbSessions: providedDbSessions,
@@ -41,12 +44,35 @@ export const useEditSessionInfo = ({
   };
 
   const handleRename = async (sessionId: string, newSessionId: string) => {
-    if (dbSessions.includes(sessionId)) {
-      await updateSessionName({
-        old_session_id: sessionId,
-        new_session_id: newSessionId,
-      });
+    // Update session name via API or sessionStorage
+    await updateSessionName({
+      old_session_id: sessionId,
+      new_session_id: newSessionId,
+    });
+
+    // Update local sessions list if this is a local session
+    if (flowId && isPlayground) {
+      try {
+        const stored = window.sessionStorage.getItem(
+          LOCAL_SESSIONS_STORAGE_KEY(flowId),
+        );
+        if (stored) {
+          const localSessions = JSON.parse(stored) as string[];
+          const index = localSessions.indexOf(sessionId);
+          if (index !== -1) {
+            localSessions[index] = newSessionId;
+            window.sessionStorage.setItem(
+              LOCAL_SESSIONS_STORAGE_KEY(flowId),
+              JSON.stringify(localSessions),
+            );
+          }
+        }
+      } catch (error) {
+        console.error("Error updating local sessions:", error);
+      }
     }
+
+    // Update selected session if the renamed session is currently selected
     if (flowId && sessionId === selectedSession) {
       setSelectedSession(newSessionId);
     }

src/frontend/src/controllers/API/queries/messages/use-rename-session.ts

@@ -1,6 +1,7 @@
 import type { UseMutationResult } from "@tanstack/react-query";
 import { useGetFlowId } from "@/modals/IOModal/hooks/useGetFlowId";
 import useFlowStore from "@/stores/flowStore";
+import { useMessagesStore } from "@/stores/messagesStore";
 import type { useMutationFunctionType } from "@/types/api";
 import type { Message } from "@/types/messages";
 import { api } from "../../api";
@@ -24,14 +25,40 @@ export const useUpdateSessionName: useMutationFunctionType<
     const isPlayground = useFlowStore.getState().playgroundPage;
     // if we are in playground we will edit the local storage instead of the API
     if (isPlayground && flowId) {
-      const messages = JSON.parse(sessionStorage.getItem(flowId) || "");
+      const messages = JSON.parse(sessionStorage.getItem(flowId) || "[]");
       const messagesWithNewSessionId = messages.map((message: Message) => {
         if (message.session_id === data.old_session_id) {
           message.session_id = data.new_session_id;
         }
         return message;
       });
       sessionStorage.setItem(flowId, JSON.stringify(messagesWithNewSessionId));
+      
+      // Update the messages store to reflect the new session_id
+      useMessagesStore.getState().renameSession(data.old_session_id, data.new_session_id);
+      
+      // Update React Query cache - move messages from old session key to new session key
+      const oldCacheKey = [
+        "useGetMessagesQuery",
+        { id: flowId, session_id: data.old_session_id },
+      ];
+      const newCacheKey = [
+        "useGetMessagesQuery",
+        { id: flowId, session_id: data.new_session_id },
+      ];
+      
+      const oldMessages = queryClient.getQueryData<Message[]>(oldCacheKey);
+      if (oldMessages) {
+        // Update session_id in cached messages and move to new cache key
+        const updatedMessages = oldMessages.map((msg) => ({
+          ...msg,
+          session_id: data.new_session_id,
+        }));
+        queryClient.setQueryData(newCacheKey, updatedMessages);
+        // Remove old cache entry
+        queryClient.removeQueries({ queryKey: oldCacheKey });
+      }
+      
       return {
         data: messagesWithNewSessionId,
       };

src/frontend/src/stores/messagesStore.ts

@@ -11,6 +11,16 @@ export const useMessagesStore = create<MessagesStoreType>((set, get) => ({
       return { messages: updatedMessages };
     });
   },
+  renameSession: (oldSessionId, newSessionId) => {
+    set((state) => {
+      const updatedMessages = state.messages.map((msg) =>
+        msg.session_id === oldSessionId
+          ? { ...msg, session_id: newSessionId }
+          : msg,
+      );
+      return { messages: updatedMessages };
+    });
+  },
   messages: [],
   setMessages: (messages) => {
     set(() => ({ messages: messages }));

src/frontend/src/types/zustand/messages/index.ts

@@ -11,5 +11,6 @@ export type MessagesStoreType = {
   clearMessages: () => void;
   removeMessages: (ids: string[]) => void;
   deleteSession: (id: string) => void;
+  renameSession: (oldSessionId: string, newSessionId: string) => void;
   displayLoadingMessage: boolean;
 };