232 lines
8.4 KiB
TypeScript
232 lines
8.4 KiB
TypeScript
import { JSDOM } from "jsdom";
|
|
import DOMPurify from "dompurify";
|
|
import { Resvg } from "@resvg/resvg-js";
|
|
import { chromium } from "playwright-core";
|
|
import { puterChatComplete } from "../puter-inference.js";
|
|
import type { RenderResult, DiagramBundle } from "./types.js";
|
|
|
|
// ─── Security stripping ────────────────────────────────────────────────────────
|
|
|
|
const INIT_BLOCK_RE = /%%\{[\s\S]*?\}%%/g;
|
|
const CLICK_LINE_RE = /^\s*click\s+.*/gim;
|
|
|
|
/**
|
|
* Remove unsafe Mermaid directives (%%{init}%% blocks and click directives)
|
|
* from a Mermaid source string.
|
|
*/
|
|
export function stripUnsafeDirectives(source: string): {
|
|
cleaned: string;
|
|
stripped: boolean;
|
|
} {
|
|
const cleaned = source
|
|
.replace(INIT_BLOCK_RE, "")
|
|
.replace(CLICK_LINE_RE, "")
|
|
.trim();
|
|
|
|
const stripped = cleaned !== source.trim();
|
|
return { cleaned, stripped };
|
|
}
|
|
|
|
// ─── LLM prompt building ────────────────────────────────────────────────────────
|
|
|
|
const DIAGRAM_TYPE_HINTS: Record<string, string> = {
|
|
flowchart:
|
|
"Use graph TD or graph LR syntax for a flowchart.",
|
|
sequence:
|
|
"Use sequenceDiagram syntax with participant declarations.",
|
|
erd:
|
|
"Use erDiagram syntax with entity relationships.",
|
|
architecture:
|
|
"Use architecture-beta syntax with groups and services.",
|
|
mindmap:
|
|
"Use mindmap syntax with indented hierarchy.",
|
|
};
|
|
|
|
/**
|
|
* Build the system + user prompt for LLM-based Mermaid diagram synthesis.
|
|
*/
|
|
export function buildDiagramPrompt(
|
|
description: string,
|
|
diagramType: string,
|
|
): { system: string; user: string } {
|
|
const typeHint =
|
|
DIAGRAM_TYPE_HINTS[diagramType] ??
|
|
`Use ${diagramType} syntax appropriate for Mermaid.`;
|
|
|
|
const system = [
|
|
"You are a Mermaid diagram generator.",
|
|
"Output ONLY valid Mermaid syntax.",
|
|
"No markdown fences, no explanation, no comments.",
|
|
"The output must be directly parseable by mermaid.render().",
|
|
"",
|
|
`Diagram type: ${diagramType}`,
|
|
typeHint,
|
|
].join("\n");
|
|
|
|
const user = description;
|
|
|
|
return { system, user };
|
|
}
|
|
|
|
// ─── Playwright browser path resolution ────────────────────────────────────────
|
|
|
|
import fs from "fs";
|
|
import path from "path";
|
|
import os from "os";
|
|
|
|
function globSync(dir: string, filename: string): string[] {
|
|
// Simple two-level glob: dir/*/filename (no deep recursion needed)
|
|
const results: string[] = [];
|
|
if (!fs.existsSync(dir)) return results;
|
|
for (const entry of fs.readdirSync(dir)) {
|
|
const candidate = path.join(dir, entry, filename);
|
|
if (fs.existsSync(candidate)) {
|
|
results.push(candidate);
|
|
}
|
|
}
|
|
return results;
|
|
}
|
|
|
|
export function resolveBrowserPath(): string {
|
|
// Check environment first
|
|
const envPath = process.env.PLAYWRIGHT_BROWSERS_PATH;
|
|
if (envPath) return envPath;
|
|
|
|
// Common installed Playwright Chromium location
|
|
const homeDir = os.homedir();
|
|
const playwrightDir = path.join(homeDir, ".cache", "ms-playwright");
|
|
const matches = globSync(playwrightDir, path.join("chrome-linux64", "chrome"));
|
|
|
|
if (matches.length > 0) {
|
|
return matches[0]!;
|
|
}
|
|
|
|
throw new Error(
|
|
"Playwright Chromium not found. Run: npx playwright install chromium",
|
|
);
|
|
}
|
|
|
|
// ─── Mermaid HTML scaffold ──────────────────────────────────────────────────────
|
|
|
|
export function buildMermaidHtml(source: string, darkMode: boolean): string {
|
|
const escaped = source
|
|
.replace(/\\/g, "\\\\")
|
|
.replace(/`/g, "\\`")
|
|
.replace(/\$/g, "\\$");
|
|
|
|
return `<!DOCTYPE html>
|
|
<html>
|
|
<head>
|
|
<meta charset="utf-8" />
|
|
<style>
|
|
body { margin: 0; background: transparent; }
|
|
#render { display: inline-block; }
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div id="render"></div>
|
|
<script src="https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.min.js"></script>
|
|
<script>
|
|
mermaid.initialize({
|
|
startOnLoad: false,
|
|
securityLevel: "strict",
|
|
theme: ${darkMode ? '"dark"' : '"default"'},
|
|
});
|
|
const source = \`${escaped}\`;
|
|
mermaid.render("render", source).then(({ svg }) => {
|
|
document.getElementById("render").innerHTML = svg;
|
|
});
|
|
</script>
|
|
</body>
|
|
</html>`;
|
|
}
|
|
|
|
// ─── SVG extractor (used by Playwright page query) ────────────────────────────
|
|
|
|
function extractInnerHtml(el: Element): string {
|
|
return el.innerHTML;
|
|
}
|
|
|
|
// ─── Main renderer ──────────────────────────────────────────────────────────────
|
|
|
|
/**
|
|
* Render a diagram from a natural language prompt or raw Mermaid source.
|
|
*
|
|
* When `input.prompt` is provided: calls the LLM to synthesize Mermaid syntax first.
|
|
* When `input.source` is provided: uses the Mermaid source directly (re-render path).
|
|
*/
|
|
export async function renderDiagram(
|
|
input: Record<string, unknown>,
|
|
): Promise<RenderResult> {
|
|
const diagramType =
|
|
typeof input.diagramType === "string" ? input.diagramType : "flowchart";
|
|
const darkMode =
|
|
typeof input.darkMode === "boolean" ? input.darkMode : false;
|
|
|
|
// ── Determine Mermaid source ───────────────────────────────────────────────
|
|
let mermaidSource: string;
|
|
|
|
if (typeof input.source === "string" && input.source.trim()) {
|
|
// Re-render path: user provided Mermaid source directly
|
|
mermaidSource = input.source;
|
|
} else {
|
|
// LLM synthesis path (DIAG-01)
|
|
const prompt = typeof input.prompt === "string" ? input.prompt : "";
|
|
if (!prompt.trim()) {
|
|
throw new Error("renderDiagram: either 'prompt' or 'source' must be provided");
|
|
}
|
|
const { system, user } = buildDiagramPrompt(prompt, diagramType);
|
|
mermaidSource = await puterChatComplete([
|
|
{ role: "system", content: system },
|
|
{ role: "user", content: user },
|
|
]);
|
|
}
|
|
|
|
// ── Strip unsafe directives (DIAG-05) ─────────────────────────────────────
|
|
const { cleaned, stripped } = stripUnsafeDirectives(mermaidSource);
|
|
|
|
// ── Playwright render ──────────────────────────────────────────────────────
|
|
const executablePath = resolveBrowserPath();
|
|
const browser = await chromium.launch({
|
|
executablePath,
|
|
headless: true,
|
|
args: ["--no-sandbox", "--disable-setuid-sandbox"],
|
|
});
|
|
|
|
let rawSvg: string;
|
|
try {
|
|
const page = await browser.newPage();
|
|
const html = buildMermaidHtml(cleaned, darkMode);
|
|
await page.setContent(html, { waitUntil: "domcontentloaded" });
|
|
await page.waitForSelector("#render svg", { timeout: 15_000 });
|
|
rawSvg = await page.$eval("#render", extractInnerHtml);
|
|
} finally {
|
|
await browser.close();
|
|
}
|
|
|
|
// ── DOMPurify sanitize ─────────────────────────────────────────────────────
|
|
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
|
const { window } = new JSDOM("");
|
|
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
|
const purify = DOMPurify(window as any);
|
|
const cleanSvg = purify.sanitize(rawSvg, { USE_PROFILES: { svg: true } });
|
|
|
|
// ── Rasterize to PNG via Resvg ─────────────────────────────────────────────
|
|
const resvg = new Resvg(cleanSvg, { dpi: 144 });
|
|
const pngData = resvg.render().asPng();
|
|
|
|
// ── Build bundle ───────────────────────────────────────────────────────────
|
|
const bundle: DiagramBundle = {
|
|
type: "diagram-bundle",
|
|
svgBase64: Buffer.from(cleanSvg).toString("base64"),
|
|
pngBase64: Buffer.from(pngData).toString("base64"),
|
|
mermaidSource: cleaned,
|
|
stripped,
|
|
};
|
|
|
|
return {
|
|
filename: "diagram-bundle.json",
|
|
contentType: "application/json",
|
|
buffer: Buffer.from(JSON.stringify(bundle)),
|
|
};
|
|
}
|