LLeMbas/src/lembas/services/markdown.py

"""Render assistant messages from Markdown to sanitised HTML.

Rendering happens on the server, in Python, so there is no JavaScript Markdown
library to vendor and the streamed and final views cannot disagree about how
something should look.

The output is sanitised with nh3 (Rust ammonia). Model output is untrusted
input: it routinely contains HTML, and a model can be talked into emitting a
script tag, so this is a real boundary and not a formality.
"""

from __future__ import annotations

import functools
import html

import nh3
from markdown_it import MarkdownIt
from pygments import highlight
from pygments.formatters import HtmlFormatter
from pygments.lexers import get_lexer_by_name, guess_lexer
from pygments.util import ClassNotFound

# Class-based highlighting; the colours come from theme tokens in chat.css, so
# code blocks follow the active theme instead of carrying their own palette.
_FORMATTER = HtmlFormatter(nowrap=True, classprefix="pg-")

ALLOWED_TAGS = {
    "p", "br", "hr", "div", "span",
    "strong", "em", "del", "sub", "sup", "mark",
    "h1", "h2", "h3", "h4", "h5", "h6",
    "ul", "ol", "li",
    "blockquote", "pre", "code",
    "table", "thead", "tbody", "tr", "th", "td",
    "a", "img",
}

ALLOWED_ATTRIBUTES = {
    # "rel" is intentionally absent: nh3 rejects it here when link_rel is set,
    # because link_rel below is what writes it.
    "a": {"href", "title", "target"},
    "img": {"src", "alt", "title"},
    "code": {"class"},
    "pre": {"class"},
    "span": {"class"},
    "div": {"class"},
    "td": {"align"},
    "th": {"align"},
}

# javascript: and data: URLs are the obvious injection route through a link.
ALLOWED_URL_SCHEMES = {"http", "https", "mailto"}


def _render_fence(tokens, idx, _options, _env) -> str:
    """Render a fenced code block.

    This replaces the renderer's `fence` rule outright rather than using
    markdown-it's `highlight` option, because that option re-wraps whatever it
    is given in <pre><code> unless the string already starts with "<pre" --
    which would nest a second <pre> inside the wrapper this returns.
    """
    token = tokens[idx]
    code = token.content
    language = (token.info or "").strip().split()[0] if token.info else ""

    lexer = None
    if language:
        try:
            lexer = get_lexer_by_name(language, stripall=False)
        except (ClassNotFound, ValueError):
            lexer = None
    elif code.strip():
        # Guessing is only worth it for a decent sample; on two lines of text
        # Pygments guesses confidently and wrongly.
        try:
            lexer = guess_lexer(code) if len(code) > 80 else None
        except (ClassNotFound, ValueError):
            lexer = None

    if lexer is None:
        body = nh3.clean_text(code)
        label = language
    else:
        body = highlight(code, lexer, _FORMATTER)
        label = language or (lexer.aliases[0] if lexer.aliases else "")

    label_html = (
        f'<div class="code-block__label">{nh3.clean_text(label)}</div>' if label else ""
    )
    return (
        f'<div class="code-block">{label_html}'
        f'<pre class="code-block__pre"><code>{body}</code></pre></div>'
    )


@functools.lru_cache(maxsize=1)
def _parser() -> MarkdownIt:
    md = MarkdownIt("commonmark", {"linkify": True, "typographer": False})
    md.enable(["table", "strikethrough", "linkify"])
    md.renderer.rules["fence"] = _render_fence
    return md


def render_markdown(text: str) -> str:
    """Markdown to safe HTML, ready to drop into a message bubble."""
    if not text:
        return ""

    html = _parser().render(text)
    return nh3.clean(
        html,
        tags=ALLOWED_TAGS,
        attributes=ALLOWED_ATTRIBUTES,
        url_schemes=ALLOWED_URL_SCHEMES,
        # Anything opened from a model's output is untrusted; noopener stops it
        # reaching back through window.opener.
        link_rel="nofollow noopener noreferrer",
    )


def escape_text(text: str) -> str:
    """Escape a plain-text run for insertion as HTML element content.

    Used for user messages and for partial assistant text mid-stream, where the
    content is not yet complete enough to parse as Markdown.

    html.escape rather than nh3.clean_text: escaping the three structural
    characters is all that is needed for a text node, and it escapes character
    by character, so escaping a stream chunk-by-chunk gives the same result as
    escaping the whole string at once. nh3.clean_text also escapes spaces and
    slashes, which triples the size of a streamed token for no benefit.
    """
    return html.escape(text, quote=False)