feat: sanitize user messages (stackblitz#42)

packages/bolt/app/components/chat/AssistantMessage.tsx

@@ -8,7 +8,7 @@ interface AssistantMessageProps {
 export const AssistantMessage = memo(({ content }: AssistantMessageProps) => {
   return (
     <div className="overflow-hidden w-full">
-      <Markdown>{content}</Markdown>
+      <Markdown html>{content}</Markdown>
     </div>
   );
 });

packages/bolt/app/components/chat/Markdown.module.scss

@@ -58,9 +58,13 @@ $code-font-size: 13px;
     color: #6a737d;
   }
 
-  p:not(:last-of-type) {
-    margin-block-start: 0;
-    margin-block-end: 16px;
+  p {
+    white-space: pre-wrap;
+
+    &:not(:last-of-type) {
+      margin-block-start: 0;
+      margin-block-end: 16px;
+    }
   }
 
   a {

packages/bolt/app/components/chat/Markdown.tsx

@@ -2,7 +2,7 @@ import { memo, useMemo } from 'react';
 import ReactMarkdown, { type Components } from 'react-markdown';
 import type { BundledLanguage } from 'shiki';
 import { createScopedLogger } from '~/utils/logger';
-import { rehypePlugins, remarkPlugins } from '~/utils/markdown';
+import { rehypePlugins, remarkPlugins, allowedHTMLElements } from '~/utils/markdown';
 import { Artifact } from './Artifact';
 import { CodeBlock } from './CodeBlock';
 
@@ -12,12 +12,14 @@ const logger = createScopedLogger('MarkdownComponent');
 
 interface MarkdownProps {
   children: string;
+  html?: boolean;
+  limitedMarkdown?: boolean;
 }
 
-export const Markdown = memo(({ children }: MarkdownProps) => {
+export const Markdown = memo(({ children, html = false, limitedMarkdown = false }: MarkdownProps) => {
   logger.trace('Render');
 
-  const components = useMemo<Components>(() => {
+  const components = useMemo(() => {
     return {
       div: ({ className, children, node, ...props }) => {
         if (className?.includes('__boltArtifact__')) {
@@ -55,15 +57,16 @@ export const Markdown = memo(({ children }: MarkdownProps) => {
 
         return <pre {...rest}>{children}</pre>;
       },
-    };
+    } satisfies Components;
   }, []);
 
   return (
     <ReactMarkdown
+      allowedElements={allowedHTMLElements}
       className={styles.MarkdownContent}
       components={components}
-      remarkPlugins={remarkPlugins}
-      rehypePlugins={rehypePlugins}
+      remarkPlugins={remarkPlugins(limitedMarkdown)}
+      rehypePlugins={rehypePlugins(html)}
     >
       {children}
     </ReactMarkdown>

packages/bolt/app/components/chat/UserMessage.tsx

@@ -8,7 +8,7 @@ interface UserMessageProps {
 export function UserMessage({ content }: UserMessageProps) {
   return (
     <div className="overflow-hidden pt-[4px]">
-      <Markdown>{sanitizeUserMessage(content)}</Markdown>
+      <Markdown limitedMarkdown>{sanitizeUserMessage(content)}</Markdown>
     </div>
   );
 }

packages/bolt/app/lib/.server/llm/prompts.ts

@@ -1,4 +1,5 @@
 import { MODIFICATIONS_TAG_NAME, WORK_DIR } from '~/utils/constants';
+import { allowedHTMLElements } from '~/utils/markdown';
 import { stripIndents } from '~/utils/stripIndent';
 
 export const getSystemPrompt = (cwd: string = WORK_DIR) => `
@@ -35,6 +36,10 @@ You are Bolt, an expert AI assistant and exceptional senior software developer w
   Use 2 spaces for code indentation
 </code_formatting_info>
 
+<message_formatting_info>
+  You can make the output pretty by using only the following available HTML elements: ${allowedHTMLElements.map((tagName) => `<${tagName}>`).join(', ')}
+</message_formatting_info>
+
 <diff_spec>
   For user-made file modifications, a \`<${MODIFICATIONS_TAG_NAME}>\` section will appear at the start of the user message. It will contain either \`<diff>\` or \`<file>\` elements for each modified file:
 

packages/bolt/app/lib/runtime/message-parser.ts

@@ -272,7 +272,7 @@ export class StreamingMessageParser {
 const createArtifactElement: ElementFactory = (props) => {
   const elementProps = [
     'class="__boltArtifact__"',
-    Object.entries(props).map(([key, value]) => {
+    ...Object.entries(props).map(([key, value]) => {
       return `data-${camelToDashCase(key)}=${JSON.stringify(value)}`;
     }),
   ];

packages/bolt/app/utils/markdown.ts

@@ -1,6 +1,116 @@
 import rehypeRaw from 'rehype-raw';
 import remarkGfm from 'remark-gfm';
-import type { PluggableList } from 'unified';
+import type { PluggableList, Plugin } from 'unified';
+import rehypeSanitize, { defaultSchema, type Options as RehypeSanitizeOptions } from 'rehype-sanitize';
+import { SKIP, visit } from 'unist-util-visit';
+import type { UnistNode, UnistParent } from 'node_modules/unist-util-visit/lib';
 
-export const remarkPlugins = [remarkGfm] satisfies PluggableList;
-export const rehypePlugins = [rehypeRaw] satisfies PluggableList;
+export const allowedHTMLElements = [
+  'a',
+  'b',
+  'blockquote',
+  'br',
+  'code',
+  'dd',
+  'del',
+  'details',
+  'div',
+  'dl',
+  'dt',
+  'em',
+  'h1',
+  'h2',
+  'h3',
+  'h4',
+  'h5',
+  'h6',
+  'hr',
+  'i',
+  'ins',
+  'kbd',
+  'li',
+  'ol',
+  'p',
+  'pre',
+  'q',
+  'rp',
+  'rt',
+  'ruby',
+  's',
+  'samp',
+  'source',
+  'span',
+  'strike',
+  'strong',
+  'sub',
+  'summary',
+  'sup',
+  'table',
+  'tbody',
+  'td',
+  'tfoot',
+  'th',
+  'thead',
+  'tr',
+  'ul',
+  'var',
+];
+
+const rehypeSanitizeOptions: RehypeSanitizeOptions = {
+  ...defaultSchema,
+  tagNames: allowedHTMLElements,
+  attributes: {
+    ...defaultSchema.attributes,
+    div: [...(defaultSchema.attributes?.div ?? []), 'data*', ['className', '__boltArtifact__']],
+  },
+  strip: [],
+};
+
+export function remarkPlugins(limitedMarkdown: boolean) {
+  const plugins: PluggableList = [remarkGfm];
+
+  if (limitedMarkdown) {
+    plugins.unshift(limitedMarkdownPlugin);
+  }
+
+  return plugins;
+}
+
+export function rehypePlugins(html: boolean) {
+  const plugins: PluggableList = [];
+
+  if (html) {
+    plugins.push(rehypeRaw, [rehypeSanitize, rehypeSanitizeOptions]);
+  }
+
+  return plugins;
+}
+
+const limitedMarkdownPlugin: Plugin = () => {
+  return (tree, file) => {
+    const contents = file.toString();
+
+    visit(tree, (node: UnistNode, index, parent: UnistParent) => {
+      if (
+        index == null ||
+        ['paragraph', 'text', 'inlineCode', 'code', 'strong', 'emphasis'].includes(node.type) ||
+        !node.position
+      ) {
+        return true;
+      }
+
+      let value = contents.slice(node.position.start.offset, node.position.end.offset);
+
+      if (node.type === 'heading') {
+        value = `\n${value}`;
+      }
+
+      parent.children[index] = {
+        type: 'text',
+        value,
+      } as any;
+
+      return [SKIP, index] as const;
+    });
+  };
+};

packages/bolt/package.json

@@ -64,9 +64,11 @@
     "react-resizable-panels": "^2.0.20",
     "react-toastify": "^10.0.5",
     "rehype-raw": "^7.0.0",
+    "rehype-sanitize": "^6.0.0",
     "remark-gfm": "^4.0.0",
     "remix-utils": "^7.6.0",
-    "shiki": "^1.9.1"
+    "shiki": "^1.9.1",
+    "unist-util-visit": "^5.0.0"
   },
   "devDependencies": {
     "@cloudflare/workers-types": "^4.20240620.0",