Fix syntax highlighting for files

Author: kylecarbs

src/api/diff.ts

@@ -157,6 +157,15 @@ function guessLang(filename?: string): string {
 // Syntax Highlighting
 // ============================================================================
 
+function escapeHtml(text: string): string {
+  return text
+    .replace(/&/g, "&")
+    .replace(/</g, "&lt;")
+    .replace(/>/g, "&gt;")
+    .replace(/"/g, "&quot;")
+    .replace(/'/g, "&#039;");
+}
+
 function hastToHtml(node: any): string {
   if (node.type === "text") {
     return escapeHtml(node.value);
@@ -171,15 +180,6 @@ function hastToHtml(node: any): string {
   return "";
 }
 
-function escapeHtml(text: string): string {
-  return text
-    .replace(/&/g, "&amp;")
-    .replace(/</g, "&lt;")
-    .replace(/>/g, "&gt;")
-    .replace(/"/g, "&quot;")
-    .replace(/'/g, "&#039;");
-}
-
 function highlight(code: string, lang: string): string {
   try {
     const tree = refractor.highlight(code, lang);
@@ -189,6 +189,95 @@ function highlight(code: string, lang: string): string {
   }
 }
 
+/**
+ * Highlight an entire file and return an array of HTML strings, one per line.
+ * This handles multi-line constructs (strings, comments) correctly by
+ * closing and reopening tags at line boundaries.
+ */
+interface OpenTag {
+  tagName: string;
+  className?: string;
+}
+
+function highlightFileByLines(content: string, lang: string): string[] {
+  if (!content) return [];
+
+  try {
+    const tree = refractor.highlight(content, lang);
+    const lines: string[] = [];
+    let currentLine: string[] = [];
+    const openTags: OpenTag[] = [];
+
+    function closeAllTags(): string {
+      return [...openTags]
+        .reverse()
+        .map((t) => `</${t.tagName}>`)
+        .join("");
+    }
+
+    function openAllTags(): string {
+      return openTags
+        .map((t) => {
+          const cls = t.className ? ` class="${t.className}"` : "";
+          return `<${t.tagName}${cls}>`;
+        })
+        .join("");
+    }
+
+    function processText(text: string) {
+      const parts = text.split("\n");
+      for (let i = 0; i < parts.length; i++) {
+        if (i > 0) {
+          // End current line with closing tags
+          currentLine.push(closeAllTags());
+          lines.push(currentLine.join(""));
+          // Start new line with opening tags
+          currentLine = [openAllTags()];
+        }
+        if (parts[i]) {
+          currentLine.push(escapeHtml(parts[i]));
+        }
+      }
+    }
+
+    function walkNode(node: any) {
+      if (node.type === "text") {
+        processText(node.value);
+      } else if (node.type === "element") {
+        const { tagName, properties, children } = node;
+        const className = (properties?.className as string[] | undefined)?.join(
+          " "
+        );
+        const tag: OpenTag = { tagName, className };
+
+        // Open tag
+        const cls = className ? ` class="${className}"` : "";
+        currentLine.push(`<${tagName}${cls}>`);
+        openTags.push(tag);
+
+        // Process children
+        children.forEach(walkNode);
+
+        // Close tag
+        openTags.pop();
+        currentLine.push(`</${tagName}>`);
+      }
+    }
+
+    tree.children.forEach(walkNode);
+
+    // Don't forget the last line
+    if (currentLine.length > 0) {
+      lines.push(currentLine.join(""));
+    }
+
+    return lines;
+  } catch {
+    // Fallback: escape each line
+    return content.split("\n").map(escapeHtml);
+  }
+}
+
 /**
  * Highlight a range of lines from file content.
  * Returns an array of DiffLine objects with syntax highlighting.
@@ -201,12 +290,18 @@ export function highlightFileLines(
 ): DiffLine[] {
   const language = guessLang(filename);
   const allLines = content.split("\n");
+
+  // Pre-highlight the entire file for proper context
+  const highlightedLines = highlightFileByLines(content, language);
+
   const result: DiffLine[] = [];
 
   for (let i = 0; i < count; i++) {
     const lineNum = startLine + i;
     const lineContent = allLines[lineNum - 1] ?? "";
-    const highlighted = highlight(lineContent, language);
+    // Use pre-highlighted HTML, fallback to individual highlighting
+    const highlighted =
+      highlightedLines[lineNum - 1] ?? highlight(lineContent, language);
 
     result.push({
       type: "normal",
@@ -549,7 +644,9 @@ export function parseDiffWithHighlighting(
   patch: string,
   filename: string,
   previousFilename?: string,
-  cacheKey?: string
+  cacheKey?: string,
+  oldContent?: string,
+  newContent?: string
 ): ParsedDiff {
   // Check cache
   if (cacheKey && diffCache.has(cacheKey)) {
@@ -570,6 +667,19 @@ ${patch}`;
   }
 
   const language = guessLang(filename);
+  const prevLanguage = previousFilename
+    ? guessLang(previousFilename)
+    : language;
+
+  // Pre-highlight full files if content is provided
+  // This ensures proper highlighting for multi-line constructs (strings, comments, etc.)
+  const oldHighlightedLines = oldContent
+    ? highlightFileByLines(oldContent, prevLanguage)
+    : null;
+  const newHighlightedLines = newContent
+    ? highlightFileByLines(newContent, language)
+    : null;
+
   const rawHunks = insertSkipBlocks(
     file.hunks.map((hunk) => parseHunk(hunk, opts))
   );
@@ -598,15 +708,63 @@ ${patch}`;
             ? (line as any).lineNumber
             : undefined;
 
+        // For lines with a single segment (no inline diff), use pre-highlighted content
+        // For lines with multiple segments (inline diff), highlight each segment
+        const hasSingleSegment = line.content.length === 1;
+        const singleSegmentIsNormal =
+          hasSingleSegment && line.content[0].type === "normal";
+
         return {
           type: line.type,
           oldLineNumber: oldNum,
           newLineNumber: newNum,
-          content: line.content.map((seg) => ({
-            value: seg.value,
-            html: highlight(seg.value, language),
-            type: seg.type,
-          })),
+          content: line.content.map((seg) => {
+            let html: string;
+
+            // Try to use pre-highlighted content for better context
+            if (singleSegmentIsNormal) {
+              // Use pre-highlighted line if available
+              if (
+                line.type === "delete" &&
+                oldHighlightedLines &&
+                oldNum !== undefined
+              ) {
+                html =
+                  oldHighlightedLines[oldNum - 1] ??
+                  highlight(seg.value, prevLanguage);
+              } else if (
+                line.type === "insert" &&
+                newHighlightedLines &&
+                newNum !== undefined
+              ) {
+                html =
+                  newHighlightedLines[newNum - 1] ??
+                  highlight(seg.value, language);
+              } else if (
+                line.type === "normal" &&
+                newHighlightedLines &&
+                newNum !== undefined
+              ) {
+                // For normal lines, prefer new file highlighting (same content in both)
+                html =
+                  newHighlightedLines[newNum - 1] ??
+                  highlight(seg.value, language);
+              } else {
+                html = highlight(seg.value, language);
+              }
+            } else {
+              // Multiple segments (inline diff) - highlight each segment individually
+              // This is acceptable since inline diffs are usually small
+              const segLang = seg.type === "delete" ? prevLanguage : language;
+              html = highlight(seg.value, segLang);
+            }
+
+            return {
+              value: seg.value,
+              html,
+              type: seg.type,
+            };
+          }),
         };
       }),
     };

src/browser/contexts/pr-review/useDiffLoader.ts

@@ -1,6 +1,7 @@
 import { useEffect } from "react";
 import type { PullRequestFile } from "@/api/types";
 import { diffService } from "@/browser/lib/diff";
+import { useGitHub } from "@/browser/contexts/github";
 import { usePRReviewStore, usePRReviewSelector, type ParsedDiff } from ".";
 
 const diffCache = new Map<string, ParsedDiff>();
@@ -26,15 +27,23 @@ function abortAllPendingFetches() {
   }
 }
 
+/** Content getter function type for fetching file content */
+type FileContentGetter = (path: string, ref: string) => Promise<string>;
+
 async function fetchParsedDiff(
   file: PullRequestFile,
-  signal?: AbortSignal
+  signal?: AbortSignal,
+  getFileContent?: FileContentGetter,
+  baseRef?: string,
+  headRef?: string
 ): Promise<ParsedDiff> {
   if (!file.patch || !file.sha) {
     return { hunks: [] };
   }
 
-  const cacheKey = file.sha;
+  // Cache key includes whether we have file content (for better highlighting)
+  const hasContent = !!(getFileContent && baseRef && headRef);
+  const cacheKey = hasContent ? `${file.sha}:full` : file.sha;
 
   // Check cache first
   if (diffCache.has(cacheKey)) {
@@ -72,11 +81,39 @@ async function fetchParsedDiff(
   }
 
   const fetchPromise = (async () => {
+    let oldContent: string | undefined;
+    let newContent: string | undefined;
+
+    // Fetch file content for better syntax highlighting if getter is provided
+    if (getFileContent && baseRef && headRef) {
+      try {
+        const [oldResult, newResult] = await Promise.all([
+          // For deleted files or renames, use previous_filename for base
+          file.status === "added"
+            ? Promise.resolve("")
+            : getFileContent(
+                file.previous_filename || file.filename,
+                baseRef
+              ).catch(() => ""),
+          // For deleted files, new content is empty
+          file.status === "removed"
+            ? Promise.resolve("")
+            : getFileContent(file.filename, headRef).catch(() => ""),
+        ]);
+        oldContent = oldResult;
+        newContent = newResult;
+      } catch {
+        // If fetching content fails, continue without it
+      }
+    }
+
     // Use WebWorker for diff parsing (off main thread)
     const parsed = await diffService.parseDiff(
       file.patch!,
       file.filename,
-      file.previous_filename
+      file.previous_filename,
+      oldContent,
+      newContent
     );
 
     // Clean up pending entry
@@ -108,6 +145,10 @@ async function fetchParsedDiff(
 
 export function useDiffLoader() {
   const store = usePRReviewStore();
+  const github = useGitHub();
+  const owner = usePRReviewSelector((s) => s.owner);
+  const repo = usePRReviewSelector((s) => s.repo);
+  const pr = usePRReviewSelector((s) => s.pr);
   const selectedFile = usePRReviewSelector((s) => s.selectedFile);
   const files = usePRReviewSelector((s) => s.files);
   const loadedDiffs = usePRReviewSelector((s) => s.loadedDiffs);
@@ -120,7 +161,7 @@ export function useDiffLoader() {
 
     const currentFile = selectedFile;
 
-    // Check cache synchronously - instant if cached
+    // Check cache synchronously - instant if cached (with full file content)
     const cached = getDiffFromCache(file);
     if (cached) {
       if (!loadedDiffs[currentFile]) {
@@ -146,14 +187,19 @@ export function useDiffLoader() {
       }
     }, 50);
 
-    // Fetch immediately
-    fetchParsedDiff(file)
+    // Create file content getter for better syntax highlighting
+    const getFileContent: FileContentGetter = (path, ref) =>
+      github.getFileContent(owner, repo, path, ref);
+
+    // Fetch immediately with full file content for better highlighting
+    fetchParsedDiff(file, undefined, getFileContent, pr.base.sha, pr.head.sha)
       .then((diff) => {
         if (store.getSnapshot().selectedFile === currentFile) {
           store.setLoadedDiff(currentFile, diff);
           store.setDiffLoading(currentFile, false);
 
           // Prefetch next files aggressively (5 ahead, 2 behind)
+          // Note: prefetch without file content for speed
           const currentIndex = files.findIndex(
             (f) => f.filename === currentFile
           );
@@ -166,7 +212,7 @@ export function useDiffLoader() {
               !getDiffFromCache(f)
           );
 
-          // Prefetch all in parallel
+          // Prefetch all in parallel (without file content for speed)
           Promise.all(
             filesToPrefetch.map((pfile) =>
               fetchParsedDiff(pfile)
@@ -191,5 +237,15 @@ export function useDiffLoader() {
       clearTimeout(loadingTimeoutId);
       store.setDiffLoading(currentFile, false);
     };
-  }, [selectedFile, files, loadedDiffs, store]);
+  }, [
+    selectedFile,
+    files,
+    loadedDiffs,
+    store,
+    github,
+    owner,
+    repo,
+    pr.base.sha,
+    pr.head.sha,
+  ]);
 }