Simplify deferred codeblock highlighting and fix streaming

.changeset/old-dragons-crash.md

@@ -0,0 +1,5 @@
+---
+"gitbook": patch
+---
+
+Fix rendering of code blocks in Ask AI when being streamed

packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx

@@ -1,11 +1,9 @@
 'use client';
 
 import type { DocumentBlockCode } from '@gitbook/api';
-import { useEffect, useRef, useState } from 'react';
+import { useEffect, useMemo, useRef, useState } from 'react';
 
 import { useInViewportListener } from '@/components/hooks/useInViewportListener';
-import { useScrollListener } from '@/components/hooks/useScrollListener';
-import { useDebounceCallback, useEventCallback } from 'usehooks-ts';
 import type { BlockProps } from '../Block';
 import { CodeBlockRenderer } from './CodeBlockRenderer';
 import type { HighlightLine, RenderedInline } from './highlight';
@@ -22,53 +20,56 @@ type ClientBlockProps = Pick<BlockProps<DocumentBlockCode>, 'block' | 'style'> &
 export function ClientCodeBlock(props: ClientBlockProps) {
     const { block, style, inlines } = props;
     const blockRef = useRef<HTMLDivElement>(null);
-    const processedRef = useRef(false);
-    const isInViewportRef = useRef<boolean | null>(null);
-    const [lines, setLines] = useState<HighlightLine[]>(() => plainHighlight(block, []));
+    const [isInViewport, setIsInViewport] = useState(false);
+    const plainLines = useMemo(() => plainHighlight(block, []), [block]);
+    const [lines, setLines] = useState<null | HighlightLine[]>(null);
 
     // Preload the highlighter when the block is mounted.
     useEffect(() => {
         import('./highlight').then(({ preloadHighlight }) => preloadHighlight(block));
     }, [block]);
 
-    const runHighlight = useEventCallback(() => {
-        if (processedRef.current) {
-            return;
-        }
-        if (typeof window !== 'undefined') {
-            import('./highlight').then(({ highlight }) => {
-                highlight(block, inlines).then((lines) => {
-                    setLines(lines);
-                    processedRef.current = true;
-                });
-            });
-        }
-    });
-    const debouncedRunHighlight = useDebounceCallback(runHighlight, 1000);
-
+    // Detect when the block is in viewport
     useInViewportListener(
         blockRef,
         (isInViewport, disconnect) => {
-            // Disconnect once in viewport
             if (isInViewport) {
+                // Disconnect once in viewport
                 disconnect();
-                // If it's initially in viewport, we need to run the highlight
-                if (isInViewportRef.current === null) {
-                    runHighlight();
-                }
+                setIsInViewport(true);
             }
-            isInViewportRef.current = isInViewport;
         },
         { rootMargin: '200px' }
     );
 
-    const handleScroll = useDebounceCallback(() => {
-        if (isInViewportRef.current) {
-            debouncedRunHighlight();
+    // When the block is visible or updated, we need to re-run the highlight
+    useEffect(() => {
+        if (isInViewport) {
+            // If the block is in viewport, we need to run the highlight
+            let cancelled = false;
+
+            if (typeof window !== 'undefined') {
+                import('./highlight').then(({ highlight }) => {
+                    highlight(block, inlines).then((lines) => {
+                        if (cancelled) {
+                            return;
+                        }
+
+                        setLines(lines);
+                    });
+                });
+            }
+
+            return () => {
+                cancelled = true;
+            };
         }
-    }, 80);
 
-    useScrollListener(handleScroll, useRef(typeof window !== 'undefined' ? window : null));
+        // Otherwise if the block is not in viewport, we reset to the plain lines
+        setLines(null);
+    }, [isInViewport, block, inlines]);
 
-    return <CodeBlockRenderer ref={blockRef} block={block} style={style} lines={lines} />;
+    return (
+        <CodeBlockRenderer ref={blockRef} block={block} style={style} lines={lines ?? plainLines} />
+    );
 }