lam-agent-vicinae/src/lib/render.ts

import { linkCitations } from "./citations";
import { CLOSING_FENCE, OPENING_FENCE, mapOutsideCodeSpans, readCodeSpan, repairCodeBlocks } from "./code";
import { sanitizeInlineHtml } from "./html";
import { looksLikeMath, renderLatex } from "./latex";
import type { Source } from "./types";

/**
 * The rendering pass for model output.
 *
 * `markdown.ts` builds the transcript, but an answer is not markdown that the
 * renderer can draw as written. Models write LaTeX, raw HTML, and diagrams made
 * of spaces, and Vicinae draws none of those: it has no maths block, block level
 * HTML is reduced to its text, and a paragraph wraps. Each of those is turned
 * into something the renderer already knows how to draw here, in one place, on
 * the way to the pane.
 *
 * The order is the whole design:
 *
 * 1. {@link repairCodeBlocks} runs first, because it decides where the fences
 *    are. Everything after this pass is fence aware and skips code.
 * 2. Maths runs next, so a formula becomes inline markup before the HTML pass
 *    looks at markup at all.
 * 3. HTML runs last, so the tags the maths pass adds are checked by the same
 *    allowlist as the ones the model wrote. It is a second opinion on this
 *    module's own output, and it costs nothing to take.
 *
 * Nothing here reads preferences or touches Vicinae, so the whole pipeline is
 * unit testable and the clipboard copy in `markdown.ts` keeps the source text
 * rather than the rendering.
 */

/** Fence labels that mean "this body is maths". */
const MATH_LANGUAGES = new Set(["math", "maths", "latex", "tex", "stex", "katex"]);

interface Segment {
  /** Fenced code, which every pass here leaves alone unless it is maths. */
  fence: boolean;
  /** Lower-cased language on the opening fence, empty when there is none. */
  language: string;
  /** For a fence, the body lines only. For text, the lines themselves. */
  lines: string[];
  /** A fence exactly as written, so it can be copied through untouched. */
  raw?: string[];
}

/**
 * Splits a document into text and fences.
 *
 * Both passes need the same answer to "is this code", and this is the one place
 * that decides. A fence with no closing marker runs to the end of the document,
 * which is how the parser reads it too.
 */
function splitSegments(markdown: string): Segment[] {
  const lines = markdown.split("\n");
  const segments: Segment[] = [];

  let text: string[] = [];
  const flushText = () => {
    if (text.length > 0) segments.push({ fence: false, language: "", lines: text });
    text = [];
  };

  let index = 0;
  while (index < lines.length) {
    const opener = OPENING_FENCE.exec(lines[index]);

    if (!opener) {
      text.push(lines[index]);
      index += 1;
      continue;
    }

    flushText();

    const raw = [lines[index]];
    index += 1;
    while (index < lines.length && !CLOSING_FENCE.test(lines[index])) {
      raw.push(lines[index]);
      index += 1;
    }
    if (index < lines.length) {
      raw.push(lines[index]);
      index += 1;
    }

    segments.push({
      fence: true,
      language: opener[3].toLowerCase(),
      lines: raw.slice(1, raw.length - (CLOSING_FENCE.test(raw[raw.length - 1]) ? 1 : 0)),
      raw,
    });
  }

  flushText();
  return segments;
}

const DISPLAY_MARKERS = ["$$", "\\["];

const DISPLAY_CLOSERS: Record<string, string> = { $$: "$$", "\\[": "\\]" };

interface DisplayOpen {
  marker: string;
  /** The formula on this line, whether or not the marker closed it here. */
  body: string;
  closed: boolean;
}

/**
 * Reads a display maths opener, if this line has one.
 *
 * A marker is only read as a block opener when it is the first thing on the
 * line, so `where $$x$$ is` stays a sentence. What follows the marker is the
 * formula whether or not the closer is on the same line, because a long formula
 * written across lines puts its first terms beside the opening `$$`.
 */
function openDisplayMath(line: string): DisplayOpen | undefined {
  const trimmed = line.trimStart();

  for (const marker of DISPLAY_MARKERS) {
    if (!trimmed.startsWith(marker)) continue;

    const closer = DISPLAY_CLOSERS[marker];
    const inner = trimmed.slice(marker.length).trim();

    if (inner.length >= closer.length && inner.endsWith(closer)) {
      return { marker, body: inner.slice(0, inner.length - closer.length), closed: true };
    }

    return { marker, body: inner, closed: false };
  }

  return undefined;
}

/** The part of a line before the closer, or `undefined` when the line has none. */
function closeDisplayMath(line: string, marker: string): string | undefined {
  const closer = DISPLAY_CLOSERS[marker];
  const trimmed = line.trimEnd();
  return trimmed.endsWith(closer) ? trimmed.slice(0, trimmed.length - closer.length) : undefined;
}

/**
 * Inline maths, in the order the delimiters have to be tried.
 *
 * `$$` has to be read before `$`, and both have to be read after the backslash
 * forms, which are unambiguous. The single `$` pattern is the fussy one and the
 * guard around it is deliberate: an opening `$` cannot be followed by a space
 * and a closing `$` cannot be preceded by one, which is what keeps `$5 and $10`
 * and `cd $A; echo $B` out of the renderer.
 */
const INLINE_MATH: Array<{ pattern: RegExp; guarded: boolean }> = [
  { pattern: /\\\((.+?)\\\)/g, guarded: false },
  { pattern: /\\\[(.+?)\\\]/g, guarded: false },
  { pattern: /\$\$(.+?)\$\$/g, guarded: false },
  { pattern: /(?<![$\\])\$(?![\s$])((?:\\.|[^\n$\\])+?)(?<!\s)\$(?!\$|\d)/g, guarded: true },
];

function convertInlineMath(text: string): string {
  return mapOutsideCodeSpans(text, (chunk) => {
    let out = chunk;

    for (const { pattern, guarded } of INLINE_MATH) {
      out = out.replace(pattern, (match: string, body: string) =>
        !guarded || looksLikeMath(body) ? renderLatex(body) : match,
      );
    }

    return out;
  });
}

/**
 * The line with every code span blanked out, at the same length.
 *
 * Used to ask whether a display marker is really at the start of the line. The
 * replacement is a non-space character on purpose: spaces would be trimmed away
 * again and a marker hidden behind a code span would look like a block opener.
 */
function maskCodeSpans(line: string): string {
  let out = "";
  let index = 0;

  while (index < line.length) {
    if (line[index] !== "`") {
      out += line[index];
      index += 1;
      continue;
    }

    const span = readCodeSpan(line, index);
    out += "\u0000".repeat(span.raw.length);
    index = span.end;
  }

  return out;
}

/**
 * Renders the maths in a run of text.
 *
 * A display formula is pulled onto its own paragraph with a blank line either
 * side, which is the only way to keep it out of the sentence around it and out
 * of the paragraph that follows.
 */
function renderMathLines(lines: string[]): string[] {
  const out: string[] = [];
  let index = 0;

  while (index < lines.length) {
    const line = lines[index];
    const display = openDisplayMath(line);

    // A marker inside a code span is being shown, not used, and it is the
    // opening one that decides whether this is a block at all.
    if (!display || !openDisplayMath(maskCodeSpans(line))) {
      out.push(convertInlineMath(line));
      index += 1;
      continue;
    }

    if (display.closed) {
      const formula = renderLatex(display.body, { display: true });
      if (formula.length > 0) out.push("", formula, "");
      index += 1;
      continue;
    }

    const body = [display.body];
    let cursor = index + 1;
    let closed = false;

    while (cursor < lines.length) {
      const rest = closeDisplayMath(lines[cursor], display.marker);
      if (rest !== undefined) {
        body.push(rest);
        closed = true;
        break;
      }

      body.push(lines[cursor]);
      cursor += 1;
    }

    // A marker with no closer is only a formula if what it swallowed looks like
    // one. A stray `$$5 per unit` would otherwise pull the rest of the answer
    // into a single paragraph, which is a lot of damage for a typo.
    if (!closed && !looksLikeMath(body.join("\n"))) {
      out.push(convertInlineMath(line));
      index += 1;
      continue;
    }

    const formula = renderLatex(body.join("\n"), { display: true });
    if (formula.length > 0) out.push("", formula, "");

    // An unclosed display marker swallows the rest of the run, which is what the
    // formula the model started to write actually covers.
    index = closed ? cursor + 1 : lines.length;
  }

  return out;
}

function mathPass(segments: Segment[]): Segment[] {
  return segments.map((segment) => {
    if (segment.fence) {
      if (!MATH_LANGUAGES.has(segment.language)) return segment;

      const formula = renderLatex(segment.lines.join("\n"), { display: true });
      return { fence: false, language: "", lines: formula.length > 0 ? ["", formula, ""] : [] };
    }

    return { fence: false, language: "", lines: renderMathLines(segment.lines) };
  });
}

/**
 * The tags a parser reads as an HTML *block* when a line starts with one.
 *
 * This is CommonMark's type 6 list. An HTML block is not a rendering of its own
 * in Vicinae: `processHtmlNodes` in `markdown-model.cpp` walks it, pulls out
 * every `<img>` as a block image, and keeps the trimmed text of everything else
 * with the markup thrown away. So `<table><tr><td>a</td><td>b</td></tr></table>`
 * arrives as the word `ab`, and a `<div>` full of styled spans arrives as one
 * run-on sentence.
 */
const HTML_BLOCK_TAGS = new Set([
  "address",
  "article",
  "aside",
  "base",
  "basefont",
  "blockquote",
  "body",
  "caption",
  "center",
  "col",
  "colgroup",
  "dd",
  "details",
  "dialog",
  "dir",
  "div",
  "dl",
  "dt",
  "fieldset",
  "figcaption",
  "figure",
  "footer",
  "form",
  "frame",
  "frameset",
  "h1",
  "h2",
  "h3",
  "h4",
  "h5",
  "h6",
  "head",
  "header",
  "hr",
  "html",
  "iframe",
  "legend",
  "li",
  "link",
  "main",
  "menu",
  "menuitem",
  "nav",
  "noframes",
  "ol",
  "optgroup",
  "option",
  "p",
  "param",
  "search",
  "section",
  "summary",
  "table",
  "tbody",
  "td",
  "tfoot",
  "th",
  "thead",
  "title",
  "tr",
  "track",
  "ul",
]);

/** Tags that never need a closing tag, which the balance check has to skip. */
const VOID_TAGS = new Set([
  "area",
  "base",
  "br",
  "col",
  "embed",
  "hr",
  "img",
  "input",
  "link",
  "meta",
  "param",
  "source",
  "track",
  "wbr",
]);

/** Whether this line would open an HTML block rather than continue a paragraph. */
function startsHtmlBlock(line: string): boolean {
  const trimmed = line.trim();
  const name = /^<([a-zA-Z][a-zA-Z0-9-]*)[\s/>]/.exec(trimmed)?.[1]?.toLowerCase();
  if (name && HTML_BLOCK_TAGS.has(name)) return true;

  // A single complete tag with nothing but whitespace after it opens a block
  // too, which is the rule that catches a standalone `<img>` or `<span>`.
  return /^<\/?[a-zA-Z][a-zA-Z0-9-]*(\s[^<>]*)?\/?>\s*$/.test(trimmed);
}

/**
 * Whether every tag in this run of HTML is closed in the right order.
 *
 * Only well-formed markup is handed to the rich text engine, because markup that
 * does not balance is markup nothing can be relied on to draw, and the text the
 * block path keeps is the better of two broken outcomes.
 */
function isBalancedHtml(text: string): boolean {
  const stack: string[] = [];

  for (const match of text.matchAll(/<(\/?)([a-zA-Z][a-zA-Z0-9-]*)[^<>]*?(\/?)>/g)) {
    const name = match[2].toLowerCase();
    if (match[3] === "/" || VOID_TAGS.has(name)) continue;

    if (match[1] === "/") {
      if (stack.pop() !== name) return false;
    } else {
      stack.push(name);
    }
  }

  return stack.length === 0;
}

/**
 * Whether a paragraph should be handed to the engine as inline HTML.
 *
 * The `<img>` exception is deliberate: an image inside an HTML block is the one
 * thing that path does better than this one, because it becomes a block image
 * the renderer actually loads.
 */
function shouldUnblock(paragraph: string[]): boolean {
  if (paragraph.length === 0) return false;
  if (!startsHtmlBlock(paragraph[0])) return false;

  const text = paragraph.join("\n");
  if (/<img\b/i.test(text)) return false;

  return isBalancedHtml(text);
}

/**
 * Moves a block of HTML into the paragraph the engine can see.
 *
 * Wrapping it in a `<span>` and collapsing it onto one line is what makes the
 * difference: a line that begins with `<span>` is a paragraph rather than an
 * HTML block, because `span` is not on that list, and a paragraph keeps its raw
 * HTML (`CMARK_NODE_HTML_INLINE` is passed through untouched). The same markup
 * therefore reaches the rich text engine, which is the only component in the
 * chain that can draw a table or a styled panel at all.
 *
 * Lines are joined because a continuation line starting with a block tag would
 * open a new HTML block of its own, which is how a multi-line HTML table loses
 * its markup today.
 */
function unblockHtml(lines: string[]): string[] {
  const out: string[] = [];
  let index = 0;

  while (index < lines.length) {
    let end = index;
    while (end < lines.length && lines[end].trim().length > 0) end += 1;

    const paragraph = lines.slice(index, end);
    if (shouldUnblock(paragraph)) {
      out.push(`<span>${paragraph.map((line) => line.trim()).join(" ")}</span>`);
    } else {
      out.push(...paragraph);
    }

    if (end === index) {
      out.push(lines[end]);
      index = end + 1;
      continue;
    }

    index = end;
  }

  return out;
}

function htmlPass(segments: Segment[]): Segment[] {
  return segments.map((segment) => {
    if (segment.fence) return segment;

    const unblocked = unblockHtml(segment.lines);
    return {
      fence: false,
      language: "",
      lines: sanitizeInlineHtml(unblocked.join("\n")).split("\n"),
    };
  });
}

function flatten(segments: Segment[]): string {
  const lines: string[] = [];
  for (const segment of segments) lines.push(...(segment.fence ? (segment.raw ?? []) : segment.lines));
  return lines.join("\n");
}

/**
 * Maths and markup only, for text a person wrote.
 *
 * A question goes through this rather than the full pass: a person pastes LaTeX
 * and writes HTML, but their typing is not the thing the code-fence repair
 * exists for, so a question is left as it was typed wherever the repair would
 * have changed it.
 */
export function renderInlineContent(markdown: string): string {
  return flatten(htmlPass(mathPass(splitSegments(markdown))));
}

/** What the pane will draw, and what the answer cited so the list can filter. */
export interface RenderedMarkdown {
  text: string;
  /**
   * The numbers this answer cited that became links.
   *
   * Reported rather than worked out again by the caller: the pass already had to
   * decide which markers were real and which were inside code, and a second reader
   * of the same text could only disagree with it.
   */
  cited: Set<number>;
}

/** Optional input to the pass. */
export interface RenderOptions {
  /**
   * The turn's sources, so a marker can become a link to the page it names.
   *
   * Absent for a question, and absent for an answer in a turn that consulted
   * nothing, in which case the citation pass is the identity function.
   */
  sources?: Source[];
}

/**
 * The full rendering pass, from what the model wrote to what the pane draws.
 *
 * Takes one message's text, not a whole transcript: the transcript adds its own
 * markdown around each message (a blockquote for a question, a rule between
 * exchanges) and none of that should be read as model output.
 */
export function renderModelOutput(markdown: string, options: RenderOptions = {}): RenderedMarkdown {
  const sources = options.sources ?? [];
  const cited = new Set<number>();

  /**
   * Maths first, then citations, then markup.
   *
   * The order is forced by one collision. A link label has to escape its brackets
   * or the parser ends the link text early, and `[\[1\]]` is exactly the shape of
   * LaTeX's own `\[...\]` display delimiter, which the maths pass would eat and
   * hand back as a bare `1`. So the model's maths is converted before any marker
   * becomes a link: the citation pass writes final form markdown, which is not
   * input for anything but the markup pass, and that one only reads angle brackets.
   */
  const segments = mathPass(splitSegments(repairCodeBlocks(markdown))).map((segment) => {
    if (segment.fence || sources.length === 0) return segment;

    const result = linkCitations(segment.lines.join("\n"), sources);
    for (const index of result.cited) cited.add(index);

    return { fence: false, language: "", lines: result.text.split("\n") };
  });

  return { text: flatten(htmlPass(segments)), cited };
}