Merge pull request #1260 from kubet/fix/stream-restore-on-thread-navigation

Fix/stream restore on thread navigation

Author: kubet

frontend/src/app/(dashboard)/projects/[projectId]/thread/[threadId]/page.tsx

@@ -137,6 +137,12 @@ export default function ThreadPage({
   const agent = threadAgentData?.agent;
   const workflowId = threadQuery.data?.metadata?.workflow_id;
 
+  // Invalidate cache when thread changes to get fresh data
+  useEffect(() => {
+    queryClient.invalidateQueries({ queryKey: threadKeys.agentRuns(threadId) });
+    queryClient.invalidateQueries({ queryKey: threadKeys.messages(threadId) });
+  }, [threadId, queryClient]);
+
   // Set initial selected agent from thread data
   useEffect(() => {
     if (threadAgentData?.agent && !selectedAgentId) {

frontend/src/app/(dashboard)/projects/[projectId]/thread/_hooks/useThreadData.ts

@@ -40,12 +40,14 @@ export function useThreadData(threadId: string, projectId: string): UseThreadDat
   const messagesLoadedRef = useRef(false);
   const agentRunsCheckedRef = useRef(false);
   const hasInitiallyScrolled = useRef<boolean>(false);
+  
 
   const threadQuery = useThreadQuery(threadId);
   const messagesQuery = useMessagesQuery(threadId);
   const projectQuery = useProjectQuery(projectId);
   const agentRunsQuery = useAgentRunsQuery(threadId);
-
+  
+  // (debug logs removed)
 
   useEffect(() => {
     let isMounted = true;
@@ -54,6 +56,9 @@ export function useThreadData(threadId: string, projectId: string): UseThreadDat
     agentRunsCheckedRef.current = false;
     messagesLoadedRef.current = false;
     initialLoadCompleted.current = false;
+    
+    // Clear messages on thread change; fresh data will set messages
+    setMessages([]);
 
     async function initializeData() {
       if (!initialLoadCompleted.current) setIsLoading(true);
@@ -78,7 +83,7 @@ export function useThreadData(threadId: string, projectId: string): UseThreadDat
         }
 
         if (messagesQuery.data && !messagesLoadedRef.current) {
-         
+          // (debug logs removed)
 
           const unifiedMessages = (messagesQuery.data || [])
             .filter((msg) => msg.type !== 'status')
@@ -91,9 +96,28 @@ export function useThreadData(threadId: string, projectId: string): UseThreadDat
               metadata: msg.metadata || '{}',
               created_at: msg.created_at || new Date().toISOString(),
               updated_at: msg.updated_at || new Date().toISOString(),
+              agent_id: (msg as any).agent_id,
+              agents: (msg as any).agents,
             }));
 
-          setMessages(unifiedMessages);
+          // Merge with any local messages that are not present in server data yet
+          const serverIds = new Set(
+            unifiedMessages.map((m) => m.message_id).filter(Boolean) as string[],
+          );
+          const localExtras = (messages || []).filter(
+            (m) =>
+              !m.message_id ||
+              (typeof m.message_id === 'string' && m.message_id.startsWith('temp-')) ||
+              !serverIds.has(m.message_id as string),
+          );
+          const mergedMessages = [...unifiedMessages, ...localExtras].sort((a, b) => {
+            const aTime = a.created_at ? new Date(a.created_at).getTime() : 0;
+            const bTime = b.created_at ? new Date(b.created_at).getTime() : 0;
+            return aTime - bTime;
+          });
+
+          setMessages(mergedMessages);
+          // Messages set only from server merge; no cross-thread cache
           messagesLoadedRef.current = true;
 
           if (!hasInitiallyScrolled.current) {
@@ -102,7 +126,7 @@ export function useThreadData(threadId: string, projectId: string): UseThreadDat
         }
 
         if (agentRunsQuery.data && !agentRunsCheckedRef.current && isMounted) {
-         
+          // (debug logs removed)
 
           agentRunsCheckedRef.current = true;
 
@@ -169,12 +193,17 @@ export function useThreadData(threadId: string, projectId: string): UseThreadDat
     agentRunsQuery.data
   ]);
 
-  // Disabled automatic message replacement to prevent optimistic message deletion
-  // Messages are now only loaded on initial page load and updated via streaming
+  // Force message reload when thread changes or new data arrives
   useEffect(() => {
-    if (messagesQuery.data && messagesQuery.status === 'success') {
-      // Only load messages on initial load, not when agent status changes
-      if (!isLoading && messages.length === 0) {
+    if (messagesQuery.data && messagesQuery.status === 'success' && !isLoading) {
+      // (debug logs removed)
+      
+      // Always reload messages when thread data changes or we have more raw messages than processed
+      const shouldReload = messages.length === 0 || messagesQuery.data.length > messages.length + 50; // Allow for status messages
+      
+      if (shouldReload) {
+        // (debug logs removed)
+        
         const unifiedMessages = (messagesQuery.data || [])
           .filter((msg) => msg.type !== 'status')
           .map((msg: ApiMessageType) => ({
@@ -190,7 +219,28 @@ export function useThreadData(threadId: string, projectId: string): UseThreadDat
             agents: (msg as any).agents,
           }));
 
-        setMessages(unifiedMessages);
+        // Merge strategy: preserve any local (optimistic/streamed) messages not in server yet
+        setMessages((prev) => {
+          const serverIds = new Set(
+            unifiedMessages.map((m) => m.message_id).filter(Boolean) as string[],
+          );
+          const localExtras = (prev || []).filter(
+            (m) =>
+              !m.message_id ||
+              (typeof m.message_id === 'string' && m.message_id.startsWith('temp-')) ||
+              !serverIds.has(m.message_id as string),
+          );
+          const merged = [...unifiedMessages, ...localExtras].sort((a, b) => {
+            const aTime = a.created_at ? new Date(a.created_at).getTime() : 0;
+            const bTime = b.created_at ? new Date(b.created_at).getTime() : 0;
+            return aTime - bTime;
+          });
+          
+          // Messages set only from server merge; no cross-thread cache
+          return merged;
+        });
+      } else {
+        // (debug logs removed)
       }
     }
   }, [messagesQuery.data, messagesQuery.status, isLoading, messages.length, threadId]);

frontend/src/hooks/useAgentStream.ts

@@ -98,7 +98,25 @@ export function useAgentStream(
       .reduce((acc, curr) => acc + curr.content, '');
   }, [textContent]);
 
-  // Update refs if threadId or setMessages changes
+  // Refs to capture current state for persistence
+  const statusRef = useRef(status);
+  const agentRunIdRef = useRef(agentRunId);
+  const textContentRef = useRef(textContent);
+  
+  // Update refs whenever state changes
+  useEffect(() => {
+    statusRef.current = status;
+  }, [status]);
+  
+  useEffect(() => {
+    agentRunIdRef.current = agentRunId;
+  }, [agentRunId]);
+  
+  useEffect(() => {
+    textContentRef.current = textContent;
+  }, [textContent]);
+
+  // Update refs if threadId changes (no persistence across navigation)
   useEffect(() => {
     threadIdRef.current = threadId;
   }, [threadId]);
@@ -171,7 +189,7 @@ export function useAgentStream(
       updateStatus(finalStatus);
       setAgentRunId(null);
       currentRunIdRef.current = null;
-
+      
       // Message refetch disabled - optimistic messages will handle updates
 
       // If the run was stopped or completed, try to get final status to update nonRunning set (keep this)
@@ -470,23 +488,14 @@ export function useAgentStream(
   useEffect(() => {
     isMountedRef.current = true;
 
-    // Cleanup function for when the component unmounts or agentRunId changes
+    // Cleanup function - be more conservative about stream cleanup
     return () => {
       isMountedRef.current = false;
-      console.log(
-        '[useAgentStream] Unmounting or agentRunId changing. Cleaning up stream.',
-      );
-      if (streamCleanupRef.current) {
-        streamCleanupRef.current();
-        streamCleanupRef.current = null;
-      }
-      // Reset state on unmount if needed, though finalizeStream should handle most cases
-      setStatus('idle');
-      setTextContent([]);
-      setToolCall(null);
-      setError(null);
-      setAgentRunId(null);
-      currentRunIdRef.current = null;
+      console.log('[useAgentStream] Component unmounting/navigation.');
+      
+      // Don't automatically cleanup streams on navigation
+      // Only set mounted flag to false to prevent new operations
+      // Streams will be cleaned up when they naturally complete or on explicit stop
     };
   }, []); // Empty dependency array for mount/unmount effect
 
@@ -496,13 +505,13 @@ export function useAgentStream(
     async (runId: string) => {
       if (!isMountedRef.current) return;
       console.log(
-        `[useAgentStream] Received request to start streaming for ${runId}`,
+        `[STREAM_FIX] Thread ${threadId} received request to start streaming for ${runId}`,
       );
 
       // Clean up any previous stream
       if (streamCleanupRef.current) {
         console.log(
-          '[useAgentStream] Cleaning up existing stream before starting new one.',
+          `[STREAM_FIX] Thread ${threadId} cleaning up existing stream before starting new one.`,
         );
         streamCleanupRef.current();
         streamCleanupRef.current = null;