fix: prevent memory leaks in animated markdown components

- Memoize context values in AppMarkdownStreaming to prevent unnecessary re-renders
- Extract component functions in useMarkdownComponents using useCallback to prevent recreation
- Add cleanup logic in TokenizedText to clear animated chunks when streaming finishes
- Fix missing isStreamFinished dependency in TokenizedText useMemo
- Memoize animation styles in animation-helpers to prevent object recreation

These changes address critical memory leak sources that were causing 2GB+ memory usage growth.

Co-Authored-By: nate@buster.so <nate@buster.so>

apps/web/src/components/ui/streaming/AppMarkdownStreaming/AppMarkdownStreaming.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import React, { useContext } from 'react';
+import React, { useContext, useMemo } from 'react';
 import { createContext } from 'react';
 import { cn } from '@/lib/classMerge';
 import { useMarkdownStreaming } from './useMarkdownStreaming';
@@ -31,16 +31,20 @@ const AppMarkdownStreaming = ({
     isStreamFinished
   });
 
+  const contextValue = useMemo(
+    () => ({
+      animation,
+      animationDuration,
+      animationTimingFunction,
+      isStreamFinished,
+      isThrottleStreamingFinished: isFinished,
+      stripFormatting
+    }),
+    [animation, animationDuration, animationTimingFunction, isStreamFinished, isFinished, stripFormatting]
+  );
+
   return (
-    <AppMarkdownStreamingContext.Provider
-      value={{
-        animation,
-        animationDuration,
-        animationTimingFunction,
-        isStreamFinished,
-        isThrottleStreamingFinished: isFinished,
-        stripFormatting
-      }}>
+    <AppMarkdownStreamingContext.Provider value={contextValue}>
       <div className={cn('flex flex-col space-y-2.5', className)}>
         {blockMatches.map((blockMatch, index) => {
           const Component = blockMatch.block.component;

apps/web/src/components/ui/typography/AnimatedMarkdown/TokenizedText.tsx

@@ -1,4 +1,4 @@
-import React, { useMemo, useRef, useCallback } from 'react';
+import React, { useMemo, useRef, useCallback, useEffect } from 'react';
 import { type MarkdownAnimation } from '../animation-common';
 import { createAnimationStyle } from '../animation-common';
 
@@ -31,7 +31,7 @@ const TokenizedText: React.FC<TokenizedTextProps> = React.memo(
           animationTimingFunction,
           isStreamFinished
         }),
-      [animation, animationDuration, animationTimingFunction]
+      [animation, animationDuration, animationTimingFunction, isStreamFinished]
     );
 
     // Memoize the span creation function to avoid recreating it on every render
@@ -99,6 +99,13 @@ const TokenizedText: React.FC<TokenizedTextProps> = React.memo(
       return [baseText, ...animatedSpans];
     }, [text, createSpan]);
 
+    useEffect(() => {
+      if (isStreamFinished) {
+        animatedChunksRef.current = [];
+        previousTextRef.current = '';
+      }
+    }, [isStreamFinished]);
+
     return <>{content}</>;
   }
 );

apps/web/src/components/ui/typography/AnimatedMarkdown/animation-helpers.tsx

@@ -1,5 +1,5 @@
 import type { AnimatedMarkdownProps } from './AnimatedMarkdown';
-import React from 'react';
+import React, { useMemo } from 'react';
 import TokenizedText from './TokenizedText';
 import { createAnimationStyle } from '../animation-common';
 
@@ -52,12 +52,17 @@ export const animateTokenizedText = (
       // This else block might still wrap elements if they are not explicitly handled
       // by the ReactMarkdown components mapping (e.g. custom components not passing animateText to children)
       // For standard HTML elements, the `components` mapping should handle animation of children.
+      const animationStyle = useMemo(
+        () => createAnimationStyle(animationsProps),
+        [animationsProps.animation, animationsProps.animationDuration, animationsProps.animationTimingFunction, animationsProps.isStreamFinished]
+      );
+
       return (
         <span
           key={`react-element-${index}`}
           data-testid="other-markdown-element"
           style={{
-            ...createAnimationStyle(animationsProps),
+            ...animationStyle,
             whiteSpace: 'pre-wrap',
             display: isInlineElement ? 'inline' : 'inline-block'
           }}>
@@ -65,12 +70,17 @@ export const animateTokenizedText = (
         </span>
       );
     }
+    const animationStyle = useMemo(
+      () => createAnimationStyle(animationsProps),
+      [animationsProps.animation, animationsProps.animationDuration, animationsProps.animationTimingFunction, animationsProps.isStreamFinished]
+    );
+
     return (
       <span
         key={`unknown-element-${index}`}
         data-testid="animated-markdown-element"
         style={{
-          ...createAnimationStyle(animationsProps),
+          ...animationStyle,
           whiteSpace: 'pre-wrap',
           display: 'inline'
         }}>

apps/web/src/components/ui/typography/AnimatedMarkdown/useMarkdownComponents.tsx

@@ -1,4 +1,4 @@
-import { useMemo } from 'react';
+import { useMemo, useCallback } from 'react';
 import { type MarkdownAnimation, type MarkdownAnimationTimingFunction } from '../animation-common';
 import { type Components } from 'react-markdown';
 import {
@@ -49,132 +49,278 @@ export const useMarkdownComponents = ({
     };
   }, [animation, animationDuration, animationTimingFunction, isStreamFinished, stripFormatting]);
 
+  const ParagraphWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <ParagraphComponent {...commonProps} className={className} style={style}>
+        {children}
+      </ParagraphComponent>
+    ),
+    [commonProps]
+  );
+
+  const H1Wrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <HeaderComponent {...commonProps} tag="h1" className={className} style={style}>
+        {children}
+      </HeaderComponent>
+    ),
+    [commonProps]
+  );
+
+  const H2Wrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <HeaderComponent {...commonProps} tag="h2" className={className} style={style}>
+        {children}
+      </HeaderComponent>
+    ),
+    [commonProps]
+  );
+
+  const H3Wrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <HeaderComponent {...commonProps} tag="h3" className={className} style={style}>
+        {children}
+      </HeaderComponent>
+    ),
+    [commonProps]
+  );
+
+  const H4Wrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <HeaderComponent {...commonProps} tag="h4" className={className} style={style}>
+        {children}
+      </HeaderComponent>
+    ),
+    [commonProps]
+  );
+
+  const H5Wrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <HeaderComponent {...commonProps} tag="h5" className={className} style={style}>
+        {children}
+      </HeaderComponent>
+    ),
+    [commonProps]
+  );
+
+  const H6Wrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <HeaderComponent {...commonProps} tag="h6" className={className} style={style}>
+        {children}
+      </HeaderComponent>
+    ),
+    [commonProps]
+  );
+
+  const BlockquoteWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <BlockquoteComponent {...commonProps} className={className} style={style}>
+        {children}
+      </BlockquoteComponent>
+    ),
+    [commonProps]
+  );
+
+  const StrongWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <StrongComponent {...commonProps} className={className} style={style}>
+        {children}
+      </StrongComponent>
+    ),
+    [commonProps]
+  );
+
+  const EmphasisWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <EmphasisComponent {...commonProps} className={className} style={style}>
+        {children}
+      </EmphasisComponent>
+    ),
+    [commonProps]
+  );
+
+  const DeleteWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <DeleteComponent {...commonProps} className={className} style={style}>
+        {children}
+      </DeleteComponent>
+    ),
+    [commonProps]
+  );
+
+  const LinkWrapper = useCallback(
+    ({ children, className, style, href }: any) => (
+      <LinkComponent {...commonProps} className={className} style={style} href={href}>
+        {children}
+      </LinkComponent>
+    ),
+    [commonProps]
+  );
+
+  const ImageWrapper = useCallback(
+    ({ className, style, src, alt }: any) => (
+      <ImageComponent {...commonProps} className={className} style={style} src={src} alt={alt} />
+    ),
+    [commonProps]
+  );
+
+  const HorizontalRuleWrapper = useCallback(
+    ({ className, style }: any) => (
+      <HorizontalRuleComponent {...commonProps} className={className} style={style} />
+    ),
+    [commonProps]
+  );
+
+  const TableWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <TableComponent {...commonProps} className={className} style={style}>
+        {children}
+      </TableComponent>
+    ),
+    [commonProps]
+  );
+
+  const TableHeadWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <TableHeadComponent {...commonProps} className={className} style={style}>
+        {children}
+      </TableHeadComponent>
+    ),
+    [commonProps]
+  );
+
+  const TableBodyWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <TableBodyComponent {...commonProps} className={className} style={style}>
+        {children}
+      </TableBodyComponent>
+    ),
+    [commonProps]
+  );
+
+  const TableRowWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <TableRowComponent {...commonProps} className={className} style={style}>
+        {children}
+      </TableRowComponent>
+    ),
+    [commonProps]
+  );
+
+  const UnorderedListWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <UnorderedListComponent {...commonProps} className={className} style={style}>
+        {children}
+      </UnorderedListComponent>
+    ),
+    [commonProps]
+  );
+
+  const OrderedListWrapper = useCallback(
+    ({ children, className, style, ...rest }: any) => (
+      <OrderedListComponent {...commonProps} className={className} style={style} {...rest}>
+        {children}
+      </OrderedListComponent>
+    ),
+    [commonProps]
+  );
+
+  const ListItemWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <ListItemComponent {...commonProps} className={className} style={style}>
+        {children}
+      </ListItemComponent>
+    ),
+    [commonProps]
+  );
+
+  const TableCellWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <TableCellComponent className={className} style={style}>
+        {children}
+      </TableCellComponent>
+    ),
+    []
+  );
+
+  const TableHeaderCellWrapper = useCallback(
+    ({ children, className, style }: any) => (
+      <TableHeaderCellComponent className={className} style={style}>
+        {children}
+      </TableHeaderCellComponent>
+    ),
+    []
+  );
+
+  const BreakWrapper = useCallback(
+    ({ className, style }: any) => <BreakComponent className={className} style={style} />,
+    []
+  );
+
+  const CodeWrapper = useCallback(
+    ({ children, className, style, ...rest }: any) => (
+      <CodeComponent {...commonProps} className={className} style={style} isInline={true}>
+        {children}
+      </CodeComponent>
+    ),
+    [commonProps]
+  );
+
   const components: Components = useMemo(() => {
     return {
-      // Components with animations
-      p: ({ children, className, style }) => (
-        <ParagraphComponent {...commonProps} className={className} style={style}>
-          {children}
-        </ParagraphComponent>
-      ),
-      h1: ({ children, className, style }) => (
-        <HeaderComponent {...commonProps} tag="h1" className={className} style={style}>
-          {children}
-        </HeaderComponent>
-      ),
-      h2: ({ children, className, style }) => (
-        <HeaderComponent {...commonProps} tag="h2" className={className} style={style}>
-          {children}
-        </HeaderComponent>
-      ),
-      h3: ({ children, className, style }) => (
-        <HeaderComponent {...commonProps} tag="h3" className={className} style={style}>
-          {children}
-        </HeaderComponent>
-      ),
-      h4: ({ children, className, style }) => (
-        <HeaderComponent {...commonProps} tag="h4" className={className} style={style}>
-          {children}
-        </HeaderComponent>
-      ),
-      h5: ({ children, className, style }) => (
-        <HeaderComponent {...commonProps} tag="h5" className={className} style={style}>
-          {children}
-        </HeaderComponent>
-      ),
-      h6: ({ children, className, style }) => (
-        <HeaderComponent {...commonProps} tag="h6" className={className} style={style}>
-          {children}
-        </HeaderComponent>
-      ),
-      blockquote: ({ children, className, style }) => (
-        <BlockquoteComponent {...commonProps} className={className} style={style}>
-          {children}
-        </BlockquoteComponent>
-      ),
-      strong: ({ children, className, style }) => (
-        <StrongComponent {...commonProps} className={className} style={style}>
-          {children}
-        </StrongComponent>
-      ),
-      em: ({ children, className, style }) => (
-        <EmphasisComponent {...commonProps} className={className} style={style}>
-          {children}
-        </EmphasisComponent>
-      ),
-      del: ({ children, className, style }) => (
-        <DeleteComponent {...commonProps} className={className} style={style}>
-          {children}
-        </DeleteComponent>
-      ),
-      a: ({ children, className, style, href }) => (
-        <LinkComponent {...commonProps} className={className} style={style} href={href}>
-          {children}
-        </LinkComponent>
-      ),
-      img: ({ className, style, src, alt }) => (
-        <ImageComponent {...commonProps} className={className} style={style} src={src} alt={alt} />
-      ),
-      hr: ({ className, style }) => (
-        <HorizontalRuleComponent {...commonProps} className={className} style={style} />
-      ),
-      table: ({ children, className, style }) => (
-        <TableComponent {...commonProps} className={className} style={style}>
-          {children}
-        </TableComponent>
-      ),
-      thead: ({ children, className, style }) => (
-        <TableHeadComponent {...commonProps} className={className} style={style}>
-          {children}
-        </TableHeadComponent>
-      ),
-      tbody: ({ children, className, style }) => (
-        <TableBodyComponent {...commonProps} className={className} style={style}>
-          {children}
-        </TableBodyComponent>
-      ),
-      tr: ({ children, className, style }) => (
-        <TableRowComponent {...commonProps} className={className} style={style}>
-          {children}
-        </TableRowComponent>
-      ),
-
-      ul: ({ children, className, style }) => (
-        <UnorderedListComponent {...commonProps} className={className} style={style}>
-          {children}
-        </UnorderedListComponent>
-      ),
-      ol: ({ children, className, style, ...rest }) => (
-        <OrderedListComponent {...commonProps} className={className} style={style} {...rest}>
-          {children}
-        </OrderedListComponent>
-      ),
-      li: ({ children, className, style }) => {
-        return (
-          <ListItemComponent {...commonProps} className={className} style={style}>
-            {children}
-          </ListItemComponent>
-        );
-      },
-      td: ({ children, className, style }) => (
-        <TableCellComponent className={className} style={style}>
-          {children}
-        </TableCellComponent>
-      ),
-      th: ({ children, className, style }) => (
-        <TableHeaderCellComponent className={className} style={style}>
-          {children}
-        </TableHeaderCellComponent>
-      ),
-      br: ({ className, style }) => <BreakComponent className={className} style={style} />,
-      code: ({ children, className, style, ...rest }) => (
-        //we can assume that code is inline if it reach to this point
-        <CodeComponent {...commonProps} className={className} style={style} isInline={true}>
-          {children}
-        </CodeComponent>
-      )
+      p: ParagraphWrapper,
+      h1: H1Wrapper,
+      h2: H2Wrapper,
+      h3: H3Wrapper,
+      h4: H4Wrapper,
+      h5: H5Wrapper,
+      h6: H6Wrapper,
+      blockquote: BlockquoteWrapper,
+      strong: StrongWrapper,
+      em: EmphasisWrapper,
+      del: DeleteWrapper,
+      a: LinkWrapper,
+      img: ImageWrapper,
+      hr: HorizontalRuleWrapper,
+      table: TableWrapper,
+      thead: TableHeadWrapper,
+      tbody: TableBodyWrapper,
+      tr: TableRowWrapper,
+      ul: UnorderedListWrapper,
+      ol: OrderedListWrapper,
+      li: ListItemWrapper,
+      td: TableCellWrapper,
+      th: TableHeaderCellWrapper,
+      br: BreakWrapper,
+      code: CodeWrapper
     };
-  }, [commonProps]);
+  }, [
+    ParagraphWrapper,
+    H1Wrapper,
+    H2Wrapper,
+    H3Wrapper,
+    H4Wrapper,
+    H5Wrapper,
+    H6Wrapper,
+    BlockquoteWrapper,
+    StrongWrapper,
+    EmphasisWrapper,
+    DeleteWrapper,
+    LinkWrapper,
+    ImageWrapper,
+    HorizontalRuleWrapper,
+    TableWrapper,
+    TableHeadWrapper,
+    TableBodyWrapper,
+    TableRowWrapper,
+    UnorderedListWrapper,
+    OrderedListWrapper,
+    ListItemWrapper,
+    TableCellWrapper,
+    TableHeaderCellWrapper,
+    BreakWrapper,
+    CodeWrapper
+  ]);
 
   return { components, commonProps };
 };