chore: fix chat charts persistance

Author: epetrow

client/src/hooks/useChatBot.tsx

@@ -17,6 +17,7 @@ export interface StreamingResponse {
   json?: Record<string, unknown>; // For chart data or other structured responses
   error?: string;
   incomplete?: boolean;
+  messageId?: string | number; // ID of the message associated with this response
 }
 
 export interface UseChatBotReturn {
@@ -174,25 +175,27 @@ export const useChatBot = (config: ChatBotConfig): UseChatBotReturn => {
         setMessages(prev => prev.filter(msg => msg.id !== typingMessageId));
 
         // Then add complete response after a tiny delay
+        const botMessageId = Date.now() + 2;
         setTimeout(() => {
           const botMessage: Message = {
-            id: Date.now() + 2,
+            id: botMessageId,
             author: bot,
             timestamp: new Date(),
             text: currentAnswer
           };
 
           setMessages(prev => [...prev, botMessage]);
         }, 10);
+        
+        // Set the final response with the message ID for any additional processing
+        if (finalResponse) {
+          finalResponse.messageId = botMessageId;
+          setLatestResponse(finalResponse);
+        }
       } else {
         // If no response, just remove typing indicator
         setMessages(prev => prev.filter(msg => msg.id !== typingMessageId));
       }
-      
-      // Set the final response for any additional processing
-      if (finalResponse) {
-        setLatestResponse(finalResponse);
-      }
     } catch (error) {
       const errorMessage = error instanceof Error ? error.message : 'Network error';
 

client/src/pages/FinanceAnalysis.tsx

@@ -147,6 +147,8 @@ export default function FinanceAnalysis() {
 
   const [selectedCharts, setSelectedCharts] = React.useState<BarChartDef[]>([]);
   const [isChartsExpanded, setIsChartsExpanded] = React.useState(false);
+  // Map message IDs to their associated charts
+  const [messageCharts, setMessageCharts] = React.useState<Map<string | number, BarChartDef[]>>(new Map());
 
   // Predefined suggestions related to financial data
   const financialSuggestions: ChatSuggestion[] = [
@@ -178,36 +180,48 @@ export default function FinanceAnalysis() {
     suggestions: financialSuggestions,
   });
 
-  // Watch for changes in the latest response to update charts
+  // Associate charts with the latest bot message when response arrives
   React.useEffect(() => {
     if (
       chatBot.latestResponse?.json?.charts &&
-      Array.isArray(chatBot.latestResponse.json.charts)
+      Array.isArray(chatBot.latestResponse.json.charts) &&
+      chatBot.latestResponse.messageId
     ) {
       const validCharts: BarChartDef[] = chatBot.latestResponse.json.charts
         .filter(isBarChartDef)
         .slice(0, 3);
-      setSelectedCharts(validCharts);
-    } else {
-      setSelectedCharts([]);
+      
+      if (validCharts.length > 0) {
+        setMessageCharts(prev => {
+          const newMap = new Map(prev);
+          newMap.set(chatBot.latestResponse!.messageId!, validCharts);
+          return newMap;
+        });
+        setSelectedCharts(validCharts);
+      }
     }
   }, [chatBot.latestResponse, isBarChartDef]);
 
   // Custom message template that includes thumbnail when charts are available
   const customMessageTemplate = React.useCallback((props: ChatMessageTemplateProps) => {
     const isBot = props.item.author.id !== chatBot.user.id;
-    const isLatestBotMessage = isBot && props.item.id === chatBot.messages[chatBot.messages.length - 1]?.id;
-    const hasCharts = selectedCharts.length > 0;
+    const chartsForThisMessage = messageCharts.get(props.item.id);
 
     return (
       <div>
         <ChatMessage {...props} />
-        {isLatestBotMessage && hasCharts && (
-          <ChartThumbnail onClick={() => setIsChartsExpanded(true)} />
+        {isBot && chartsForThisMessage && chartsForThisMessage.length > 0 && (
+          <ChartThumbnail onClick={() => {
+            const charts = messageCharts.get(props.item.id);
+            if (charts) {
+              setSelectedCharts(charts);
+              setIsChartsExpanded(true);
+            }
+          }} />
         )}
       </div>
     );
-  }, [chatBot.user.id, chatBot.messages, selectedCharts]);
+  }, [chatBot.user.id, messageCharts]);
 
   // Memoized callback for sending messages
   const handleSendMessage = React.useCallback((text: string) => {