diff --git a/frontend/src/components/workspace/messages/markdown-content.tsx b/frontend/src/components/workspace/messages/markdown-content.tsx index 6772f5f61..25aee1745 100644 --- a/frontend/src/components/workspace/messages/markdown-content.tsx +++ b/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 ( +
+ {language && ( +
+ {language} +
+ )} +
+        
+          {children}
+        
+      
+
+ ); +} + +function StreamingCode({ + children, + className, + node: _node, + ...props +}: StreamingCodeProps) { + const isBlock = useContext(StreamingCodeBlockContext); + + if (!isBlock) { + return ( + + {children} + + ); + } + + return ( + + {children} + + ); +} + /** 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; diff --git a/frontend/tests/unit/components/workspace/messages/markdown-content.test.ts b/frontend/tests/unit/components/workspace/messages/markdown-content.test.ts new file mode 100644 index 000000000..7c29c7539 --- /dev/null +++ b/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[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", '
Hello
', "```"].join("\n"), + true, + ); + + expect(html).toContain("data-streaming-code-block"); + expect(html).toContain('data-language="html"'); + expect(html).toContain( + "<main class="report">Hello</main>", + ); + 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", '
Hello
', "```"].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) => + 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", "
", "```"].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"); + }); +});