Feature/add realtime updates (#333)

* Add realtime updates

* Show realtime updates in planning messages

* Implement streaming text renderer

* Implement multi-comment streaming

* Use ChatContent component to render streams

Author: ppsreejith

web/src/components/common/ActionStack.tsx

@@ -313,7 +313,10 @@ const PlanningActionStack: React.FC = () => {
   const metadataProcessingCache = useSelector((state: RootState) => state.settings.metadataProcessingCache)
   const isAnalystmode = useSelector((state: RootState) => state.settings.analystMode) || false;
   const dbMetadata = get(metadataProcessingCache, [dbId, 'result'], null);
-  
+  const thread = useSelector((state: RootState) => state.chat.activeThread)
+  const planningMessage = useSelector((state: RootState) => state.chat.threads[thread].planningMessage)
+  const streamingContents = useSelector((state: RootState) => state.chat.threads[thread].streamingContents)
+
   const planningActions = isEmpty(dbMetadata) && isAnalystmode
       ? ['One time optimizations underway']
       : ['Planning next steps', 'Thinking about the question', 'Understanding App state', 'Finalizing Actions', 'Validating Answers']
@@ -325,8 +328,11 @@ const PlanningActionStack: React.FC = () => {
     return () => clearInterval(intervalId);
   }, []);
 
+  // Use socket planning message if available, otherwise cycle through default messages
+  const displayMessage = planningMessage || planningActions[currentTitleIndex % planningActions.length]
+
   return (
-  <HStack aria-label={"planning"} className={'action-stack'} justifyContent={'start'} width={"100%"}>
+  <VStack aria-label={"planning"} className={'action-stack'} justifyContent={'start'} width={"100%"} spacing={2}>
     <Box
       bg={'minusxGreen.800'}
       p={2}
@@ -339,10 +345,31 @@ const PlanningActionStack: React.FC = () => {
     >
       <HStack>
         <Box>
-          <Text key={currentTitleIndex} animation={currentTitleIndex > 0 ? `${scrollUp} 0.5s ease-in-out` : ""} >{planningActions[currentTitleIndex % planningActions.length]}</Text>
+          <Text key={planningMessage || currentTitleIndex} animation={currentTitleIndex > 0 ? `${scrollUp} 0.5s ease-in-out` : ""} >{displayMessage}</Text>
         </Box>
         <Spinner size="xs" speed={'0.75s'} color="minusxBW.100" aria-label={"planning-spinner"}/>
       </HStack>
     </Box>
-  </HStack>
+    {streamingContents?.map((streamingContent) => (
+      <Box
+        key={streamingContent.id}
+        bg={'minusxGreen.800'}
+        p={3}
+        borderRadius={'10px 10px 10px 0'}
+        color={'minusxBW.50'}
+        width={"90%"}
+        alignSelf={'flex-start'}
+        aria-label={"streaming-content-bubble"}
+      >
+        <ChatContent
+          content={{
+            type: 'DEFAULT',
+            text: streamingContent.text,
+            images: [],
+          }}
+          role="assistant"
+        />
+      </Box>
+    ))}
+  </VStack>
 )}

web/src/components/common/App.tsx

@@ -46,7 +46,7 @@ import { SupportButton } from './Support'
 import { Markdown } from './Markdown'
 import { getMXToken, setMinusxMode, toggleMinusXRoot } from '../../app/rpc'
 import { configs } from '../../constants'
-import { abortPlan, startNewThread, updateThreadID } from '../../state/chat/reducer'
+import { abortPlan, startNewThread, updateThreadID, setPlanningMessage, appendStreamingContent } from '../../state/chat/reducer'
 import { intelligentThreadStart } from '../../helpers/threadHistory'
 
 // Agent constants
@@ -192,22 +192,31 @@ const AppLoggedIn = forwardRef((_props, ref) => {
   }, [atlasData, atlasLoading, atlasError])
 
   // Disabling sockets for now
-  // useSocketIO({
-  //   sessionToken: sessionJwt,
-  //   onMessage: (message) => {
-  //     console.log('Socket.io message received:', message);
-  //   },
-  //   onConnect: () => {
-  //     console.log('Socket.io connected successfully');
-  //   },
-  //   onDisconnect: (reason) => {
-  //     console.log('Socket.io disconnected:', reason);
-  //   },
-  //   onError: (error) => {
-  //     console.log(error.message)
-  //     console.error('Socket.io connection error:', error);
-  //   }
-  // });
+  useSocketIO({
+    sessionToken: sessionJwt,
+    onMessage: (message) => {
+      console.log('Socket.io message received:', message);
+      // Handle planning messages
+      if (message?.type === 'message' && message?.content?.agent) {
+        const agentName = message.content.agent;
+        dispatch(setPlanningMessage(`Running ${agentName}`));
+      }
+      // Handle streaming content chunks
+      if (message?.type === 'content' && message?.id && message?.content) {
+        dispatch(appendStreamingContent({ id: message.id, chunk: message.content }));
+      }
+    },
+    onConnect: () => {
+      console.log('Socket.io connected successfully');
+    },
+    onDisconnect: (reason) => {
+      console.log('Socket.io disconnected:', reason);
+    },
+    onError: (error) => {
+      console.log(error.message)
+      console.error('Socket.io connection error:', error);
+    }
+  });
 
   // Update thread id on start
   useEffect(() => {

web/src/state/chat/reducer.ts

@@ -153,6 +153,8 @@ export interface ChatThread {
   interrupted: boolean
   tasks: Tasks
   id: string
+  planningMessage?: string
+  streamingContents?: Array<{ id: string, text: string }>
 }
 
 interface ChatState {
@@ -735,7 +737,14 @@ export const chatSlice = createSlice({
       activeThread.debugChatIndex = action.payload
     },
     setActiveThreadStatus: (state, action: PayloadAction<ChatThreadStatus>) => {
-      state.threads[state.activeThread].status = action.payload
+      const activeThread = state.threads[state.activeThread]
+      const oldStatus = activeThread.status
+      activeThread.status = action.payload
+      // Clear planning data when leaving PLANNING status
+      if (oldStatus === 'PLANNING' && action.payload !== 'PLANNING') {
+        activeThread.planningMessage = undefined
+        activeThread.streamingContents = undefined
+      }
     },
     toggleUserConfirmation: (state, action: PayloadAction<{
       show: boolean
@@ -884,11 +893,43 @@ export const chatSlice = createSlice({
         console.error('Error cloning thread from history:', error)
         // Don't change state on error - let existing thread remain active
       }
+    },
+    setPlanningMessage: (state, action: PayloadAction<string>) => {
+      const activeThread = getActiveThread(state)
+      activeThread.planningMessage = action.payload
+    },
+    clearPlanningMessage: (state) => {
+      const activeThread = getActiveThread(state)
+      activeThread.planningMessage = undefined
+    },
+    appendStreamingContent: (state, action: PayloadAction<{ id: string, chunk: string }>) => {
+      const activeThread = getActiveThread(state)
+      const { id, chunk } = action.payload
+
+      // Initialize array if needed
+      if (!activeThread.streamingContents) {
+        activeThread.streamingContents = []
+      }
+
+      // Find existing entry by id
+      const existingEntry = activeThread.streamingContents.find(entry => entry.id === id)
+
+      if (existingEntry) {
+        // Append chunk to existing entry
+        existingEntry.text += chunk
+      } else {
+        // Add new entry
+        activeThread.streamingContents.push({ id, text: chunk })
+      }
+    },
+    clearStreamingContent: (state) => {
+      const activeThread = getActiveThread(state)
+      activeThread.streamingContents = undefined
     }
   },
 })
 
 // Action creators are generated for each case reducer function
-export const { addUserMessage, deleteUserMessage, addActionPlanMessage, addActionPlanMessageV2, startAction, finishAction, interruptPlan, startNewThread, addReaction, removeReaction, updateDebugChatIndex, setActiveThreadStatus, toggleUserConfirmation, setUserConfirmationInput, toggleClarification, setClarificationAnswer, switchToThread, abortPlan, updateThreadID, updateLastWarmedOn, clearTasks, cloneThreadFromHistory, setUserConfirmationFeedback } = chatSlice.actions
+export const { addUserMessage, deleteUserMessage, addActionPlanMessage, addActionPlanMessageV2, startAction, finishAction, interruptPlan, startNewThread, addReaction, removeReaction, updateDebugChatIndex, setActiveThreadStatus, toggleUserConfirmation, setUserConfirmationInput, toggleClarification, setClarificationAnswer, switchToThread, abortPlan, updateThreadID, updateLastWarmedOn, clearTasks, cloneThreadFromHistory, setUserConfirmationFeedback, setPlanningMessage, clearPlanningMessage, appendStreamingContent, clearStreamingContent } = chatSlice.actions
 
 export default chatSlice.reducer