From ca57a0bcabf399b222d18ffd151edc6895830c74 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Thu, 19 Mar 2026 17:34:45 +0200
Subject: [PATCH] fix(site): prevent rehype-raw from swallowing JSX in chat
 output (#23293)

Omit rehype-raw from the Streamdown rehype plugin list so
HTML-like syntax in LLM output is escaped as text instead of
being parsed by the HTML5 engine and stripped by rehype-sanitize.

When the LLM writes JSX fragments like <Component prop={val} />
outside code fences, remark-parse tags them as html nodes.
rehype-raw then feeds them to parse5, and rehype-sanitize strips
the unknown elements, silently destroying content. Without
rehype-raw, Streamdown auto-injects a remark plugin that converts
html nodes to text, preserving them as visible escaped text.

Markdown formatting (bold, italic, links, code blocks, tables)
is unaffected since those go through remark/rehype directly.
---
 .../ai-elements/response.stories.tsx          | 37 +++++++++++++++++++
 site/src/components/ai-elements/response.tsx  | 18 ++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/site/src/components/ai-elements/response.stories.tsx b/site/src/components/ai-elements/response.stories.tsx
index cef2259010..5a2392567c 100644
--- a/site/src/components/ai-elements/response.stories.tsx
+++ b/site/src/components/ai-elements/response.stories.tsx
@@ -1,4 +1,5 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
+import { expect, within } from "storybook/test";
 import { Response } from "./response";
 
 const sampleMarkdown = `
@@ -66,3 +67,39 @@ export const MarkdownAndLinksLight: Story = {
 		theme: "light",
 	},
 };
+
+// Verifies that JSX-like syntax in LLM output is preserved as
+// escaped text rather than being swallowed by the HTML pipeline.
+const jsxProseMarkdown = `
+\`getLineAnnotations\` depends on \`activeCommentBox\` which could shift.
+
+<RemoteDiffPanel
+  commentBox={commentBox}
+  scrollToFile={scrollTarget}
+  onScrollToFileComplete={handleScrollComplete}
+/>
+
+The props that might change on every \`RemoteDiffPanel\` re-render:
+- \`isLoading\` only during refetch
+- \`getLineAnnotations\` only when \`activeCommentBox\` changes
+`;
+
+export const JsxInProse: Story = {
+	args: {
+		children: jsxProseMarkdown,
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		// These strings live inside the <RemoteDiffPanel .../> JSX block.
+		// Without the rehype-raw fix they are silently eaten by the
+		// HTML sanitizer and never reach the DOM.
+		// The tag name itself is the token most likely to be consumed
+		// by HTML parsing, so assert it explicitly.
+		const tagName = await canvas.findByText(/<RemoteDiffPanel/);
+		expect(tagName).toBeInTheDocument();
+		const marker = await canvas.findByText(/scrollToFile=\{scrollTarget\}/);
+		expect(marker).toBeInTheDocument();
+		const marker2 = await canvas.findByText(/commentBox=\{commentBox\}/);
+		expect(marker2).toBeInTheDocument();
+	},
+};
diff --git a/site/src/components/ai-elements/response.tsx b/site/src/components/ai-elements/response.tsx
index 626f7fe600..fea04d1064 100644
--- a/site/src/components/ai-elements/response.tsx
+++ b/site/src/components/ai-elements/response.tsx
@@ -5,7 +5,12 @@ import {
 } from "@pierre/diffs/react";
 import type { ComponentPropsWithRef, ReactNode } from "react";
 import { useMemo } from "react";
-import { type Components, Streamdown, type UrlTransform } from "streamdown";
+import {
+	type Components,
+	defaultRehypePlugins,
+	Streamdown,
+	type UrlTransform,
+} from "streamdown";
 import { cn } from "utils/cn";
 
 interface ResponseProps extends Omit<ComponentPropsWithRef<"div">, "children"> {
@@ -13,6 +18,16 @@ interface ResponseProps extends Omit<ComponentPropsWithRef<"div">, "children"> {
 	urlTransform?: UrlTransform;
 }
 
+// Omit rehype-raw so HTML-like syntax in LLM output is rendered as
+// escaped text instead of being parsed by the HTML5 engine. Without
+// this, JSX fragments such as <ComponentName prop={value} /> are
+// consumed by rehype-raw and then stripped by rehype-sanitize,
+// silently destroying content mid-stream.
+const chatRehypePlugins = [
+	defaultRehypePlugins.sanitize,
+	defaultRehypePlugins.harden,
+];
+
 const fileViewerCSS =
 	"pre, [data-line], [data-diffs-header] { background-color: transparent !important; }";
 
@@ -241,6 +256,7 @@ export const Response = ({
 				controls={false}
 				components={components}
 				urlTransform={urlTransform}
+				rehypePlugins={chatRehypePlugins}
 			>
 				{children}
 			</Streamdown>