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 ( +
+++ {children} + +
+ {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