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 re

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",
    )


# A mention is `@` followed by a run of non-space, claimed only at the start of
# the text or after whitespace. That last part is the whole rule: without it
# every email address in a message becomes a highlighted file reference, which
# is both wrong and ugly. It matches what composer.js recognises while typing,
# and the two must stay in step or the box and the transcript disagree.
_MENTION = re.compile(r"(?:(?<=\s)|^)@([^\s@]+)")


def highlight_tokens(text: str) -> str:
    """A user's own message, escaped, with `@mentions` marked.

    User turns have no render step at all -- the template prints the column and
    relies on `white-space: pre-wrap` -- so this is it, and it must escape
    before it injects or it is an XSS hole in the one place a person controls
    the bytes exactly.

    Only mentions. A `/command` never survives to a message: commands are
    intercepted in the composer and never posted, so anything beginning with a
    slash in a transcript is text somebody meant as text, and marking it as a
    command would be marking it as something it is not.
    """
    if not text:
        return ""
    escaped = html.escape(text, quote=False)
    # Applied to the *escaped* string, so the span is the only markup that can
    # exist. `@` and the path characters are untouched by html.escape, and a
    # `&amp;` it produced contains no whitespace -- which is why the pattern is
    # anchored on whitespace rather than on a character class.
    return _MENTION.sub(r'<span class="tok-mention">@\1</span>', escaped)


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)


# Code blocks are dropped whole rather than read out. A speech model given a
# code fence pronounces every bracket and underscore, which is unlistenable and
# takes longer than the prose it was buried in.
#
# Matched on <pre> rather than on the .code-block wrapper: the wrapper also
# contains a label div, so a non-greedy match for its closing tag stops at the
# label's and leaves the code behind. <pre> cannot nest, so this is exact.
_CODE_BLOCK = re.compile(r"<pre\b[^>]*>.*?</pre>", re.DOTALL)
_CODE_LABEL = re.compile(r"<div class=\"code-block__label\">.*?</div>", re.DOTALL)
_TAG = re.compile(r"<[^>]+>")
_WHITESPACE = re.compile(r"[ \t]*\n\s*\n\s*")

# Speech endpoints reject or truncate very long inputs, and a reply long enough
# to hit this is not one anybody is listening to in full.
MAX_SPEAKABLE = 8000


def speakable_text(text: str) -> str:
    """Markdown reduced to something worth reading aloud.

    Goes through the renderer rather than stripping the Markdown source
    directly, so tables, lists and links come out as their text instead of as
    punctuation, and there is one definition of what a message *says*.
    """
    if not text:
        return ""

    rendered = _CODE_LABEL.sub(" ", _CODE_BLOCK.sub("\n", render_markdown(text)))
    stripped = html.unescape(_TAG.sub(" ", rendered))

    # Paragraph breaks survive as a single newline: speech models use them as a
    # pause, and a wall of one line is read without any.
    stripped = _WHITESPACE.sub("\n", stripped)
    lines = [" ".join(line.split()) for line in stripped.splitlines()]
    return "\n".join(line for line in lines if line)[:MAX_SPEAKABLE]