feat(ui): performance fixes

surfsense_web/app/dashboard/[search_space_id]/researcher/[chat_id]/page.tsx

@@ -240,7 +240,7 @@ const SourcesDialogContent = ({
 const ChatPage = () => {
   const [token, setToken] = React.useState<string | null>(null);
   const [activeTab, setActiveTab] = useState("");
-  const [dialogOpen, setDialogOpen] = useState(false);
+  const [dialogOpenId, setDialogOpenId] = useState<number | null>(null);
   const [sourcesPage, setSourcesPage] = useState(1);
   const [expandedSources, setExpandedSources] = useState(false);
   const [canScrollLeft, setCanScrollLeft] = useState(false);
@@ -260,6 +260,13 @@ const ChatPage = () => {
 
   const { search_space_id, chat_id } = useParams();
 
+  // Function to scroll terminal to bottom
+  const scrollTerminalToBottom = () => {
+    if (terminalMessagesRef.current) {
+      terminalMessagesRef.current.scrollTop = terminalMessagesRef.current.scrollHeight;
+    }
+  };
+
   // Get token from localStorage on client side only
   React.useEffect(() => {
     setToken(localStorage.getItem('surfsense_bearer_token'));
@@ -469,54 +476,60 @@ const ChatPage = () => {
     updateChat();
   }, [messages, status, chat_id, researchMode, selectedConnectors, search_space_id]);
 
-  // Log messages whenever they update and extract annotations from the latest assistant message if available
+  // Memoize connector sources to prevent excessive re-renders
-  useEffect(() => {
+  const processedConnectorSources = React.useMemo(() => {
-    console.log('Messages updated:', messages);
+    if (messages.length === 0) return connectorSources;
     
-    // Extract annotations from the latest assistant message if available
+    // Only process when we have a complete message (not streaming)
+    if (status !== 'ready') return connectorSources;
+    
+    // Find the latest assistant message
     const assistantMessages = messages.filter(msg => msg.role === 'assistant');
-    if (assistantMessages.length > 0) {
+    if (assistantMessages.length === 0) return connectorSources;
+    
     const latestAssistantMessage = assistantMessages[assistantMessages.length - 1];
-      if (latestAssistantMessage?.annotations) {
+    if (!latestAssistantMessage?.annotations) return connectorSources;
+    
+    // Find the latest SOURCES annotation
     const annotations = latestAssistantMessage.annotations as any[];
+    const sourcesAnnotations = annotations.filter(a => a.type === 'SOURCES');
     
-        // Debug log to track streaming annotations
+    if (sourcesAnnotations.length === 0) return connectorSources;
-        if (process.env.NODE_ENV === 'development') {
-          console.log('Streaming annotations:', annotations);
     
-          // Log counts of each annotation type
-          const terminalInfoCount = annotations.filter(a => a.type === 'TERMINAL_INFO').length;
-          const sourcesCount = annotations.filter(a => a.type === 'SOURCES').length;
-          const answerCount = annotations.filter(a => a.type === 'ANSWER').length;
-
-          console.log(`Annotation counts - Terminal: ${terminalInfoCount}, Sources: ${sourcesCount}, Answer: ${answerCount}`);
-        }
-
-        // Process SOURCES annotation - get the last one to ensure we have the latest
-        const sourcesAnnotations = annotations.filter(
-          (annotation) => annotation.type === 'SOURCES'
-        );
-
-        if (sourcesAnnotations.length > 0) {
-          // Get the last SOURCES annotation to ensure we have the most recent one
     const latestSourcesAnnotation = sourcesAnnotations[sourcesAnnotations.length - 1];
-          if (latestSourcesAnnotation.content) {
+    if (!latestSourcesAnnotation.content) return connectorSources;
-            setConnectorSources(latestSourcesAnnotation.content);
-          }
-        }
     
-        // Check for terminal info annotations and scroll terminal to bottom if they exist
+    // Use this content if it differs from current
-        const terminalInfoAnnotations = annotations.filter(
+    return latestSourcesAnnotation.content;
-          (annotation) => annotation.type === 'TERMINAL_INFO'
+  }, [messages, status, connectorSources]);
-        );
+  
+  // Update connector sources when processed value changes
+  useEffect(() => {
+    if (processedConnectorSources !== connectorSources) {
+      setConnectorSources(processedConnectorSources);
+    }
+  }, [processedConnectorSources, connectorSources]);
+  
+  // Check and scroll terminal when terminal info is available
+  useEffect(() => {
+    if (messages.length === 0 || status !== 'ready') return;
+    
+    // Find the latest assistant message
+    const assistantMessages = messages.filter(msg => msg.role === 'assistant');
+    if (assistantMessages.length === 0) return;
+    
+    const latestAssistantMessage = assistantMessages[assistantMessages.length - 1];
+    if (!latestAssistantMessage?.annotations) return;
+    
+    // Check for terminal info annotations
+    const annotations = latestAssistantMessage.annotations as any[];
+    const terminalInfoAnnotations = annotations.filter(a => a.type === 'TERMINAL_INFO');
     
     if (terminalInfoAnnotations.length > 0) {
       // Schedule scrolling after the DOM has been updated
       setTimeout(scrollTerminalToBottom, 100);
     }
-      }
+  }, [messages, status]);
-    }
-  }, [messages]);
 
   // Custom handleSubmit function to include selected connectors and answer type
   const handleSubmit = (e: React.FormEvent) => {
@@ -543,25 +556,23 @@ const ChatPage = () => {
     messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
   };
 
-  // Function to scroll terminal to bottom
-  const scrollTerminalToBottom = () => {
-    if (terminalMessagesRef.current) {
-      terminalMessagesRef.current.scrollTop = terminalMessagesRef.current.scrollHeight;
-    }
-  };
-
   // Scroll to bottom when messages change
   useEffect(() => {
     scrollToBottom();
   }, [messages]);
 
-  // Set activeTab when connectorSources change
+  // Set activeTab when connectorSources change using a memoized value
-  useEffect(() => {
+  const activeTabValue = React.useMemo(() => {
-    if (connectorSources.length > 0) {
+    return connectorSources.length > 0 ? connectorSources[0].type : "";
-      setActiveTab(connectorSources[0].type);
-    }
   }, [connectorSources]);
   
+  // Update activeTab when the memoized value changes
+  useEffect(() => {
+    if (activeTabValue && activeTabValue !== activeTab) {
+      setActiveTab(activeTabValue);
+    }
+  }, [activeTabValue, activeTab]);
+
   // Scroll terminal to bottom when expanded
   useEffect(() => {
     if (terminalExpanded) {
@@ -617,7 +628,7 @@ const ChatPage = () => {
   };
 
   // Function to get a citation source by ID
-  const getCitationSource = (citationId: number, messageIndex?: number): Source | null => {
+  const getCitationSource = React.useCallback((citationId: number, messageIndex?: number): Source | null => {
     if (!messages || messages.length === 0) return null;
 
     // If no specific message index is provided, use the latest assistant message
@@ -699,7 +710,7 @@ const ChatPage = () => {
 
       return foundSource || null;
     }
-  };
+  }, [messages]);
 
   return (
     <>
@@ -900,7 +911,7 @@ const ChatPage = () => {
                               ))}
 
                               {connector.sources.length > INITIAL_SOURCES_DISPLAY && (
-                                <Dialog open={dialogOpen && activeTab === connector.type} onOpenChange={(open) => setDialogOpen(open)}>
+                                <Dialog open={dialogOpenId === connector.id} onOpenChange={(open) => setDialogOpenId(open ? connector.id : null)}>
                                   <DialogTrigger asChild>
                                     <Button variant="ghost" className="w-full text-sm text-gray-500 dark:text-gray-400">
                                       Show {connector.sources.length - INITIAL_SOURCES_DISPLAY} More Sources

surfsense_web/components/chat/Citation.tsx

@@ -20,7 +20,7 @@ type CitationProps = {
 /**
  * Citation component to handle individual citations
  */
-export const Citation = ({ citationId, citationText, position, source }: CitationProps) => {
+export const Citation = React.memo(({ citationId, citationText, position, source }: CitationProps) => {
   const [open, setOpen] = useState(false);
   const citationKey = `citation-${citationId}-${position}`;
   
@@ -38,7 +38,8 @@ export const Citation = ({ citationId, citationText, position, source }: Citatio
             </span>
           </sup>
         </DropdownMenuTrigger>
-        <DropdownMenuContent align="start" className="w-80 p-0">
+        {open && (
+          <DropdownMenuContent align="start" className="w-80 p-0" forceMount>
             <Card className="border-0 shadow-none">
               <div className="p-3 flex items-start gap-3">
                 <div className="flex-shrink-0 w-7 h-7 flex items-center justify-center bg-muted rounded-full">
@@ -65,10 +66,13 @@ export const Citation = ({ citationId, citationText, position, source }: Citatio
               </div>
             </Card>
           </DropdownMenuContent>
+        )}
       </DropdownMenu>
     </span>
   );
-};
+});
+
+Citation.displayName = 'Citation';
 
 /**
  * Function to render text with citations

surfsense_web/components/markdown-viewer.tsx

@@ -1,4 +1,4 @@
-import React from "react";
+import React, { useMemo } from "react";
 import ReactMarkdown from "react-markdown";
 import rehypeRaw from "rehype-raw";
 import rehypeSanitize from "rehype-sanitize";
@@ -14,14 +14,11 @@ interface MarkdownViewerProps {
 }
 
 export function MarkdownViewer({ content, className, getCitationSource }: MarkdownViewerProps) {
-  return (
+  // Memoize the markdown components to prevent unnecessary re-renders
-    <div className={cn("prose prose-sm dark:prose-invert max-w-none", className)}>
+  const components = useMemo(() => {
-      <ReactMarkdown
+    return {
-        rehypePlugins={[rehypeRaw, rehypeSanitize]}
-        remarkPlugins={[remarkGfm]}
-        components={{
       // Define custom components for markdown elements
-          p: ({node, children, ...props}) => {
+      p: ({node, children, ...props}: any) => {
         // If there's no getCitationSource function, just render normally
         if (!getCitationSource) {
           return <p className="my-2" {...props}>{children}</p>;
@@ -30,52 +27,52 @@ export function MarkdownViewer({ content, className, getCitationSource }: Markdo
         // Process citations within paragraph content
         return <p className="my-2" {...props}>{processCitationsInReactChildren(children, getCitationSource)}</p>;
       },
-          a: ({node, children, ...props}) => {
+      a: ({node, children, ...props}: any) => {
         // Process citations within link content if needed
         const processedChildren = getCitationSource 
           ? processCitationsInReactChildren(children, getCitationSource) 
           : children;
         return <a className="text-primary hover:underline" {...props}>{processedChildren}</a>;
       },
-          li: ({node, children, ...props}) => {
+      li: ({node, children, ...props}: any) => {
         // Process citations within list item content
         const processedChildren = getCitationSource 
           ? processCitationsInReactChildren(children, getCitationSource) 
           : children;
         return <li {...props}>{processedChildren}</li>;
       },
-          ul: ({node, ...props}) => <ul className="list-disc pl-5 my-2" {...props} />,
+      ul: ({node, ...props}: any) => <ul className="list-disc pl-5 my-2" {...props} />,
-          ol: ({node, ...props}) => <ol className="list-decimal pl-5 my-2" {...props} />,
+      ol: ({node, ...props}: any) => <ol className="list-decimal pl-5 my-2" {...props} />,
-          h1: ({node, children, ...props}) => {
+      h1: ({node, children, ...props}: any) => {
         const processedChildren = getCitationSource 
           ? processCitationsInReactChildren(children, getCitationSource) 
           : children;
         return <h1 className="text-2xl font-bold mt-6 mb-2" {...props}>{processedChildren}</h1>;
       },
-          h2: ({node, children, ...props}) => {
+      h2: ({node, children, ...props}: any) => {
         const processedChildren = getCitationSource 
           ? processCitationsInReactChildren(children, getCitationSource) 
           : children;
         return <h2 className="text-xl font-bold mt-5 mb-2" {...props}>{processedChildren}</h2>;
       },
-          h3: ({node, children, ...props}) => {
+      h3: ({node, children, ...props}: any) => {
         const processedChildren = getCitationSource 
           ? processCitationsInReactChildren(children, getCitationSource) 
           : children;
         return <h3 className="text-lg font-bold mt-4 mb-2" {...props}>{processedChildren}</h3>;
       },
-          h4: ({node, children, ...props}) => {
+      h4: ({node, children, ...props}: any) => {
         const processedChildren = getCitationSource 
           ? processCitationsInReactChildren(children, getCitationSource) 
           : children;
         return <h4 className="text-base font-bold mt-3 mb-1" {...props}>{processedChildren}</h4>;
       },
-          blockquote: ({node, ...props}) => <blockquote className="border-l-4 border-muted pl-4 italic my-2" {...props} />,
+      blockquote: ({node, ...props}: any) => <blockquote className="border-l-4 border-muted pl-4 italic my-2" {...props} />,
-          hr: ({node, ...props}) => <hr className="my-4 border-muted" {...props} />,
+      hr: ({node, ...props}: any) => <hr className="my-4 border-muted" {...props} />,
-          img: ({node, ...props}) => <img className="max-w-full h-auto my-4 rounded" {...props} />,
+      img: ({node, ...props}: any) => <img className="max-w-full h-auto my-4 rounded" {...props} />,
-          table: ({node, ...props}) => <div className="overflow-x-auto my-4"><table className="min-w-full divide-y divide-border" {...props} /></div>,
+      table: ({node, ...props}: any) => <div className="overflow-x-auto my-4"><table className="min-w-full divide-y divide-border" {...props} /></div>,
-          th: ({node, ...props}) => <th className="px-3 py-2 text-left font-medium bg-muted" {...props} />,
+      th: ({node, ...props}: any) => <th className="px-3 py-2 text-left font-medium bg-muted" {...props} />,
-          td: ({node, ...props}) => <td className="px-3 py-2 border-t border-border" {...props} />,
+      td: ({node, ...props}: any) => <td className="px-3 py-2 border-t border-border" {...props} />,
       code: ({node, className, children, ...props}: any) => {
         const match = /language-(\w+)/.exec(className || '');
         const isInline = !match;
@@ -89,7 +86,15 @@ export function MarkdownViewer({ content, className, getCitationSource }: Markdo
             </div>
           );
       }
-        }}
+    };
+  }, [getCitationSource]);
+
+  return (
+    <div className={cn("prose prose-sm dark:prose-invert max-w-none", className)}>
+      <ReactMarkdown
+        rehypePlugins={[rehypeRaw, rehypeSanitize]}
+        remarkPlugins={[remarkGfm]}
+        components={components}
       >
         {content}
       </ReactMarkdown>
@@ -98,7 +103,7 @@ export function MarkdownViewer({ content, className, getCitationSource }: Markdo
 }
 
 // Helper function to process citations within React children
-function processCitationsInReactChildren(children: React.ReactNode, getCitationSource: (id: number) => Source | null): React.ReactNode {
+const processCitationsInReactChildren = (children: React.ReactNode, getCitationSource: (id: number) => Source | null): React.ReactNode => {
   // If children is not an array or string, just return it
   if (!children || (typeof children !== 'string' && !Array.isArray(children))) {
     return children;
@@ -120,10 +125,10 @@ function processCitationsInReactChildren(children: React.ReactNode, getCitationS
   }
   
   return children;
-}
+};
 
 // Process citation references in text content
-function processCitationsInText(text: string, getCitationSource: (id: number) => Source | null): React.ReactNode[] {
+const processCitationsInText = (text: string, getCitationSource: (id: number) => Source | null): React.ReactNode[] => {
   // Use improved regex to catch citation numbers more reliably
   // This will match patterns like [1], [42], etc. including when they appear at the end of a line or sentence
   const citationRegex = /\[(\d+)\]/g;
@@ -132,13 +137,7 @@ function processCitationsInText(text: string, getCitationSource: (id: number) =>
   let match;
   let position = 0;
   
-  // Debug log for troubleshooting
-  console.log("Processing citations in text:", text);
-  
   while ((match = citationRegex.exec(text)) !== null) {
-    // Log each match for debugging
-    console.log("Citation match found:", match[0], "at index", match.index);
-    
     // Add text before the citation
     if (match.index > lastIndex) {
       parts.push(text.substring(lastIndex, match.index));
@@ -148,9 +147,6 @@ function processCitationsInText(text: string, getCitationSource: (id: number) =>
     const citationId = parseInt(match[1], 10);
     const source = getCitationSource(citationId);
     
-    // Log the citation details
-    console.log("Citation ID:", citationId, "Source:", source ? "found" : "not found");
-    
     parts.push(
       <Citation 
         key={`citation-${citationId}-${position}`}
@@ -171,4 +167,4 @@ function processCitationsInText(text: string, getCitationSource: (id: number) =>
   }
   
   return parts;
-} 
+};