hearthforge

hearthforge/src/services/highlight.ts
highlight.ts
⎇
Raw
1import { fileTypeFromBuffer } from "file-type";
2import type { Language } from "linguist-languages";
3import * as linguistLangs from "linguist-languages";
4import {
5 type BundledLanguage,
6 bundledLanguagesInfo,
7 createHighlighter,
8 type Highlighter,
9} from "shiki";
10import config from "../config.ts";
11import { BINARY_DETECT_BYTES, MAX_FILE_CACHE } from "../constants.ts";
12
13let highlighter: Highlighter | null = null;
14let extToLangId: Map<string, string> | null = null;
15let filenameToLangId: Map<string, string> | null = null;
16
17export async function highlightStartup(): Promise<void> {
18 console.log(
19 "[highlight] initializing highlighter and language detection maps...",
20 );
21
22 // Index linguist languages by lowercase name and aliases for precise lookup.
23 // Many shiki languages share the same tmScope (e.g. source.js has 31 entries),
24 // so scope-based matching is unreliable — name/alias matching is used instead.
25 const linguistByName = new Map<string, Language>();
26 for (const lang of Object.values(
27 linguistLangs as Record<string, Language>,
28 )) {
29 linguistByName.set(lang.name.toLowerCase(), lang);
30 for (const alias of lang.aliases ?? []) {
31 if (!linguistByName.has(alias.toLowerCase())) {
32 linguistByName.set(alias.toLowerCase(), lang);
33 }
34 }
35 }
36
37 const extMap = new Map<string, string>();
38 const fnMap = new Map<string, string>();
39
40 for (const { id, aliases } of bundledLanguagesInfo) {
41 // Match shiki lang → linguist language by ID then aliases (no scope fallback)
42 const linguistLang =
43 linguistByName.get(id) ??
44 aliases?.reduce<Language | undefined>(
45 (found, a) => found ?? linguistByName.get(a.toLowerCase()),
46 undefined,
47 );
48 if (!linguistLang) continue;
49 for (const ext of linguistLang.extensions ?? []) {
50 if (!extMap.has(ext)) extMap.set(ext, id);
51 }
52 for (const fn of linguistLang.filenames ?? []) {
53 if (!fnMap.has(fn)) fnMap.set(fn, id);
54 }
55 }
56
57 extToLangId = extMap;
58 filenameToLangId = fnMap;
59 // Grammars are loaded on first use by ensureLang(). Preloading all ~350
60 // bundled languages costs ~4s and ~265MB per worker, and a given repo
61 // only ever touches a handful.
62 highlighter = await createHighlighter({
63 themes: ["github-light", "github-dark"],
64 langs: [],
65 });
66
67 console.log(
68 `[highlight] ready: ${extMap.size} extensions, ${fnMap.size} filenames`,
69 );
70}
71
72export function getHighlighter(): Highlighter {
73 if (!highlighter) throw new Error("highlightStartup() has not been called");
74 return highlighter;
75}
76
77/**
78 * Load `lang`'s grammar if it is not loaded yet, and return the language to
79 * actually highlight with — falling back to "text" if the grammar is missing
80 * or fails to load.
81 *
82 * Callers MUST await this before codeToHtml/codeToTokensWithThemes. Those
83 * throw on an unloaded language, and both call sites catch that and degrade
84 * to unhighlighted plain text — so skipping this step turns highlighting off
85 * silently rather than loudly.
86 */
87export async function ensureLang(lang: string): Promise<string> {
88 const h = getHighlighter();
89 if (h.getLoadedLanguages().includes(lang)) return lang;
90 try {
91 await h.loadLanguage(lang as BundledLanguage);
92 return lang;
93 } catch {
94 return "text";
95 }
96}
97
98export function hasBinaryContent(buf: Buffer): boolean {
99 return buf.subarray(0, BINARY_DETECT_BYTES).includes(0);
100}
101
102export function detectLang(filename: string): string {
103 if (!extToLangId || !filenameToLangId) return "text";
104 const base = filename.split("/").pop() ?? filename;
105 const ext = base.includes(".") ? base.slice(base.lastIndexOf(".")) : "";
106 return (
107 filenameToLangId.get(base) ??
108 (ext ? extToLangId.get(ext) : undefined) ??
109 "text"
110 );
111}
112
113export type FileView =
114 | { type: "inline"; html: string; lines: number }
115 | { type: "download"; size: number }
116 | { type: "binary"; size: number }
117 | { type: "media"; mimeType: string; size: number };
118
119const fileCache = new Map<string, FileView>();
120
121function shikiHtmlToTable(html: string): string {
122 const styleMatch = html.match(/<pre[^>]*style="([^"]*)"/);
123 const rawStyle = styleMatch?.[1] ?? "";
124 // Convert Shiki's direct background-color/color to CSS custom properties so
125 // the CSS (not the inline style) controls the background in dark mode.
126 const lightBg =
127 rawStyle.match(/(?:^|;)background-color:([^;]+)/)?.[1] ?? "#fff";
128 const lightColor = rawStyle.match(/(?:^|;)color:([^;]+)/)?.[1] ?? "";
129 const darkBg = rawStyle.match(/--shiki-dark-bg:([^;]+)/)?.[1] ?? "";
130 const darkColor = rawStyle.match(/--shiki-dark:([^;]+)/)?.[1] ?? "";
131 const varParts = [`--shiki-light-bg:${lightBg}`];
132 if (lightColor) varParts.push(`--shiki-light:${lightColor}`);
133 if (darkBg) varParts.push(`--shiki-dark-bg:${darkBg}`);
134 if (darkColor) varParts.push(`--shiki-dark:${darkColor}`);
135 const tableStyle = varParts.join(";");
136
137 const codeMatch = html.match(/<code>([\s\S]*?)<\/code>/);
138 const codeContent = codeMatch?.[1] ?? "";
139 const rawLines = codeContent.split("\n");
140 // Shiki adds a trailing newline, so drop the last empty entry
141 const lines =
142 rawLines.length > 0 && rawLines[rawLines.length - 1] === ""
143 ? rawLines.slice(0, -1)
144 : rawLines;
145 const rows = lines
146 .map((line, i) => {
147 const n = i + 1;
148 return `<tr id="L${n}"><td class="blob-ln"><a href="#L${n}">${n}</a></td><td class="blob-code">${line}</td></tr>`;
149 })
150 .join("");
151 return `<table class="blob-table" style="${tableStyle}"><tbody>${rows}</tbody></table>`;
152}
153
154function plainTextToTable(text: string): string {
155 const rawLines = text.split("\n");
156 const lines =
157 rawLines.length > 0 && rawLines[rawLines.length - 1] === ""
158 ? rawLines.slice(0, -1)
159 : rawLines;
160 const rows = lines
161 .map((line, i) => {
162 const n = i + 1;
163 const escaped = line
164 .replace(/&/g, "&amp;")
165 .replace(/</g, "&lt;")
166 .replace(/>/g, "&gt;");
167 return `<tr id="L${n}"><td class="blob-ln"><a href="#L${n}">${n}</a></td><td class="blob-code">${escaped}</td></tr>`;
168 })
169 .join("");
170 return `<table class="blob-table"><tbody>${rows}</tbody></table>`;
171}
172
173export async function serveFile(
174 content: Buffer,
175 filename: string,
176 cacheKey: string,
177): Promise<FileView> {
178 const cached = fileCache.get(cacheKey);
179 if (cached) return cached;
180
181 const fileType = await fileTypeFromBuffer(content);
182 if (fileType) {
183 const group = fileType.mime.split("/")[0];
184 if (group === "image" || group === "audio" || group === "video") {
185 return {
186 type: "media",
187 mimeType: fileType.mime,
188 size: content.length,
189 };
190 }
191 }
192
193 const isBinary = hasBinaryContent(content);
194 if (isBinary) {
195 return { type: "binary", size: content.length };
196 }
197 if (content.length > config.INLINE_MAX_BYTES) {
198 return { type: "download", size: content.length };
199 }
200 const text = content.toString("utf-8");
201 const lineCount = text.split("\n").length;
202 // Outside the try: a load failure must fall back to "text", not to the
203 // unhighlighted plainTextToTable path below.
204 const lang = await ensureLang(detectLang(filename));
205 let view: FileView;
206 try {
207 const h = getHighlighter();
208 const shikiHtml = h.codeToHtml(text, {
209 lang,
210 themes: { light: "github-light", dark: "github-dark" },
211 });
212 view = {
213 type: "inline",
214 html: shikiHtmlToTable(shikiHtml),
215 lines: lineCount,
216 };
217 } catch {
218 view = {
219 type: "inline",
220 html: plainTextToTable(text),
221 lines: lineCount,
222 };
223 }
224 if (cacheKey) {
225 if (fileCache.size >= MAX_FILE_CACHE)
226 fileCache.delete(fileCache.keys().next().value!);
227 fileCache.set(cacheKey, view);
228 }
229 return view;
230}
231