Feat/code block cls lower lazy boundary

.changeset/mighty-bikes-eat.md

@@ -0,0 +1,5 @@
+---
+"streamdown": minor
+---
+
+Move code block lazy loading to the highlighting layer so block shells render immediately with plain text content before syntax colors resolve. This improves visual stability and removes the spinner fallback for standard code blocks.

apps/website/content/docs/code-blocks.mdx

@@ -211,6 +211,12 @@ function example() {
 
 The unterminated block parser ensures the code block renders properly even without the closing backticks.
 
+### Loading Behavior
+
+Code block shells render immediately with plain text content, then syntax colors are applied when highlighting resolves.
+
+This keeps code readable on first paint and improves visual stability during lazy highlight loading.
+
 ### Disabling Interactions During Streaming
 
 Use the `isAnimating` prop to disable copy buttons while streaming:

packages/streamdown/__tests__/code-block-loading.test.tsx

@@ -0,0 +1,173 @@
+import { act, render, waitFor } from "@testing-library/react";
+import type { ComponentType } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { HighlightOptions, HighlightResult } from "../lib/plugin-types";
+
+// Helper for controllable promise (to avoid arbitrary delays)
+const createDeferred = <T,>() => {
+  let resolve!: (value: T | PromiseLike<T>) => void;
+  let reject!: (reason?: unknown) => void;
+
+  const promise = new Promise<T>((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+
+  return { promise, resolve, reject };
+};
+
+describe("Code block loading behavior", () => {
+  afterEach(() => {
+    vi.resetModules();
+    vi.clearAllMocks();
+    vi.doUnmock("../lib/code-block/highlighted-body");
+  });
+
+  it("renders readable text and no loader before lazy module resolves", async () => {
+    const lazyModule = createDeferred<{
+      HighlightedCodeBlockBody: ComponentType<{
+        code: string;
+        language: string;
+        raw: HighlightResult;
+      }>;
+    }>();
+
+    // mock the targeted part that MUST be lazyloaded
+    let lazyLoaded = false;
+    vi.doMock("../lib/code-block/highlighted-body", () =>
+      lazyModule.promise.then((mod) => {
+        lazyLoaded = true;
+        return mod;
+      })
+    );
+
+    const { StreamdownContext } = await import("../index");
+    const { CodeBlock } = await import("../lib/code-block");
+
+    const { container } = render(
+      <StreamdownContext.Provider
+        value={{
+          shikiTheme: ["github-light", "github-dark"],
+          controls: true,
+          isAnimating: false,
+          mode: "streaming",
+        }}
+      >
+        <CodeBlock code={"const x = 1;\n"} language="javascript" />
+      </StreamdownContext.Provider>
+    );
+
+    // Before Highlighter lazy component loads we already see text content and no loader
+    const body = container.querySelector('[data-streamdown="code-block-body"]');
+    expect(body?.textContent).toContain("const x = 1;");
+    expect(container.querySelector(".animate-spin")).toBeNull();
+    expect(lazyLoaded).toBe(false);
+
+    // trigger mocked lazyload resolve
+    lazyModule.resolve({
+      HighlightedCodeBlockBody: () => (
+        <pre data-streamdown="code-block-body">lazy module resolved</pre>
+      ),
+    });
+
+    await waitFor(() => {
+      expect(container.textContent).toContain("lazy module resolved");
+      expect(lazyLoaded).toBe(true);
+    });
+  });
+
+  it("applies highlight styles only after manual callback resolution", async () => {
+    const { StreamdownContext } = await import("../index");
+    const { PluginContext } = await import("../lib/plugin-context");
+    const { HighlightedCodeBlockBody } = await import(
+      "../lib/code-block/highlighted-body"
+    );
+
+    // keep external ref to trigger manually
+    let resolveHighlight: ((result: HighlightResult) => void) | null = null;
+
+    const rawResult: HighlightResult = {
+      bg: "transparent",
+      fg: "inherit",
+      tokens: [
+        [
+          {
+            content: "const x = 1;",
+            color: "inherit",
+            bgColor: "transparent",
+            htmlStyle: {},
+            offset: 0,
+          },
+        ],
+      ],
+    };
+
+    const highlightedResult: HighlightResult = {
+      ...rawResult,
+      tokens: [
+        [
+          {
+            ...rawResult.tokens[0][0],
+            color: "#ff0000",
+          },
+        ],
+      ],
+    };
+
+    const codePlugin = {
+      name: "shiki" as const,
+      type: "code-highlighter" as const,
+      highlight: vi.fn(
+        (_: HighlightOptions, callback?: (result: HighlightResult) => void) => {
+          resolveHighlight = callback ?? null;
+          return null;
+        }
+      ),
+      supportsLanguage: vi.fn().mockReturnValue(true),
+      getSupportedLanguages: vi.fn().mockReturnValue(["javascript"]),
+      getThemes: vi.fn().mockReturnValue(["github-light", "github-dark"]),
+    };
+
+    const { container } = render(
+      <PluginContext.Provider value={{ code: codePlugin as any }}>
+        <StreamdownContext.Provider
+          value={{
+            shikiTheme: ["github-light", "github-dark"],
+            controls: true,
+            isAnimating: false,
+            mode: "streaming",
+          }}
+        >
+          <HighlightedCodeBlockBody
+            code="const x = 1;"
+            language="javascript"
+            raw={rawResult}
+          />
+        </StreamdownContext.Provider>
+      </PluginContext.Provider>
+    );
+
+    const initialToken = container.querySelector(
+      '[data-streamdown="code-block-body"] code > span > span'
+    ) as HTMLElement | null;
+    expect(initialToken).toBeTruthy();
+    expect(initialToken?.style.getPropertyValue("--sdm-c")).toBe("inherit");
+
+    await waitFor(() => {
+      expect(codePlugin.highlight).toHaveBeenCalledTimes(1);
+      expect(resolveHighlight).toBeTruthy();
+    });
+
+    // Manually trigger the highlighting
+    await act(async () => {
+      resolveHighlight?.(highlightedResult);
+    });
+
+    await waitFor(() => {
+      const updatedToken = container.querySelector(
+        '[data-streamdown="code-block-body"] code > span > span'
+      ) as HTMLElement | null;
+      expect(updatedToken?.style.getPropertyValue("--sdm-c")).toBe("#ff0000");
+    });
+  });
+});

packages/streamdown/__tests__/components.test.tsx

@@ -247,7 +247,7 @@ describe("Markdown Components", () => {
         </Code>
       );
 
-      // Wait for lazy-loaded CodeBlock component
+      // Wait for code block to render
       await waitFor(() => {
         const codeBlock = container.querySelector(
           '[data-streamdown="code-block"]'
@@ -295,7 +295,7 @@ describe("Markdown Components", () => {
         </Code>
       );
 
-      // Wait for lazy-loaded CodeBlock component
+      // Wait for code block to render
       await waitFor(() => {
         const codeBlock = container.querySelector(
           '[data-streamdown="code-block"]'
@@ -332,7 +332,7 @@ describe("Markdown Components", () => {
         </Code>
       );
 
-      // Wait for Suspense boundary to resolve
+      // Wait for code block to render
       await waitFor(() => {
         const codeBlock = container.querySelector(
           '[data-streamdown="code-block"]'

packages/streamdown/lib/code-block/highlighted-body.tsx

@@ -0,0 +1,58 @@
+import { type HTMLAttributes, useContext, useEffect, useState } from "react";
+import type { BundledLanguage } from "shiki";
+import { StreamdownContext } from "../../index";
+import { useCodePlugin } from "../plugin-context";
+import type { HighlightResult } from "../plugin-types";
+import { CodeBlockBody } from "./body";
+
+type HighlightedCodeBlockBodyProps = HTMLAttributes<HTMLPreElement> & {
+  code: string;
+  language: string;
+  raw: HighlightResult;
+};
+
+export const HighlightedCodeBlockBody = ({
+  code,
+  language,
+  raw,
+  className,
+  ...rest
+}: HighlightedCodeBlockBodyProps) => {
+  const { shikiTheme } = useContext(StreamdownContext);
+  const codePlugin = useCodePlugin();
+  const [result, setResult] = useState<HighlightResult>(raw);
+
+  useEffect(() => {
+    if (!codePlugin) {
+      setResult(raw);
+      return;
+    }
+
+    const cachedResult = codePlugin.highlight(
+      {
+        code,
+        language: language as BundledLanguage,
+        themes: shikiTheme,
+      },
+      (highlightedResult) => {
+        setResult(highlightedResult);
+      }
+    );
+
+    if (cachedResult) {
+      setResult(cachedResult);
+      return;
+    }
+
+    setResult(raw);
+  }, [code, language, shikiTheme, codePlugin, raw]);
+
+  return (
+    <CodeBlockBody
+      className={className}
+      language={language}
+      result={result}
+      {...rest}
+    />
+  );
+};

packages/streamdown/lib/code-block/index.tsx

@@ -1,13 +1,4 @@
-import {
-  type HTMLAttributes,
-  useContext,
-  useEffect,
-  useMemo,
-  useState,
-} from "react";
-import type { BundledLanguage } from "shiki";
-import { StreamdownContext } from "../../index";
-import { useCodePlugin } from "../plugin-context";
+import { type HTMLAttributes, lazy, Suspense, useMemo } from "react";
 import type { HighlightResult } from "../plugin-types";
 import { CodeBlockBody } from "./body";
 import { CodeBlockContainer } from "./container";
@@ -23,6 +14,12 @@ type CodeBlockProps = HTMLAttributes<HTMLPreElement> & {
   isIncomplete?: boolean;
 };
 
+const HighlightedCodeBlockBody = lazy(() =>
+  import("./highlighted-body").then((mod) => ({
+    default: mod.HighlightedCodeBlockBody,
+  }))
+);
+
 export const CodeBlock = ({
   code,
   language,
@@ -31,9 +28,6 @@ export const CodeBlock = ({
   isIncomplete = false,
   ...rest
 }: CodeBlockProps) => {
-  const { shikiTheme } = useContext(StreamdownContext);
-  const codePlugin = useCodePlugin();
-
   // Remove trailing newlines to prevent empty line at end of code blocks
   const trimmedCode = useMemo(
     () => code.replace(TRAILING_NEWLINES_REGEX, ""),
@@ -58,49 +52,28 @@ export const CodeBlock = ({
     [trimmedCode]
   );
 
-  // Use raw as initial state
-  const [result, setResult] = useState<HighlightResult>(raw);
-
-  // Try to get cached result or subscribe to highlighting
-  useEffect(() => {
-    // If no code plugin, just use raw tokens (plain text)
-    if (!codePlugin) {
-      setResult(raw);
-      return;
-    }
-
-    const cachedResult = codePlugin.highlight(
-      {
-        code: trimmedCode,
-        language: language as BundledLanguage,
-        themes: shikiTheme,
-      },
-      (highlightedResult) => {
-        setResult(highlightedResult);
-      }
-    );
-
-    if (cachedResult) {
-      // Already cached, use it immediately
-      setResult(cachedResult);
-      return;
-    }
-
-    // Not cached - reset to raw tokens while waiting for highlighting
-    // This is critical for streaming: ensures we show current code, not stale tokens
-    setResult(raw);
-  }, [trimmedCode, language, shikiTheme, codePlugin, raw]);
-
   return (
     <CodeBlockContext.Provider value={{ code }}>
       <CodeBlockContainer isIncomplete={isIncomplete} language={language}>
         <CodeBlockHeader language={language}>{children}</CodeBlockHeader>
-        <CodeBlockBody
-          className={className}
-          language={language}
-          result={result}
-          {...rest}
-        />
+        <Suspense
+          fallback={
+            <CodeBlockBody
+              className={className}
+              language={language}
+              result={raw}
+              {...rest}
+            />
+          }
+        >
+          <HighlightedCodeBlockBody
+            className={className}
+            code={trimmedCode}
+            language={language}
+            raw={raw}
+            {...rest}
+          />
+        </Suspense>
       </CodeBlockContainer>
     </CodeBlockContext.Provider>
   );