@mui/internal-docs-infra 0.11.1-canary.21 → 0.11.1-canary.22

package/CodeHighlighter/CodeHighlighter.mjs

@@ -106,7 +106,11 @@ export function CodeHighlighter(props) {
   };
 
   // No ContentLoading: render the content/full-load directly, with no loading
-  // fallback (the client shows nothing until content is ready).
+  // fallback (the client shows nothing until content is ready). For
+  // `highlightAt: 'init'`, `createClientProps` folds each variant's highlighted-visible
+  // `fallbackCritical` over its plain `fallback`, so `<Pre>` paints the visible frames
+  // highlighted on the first render (no decompression) and decodes the full tree after
+  // paint (the `decodeAllowed` latch).
   if (process.env.NODE_ENV !== "production") renderChunk.displayName = "renderChunk";
   if (!ContentLoading) {
     if (props.highlightAfter === 'stream') {
@@ -158,7 +162,12 @@ export function CodeHighlighter(props) {
     initialFilename: initialData.initialFilename,
     initialSource: initialData.initialSource,
     initialExtraFiles: initialData.initialExtraFiles,
-    ContentLoading
+    ContentLoading,
+    // Compressing the residual fallbacks only shrinks the server→client payload. This
+    // entry is isomorphic, so when it runs on the client (e.g. a Pages-Router app that
+    // renders everything client-side) there is no wire — skip the compress and keep the
+    // fallbacks inline rather than compressing them only to decompress them right back.
+    compressResidual: typeof window === 'undefined'
   });
   return renderChunk({
     preloaded: codeForClient,

package/CodeHighlighter/CodeHighlighterClient.mjs

@@ -8,6 +8,7 @@ import { hasAllVariants } from "../pipeline/loadIsomorphicCodeVariant/hasAllCode
 import { CodeHighlighterFallbackContext } from "./CodeHighlighterFallbackContext.mjs";
 import { useControlledCode } from "../CodeControllerContext/index.mjs";
 import { codeToFallbackProps, deriveFallbacksFromCode, stripFallbackHastsFromCode } from "./codeToFallbackProps.mjs";
+import { resolveFallbackCritical } from "./resolveFallbackCritical.mjs";
 import { decompressResidualFallbacks, residualDictionaryText, scatterResidualFallbacks } from "./fallbackCompression.mjs";
 import { mergeCodeMetadata } from "../pipeline/loadIsomorphicCodeVariant/mergeCodeMetadata.mjs";
 import { getAvailableTransforms } from "../pipeline/loadIsomorphicCodeVariant/getAvailableTransforms.mjs";
@@ -21,6 +22,7 @@ import { useChunk } from "../CoordinatedLazy/useChunk.mjs";
 import { useCoordinatedSwap } from "../CoordinatedLazy/useCoordinatedSwap.mjs";
 import { CoordinatedFallbackContext } from "../CoordinatedLazy/CoordinatedFallbackContext.mjs";
 import { CoordinatedContentContext } from "../CoordinatedLazy/CoordinatedContentContext.mjs";
+import { requestIdle } from "../useCoordinated/scheduleTasks.mjs";
 import * as Errors from "./errors.mjs";
 import { jsx as _jsx } from "react/jsx-runtime";
 const DEBUG = false; // Set to true for debugging purposes
@@ -117,11 +119,20 @@ function useInitialData({
         return code ?? {};
       }
 
+      // Fold each variant's highlighted-visible `fallbackCritical` over its plain
+      // `fallback` (under `highlightAt: 'init'`) and strip the staging field, so the
+      // hoisted loading fallback is already highlighted and nothing leaks to the
+      // content. `collapseToEmpty` isn't threaded into the client here, so the `false`
+      // form is assumed: under collapse-to-empty this may promote a few frames that are
+      // then CSS-hidden, but that is harmless — the promoted text is byte-identical
+      // (a valid dictionary) and the frames never paint.
+      const resolved = resolveFallbackCritical(loaded.code, highlightAfter, false) ?? loaded.code;
+
       // Strip fallbacks from code and hoist them directly
       const {
         strippedCode,
         allFallbackHasts
-      } = stripFallbackHastsFromCode(loaded.code, variantName, fallbackUsesExtraFiles, fallbackUsesAllVariants);
+      } = stripFallbackHastsFromCode(resolved, variantName, fallbackUsesExtraFiles, fallbackUsesAllVariants);
       if (!signal.aborted) {
         setCode(strippedCode);
         for (const [variant, hasts] of Object.entries(allFallbackHasts)) {
@@ -279,12 +290,18 @@ function useAllVariants({
             resultCode[item.name] = item.variant.code;
           }
         }
+
+        // Strip the staging `fallbackCritical` before it enters `code` state and
+        // reaches the content. The full load runs with `disableParsing`, so the source
+        // is a raw string and no `fallbackCritical` is produced here — the strip is
+        // purely defensive, hence the strip-only `'idle'` (promotion is `'init'`-gated).
+        const resolvedResultCode = resolveFallbackCritical(resultCode, 'idle', false) ?? resultCode;
         if (errors.length > 0) {
           console.error(new Errors.ErrorCodeHighlighterClientLoadVariantsFailure(url, errors));
         } else if (!signal.aborted) {
-          setCode(resultCode);
+          setCode(resolvedResultCode);
         }
-        return resultCode;
+        return resolvedResultCode;
       } catch (error) {
         console.error(new Errors.ErrorCodeHighlighterClientLoadAllVariantsFailure(url, error));
         return code ?? {};
@@ -304,17 +321,6 @@ function useAllVariants({
     refresh: refreshAllVariants
   };
 }
-function yieldToMain() {
-  const scheduler = globalThis.scheduler;
-  if (scheduler?.yield) {
-    return scheduler.yield();
-  }
-
-  // Fall back to yielding with setTimeout.
-  return new Promise(resolve => {
-    setTimeout(resolve, 0);
-  });
-}
 function useCodeParsing({
   code,
   readyForContent,
@@ -331,22 +337,18 @@ function useCodeParsing({
   const [isHighlightAllowed, setIsHighlightAllowed] = React.useState(highlightAfter === 'init' || highlightAfter === 'hydration' && isHydrated);
   React.useEffect(() => {
     if (highlightAfter === 'idle') {
-      const requestIdleCallback = window.requestIdleCallback ?? setTimeout;
-      const cancelIdleCallback = window.cancelIdleCallback ?? clearTimeout;
-      const idleRequest = requestIdleCallback(() => {
-        setIsHighlightAllowed(true);
-      });
-      return () => cancelIdleCallback(idleRequest);
+      return requestIdle(() => setIsHighlightAllowed(true));
     }
     return undefined;
   }, [highlightAfter]);
 
-  // Update highlight allowed state when hydration completes
+  // Highlight instantly once hydrated, as a non-blocking client transition,
+  // rather than deferring to a scheduled task. (`highlightAt: 'idle'` above is
+  // the mode that deliberately keeps the unhighlighted first paint and swaps in
+  // the highlighted tree on a later idle render.)
   React.useEffect(() => {
     if (highlightAfter === 'hydration' && isHydrated) {
-      // we should ensure that each code highlighter is enhanced as a separate task
-      // this should run from top to bottom
-      yieldToMain().then(() => setIsHighlightAllowed(true));
+      React.startTransition(() => setIsHighlightAllowed(true));
     }
   }, [highlightAfter, isHydrated]);
 
@@ -475,13 +477,12 @@ function useCodeTransforms({
   // Get available transforms from the current variant (separate memo for efficiency)
   const availableTransforms = React.useMemo(() => getAvailableTransforms(parsedCode ?? loadedCode, variantName), [parsedCode, loadedCode, variantName]);
 
-  // Effect to compute transformations for all variants
+  // Effect to compute transformations for all variants. Only runs when the
+  // full async pipeline is wired (`parsedCode` + worker + deltas computer);
+  // the no-async case is derived during render below instead of being stored,
+  // so this effect never publishes a synchronous pass-through state.
   React.useEffect(() => {
     if (!parsedCode || !sourceParser || !computeHastDeltasLoader) {
-      setTransformedState({
-        input: parsedCode,
-        output: parsedCode
-      });
       return;
     }
 
@@ -507,13 +508,16 @@ function useCodeTransforms({
     })();
   }, [parsedCode, sourceParser, computeHastDeltasLoader]);
 
-  // Expose the cached output regardless of whether `parsedCode` changed since
-  // the last computation — falling back to `undefined` here would yank the
-  // currently-displayed HAST for a frame while the async pipeline catches up.
-  // Staleness is signalled via `waitingForTransformedCode` so downstream
-  // gates (e.g. `useTransformManagement` / `useVariantSelection`) hold off
-  // committing a swap until fresh deltas land.
-  const transformedCode = transformedState.output;
+  // When the full async pipeline is wired, expose the cached output regardless
+  // of whether `parsedCode` changed since the last computation — falling back
+  // to `undefined` here would yank the currently-displayed HAST for a frame
+  // while the async pipeline catches up. Staleness is signalled via
+  // `waitingForTransformedCode` so downstream gates (e.g.
+  // `useTransformManagement` / `useVariantSelection`) hold off committing a
+  // swap until fresh deltas land. Without the pipeline, `transformedCode` is a
+  // synchronous pass-through of `parsedCode` derived during render.
+  const hasAsyncPipeline = !!parsedCode && !!sourceParser && !!computeHastDeltasLoader;
+  const transformedCode = hasAsyncPipeline ? transformedState.output : parsedCode;
 
   // Async hast-deltas pipeline status. While true, consumers (notably
   // `useTransformManagement`'s `deferHighlight` gate) should treat
@@ -531,7 +535,7 @@ function useCodeTransforms({
   // `input` against the live `parsedCode` instead of just checking
   // `!transformedCode` so a freshly-arriving variant re-engages the wait
   // until its deltas land.
-  const waitingForTransformedCode = !!parsedCode && !!sourceParser && !!computeHastDeltasLoader && transformedState.input !== parsedCode;
+  const waitingForTransformedCode = hasAsyncPipeline && transformedState.input !== parsedCode;
   return {
     transformedCode,
     availableTransforms,
@@ -835,15 +839,23 @@ export function CodeHighlighterClient(props) {
   const isControlled = Boolean(props.code || controlled?.code);
   const [code, setCode] = React.useState(typeof props.precompute === 'object' ? props.precompute : undefined);
 
-  // Sync code state with precompute prop changes (for hot-reload)
-  React.useEffect(() => {
+  // Sync code state with precompute prop changes (for hot-reload). Done with
+  // the store-previous-prop render-phase derivation rather than an effect:
+  // `code` is genuinely state (also mutated by `useInitialData` via `setCode`
+  // for client fallback loading) so it can't be pure derivation, but the
+  // re-seed on a new `precompute` is a render-time setState off the previous
+  // prop value. Match the original effect's branch logic: only object values
+  // re-seed and only an explicit `undefined` clears — any other value (e.g. a
+  // loader) leaves `code` untouched.
+  const [prevPrecompute, setPrevPrecompute] = React.useState(props.precompute);
+  if (props.precompute !== prevPrecompute) {
+    setPrevPrecompute(props.precompute);
     if (typeof props.precompute === 'object') {
       setCode(props.precompute);
     } else if (props.precompute === undefined) {
-      // Only reset to undefined if precompute is explicitly undefined
       setCode(undefined);
     }
-  }, [props.precompute]);
+  }
 
   // State to store processed globalsCode to avoid duplicate loading
   const [processedGlobalsCode, setProcessedGlobalsCode] = React.useState(undefined);
@@ -1049,22 +1061,18 @@ export function CodeHighlighterClient(props) {
   const [isEnhanceAllowed, setIsEnhanceAllowed] = React.useState(enhanceAfter === 'init' || enhanceAfter === 'hydration' && isHydrated);
   React.useEffect(() => {
     if (enhanceAfter === 'idle') {
-      const requestIdleCallback = window.requestIdleCallback ?? setTimeout;
-      const cancelIdleCallback = window.cancelIdleCallback ?? clearTimeout;
-      const idleRequest = requestIdleCallback(() => {
-        setIsEnhanceAllowed(true);
-      });
-      return () => cancelIdleCallback(idleRequest);
+      return requestIdle(() => setIsEnhanceAllowed(true));
     }
     return undefined;
   }, [enhanceAfter]);
 
-  // Update enhance allowed state when hydration completes
+  // Enhance instantly once hydrated, as a non-blocking client transition,
+  // rather than deferring to a scheduled task. (`enhanceAfter: 'idle'` above is
+  // the mode that deliberately keeps the un-enhanced first paint and swaps in
+  // the enhanced tree on a later idle render.)
   React.useEffect(() => {
     if (enhanceAfter === 'hydration' && isHydrated) {
-      // we should ensure that each code highlighter is enhanced as a separate task
-      // this should run from top to bottom
-      yieldToMain().then(() => setIsEnhanceAllowed(true));
+      React.startTransition(() => setIsEnhanceAllowed(true));
     }
   }, [enhanceAfter, isHydrated]);
   const readyForContent = React.useMemo(() => {
@@ -1288,6 +1296,7 @@ export function CodeHighlighterClient(props) {
   // ContentLoading wired the hook; the child's `useCodeFallback` effect sets it
   // again before that runs.
   if (showFallback) {
+    // eslint-disable-next-line react-hooks/refs -- dev-only validation flag; reset during render so the child useCodeFallback effect (fires before this parent's validation effect) can re-set it each time the fallback shows
     hookCalledRef.current = false;
   }
 

package/CodeHighlighter/createClientProps.mjs

@@ -1,5 +1,6 @@
 import * as React from 'react';
 import { replaceUrlPrefix } from "../pipeline/loaderUtils/applyUrlPrefix.mjs";
+import { resolveFallbackCritical } from "./resolveFallbackCritical.mjs";
 import { jsx as _jsx } from "react/jsx-runtime";
 /**
  * Assemble the props for the `'use client'` `CodeHighlighterClient` from the
@@ -11,6 +12,16 @@ export function createClientProps(props) {
   const highlightAfter = props.highlightAfter === 'stream' ? 'init' : props.highlightAfter;
   const enhanceAfter = props.enhanceAfter === 'stream' ? 'init' : props.enhanceAfter;
 
+  // Resolve the staging `fallbackCritical` on every variant of both carriers before
+  // they cross to the client: under `highlightAt: 'init'` (not `collapseToEmpty`)
+  // promote it over the plain `fallback` so the first paint is highlighted with no
+  // decompression, and always strip it so it never reaches `Content`/`ContentLoading`
+  // or bloats the payload. `code` and `precompute` are forwarded separately (the
+  // client seeds its `code` state from `precompute`), so both must be resolved.
+  const collapseToEmpty = props.collapseToEmpty ?? props.contentProps?.collapseToEmpty ?? false;
+  const code = resolveFallbackCritical(props.code, highlightAfter, collapseToEmpty);
+  const precompute = resolveFallbackCritical(props.precompute, highlightAfter, collapseToEmpty);
+
   // Rewrite the top-level URL before it leaves the server. The client never
   // receives `urlPrefix` (and shouldn't deal with `file://` URLs), so any
   // local URL must be translated to its hosted form here. Variant-level URLs
@@ -20,7 +31,7 @@ export function createClientProps(props) {
   const url = props.urlPrefix && props.url ? replaceUrlPrefix(props.url, props.urlPrefix) : props.url;
   const contentProps = {
     ...props.contentProps,
-    code: props.code || props.precompute,
+    code: code || precompute,
     components: props.components,
     name: props.name,
     slug: props.slug,
@@ -38,8 +49,8 @@ export function createClientProps(props) {
   };
   return {
     url,
-    code: props.code,
-    precompute: props.precompute,
+    code,
+    precompute,
     components: props.components,
     variants: props.variants,
     variant: props.variant,

package/CodeHighlighter/fallbackFormat.d.mts

@@ -65,6 +65,44 @@ export declare function fallbackToText(nodes: FallbackNode[]): string;
  * nodes (e.g. whitespace text between frames) are preserved in place.
  */
 export declare function buildRootFallback(root: HastRoot): FallbackNode[];
+/**
+ * Whether a compact `fallback` already carries highlighting — i.e. at least one frame
+ * keeps nested token spans (array children) instead of flat plain text. True exactly for
+ * the promoted {@link promoteCriticalFallback} form; a plain {@link buildRootFallback} has
+ * a string child on every frame. `<Pre>` uses this to defer the decompressing decode ONLY
+ * when the first paint is already highlighted, so it never flashes plain → highlighted.
+ */
+export declare function fallbackIsHighlighted(fallback: FallbackNode[]): boolean;
+/**
+ * The **sparse** highlighted-visible fallback: a map from frame index to the
+ * highlighted `FallbackNode` for ONLY the frames visible on the initial collapsed
+ * render (`visibleFrames`). Off-screen frames are omitted — they flatten to exactly
+ * {@link buildRootFallback}'s plain output, so storing them would just duplicate
+ * `fallback` in the precompute. {@link promoteCriticalFallback} splices these back
+ * over the plain fallback for `highlightAt: 'init'` (paint highlighted on the first
+ * render, zero decompression). Frame indices count `span.frame` children only,
+ * matching `getInitialVisibleFrames`.
+ *
+ * The decoded `root` is shared/read-only; this only reads its frames (the synthetic
+ * visible frame reuses the frame's `children` array without mutating it, and
+ * `data-lined` is dropped via destructuring, not deletion).
+ */
+export declare function buildCriticalFallback(root: HastRoot, visibleFrames: {
+  [key: number]: boolean;
+}): {
+  [frameIndex: number]: FallbackNode;
+};
+/**
+ * Splice a sparse {@link buildCriticalFallback} diff back onto a plain `fallback`,
+ * replacing each visible frame (matched by index in document order) with its
+ * highlighted node. The result is the full highlighted-visible fallback; its text is
+ * byte-identical to `fallback` (the highlight spans only wrap the same characters), so
+ * it stays a valid DEFLATE dictionary — asserted by tests. Returns a new array; the
+ * input is not mutated.
+ */
+export declare function promoteCriticalFallback(fallback: FallbackNode[], critical: {
+  [frameIndex: number]: FallbackNode;
+}): FallbackNode[];
 /**
  * Redistributes a root fallback (built by `buildRootFallback`) back onto the
  * frames of a decoded HAST root, setting each frame's `data.fallback` to the

package/CodeHighlighter/fallbackFormat.mjs

@@ -179,6 +179,17 @@ function collectFrameText(frame) {
   return out;
 }
 
+/**
+ * The frame's precomputed plain-text fallback (`frame.data.fallback`, set by
+ * `addLineGutters`), falling back to walking the frame's text when absent. This
+ * is the exact text the highlighted children render, so it doubles as the
+ * DEFLATE dictionary contribution for the frame.
+ */
+function framePlainText(frame) {
+  const fallbackNodes = frame.data?.fallback;
+  return fallbackNodes && fallbackNodes.length > 0 ? fallbackNodes.map(node => node.type === 'text' ? node.value : '').join('') : collectFrameText(frame);
+}
+
 /**
  * Builds the variant-level root fallback from a final (post-enhancer) HAST
  * root. Each `span.frame` becomes a compact frame element whose single text
@@ -201,15 +212,13 @@ export function buildRootFallback(root) {
         dataLined,
         ...properties
       } = frame.properties || {};
-      const fallbackNodes = frame.data?.fallback;
-      const text = fallbackNodes && fallbackNodes.length > 0 ? fallbackNodes.map(node => node.type === 'text' ? node.value : '').join('') : collectFrameText(frame);
       syntheticChildren.push({
         type: 'element',
         tagName: 'span',
         properties,
         children: [{
           type: 'text',
-          value: text
+          value: framePlainText(frame)
         }]
       });
     } else {
@@ -219,6 +228,90 @@ export function buildRootFallback(root) {
   return convertChildren(syntheticChildren);
 }
 
+/** A frame node in compact `FallbackNode` form — `['span', 'frame', …]`, in any of the
+ *  3/4/5-element shapes (`className` is always the second element). */
+function isFrameFallbackNode(node) {
+  return Array.isArray(node) && node.length >= 3 && node[0] === 'span' && typeof node[1] === 'string' && node[1].split(' ').includes('frame');
+}
+
+/**
+ * Whether a compact `fallback` already carries highlighting — i.e. at least one frame
+ * keeps nested token spans (array children) instead of flat plain text. True exactly for
+ * the promoted {@link promoteCriticalFallback} form; a plain {@link buildRootFallback} has
+ * a string child on every frame. `<Pre>` uses this to defer the decompressing decode ONLY
+ * when the first paint is already highlighted, so it never flashes plain → highlighted.
+ */
+export function fallbackIsHighlighted(fallback) {
+  return fallback.some(node => {
+    if (!isFrameFallbackNode(node)) {
+      return false;
+    }
+    // `children` is the last tuple slot, except the 5-element form keeps `extra` last
+    // (matches `nodeText`). An array means nested token spans, i.e. highlighted.
+    const children = node.length === 5 ? node[3] : node[node.length - 1];
+    return Array.isArray(children);
+  });
+}
+
+/**
+ * The **sparse** highlighted-visible fallback: a map from frame index to the
+ * highlighted `FallbackNode` for ONLY the frames visible on the initial collapsed
+ * render (`visibleFrames`). Off-screen frames are omitted — they flatten to exactly
+ * {@link buildRootFallback}'s plain output, so storing them would just duplicate
+ * `fallback` in the precompute. {@link promoteCriticalFallback} splices these back
+ * over the plain fallback for `highlightAt: 'init'` (paint highlighted on the first
+ * render, zero decompression). Frame indices count `span.frame` children only,
+ * matching `getInitialVisibleFrames`.
+ *
+ * The decoded `root` is shared/read-only; this only reads its frames (the synthetic
+ * visible frame reuses the frame's `children` array without mutating it, and
+ * `data-lined` is dropped via destructuring, not deletion).
+ */
+export function buildCriticalFallback(root, visibleFrames) {
+  const critical = {};
+  let frameIndex = 0;
+  for (const child of root.children) {
+    if (child.type === 'element' && isFrameSpan(child)) {
+      if (visibleFrames[frameIndex]) {
+        const frame = child;
+        const {
+          dataLined,
+          ...properties
+        } = frame.properties || {};
+        const [node] = convertChildren([{
+          type: 'element',
+          tagName: 'span',
+          properties,
+          children: frame.children
+        }]);
+        critical[frameIndex] = node;
+      }
+      frameIndex += 1;
+    }
+  }
+  return critical;
+}
+
+/**
+ * Splice a sparse {@link buildCriticalFallback} diff back onto a plain `fallback`,
+ * replacing each visible frame (matched by index in document order) with its
+ * highlighted node. The result is the full highlighted-visible fallback; its text is
+ * byte-identical to `fallback` (the highlight spans only wrap the same characters), so
+ * it stays a valid DEFLATE dictionary — asserted by tests. Returns a new array; the
+ * input is not mutated.
+ */
+export function promoteCriticalFallback(fallback, critical) {
+  let frameIndex = 0;
+  return fallback.map(node => {
+    if (!isFrameFallbackNode(node)) {
+      return node;
+    }
+    const replacement = critical[frameIndex];
+    frameIndex += 1;
+    return replacement ?? node;
+  });
+}
+
 /**
  * Redistributes a root fallback (built by `buildRootFallback`) back onto the
  * frames of a decoded HAST root, setting each frame's `data.fallback` to the

package/CodeHighlighter/prepareInitialSource.d.mts

@@ -8,6 +8,17 @@ export interface PrepareInitialSourceOptions<T extends {}> extends CodeHighlight
   initialSource: VariantSource;
   initialExtraFiles?: VariantExtraFiles;
   ContentLoading: React.ComponentType<ContentLoadingProps<T>>;
+  /**
+   * Whether to DEFLATE-compress the consolidated residual fallbacks. This only
+   * shrinks the server→client serialized payload, so it pays off only when the
+   * result actually crosses that wire (a server render). On the client there is no
+   * wire — compressing here just to have `CodeHighlighterClient` decompress it right
+   * back is a wasted round-trip — so the isomorphic `CodeHighlighter` entry passes
+   * `typeof window === 'undefined'`. Defaults to `true` for the server-only loaders
+   * (and tests) that always produce wire output; when `false`, the fallbacks stay
+   * inline on `codeForClient`.
+   */
+  compressResidual?: boolean;
 }
 export interface PreparedInitialSource {
   /** The pre-rendered loading fallback (`<ContentLoading />`). */

package/CodeHighlighter/prepareInitialSource.mjs

@@ -1,5 +1,6 @@
 import * as React from 'react';
 import { buildStringFallback } from "./buildStringFallback.mjs";
+import { resolveFallbackCritical } from "./resolveFallbackCritical.mjs";
 import { codeToFallbackProps, stripFallbackHastsFromCode } from "./codeToFallbackProps.mjs";
 import { collapseRenderedFallbacks, compressResidualFallbacks, extractResidualFallbacks, mergeResidualFallbacks, residualDictionaryText } from "./fallbackCompression.mjs";
 import { replaceUrlPrefix } from "../pipeline/loaderUtils/applyUrlPrefix.mjs";
@@ -19,7 +20,7 @@ export function prepareInitialSource(props) {
     slug,
     name,
     initialVariant,
-    code,
+    code: initialCode,
     initialFilename,
     fallbackUsesExtraFiles,
     fallbackUsesAllVariants,
@@ -31,6 +32,18 @@ export function prepareInitialSource(props) {
   const collapseToEmptyEnabled = collapseToEmpty === true;
   const initialExpandedEnabled = initialExpanded === true;
 
+  // Fold each variant's staging `fallbackCritical` into its plain `fallback` up front
+  // (under `highlightAt: 'init'`, not `collapseToEmpty`), then strip it. The hoisted
+  // loading fallback is therefore already highlighted-visible — so the first paint is
+  // highlighted with no decompression — while the rest of this function (strip, hoist,
+  // window, compress) operates on a single `fallback` field with no awareness of the
+  // staging companion.
+  // Normalize `'stream'` → `'init'` before resolving, mirroring `createClientProps`: stream
+  // mode wants the loading fallback highlighted on first paint too (the client highlightAfter
+  // type even collapses 'stream' into 'init').
+  const highlightAfter = props.highlightAfter === 'stream' ? 'init' : props.highlightAfter;
+  const code = resolveFallbackCritical(initialCode, highlightAfter, collapseToEmptyEnabled) ?? initialCode;
+
   // When the block starts expanded, the loading UI needs the full content, so
   // the `fallbackCollapsed` window optimization (paint only the collapsed slice,
   // defer the rest) doesn't apply — treat it as off everywhere below.
@@ -125,6 +138,41 @@ export function prepareInitialSource(props) {
           focusedLines: collapseToEmptyEnabled ? 0 : counts.focusedLines,
           collapsible: collapseToEmptyEnabled ? true : counts.collapsible
         };
+        // Carry the RAW counts onto the wire `code` too, so the content component's base
+        // render (before `hast` decodes) reads the same `collapsible`/window metadata the
+        // loading fallback got. A windowed inline string — or a `hastCompressed` source
+        // whose dictionary was stripped for residual compression — otherwise has no stored
+        // counts on `codeForClient`, so `getVariantFileLineCounts` falls back to the raw
+        // (non-collapsible) string count and `data-collapsible` flashes off until the
+        // source highlights. `<Pre>` re-applies collapse-to-empty from the stored counts.
+        const wireVariant = strippedCode[variantName];
+        if (wireVariant && typeof wireVariant === 'object') {
+          const stored = {
+            totalLines: counts.totalLines,
+            focusedLines: counts.focusedLines,
+            collapsible: counts.collapsible
+          };
+          if (wireVariant.fileName === file.fileName) {
+            strippedCode[variantName] = {
+              ...wireVariant,
+              ...stored
+            };
+          } else {
+            const extra = wireVariant.extraFiles?.[file.fileName];
+            if (extra && typeof extra === 'object') {
+              strippedCode[variantName] = {
+                ...wireVariant,
+                extraFiles: {
+                  ...wireVariant.extraFiles,
+                  [file.fileName]: {
+                    ...extra,
+                    ...stored
+                  }
+                }
+              };
+            }
+          }
+        }
       }
     }
   }
@@ -182,13 +230,24 @@ export function prepareInitialSource(props) {
   // onto the code so its consumers (render and the swap line-count classifier)
   // read the dictionary off `code` regardless of which variant is active. When
   // there's nothing worth compressing, keep the plain inline fallbacks unchanged.
-  const {
-    wireCode,
-    residual
-  } = extractResidualFallbacks(strippedCode);
-  const fullResidual = effectiveFallbackCollapsed ? mergeResidualFallbacks(residual, allFallbackHasts) : residual;
-  const residualFallbacks = compressResidualFallbacks(fullResidual, residualDictionaryText(contentLoadingHasts));
-  const codeForClient = residualFallbacks ? wireCode : strippedCode;
+  // Compressing the residual only shrinks the server→client wire, so skip it entirely
+  // on the client (`compressResidual === false`): keep the fallbacks INLINE on
+  // `codeForClient` so `CodeHighlighterClient` reads them directly — no
+  // compress→decompress round-trip, and nothing to recompute on every re-render. This
+  // is the same shape the `strippedCode` branch below already produces for a payload
+  // too small to be worth compressing.
+  const compressResidual = props.compressResidual ?? true;
+  let residualFallbacks;
+  let codeForClient = strippedCode;
+  if (compressResidual) {
+    const {
+      wireCode,
+      residual
+    } = extractResidualFallbacks(strippedCode);
+    const fullResidual = effectiveFallbackCollapsed ? mergeResidualFallbacks(residual, allFallbackHasts) : residual;
+    residualFallbacks = compressResidualFallbacks(fullResidual, residualDictionaryText(contentLoadingHasts));
+    codeForClient = residualFallbacks ? wireCode : strippedCode;
+  }
 
   // Get the component for the selected variant
   const component = props.components?.[initialVariant];

package/CodeHighlighter/resolveFallbackCritical.d.mts

@@ -0,0 +1,23 @@
+import type { Code, CodeHighlighterBaseProps } from "./types.mjs";
+type HighlightAfter = CodeHighlighterBaseProps<{}>['highlightAfter'];
+/**
+ * Resolve the staging `fallbackCritical` field at a server→client (or client-load)
+ * boundary, returning a clone of `code` ready to cross to the client:
+ *
+ * - **Promote** — under `highlightAt: 'init'` (and not `collapseToEmpty`), each
+ *   variant whose `fallbackCritical` *and* plain `fallback` are both present has the
+ *   sparse `fallbackCritical` diff spliced over its `fallback` (`promoteCriticalFallback`)
+ *   — the visible frames become highlighted, the rest stay plain — so the first paint is
+ *   highlighted with no client-side decompression. The promoted `fallback` stays a valid
+ *   DEFLATE dictionary because the spliced frames have byte-identical text. With
+ *   `collapseToEmpty`, the correct critical fallback is all-plain (no frame is visible),
+ *   which already equals `fallback` — so promotion is skipped.
+ * - **Strip** — `fallbackCritical` is always deleted (even when not promoting, and
+ *   even for non-`init` modes), so it never crosses to the `Content`/`ContentLoading`
+ *   components or bloats the serialized payload.
+ *
+ * Returns `code` unchanged (same reference) when no variant carries `fallbackCritical`,
+ * and clones only the variants that do.
+ */
+export declare function resolveFallbackCritical(code: Code | undefined, highlightAfter: HighlightAfter, collapseToEmpty: boolean): Code | undefined;
+export {};

package/CodeHighlighter/resolveFallbackCritical.mjs

@@ -0,0 +1,44 @@
+import { promoteCriticalFallback } from "./fallbackFormat.mjs";
+/**
+ * Resolve the staging `fallbackCritical` field at a server→client (or client-load)
+ * boundary, returning a clone of `code` ready to cross to the client:
+ *
+ * - **Promote** — under `highlightAt: 'init'` (and not `collapseToEmpty`), each
+ *   variant whose `fallbackCritical` *and* plain `fallback` are both present has the
+ *   sparse `fallbackCritical` diff spliced over its `fallback` (`promoteCriticalFallback`)
+ *   — the visible frames become highlighted, the rest stay plain — so the first paint is
+ *   highlighted with no client-side decompression. The promoted `fallback` stays a valid
+ *   DEFLATE dictionary because the spliced frames have byte-identical text. With
+ *   `collapseToEmpty`, the correct critical fallback is all-plain (no frame is visible),
+ *   which already equals `fallback` — so promotion is skipped.
+ * - **Strip** — `fallbackCritical` is always deleted (even when not promoting, and
+ *   even for non-`init` modes), so it never crosses to the `Content`/`ContentLoading`
+ *   components or bloats the serialized payload.
+ *
+ * Returns `code` unchanged (same reference) when no variant carries `fallbackCritical`,
+ * and clones only the variants that do.
+ */
+export function resolveFallbackCritical(code, highlightAfter, collapseToEmpty) {
+  if (!code) {
+    return code;
+  }
+  const promote = highlightAfter === 'init' && !collapseToEmpty;
+  let changed = false;
+  const resolved = {};
+  for (const [key, variant] of Object.entries(code)) {
+    if (!variant || typeof variant === 'string' || variant.fallbackCritical === undefined) {
+      resolved[key] = variant;
+      continue;
+    }
+    const {
+      fallbackCritical,
+      ...rest
+    } = variant;
+    resolved[key] = promote && rest.fallback ? {
+      ...rest,
+      fallback: promoteCriticalFallback(rest.fallback, fallbackCritical)
+    } : rest;
+    changed = true;
+  }
+  return changed ? resolved : code;
+}