fix: defer markdown code line number sync during streaming

Avoids syncing code line numbers while markdown is still streaming
Keeps code block wrapping and line numbers stable after render
Updates code block layout to support deferred gutter insertion

Author: btriapitsyn

packages/ui/src/components/chat/MarkdownRendererImpl.tsx

@@ -860,6 +860,7 @@ const mermaidColorsFromTheme = (theme: Theme) => ({
 
 const useDecorateContext = (
   currentTheme: Theme,
+  deferCodeLineNumberSync: boolean,
   onPreviewLoopback?: (url: string) => void,
 ): DecorateContext => {
   const { t } = useI18n();
@@ -895,8 +896,8 @@ const useDecorateContext = (
           return {};
         }
       });
-    return { labels, codeBlockLineWrap, onToggleCodeBlockLineWrap: toggleCodeBlockLineWrap, renderMermaid, onPreviewLoopback };
-  }, [currentTheme, labels, codeBlockLineWrap, toggleCodeBlockLineWrap, onPreviewLoopback]);
+    return { labels, codeBlockLineWrap, deferCodeLineNumberSync, onToggleCodeBlockLineWrap: toggleCodeBlockLineWrap, renderMermaid, onPreviewLoopback };
+  }, [currentTheme, labels, codeBlockLineWrap, deferCodeLineNumberSync, toggleCodeBlockLineWrap, onPreviewLoopback]);
 };
 
 // Runs the async render pipeline into the container and keeps a stable
@@ -984,7 +985,9 @@ const useMorphdomMarkdown = ({
         existing[i]?.remove();
       }
 
-      scheduleMarkdownCodeLineNumberSync(target);
+      if (!ctx.deferCodeLineNumberSync) {
+        scheduleMarkdownCodeLineNumberSync(target);
+      }
     });
 
     return () => {
@@ -1012,8 +1015,9 @@ const useMorphdomMarkdown = ({
     const container = containerRef.current;
     const target = container?.querySelector<HTMLElement>('[data-markdown-content]') ?? container;
     if (!target) return;
+    if (ctx.deferCodeLineNumberSync) return;
     applyMarkdownCodeBlockWrapState(target, ctx.codeBlockLineWrap, ctx.labels);
-  }, [containerRef, ctx.codeBlockLineWrap, ctx.labels]);
+  }, [containerRef, ctx.codeBlockLineWrap, ctx.deferCodeLineNumberSync, ctx.labels]);
 
   React.useEffect(() => {
     const container = containerRef.current;
@@ -1081,7 +1085,7 @@ const MarkdownRendererImpl: React.FC<MarkdownRendererProps> = ({
   useExternalLinkInteractions({ containerRef });
 
   const syntaxVars = React.useMemo(() => getMarkdownSyntaxVars(currentTheme), [currentTheme]);
-  const ctx = useDecorateContext(currentTheme, effectiveDirectory ? handlePreviewLoopback : undefined);
+  const ctx = useDecorateContext(currentTheme, live, effectiveDirectory ? handlePreviewLoopback : undefined);
   const cacheKey = `markdown-${part?.id ? `part-${part.id}` : `message-${messageId}`}`;
 
   useMorphdomMarkdown({ containerRef, text: pacedText, streaming: live, cacheKey, syntaxVars, ctx });
@@ -1164,7 +1168,7 @@ const SimpleMarkdownRendererImpl: React.FC<{
   useExternalLinkInteractions({ containerRef, enabled: !disableLinkSafety });
 
   const syntaxVars = React.useMemo(() => getMarkdownSyntaxVars(currentTheme), [currentTheme]);
-  const ctx = useDecorateContext(currentTheme);
+  const ctx = useDecorateContext(currentTheme, false);
 
   useMorphdomMarkdown({
     containerRef,

packages/ui/src/components/chat/markdown/decorate.ts

@@ -24,6 +24,7 @@ export type DecorateLabels = {
 export type DecorateContext = {
   labels: DecorateLabels;
   codeBlockLineWrap: boolean;
+  deferCodeLineNumberSync?: boolean;
   onToggleCodeBlockLineWrap?: () => void;
   // Renders a mermaid block source to svg/ascii using current theme colors.
   renderMermaid: (source: string) => MermaidRender;
@@ -97,7 +98,7 @@ const createCodeLineNumbers = (pre: HTMLPreElement): HTMLDivElement => {
   const gutter = document.createElement('div');
   gutter.setAttribute('data-md-code-line-numbers', '');
   gutter.setAttribute('aria-hidden', 'true');
-  gutter.className = 'select-none border-r border-border/50 pr-3 text-right text-muted-foreground/45';
+  gutter.className = 'min-w-8 shrink-0 select-none border-r border-border/50 pr-3 text-right font-mono text-[13px] text-muted-foreground/45';
 
   const text = pre.textContent ?? '';
   const lineCount = Math.max(1, text.endsWith('\n') ? text.split('\n').length - 1 : text.split('\n').length);
@@ -194,6 +195,12 @@ export const scheduleMarkdownCodeLineNumberSync = (root: HTMLElement): void => {
 export const applyMarkdownCodeBlockWrapState = (root: HTMLElement, enabled: boolean, labels: DecorateLabels): void => {
   const wrappers = root.querySelectorAll<HTMLElement>('[data-component="markdown-code"]');
   for (const wrapper of Array.from(wrappers)) {
+    const body = wrapper.querySelector<HTMLElement>('[data-md-code-body]');
+    const pre = wrapper.querySelector<HTMLPreElement>('pre');
+    if (body && pre && !body.querySelector('[data-md-code-line-numbers]')) {
+      body.classList.add('flex', 'gap-3');
+      body.insertBefore(createCodeLineNumbers(pre), pre);
+    }
     applyCodeBlockWrapState(wrapper, enabled, labels);
   }
   scheduleMarkdownCodeLineNumberSync(root);
@@ -258,18 +265,22 @@ const decorateCodeBlocks = (root: HTMLElement, ctx: DecorateContext): void => {
 
     const body = document.createElement('div');
     body.setAttribute('data-md-code-body', '');
-    body.className = 'flex gap-3 px-3 py-2.5 overflow-x-auto';
+    body.className = ctx.deferCodeLineNumberSync ? 'px-3 py-2.5 overflow-x-auto' : 'flex gap-3 px-3 py-2.5 overflow-x-auto';
 
     parent.replaceChild(wrapper, pre);
     pre.style.margin = '0';
     pre.style.background = 'transparent';
     pre.classList.add('min-w-0', 'w-full', 'flex-1');
-    body.appendChild(createCodeLineNumbers(pre));
+    if (!ctx.deferCodeLineNumberSync) {
+      body.appendChild(createCodeLineNumbers(pre));
+    }
     body.appendChild(pre);
     wrapper.appendChild(header);
     wrapper.appendChild(body);
     applyCodeBlockWrapState(wrapper, ctx.codeBlockLineWrap, ctx.labels);
-    scheduleMarkdownCodeLineNumberSync(wrapper);
+    if (!ctx.deferCodeLineNumberSync) {
+      scheduleMarkdownCodeLineNumberSync(wrapper);
+    }
   }
 };