Streaming Markdown

npm install @kushagradhawan/kookie-ui react-markdown remark-gfm rehype-raw harden-react-markdown shiki @hugeicons/react @hugeicons/core-free-icons

'use client';

import * as React from 'react';
import ReactMarkdown, { type Components } from 'react-markdown';
import remarkGfm from 'remark-gfm';
import rehypeRaw from 'rehype-raw';
import hardenReactMarkdown from 'harden-react-markdown';
import { MarkdownContent, createMarkdownComponents, type MarkdownComponentOptions } from '../markdown/markdown';
import { completeUnterminatedMarkdown, parseMarkdownIntoBlocks } from './complete-markdown';

const HardenedMarkdown = hardenReactMarkdown(ReactMarkdown);

const LINK_PREFIXES = ['https://', 'http://', '/'];
const IMAGE_PREFIXES = ['https://', 'http://', '/', 'data:'];
const ALLOWED_PROTOCOLS = ['mailto:', 'tel:', 'data:', 'http:', 'https:'];
const REMARK_PLUGINS = [remarkGfm];
const REHYPE_PLUGINS = [rehypeRaw];

export type StreamingMarkdownOptions = MarkdownComponentOptions & {
  /** @default "compact" */
  spacing?: 'compact' | 'spacious';
  /** Origin used to resolve and validate relative links and images. Defaults to the page origin. */
  defaultOrigin?: string;
  /** Memoize each top-level block so only the growing tail re-renders. Needs `blockParser`. @default true */
  enableBlockMemoization?: boolean;
  /** Splits content into blocks, e.g. `(md) => marked.lexer(md, { gfm: true })`. */
  blockParser?: (content: string) => Array<{ raw?: string }>;
  /** Override individual element renderers. */
  components?: Partial<Components>;
};

export interface StreamingMarkdownProps {
  /** Markdown so far. Incomplete syntax at the end is closed before parsing. */
  content: string;
  /** Stable id for this message, used in block keys. */
  id: string;
  options?: StreamingMarkdownOptions;
}

function resolveOrigin(customOrigin?: string): string {
  if (customOrigin) return customOrigin;
  if (typeof window !== 'undefined') return window.location.origin;
  return 'http://localhost';
}

type MarkdownBlockProps = {
  content: string;
  defaultOrigin: string;
  components: Components;
};

// Each block renders straight into the shared MarkdownContent, so spacing flows across blocks.
const MarkdownBlock = React.memo(function MarkdownBlock({ content, defaultOrigin, components }: MarkdownBlockProps) {
  return (
    <HardenedMarkdown
      defaultOrigin={defaultOrigin}
      allowedLinkPrefixes={LINK_PREFIXES}
      allowedImagePrefixes={IMAGE_PREFIXES}
      allowedProtocols={ALLOWED_PROTOCOLS}
      allowDataImages
      components={components}
      remarkPlugins={REMARK_PLUGINS}
      rehypePlugins={REHYPE_PLUGINS}
    >
      {content}
    </HardenedMarkdown>
  );
});

const EMPTY_OPTIONS: StreamingMarkdownOptions = {};

/**
 * Markdown renderer for AI responses that arrive token by token. Closes unfinished syntax,
 * memoizes finished blocks, and only allows safe link and image targets.
 */
export function StreamingMarkdown({ content, id, options = EMPTY_OPTIONS }: StreamingMarkdownProps) {
  const {
    defaultOrigin: customOrigin,
    enableBlockMemoization = true,
    blockParser,
    components: customComponents,
    codeBlockCollapsible = false,
    imageComponent,
    linkComponent,
    spacing = 'compact',
  } = options;

  const defaultOrigin = resolveOrigin(customOrigin);

  // Primitive deps, so an inline `options` object does not rebuild every renderer.
  const components = React.useMemo(
    () => ({
      ...createMarkdownComponents({ codeBlockCollapsible, imageComponent, linkComponent }),
      ...customComponents,
    }),
    [codeBlockCollapsible, imageComponent, linkComponent, customComponents],
  );

  const blocks = React.useMemo(() => {
    if (enableBlockMemoization && blockParser) return parseMarkdownIntoBlocks(content, blockParser);
    const completed = completeUnterminatedMarkdown(content);
    return completed.trim() ? [completed] : [];
  }, [content, enableBlockMemoization, blockParser]);

  if (blocks.length === 0) return null;

  return (
    <MarkdownContent spacing={spacing}>
      {blocks.map((block, index) => (
        <MarkdownBlock key={`${id}-${index}`} content={block} defaultOrigin={defaultOrigin} components={components} />
      ))}
    </MarkdownContent>
  );
}

/**
 * Closes markdown that a stream has cut off mid-token, so a partial message still parses
 * the way the finished one will. Handles code fences, inline code, links, bold, italic,
 * strikethrough, and trailing headings, list items and blockquotes.
 */
export function completeUnterminatedMarkdown(content: string): string {
  if (!content.trim()) return content;

  const trimmed = content.trimEnd();
  const trailingWhitespace = content.slice(trimmed.length);
  const close = (suffix: string) => trimmed + suffix + trailingWhitespace;

  // A fence that has only just opened, e.g. "```python" with no newline yet.
  if (/```[\w-]*$/.test(trimmed) && !trimmed.includes('```\n') && !/```[\w-]+\n/.test(trimmed)) {
    return close('\n```');
  }

  if (countMatches(trimmed, /```/g) % 2 === 1) return close('\n```');

  // Inline backticks, ignoring fence markers.
  if (countMatches(trimmed.replace(/```[\w-]*/g, ''), /`/g) % 2 === 1) return close('`');

  if (/\[[^\]]*\]\([^)]*$/.test(trimmed)) return close(')');

  if (hasUnpairedMarker(trimmed, '**')) return close('**');
  if (hasUnpairedMarker(trimmed, '__')) return close('__');
  if (hasUnpairedMarker(trimmed, '~~')) return close('~~');

  // Structural endings need a newline to parse. Checked before italics so "* item" stays a list.
  const lastLine = trimmed.slice(trimmed.lastIndexOf('\n') + 1);
  if (/^#{1,6}\s+.+$/.test(lastLine) || /^(\s*)([-*+]|\d+\.)\s+.+$/.test(lastLine) || /^>\s+.+$/.test(lastLine)) {
    return close('\n');
  }

  if (hasUnpairedItalic(trimmed, '*')) return close('*');
  if (hasUnpairedItalic(trimmed, '_')) return close('_');

  return content;
}

/**
 * Splits markdown into top-level blocks so each can be memoized while a stream grows.
 * Pass a lexer such as `marked.lexer`; without one the content stays a single block.
 */
export function parseMarkdownIntoBlocks(
  content: string,
  parser?: (content: string) => Array<{ raw?: string }>,
): string[] {
  if (!content.trim()) return [];

  const completed = completeUnterminatedMarkdown(content);
  if (!parser) return [completed];

  return parser(completed)
    .map((token) => (typeof token.raw === 'string' ? token.raw : ''))
    .filter((raw) => raw.trim().length > 0);
}

function countMatches(content: string, pattern: RegExp): number {
  return content.match(pattern)?.length ?? 0;
}

/** An odd number of a double marker, with text after the last one. */
function hasUnpairedMarker(content: string, marker: '**' | '__' | '~~'): boolean {
  if (content.split(marker).length % 2 === 1) return false;
  return content.lastIndexOf(marker) + marker.length < content.length;
}

/** An odd number of single `*` or `_`, skipping list bullets and bold markers. */
function hasUnpairedItalic(content: string, marker: '*' | '_'): boolean {
  let count = 0;
  let lastIndex = -1;

  for (let i = 0; i < content.length; i++) {
    if (content[i] !== marker) continue;

    const isListBullet = marker === '*' && (i === 0 || content[i - 1] === '\n') && content[i + 1] === ' ';
    if (isListBullet) continue;

    if (content[i + 1] === marker) {
      i++; // Skip the pair — it is bold, not italic.
      continue;
    }
    if (content[i - 1] === marker) continue;

    count++;
    lastIndex = i;
  }

  if (count % 2 === 0 || lastIndex === -1) return false;

  // A trailing marker opens new italics only if something comes before it.
  if (lastIndex === content.length - 1) return lastIndex > 0;
  return true;
}

import { StreamingMarkdown } from '@/components/blocks/streaming-markdown/streaming-markdown';
 
export function AssistantMessage({ message }: { message: { id: string; text: string } }) {
  return (
    <StreamingMarkdown
      id={message.id}
      content={message.text}
      options={{ spacing: 'compact' }}
    />
  );
}

import { marked } from 'marked';
 
const blockParser = (content: string) => marked.lexer(content, { gfm: true });
 
<StreamingMarkdown
  id={message.id}
  content={message.text}
  options={{ blockParser }}
/>

import {
  completeUnterminatedMarkdown,
  parseMarkdownIntoBlocks,
} from '@/components/blocks/streaming-markdown/complete-markdown';
 
completeUnterminatedMarkdown('Use **bold');
// 'Use **bold**'