fix(playground): Fix session rename not showing messages immediately

Author: keval718

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

@@ -13,6 +13,7 @@ interface ChatSidebarProps {
   currentSessionId?: string;
   onDeleteSession?: (sessionId: string) => void;
   onOpenLogs?: (sessionId: string) => void;
+  renameLocalSession?: (oldSessionId: string, newSessionId: string) => void;
 }
 
 export function ChatSidebar({
@@ -22,9 +23,13 @@ export function ChatSidebar({
   currentSessionId,
   onDeleteSession,
   onOpenLogs,
+  renameLocalSession,
 }: ChatSidebarProps) {
   const currentFlowId = useGetFlowId();
-  const { handleDelete } = useEditSessionInfo({ flowId: currentFlowId });
+  const { handleDelete, handleRename } = useEditSessionInfo({
+    flowId: currentFlowId,
+    renameLocalSession,
+  });
 
   const sessionIds = useMemo(() => sessions, [sessions]);
 
@@ -81,6 +86,7 @@ export function ChatSidebar({
               isVisible={visibleSession === session}
               updateVisibleSession={handleSessionClick}
               inspectSession={onOpenLogs}
+              handleRename={handleRename}
               setActiveSession={() => {
                 // TODO: Implement active session
               }}

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

@@ -18,6 +18,7 @@ export interface SessionSelectorProps {
   setSelectedView?: (view: { type: string; id: string } | undefined) => void;
   playgroundPage?: boolean;
   setActiveSession?: (session: string) => void;
+  handleRename?: (oldSessionId: string, newSessionId: string) => Promise<void>;
 }
 
 export function SessionSelector({
@@ -32,6 +33,7 @@ export function SessionSelector({
   setSelectedView,
   playgroundPage = false,
   setActiveSession,
+  handleRename,
 }: SessionSelectorProps) {
   const [isEditing, setIsEditing] = useState(false);
   const { mutate: updateSessionName } = useUpdateSessionName();
@@ -43,27 +45,41 @@ export function SessionSelector({
     setIsEditing(true);
   };
 
-  const handleRenameSave = (newSessionId: string) => {
+  const handleRenameSave = async (newSessionId: string) => {
     setIsEditing(false);
     const trimmed = newSessionId.trim();
     if (!trimmed || trimmed === session) return;
-    updateSessionName(
-      { old_session_id: session, new_session_id: trimmed },
-      {
-        onSuccess: () => {
-          if (isVisible) {
+    
+    // Use handleRename if provided (from sidebar), otherwise use mutation directly (from header)
+    if (handleRename) {
+      await handleRename(session, trimmed);
+      updateVisibleSession(trimmed);
+      if (
+        selectedView?.type === "Session" &&
+        selectedView?.id === session &&
+        setSelectedView
+      ) {
+        setSelectedView({ type: "Session", id: trimmed });
+      }
+    } else {
+      // Wait for the mutation to complete before updating visible session
+      await updateSessionName(
+        { old_session_id: session, new_session_id: trimmed },
+        {
+          onSuccess: () => {
+            // Update visible session after rename is complete
             updateVisibleSession(trimmed);
-          }
-          if (
-            selectedView?.type === "Session" &&
-            selectedView?.id === session &&
-            setSelectedView
-          ) {
-            setSelectedView({ type: "Session", id: trimmed });
-          }
+            if (
+              selectedView?.type === "Session" &&
+              selectedView?.id === session &&
+              setSelectedView
+            ) {
+              setSelectedView({ type: "Session", id: trimmed });
+            }
+          },
         },
-      },
-    );
+      );
+    }
   };
 
   return (

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

@@ -9,9 +9,11 @@ const LOCAL_SESSIONS_STORAGE_KEY = (flowId: string) =>
 export const useEditSessionInfo = ({
   flowId,
   dbSessions: providedDbSessions,
+  renameLocalSession,
 }: {
   flowId?: string;
   dbSessions?: string[];
+  renameLocalSession?: (oldSessionId: string, newSessionId: string) => void;
 }) => {
   const setSelectedSession = usePlaygroundStore(
     (state) => state.setSelectedSession,
@@ -43,34 +45,17 @@ export const useEditSessionInfo = ({
     }
   };
 
-  const handleRename = async (sessionId: string, newSessionId: string) => {
+  const handleRename = async (sessionId: string, newSessionId: string) => {    
     // 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 local sessions list using the provided function
+    if (renameLocalSession) {
+      renameLocalSession(sessionId, newSessionId);
+    } 
 
     // Update selected session if the renamed session is currently selected
     if (flowId && sessionId === selectedSession) {

src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-get-add-sessions.ts

@@ -11,6 +11,7 @@ interface UseGetAddSessionsProps {
 type UseGetAddSessionsReturnType = (props: UseGetAddSessionsProps) => {
   addNewSession: (allSessions?: string[]) => string;
   removeLocalSession: (sessionId: string) => void;
+  renameLocalSession: (oldSessionId: string, newSessionId: string) => void;
   sessions: string[];
   fetchedSessions: string[];
 };
@@ -138,11 +139,22 @@ export const useGetAddSessions: UseGetAddSessionsReturnType = ({
     });
   };
 
+  const renameLocalSession = (oldSessionId: string, newSessionId: string) => {
+    setLocalSessions((prev) => {
+      const updated = new Set(prev);
+      // Remove old session name and add new one
+      const hadOld = updated.delete(oldSessionId);
+      updated.add(newSessionId);
+      return updated;
+    });
+  };
+
   const stableSessions = useMemo(() => [...sessions], [sessions]);
 
   return {
     addNewSession,
     removeLocalSession,
+    renameLocalSession,
     sessions: stableSessions,
     fetchedSessions,
   };

src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/hooks/use-chat-history.ts

@@ -29,7 +29,8 @@ export const useChatHistory = (visibleSession: string | null) => {
     queryFn: () => {
       // Return cached data immediately - this makes the query active and reactive
       // We're using useQuery purely as a subscription mechanism, not for fetching
-      return queryClient.getQueryData<Message[]>(sessionCacheKey) || [];
+      const cachedData = queryClient.getQueryData<Message[]>(sessionCacheKey) || [];
+      return cachedData;
     },
     staleTime: Infinity, // Never refetch - updates come from setQueryData, not server
     gcTime: 5 * 60 * 1000, // Keep in cache for 5 minutes (allows cleanup of old sessions)
@@ -69,12 +70,12 @@ export const useChatHistory = (visibleSession: string | null) => {
         // In the default session, we show messages that have the same session_id as the flow_id
         // OR messages that have NO session_id (legacy behavior)
         if (visibleSession === currentFlowId) {
-          return (
-            isCurrentFlow &&
-            (message.session_id === visibleSession || !message.session_id)
-          );
+          const matches = isCurrentFlow &&
+            (message.session_id === visibleSession || !message.session_id);
+          return matches;
         }
-        return isCurrentFlow && message.session_id === visibleSession;
+        const matches = isCurrentFlow && message.session_id === visibleSession;
+        return matches;
       })
       .map((message: Message) => {
         let files = message.files;
@@ -111,7 +112,8 @@ export const useChatHistory = (visibleSession: string | null) => {
         };
       });
 
-    return [...filteredMessages].sort(sortSenderMessages);
+    const sorted = [...filteredMessages].sort(sortSenderMessages);
+    return sorted;
   }, [messages, visibleSession, currentFlowId]);
 
   return chatHistory;

src/frontend/src/components/core/playgroundComponent/chat-view/hooks/use-send-message.ts

@@ -33,13 +33,6 @@ export const useSendMessage = ({ sessionId }: UseSendMessageProps = {}) => {
     }): Promise<void> => {
       // Resolve session: sessionId prop > selectedSession > flowId (default)
       const actualSession = sessionId ?? selectedSession ?? flowId;
-
-      console.debug("[useSendMessage] invoked", {
-        sessionId: actualSession,
-        chatInputId,
-        hasFiles: Boolean(files?.length),
-      });
-
       // Add placeholder user message immediately
       addUserMessage({
         id: null,

src/frontend/src/components/core/playgroundComponent/sliding-container/components/flow-page-sliding-container.tsx

@@ -37,14 +37,15 @@ export function FlowPageSlidingContainerContent({
   );
   const [openLogsModal, setOpenLogsModal] = useState(false);
 
-  const { sessions, addNewSession, removeLocalSession, fetchedSessions } =
+  const { sessions, addNewSession, removeLocalSession, renameLocalSession, fetchedSessions } =
     useGetAddSessions({
       flowId: currentFlowId,
       currentSessionId,
     });
   const { handleDelete } = useEditSessionInfo({
     flowId: currentFlowId,
     dbSessions: fetchedSessions,
+    renameLocalSession,
   });
 
   // Ensure currentFlowId is always first in sessions list
@@ -157,6 +158,7 @@ export function FlowPageSlidingContainerContent({
                 currentSessionId={currentSessionId}
                 onDeleteSession={handleDeleteSession}
                 onOpenLogs={handleOpenLogs}
+                renameLocalSession={renameLocalSession}
               />
             </div>
           </div>

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

@@ -2,6 +2,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 { usePlaygroundStore } from "@/stores/playgroundStore";
 import type { useMutationFunctionType } from "@/types/api";
 import type { Message } from "@/types/messages";
 import { api } from "../../api";
@@ -22,22 +23,23 @@ export const useUpdateSessionName: useMutationFunctionType<
   const flowId = useGetFlowId();
 
   const updateSessionApi = async (data: UpdateSessionParams) => {
-    const isPlayground = useFlowStore.getState().playgroundPage;
+    const isPlaygroundFromFlow = useFlowStore.getState().playgroundPage;
+    const isPlaygroundFromPlayground = usePlaygroundStore.getState().isPlayground;
+    const isPlayground = isPlaygroundFromFlow || isPlaygroundFromPlayground;
     // 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));
-      
+      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
+      // CRITICAL: Move messages from old cache key to new cache key
       const oldCacheKey = [
         "useGetMessagesQuery",
         { id: flowId, session_id: data.old_session_id },
@@ -47,18 +49,25 @@ export const useUpdateSessionName: useMutationFunctionType<
         { 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
+      const oldCacheData = queryClient.getQueryData<Message[]>(oldCacheKey);      
+      // Get fresh messages from sessionStorage (source of truth after rename)
+      const freshMessages = JSON.parse(sessionStorage.getItem(flowId) || "[]");
+      const messagesForNewSession = freshMessages.filter(
+        (msg: Message) => msg.session_id === data.new_session_id
+      );
+      
+      // Set the new cache with fresh messages from sessionStorage
+      queryClient.setQueryData(newCacheKey, messagesForNewSession);
+      
+      // Remove the old cache key
+      if (oldCacheData && oldCacheData.length > 0) {
         queryClient.removeQueries({ queryKey: oldCacheKey });
       }
 
+      queryClient.invalidateQueries({
+        queryKey: ["useGetSessionsFromFlowQuery"],
+      });
+      
       return {
         data: messagesWithNewSessionId,
       };
@@ -70,17 +79,45 @@ export const useUpdateSessionName: useMutationFunctionType<
           params: { new_session_id: data.new_session_id },
         },
       );
+            
+      // Update React Query cache with the renamed messages
+      if (result.data && flowId) {
+        const newCacheKey = [
+          "useGetMessagesQuery",
+          { id: flowId, session_id: data.new_session_id },
+        ];
+        
+        queryClient.setQueryData(newCacheKey, result.data);
+        
+        // Remove old cache key
+        const oldCacheKey = [
+          "useGetMessagesQuery",
+          { id: flowId, session_id: data.old_session_id },
+        ];
+        queryClient.removeQueries({ queryKey: oldCacheKey });
+      }
+      
       return result.data;
     }
   };
 
   const mutation: UseMutationResult<Message[], any, UpdateSessionParams> =
     mutate(["useUpdateSessionName"], updateSessionApi, {
-      ...options,
-      onSettled: () => {
+      onMutate: (variables) => {
+      },
+      onSuccess: (data, variables, context, ...rest) => {
+        // Call the original onSuccess if provided
+        options?.onSuccess?.(data, variables, context, ...rest);
+      },
+      onError: (error, variables, context, ...rest) => {
+        options?.onError?.(error, variables, context, ...rest);
+      },
+      onSettled: (data, error, variables, context, ...rest) => {
         queryClient.invalidateQueries({
           queryKey: ["useGetSessionsFromFlowQuery"],
         });
+        // Call the original onSettled if provided
+        options?.onSettled?.(data, error, variables, context, ...rest);
       },
     });