fix(chat): render marimo HTML styled in mo.ui.chat (#9847) (#9888)

Fixes #9847.

The `mo.ui.chat` component does not render marimo Markdown/HTML the way
the rest
of the notebook does — marimo admonitions, in particular, come out
unstyled.

When a model returns a rich object, the backend serializes it to HTML
via
`as_html(response).text` (`marimo/_plugins/ui/_impl/chat/chat.py`). For
an
admonition, `mo.md(...)` produces:

```html
<span class="markdown prose dark:prose-invert contents"><div class="admonition error">
<p class="admonition-title">Error</span>
<span class="paragraph">Nope!</span>
</div></span>
```

The frontend renders assistant text parts with `MarkdownRenderer`, which
wraps
**Streamdown**. Streamdown sanitizes with `rehype-sanitize` using the
default
GitHub schema, which strips `class` from most elements. The
`admonition`/`prose`
classes that carry all the styling are removed, so the content renders
as plain
text. The marimo CSS itself is fine and global
(`frontend/src/css/admonition.css`); it just never matches because the
class
names are gone.

## Fix

Configure Streamdown's rehype pipeline to **preserve class names**. We
extend
`rehype-sanitize`'s default schema to allow `className` on every element
and pass
it via `rehypePlugins`, reusing Streamdown's other defaults
(`rehype-raw` to
parse the HTML, `rehype-harden` to lock down links/images). Scripts,
inline
styles, and unsafe URLs are still sanitized away — we only stop
discarding class
names, which are not an execution vector.

This is a single-renderer fix: backend-rendered marimo HTML
(admonitions, prose,
etc.) now styles correctly through Streamdown, and normal streamed LLM
markdown
is unchanged. Interactive marimo UI elements (`<marimo-...>`) still
route to
marimo's HTML renderer in `chat-display.tsx`, since they're custom
elements that
Streamdown can neither keep nor hydrate.

## Fix malformed paragraph tags

`mo.md(...)` also emitted malformed HTML for admonitions — note the
`<p class="admonition-title">Error</span>` above. `_md` rewrote `<p>` to
`<span class="paragraph">` to allow nested block elements, but the
opening
replacement only matched bare `<p>` while the closing replacement
matched every
`</p>`, so an attributed paragraph (the admonition title) was left with
a
mismatched close tag. Browsers auto-correct this, but the stricter
sanitize/parse path does not. We now rewrite only attribute-less
`<p>...</p>`
pairs, so attributed paragraphs keep a balanced `</p>`.

Author: mscolnick

frontend/package.json

@@ -141,6 +141,7 @@
     "react-vega": "^8.0.0",
     "react-virtuoso": "^4.18.1",
     "reactflow": "^11.11.4",
+    "rehype-sanitize": "^6.0.0",
     "remark-gfm": "^4.0.1",
     "reveal.js": "^6.0.0",
     "rpc-anywhere": "^1.7.0",

frontend/src/components/chat/chat-display.tsx

@@ -60,8 +60,8 @@ export const renderUIMessage = ({
 
     switch (part.type) {
       case "text":
-        // Streamdown sanitizes the HTML which strips out marimo elements
-        // So instead, we render the HTML with our custom renderer.
+        // Streamdown strips marimo elements, so render them with our own HTML
+        // renderer. Other markdown (incl. mo.md HTML) goes through Streamdown.
         if (part.text.includes("<marimo-")) {
           return (
             <React.Fragment key={index}>

frontend/src/components/markdown/__tests__/markdown-renderer.test.tsx

@@ -0,0 +1,43 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { render, screen, waitFor } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { MarkdownRenderer } from "../markdown-renderer";
+
+describe("MarkdownRenderer", () => {
+  // Regression test for https://github.com/marimo-team/marimo/issues/9847
+  it("preserves class names on raw HTML (e.g. marimo admonitions)", async () => {
+    const { container } = render(
+      <MarkdownRenderer
+        content={
+          '<div class="admonition error">' +
+          '<p class="admonition-title">Error</p>' +
+          "<p>Nope!</p>" +
+          "</div>"
+        }
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Nope!")).toBeInTheDocument();
+    });
+
+    const admonition = container.querySelector(".admonition.error");
+    expect(admonition).toBeInTheDocument();
+    expect(container.querySelector(".admonition-title")).toBeInTheDocument();
+  });
+
+  it("strips unsafe tags while keeping class names", async () => {
+    const { container } = render(
+      <MarkdownRenderer
+        content={'<div class="safe"><script>alert(1)</script>hello</div>'}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("hello")).toBeInTheDocument();
+    });
+
+    expect(container.querySelector(".safe")).toBeInTheDocument();
+    expect(container.querySelector("script")).not.toBeInTheDocument();
+  });
+});

frontend/src/components/markdown/markdown-renderer.tsx

@@ -4,7 +4,12 @@ import { EditorView } from "@codemirror/view";
 import { useAtomValue } from "jotai";
 import { BetweenHorizontalStartIcon } from "lucide-react";
 import { memo, Suspense, useState } from "react";
-import { Streamdown, type StreamdownProps } from "streamdown";
+import rehypeSanitize, { defaultSchema } from "rehype-sanitize";
+import {
+  defaultRehypePlugins,
+  Streamdown,
+  type StreamdownProps,
+} from "streamdown";
 import { Button, type ButtonProps } from "@/components/ui/button";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
 import { useCellActions } from "@/core/cells/cells";
@@ -148,6 +153,23 @@ const CopyButton: React.FC<ButtonProps> = ({ onClick, ...props }) => {
   );
 };
 
+// Allow `className` on every element; the default GitHub schema strips it,
+// which drops the styling on marimo HTML like `mo.md(...)` admonitions.
+const sanitizeSchema = {
+  ...defaultSchema,
+  attributes: {
+    ...defaultSchema.attributes,
+    "*": [...(defaultSchema.attributes?.["*"] ?? []), "className"],
+  },
+};
+
+// Keep Streamdown's other rehype plugins (raw, harden) so scripts and unsafe
+// URLs are still sanitized.
+const REHYPE_PLUGINS: StreamdownProps["rehypePlugins"] = Object.values({
+  ...defaultRehypePlugins,
+  sanitize: [rehypeSanitize, sanitizeSchema],
+});
+
 type Components = StreamdownProps["components"];
 
 const COMPONENTS: Components = {
@@ -187,6 +209,7 @@ export const MarkdownRenderer = memo(({ content }: { content: string }) => {
     <Streamdown
       components={COMPONENTS}
       plugins={plugins}
+      rehypePlugins={REHYPE_PLUGINS}
       className="mo-markdown-renderer"
     >
       {content}

marimo/_output/md.py

@@ -269,10 +269,16 @@ def __init__(
 
         # markdown.markdown appends a newline, hence strip
         html_text = _render_markdown(text).strip()
-        # replace <p> tags with <span> as HTML doesn't allow nested <div>s in <p>s
-        html_text = html_text.replace(
-            "<p>", '<span class="paragraph">'
-        ).replace("</p>", "</span>")
+        # replace <p> tags with <span> as HTML doesn't allow nested <div>s in
+        # <p>s, including the admonition title's <p> so its </p> stays matched
+        html_text = (
+            html_text.replace(
+                '<p class="admonition-title">',
+                '<span class="admonition-title">',
+            )
+            .replace("<p>", '<span class="paragraph">')
+            .replace("</p>", "</span>")
+        )
 
         if apply_markdown_class:
             classes = ["markdown", "prose", "dark:prose-invert", "contents"]

pnpm-lock.yaml

@@ -278,6 +278,45 @@ def test_md_nested_list_preserves_multi_paragraph() -> None:
     )
 
 
+def test_md_admonition_paragraph_tags_balanced() -> None:
+    # Regression for https://github.com/marimo-team/marimo/issues/9847
+    input_text = """/// error
+Nope!
+///"""
+
+    result = _md(input_text, apply_markdown_class=False).text
+    assert result == snapshot(
+        """\
+<div class="admonition error">
+<span class="admonition-title">Error</span>
+<span class="paragraph">Nope!</span>
+</div>\
+"""
+    )
+    assert '<p class="admonition-title"' not in result
+    assert "<p>" not in result
+    assert "</p>" not in result
+
+
+@pytest.mark.parametrize(
+    "kind", ["note", "warning", "danger", "tip", "error", "caution"]
+)
+def test_md_admonition_well_formed(kind: str) -> None:
+    result = _md(f"/// {kind}\nBody text\n///").text
+    assert f'<div class="admonition {kind}">' in result
+    assert '<span class="admonition-title">' in result
+    assert '<span class="paragraph">Body text</span>' in result
+    assert "<p>" not in result
+    assert "</p>" not in result
+
+
+def test_md_admonition_custom_title_well_formed() -> None:
+    result = _md("/// note | My Title\nbody here\n///").text
+    assert '<span class="admonition-title">My Title</span>' in result
+    assert "<p>" not in result
+    assert "</p>" not in result
+
+
 def test_md_nested_list_with_inline_elements() -> None:
     # Inline elements like bold/italic/code should survive p-unwrapping
     input_text = """- **Bold item**