Repository: GitbookIO/gitbook
Pull request #2977: Simplify deferred codeblock highlighting and fix streaming
Author: login: SamyPesse
Created: 2025-03-12T19:19:03Z | Merged: 2025-03-12T23:25:49Z



Diff:
diff --git a/.changeset/old-dragons-crash.md b/.changeset/old-dragons-crash.md
new file mode 100644
index 0000000000..244b97ca09
--- /dev/null
+++ b/.changeset/old-dragons-crash.md
@@ -0,0 +1,5 @@
+---
+"gitbook": patch
+---
+
+Fix rendering of code blocks in Ask AI when being streamed
diff --git a/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx b/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
index 0c11112986..77d9b33c52 100644
--- a/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
+++ b/packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
@@ -1,11 +1,11 @@
 '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 { useDebounceCallback } from 'usehooks-ts';
 import type { BlockProps } from '../Block';
 import { CodeBlockRenderer } from './CodeBlockRenderer';
 import type { HighlightLine, RenderedInline } from './highlight';
@@ -22,53 +22,82 @@ 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 isInViewportRef = useRef(false);
+    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;
-                });
-            });
+    // When user scrolls, we need to wait for the scroll to finish before running the highlight
+    const isScrollingRef = useRef(false);
+    const onFinishScrolling = useDebounceCallback(() => {
+        isScrollingRef.current = false;
+
+        // If the block is in the viewport after the scroll, we need to run the highlight
+        if (isInViewportRef.current) {
+            setIsInViewport(true);
         }
-    });
-    const debouncedRunHighlight = useDebounceCallback(runHighlight, 1000);
+    }, 100);
+    useScrollListener(
+        () => {
+            isScrollingRef.current = true;
+            onFinishScrolling();
+        },
+        useRef(typeof window !== 'undefined' ? window : null)
+    );
 
+    // Detect when the block is in viewport
     useInViewportListener(
         blockRef,
         (isInViewport, disconnect) => {
-            // Disconnect once in viewport
+            isInViewportRef.current = isInViewport;
+            if (isScrollingRef.current) {
+                // If the user is scrolling, we don't want to run the highlight
+                // because it will be run when the scroll is finished
+                return;
+            }
+
             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} />
+    );
 }


Reviews:
user: login: SamyPesse
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T19:19:57Z
user: login: gregberge
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T19:44:59Z
user: login: gregberge
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T19:45:50Z
user: login: gregberge
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T19:47:03Z
user: login: gregberge
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T19:47:56Z
user: login: gregberge
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T19:54:08Z

Review comments:
diff_hunk: @@ -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();
path: packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
user: login: SamyPesse
body: @gregberge I don't understand what you tried to do here: if `isInViewportRef.current` is true, then you already ran the highlighting when it becomes visible, why defer another highlighting?
created_at: 2025-03-12T19:19:57Z
author_association: MEMBER
diff_hunk: @@ -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();
path: packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
user: login: gregberge
body: I trigger highlighting on scroll and not from the observer, so it's smoother.

But the main issue on scroll is the outline, not that so I think you can simplify it.
created_at: 2025-03-12T19:44:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -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();
path: packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
user: login: gregberge
body: The main problem is triggering compute while the user scrolls. It's not good. I wanted to highlight when the user pause the scroll.
created_at: 2025-03-12T19:45:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -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();
path: packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
user: login: gregberge
body: But are you sure it's the issue? Because the blocks are empty if it was the issue they would be not highlighted but displayed at least.
created_at: 2025-03-12T19:47:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -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();
path: packages/gitbook/src/components/DocumentView/CodeBlock/ClientCodeBlock.tsx
user: login: gregberge
body: Because the initial presence in viewport was working so, I have some doubts on this fix.
created_at: 2025-03-12T19:47:56Z
author_association: CONTRIBUTOR

PR comments:
user: login: changeset-bot[bot]
created_at: 2025-03-12T19:19:06Z
body: ###  ⚠️  No Changeset found

Latest commit: 2f8f41c718b2dd33ebdccf345d4e8ad5f081f302

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. **If these changes should result in a version bump, you need to add a changeset.**

<details><summary>This PR includes no changesets</summary>

  When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

</details>

[Click here to learn what changesets are, and how to add one](https://github.com/changesets/changesets/blob/main/docs/adding-a-changeset.md).

[Click here if you're a maintainer who wants to add a changeset to this PR](https://github.com/GitbookIO/gitbook/new/fix-code-highlighting-streaming?filename=.changeset/calm-pianos-crash.md&value=---%0A%22gitbook%22%3A%20patch%0A---%0A%0ASImplify%20deferred%20codeblock%20highlighting%20and%20fix%20streaming%0A)


author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-12T19:24:42Z
body: Summary of the deployments:

### Version 1 (production)

| Version | URL | Status |
| --- | --- | --- |
| Latest commit | [https://1559d9d2.gitbook-open.pages.dev](https://1559d9d2.gitbook-open.pages.dev) | ✅ |
| PR | [https://pr2977.gitbook-open.pages.dev](https://pr2977.gitbook-open.pages.dev) | ✅ |

### Version 2 (experimental)

| Version | URL | Status |
| --- | --- | --- |
| Vercel | [https://gitbook-v2-mzg4mlren-gitbook.vercel.app](https://gitbook-v2-mzg4mlren-gitbook.vercel.app) | ✅ |
| Cloudflare | [https://bcc97d96-gitbook-open-v2-preview.gitbook.workers.dev](https://bcc97d96-gitbook-open-v2-preview.gitbook.workers.dev) | ✅ |

### Test content

| Site | v1 | v2 |
| --- | --- | --- |
| GitBook | [https://1559d9d2.gitbook-open.pages.dev/docs.gitbook.com](https://1559d9d2.gitbook-open.pages.dev/docs.gitbook.com) | [https://gitbook-v2-mzg4mlren-gitbook.vercel.app/url/docs.gitbook.com](https://gitbook-v2-mzg4mlren-gitbook.vercel.app/url/docs.gitbook.com) |
| E2E | [https://1559d9d2.gitbook-open.pages.dev/gitbook.gitbook.io/test-gitbook-open](https://1559d9d2.gitbook-open.pages.dev/gitbook.gitbook.io/test-gitbook-open) | [https://gitbook-v2-mzg4mlren-gitbook.vercel.app/url/gitbook.gitbook.io/test-gitbook-open](https://gitbook-v2-mzg4mlren-gitbook.vercel.app/url/gitbook.gitbook.io/test-gitbook-open) |
author_association: CONTRIBUTOR
user: login: argos-ci[bot]
created_at: 2025-03-12T19:26:39Z
body: **The latest updates on your projects.** Learn more about [Argos notifications ↗︎](https://argos-ci.com/docs/notifications)

| Build | Status | Details | Updated (UTC) |
| :---- | :----- | :------ | :------------ |
| **customers-v1** ([Inspect](https://app.argos-ci.com/GitbookIO/gitbook/builds/4174)) | 👍 Changes approved | 9 changed | Mar 12, 2025, 10:20 PM |
| **customers-v2** ([Inspect](https://app.argos-ci.com/GitbookIO/gitbook/builds/4175)) | 👍 Changes approved | 8 changed | Mar 12, 2025, 10:21 PM |
| **default** ([Inspect](https://app.argos-ci.com/GitbookIO/gitbook/builds/4176)) | 👍 Changes approved | 8 changed | Mar 12, 2025, 10:22 PM |
| **v2-vercel** ([Inspect](https://app.argos-ci.com/GitbookIO/gitbook/builds/4177)) | 👍 Changes approved | 2 changed | Mar 12, 2025, 10:24 PM |
author_association: NONE