fix(frontend): avoid highlighting code while streaming (#4068)

* fix(frontend): avoid highlighting code while streaming

* fix(frontend): harden streaming code rendering

frontend/src/components/workspace/messages/markdown-content.tsx

@@ -1,6 +1,13 @@
 "use client";
 
-import { useMemo } from "react";
+import {
+  createContext,
+  type ComponentProps,
+  isValidElement,
+  type ReactNode,
+  useContext,
+  useMemo,
+} from "react";
 
 import { type ClipboardSafeStreamdownProps } from "@/components/ai-elements/streamdown";
 import {
@@ -8,6 +15,7 @@ import {
   streamdownPluginsWithoutRawHtml,
 } from "@/core/streamdown";
 import { SafeMessageResponse } from "@/core/streamdown/components";
+import { cn } from "@/lib/utils";
 
 import { createMarkdownLinkComponent } from "./markdown-link";
 
@@ -20,6 +28,70 @@ export type MarkdownContentProps = {
   components?: ClipboardSafeStreamdownProps["components"];
 };
 
+type StreamingCodeProps = ComponentProps<"code"> & {
+  node?: unknown;
+  children?: ReactNode;
+};
+
+const StreamingCodeBlockContext = createContext(false);
+
+function StreamingPre({ children }: ComponentProps<"pre">) {
+  const childClassName = isValidElement<{ className?: string }>(children)
+    ? children.props.className
+    : undefined;
+  const language =
+    /(?:^|\s)language-([^\s]+)/.exec(childClassName ?? "")?.[1] ?? "";
+
+  return (
+    <div
+      className="my-4 w-full overflow-hidden rounded-xl border"
+      data-language={language}
+      data-streaming-code-block="true"
+    >
+      {language && (
+        <div className="bg-muted/80 text-muted-foreground p-3 text-xs">
+          <span className="ml-1 font-mono lowercase">{language}</span>
+        </div>
+      )}
+      <pre className="bg-muted/40 overflow-x-auto border-t p-4 font-mono text-xs">
+        <StreamingCodeBlockContext.Provider value={true}>
+          {children}
+        </StreamingCodeBlockContext.Provider>
+      </pre>
+    </div>
+  );
+}
+
+function StreamingCode({
+  children,
+  className,
+  node: _node,
+  ...props
+}: StreamingCodeProps) {
+  const isBlock = useContext(StreamingCodeBlockContext);
+
+  if (!isBlock) {
+    return (
+      <code
+        {...props}
+        className={cn(
+          "bg-muted rounded px-1.5 py-0.5 font-mono text-sm",
+          className,
+        )}
+        data-streaming-inline-code="true"
+      >
+        {children}
+      </code>
+    );
+  }
+
+  return (
+    <code {...props} className={className}>
+      {children}
+    </code>
+  );
+}
+
 /** Renders markdown content. */
 export function MarkdownContent({
   content,
@@ -39,11 +111,19 @@ export function MarkdownContent({
     return [...base, ...extra] as ClipboardSafeStreamdownProps["rehypePlugins"];
   }, [rehypePlugins]);
   const components = useMemo(() => {
-    return {
+    const baseComponents = {
       a: createMarkdownLinkComponent(),
       ...componentsFromProps,
     };
-  }, [componentsFromProps]);
+    if (!isLoading) {
+      return baseComponents;
+    }
+    return {
+      ...baseComponents,
+      code: componentsFromProps?.code ?? StreamingCode,
+      pre: componentsFromProps?.pre ?? StreamingPre,
+    };
+  }, [componentsFromProps, isLoading]);
 
   if (!content) return null;
 

frontend/tests/unit/components/workspace/messages/markdown-content.test.ts

@@ -0,0 +1,85 @@
+import { describe, expect, it } from "@rstest/core";
+import { createElement, type ImgHTMLAttributes } from "react";
+import { renderToStaticMarkup } from "react-dom/server";
+
+import { MarkdownContent } from "@/components/workspace/messages/markdown-content";
+
+function renderMarkdown(
+  content: string,
+  isLoading: boolean,
+  components?: Parameters<typeof MarkdownContent>[0]["components"],
+) {
+  return renderToStaticMarkup(
+    createElement(MarkdownContent, { content, isLoading, components }),
+  );
+}
+
+describe("MarkdownContent streaming code blocks", () => {
+  it("renders fenced code without Streamdown highlighting while streaming", () => {
+    const html = renderMarkdown(
+      ["```html", '<main class="report">Hello</main>', "```"].join("\n"),
+      true,
+    );
+
+    expect(html).toContain("data-streaming-code-block");
+    expect(html).toContain('data-language="html"');
+    expect(html).toContain(
+      "&lt;main class=&quot;report&quot;&gt;Hello&lt;/main&gt;",
+    );
+    expect(html).not.toContain('data-streamdown="code-block"');
+  });
+
+  it("keeps inline code inline while streaming", () => {
+    const html = renderMarkdown("Use `const answer = 42` here.", true);
+
+    expect(html).toContain('data-streaming-inline-code="true"');
+    expect(html).not.toContain("data-streaming-code-block");
+  });
+
+  it("keeps an unlabeled single-line fence as a block while streaming", () => {
+    const html = renderMarkdown(["```", "x", "```"].join("\n"), true);
+
+    expect(html).toContain("data-streaming-code-block");
+    expect(html).not.toContain('data-streaming-inline-code="true"');
+  });
+
+  it("restores Streamdown highlighting after streaming finishes", () => {
+    const html = renderMarkdown(
+      ["```html", '<main class="report">Hello</main>', "```"].join("\n"),
+      false,
+    );
+
+    expect(html).toContain('data-streamdown="code-block"');
+    expect(html).not.toContain("data-streaming-code-block");
+  });
+
+  it("preserves custom non-code renderers while streaming", () => {
+    const html = renderMarkdown(
+      "[Docs](https://example.com)\n\n![Chart](chart.png)",
+      true,
+      {
+        a: ({ children, href }) =>
+          createElement("a", { "data-custom-link": true, href }, children),
+        img: (props: ImgHTMLAttributes<HTMLImageElement>) =>
+          createElement("img", { ...props, "data-custom-image": true }),
+      },
+    );
+
+    expect(html).toContain('data-custom-link="true"');
+    expect(html).toContain('data-custom-image="true"');
+  });
+
+  it("preserves a caller-provided code renderer while streaming", () => {
+    const html = renderMarkdown(
+      ["```html", "<main />", "```"].join("\n"),
+      true,
+      {
+        code: ({ children }) =>
+          createElement("code", { "data-custom-code": true }, children),
+      },
+    );
+
+    expect(html).toContain('data-custom-code="true"');
+    expect(html).toContain("data-streaming-code-block");
+  });
+});