pi-diff-review 0.1.9 → 0.1.11

package/README.md

@@ -36,6 +36,7 @@ pi install https://github.com/cmpadden/pi-diff-review
 - `C` to add or edit an overall diff comment
 - `x` to delete a comment for the current line or selected range
 - `Enter` to submit comments back to pi
+- Comments are cached per session and restored when reopening the same diff
 - `q` to exit
 
 ## Release

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-diff-review",
-  "version": "0.1.9",
+  "version": "0.1.11",
   "description": "Local diff review TUI extension for pi",
   "license": "MIT",
   "repository": {

package/src/index.ts

@@ -1,3 +1,4 @@
+import { createHash } from "node:crypto";
 import type {
   ExtensionAPI,
   ExtensionCommandContext,
@@ -9,6 +10,67 @@ import { buildReviewPrompt } from "./prompt.ts";
 import { ReviewComponent } from "./review-component.ts";
 import type { DiffSource, ReviewComment, ReviewResult } from "./types.ts";
 
+const DIFF_REVIEW_CACHE_ENTRY = "pi-diff-review-cache";
+
+type DiffReviewCacheEntry = {
+  cacheKey: string;
+  comments: ReviewComment[];
+  updatedAt: number;
+};
+
+function getDiffCacheKey(
+  cwd: string,
+  source: DiffSource,
+  diffText: string,
+): string {
+  const hash = createHash("sha256").update(diffText).digest("hex");
+  return `${cwd}\0${source.label}\0${hash}`;
+}
+
+function getCachedComments(
+  ctx: ExtensionCommandContext,
+  cacheKey: string,
+): Map<string, ReviewComment> {
+  let latest: DiffReviewCacheEntry | undefined;
+  for (const entry of ctx.sessionManager.getEntries()) {
+    if (
+      entry.type !== "custom" ||
+      entry.customType !== DIFF_REVIEW_CACHE_ENTRY
+    ) {
+      continue;
+    }
+
+    const data = entry.data as Partial<DiffReviewCacheEntry> | undefined;
+    if (data?.cacheKey !== cacheKey || !Array.isArray(data.comments)) {
+      continue;
+    }
+
+    if (!latest || (data.updatedAt ?? 0) >= latest.updatedAt) {
+      latest = {
+        cacheKey: data.cacheKey,
+        comments: data.comments,
+        updatedAt: data.updatedAt ?? 0,
+      };
+    }
+  }
+
+  return new Map(
+    (latest?.comments ?? []).map((comment) => [comment.id, comment]),
+  );
+}
+
+function persistCachedComments(
+  pi: ExtensionAPI,
+  cacheKey: string,
+  comments: Iterable<ReviewComment>,
+): void {
+  pi.appendEntry(DIFF_REVIEW_CACHE_ENTRY, {
+    cacheKey,
+    comments: [...comments],
+    updatedAt: Date.now(),
+  } satisfies DiffReviewCacheEntry);
+}
+
 export function registerDiffReviewCommand(pi: ExtensionAPI): void {
   pi.registerCommand("diff", {
     description: "Review a git diff in a custom TUI (/diff [git diff args])",
@@ -30,9 +92,17 @@ export function registerDiffReviewCommand(pi: ExtensionAPI): void {
       }
 
       const reviewLines = parseDiff(diffText);
+      const cacheKey = getDiffCacheKey(ctx.cwd, source, diffText);
+      const comments = getCachedComments(ctx, cacheKey);
+      if (comments.size > 0) {
+        ctx.ui.notify(
+          `Restored ${comments.size} cached diff comment${comments.size === 1 ? "" : "s"}.`,
+          "info",
+        );
+      }
+
       const result = await ctx.ui.custom<ReviewResult>(
         (tui, theme, _keybindings, done) => {
-          const comments = new Map<string, ReviewComment>();
           return new ReviewComponent(
             tui,
             theme,
@@ -41,6 +111,9 @@ export function registerDiffReviewCommand(pi: ExtensionAPI): void {
             comments,
             done,
             new PiModelDiffExplainer(ctx),
+            (updatedComments) => {
+              persistCachedComments(pi, cacheKey, updatedComments.values());
+            },
           );
         },
       );
@@ -48,9 +121,11 @@ export function registerDiffReviewCommand(pi: ExtensionAPI): void {
       if (!result || result.action !== "submit") return;
       if (result.comments.length === 0) {
         ctx.ui.notify("No review comments to send.", "info");
+        persistCachedComments(pi, cacheKey, []);
         return;
       }
 
+      persistCachedComments(pi, cacheKey, []);
       pi.sendUserMessage(
         buildReviewPrompt(result.comments, source.promptLabel),
       );

package/src/review-component.ts

@@ -1,3 +1,7 @@
+import {
+  getLanguageFromPath,
+  highlightCode,
+} from "@earendil-works/pi-coding-agent";
 import {
   Editor,
   matchesKey,
@@ -58,6 +62,7 @@ export class ReviewComponent {
   private commentLineKeys = new Map<number, string[]>();
   private commentsRevision = 0;
   private commentLineKeysRevision = -1;
+  private highlightedLineCache = new Map<string, string>();
 
   constructor(
     private tui: ReviewTui,
@@ -67,6 +72,7 @@ export class ReviewComponent {
     private comments: Map<string, ReviewComment>,
     private done: (result: ReviewResult) => void,
     private explainer?: DiffExplainer,
+    private onCommentsChanged?: (comments: Map<string, ReviewComment>) => void,
   ) {
     const firstCommentable = this.lines.findIndex((line) => line.commentable);
     this.selected = firstCommentable >= 0 ? firstCommentable : 0;
@@ -388,7 +394,9 @@ export class ReviewComponent {
     return { diffHeight, commentsHeight };
   }
 
-  invalidate(): void {}
+  invalidate(): void {
+    this.highlightedLineCache.clear();
+  }
 
   dispose(): void {
     this.explanationAbort?.abort();
@@ -508,6 +516,7 @@ export class ReviewComponent {
 
   private markCommentsChanged(): void {
     this.commentsRevision++;
+    this.onCommentsChanged?.(this.comments);
   }
 
   private ensureCommentLineKeys(): void {
@@ -726,22 +735,8 @@ export class ReviewComponent {
     const commentMark = hasComment ? this.theme.fg("warning", "●") : " ";
     const lineNumber =
       side === "left" ? line.oldLineNumber : line.newLineNumber;
-    const raw = `${commentMark} ${lineNumberCell(lineNumber)} ${this.getDisplayText(line)}`;
-
-    let styled: string;
-    switch (line.kind) {
-      case "add":
-        styled = this.theme.fg("toolDiffAdded", raw);
-        break;
-      case "remove":
-        styled = this.theme.fg("toolDiffRemoved", raw);
-        break;
-      case "context":
-        styled = this.theme.fg("toolDiffContext", raw);
-        break;
-      default:
-        styled = this.theme.fg("muted", raw);
-    }
+    const prefix = `${commentMark} ${lineNumberCell(lineNumber)} `;
+    let styled = this.renderDiffRowContent(line, prefix);
 
     styled = truncateToWidth(styled, width);
     const selection = this.getSelectionBounds();
@@ -750,7 +745,7 @@ export class ReviewComponent {
     if (index === this.selected || inSelection) {
       return this.theme.bg("selectedBg", padToWidth(styled, width));
     }
-    return styled;
+    return this.applyDiffBackground(line, styled, width);
   }
 
   private ensureScroll(viewportHeight: number): void {
@@ -777,43 +772,75 @@ export class ReviewComponent {
       : line.text;
   }
 
-  private renderDiffLine(
+  private applyDiffBackground(
     line: ReviewLine,
-    index: number,
+    styled: string,
     width: number,
-    selected: boolean,
-    selection?: SelectionBounds,
   ): string {
-    const hasComment = this.getCommentKeysForLine(index).length > 0;
-    const commentMark = hasComment ? this.theme.fg("warning", "●") : " ";
-    const numbers = `${lineNumberCell(line.oldLineNumber)} ${lineNumberCell(line.newLineNumber)}`;
-    const raw = `${commentMark} ${numbers} ${this.getDisplayText(line)}`;
+    if (line.kind === "add") {
+      return this.theme.bg("toolSuccessBg", padToWidth(styled, width));
+    }
+    if (line.kind === "remove") {
+      return this.theme.bg("toolErrorBg", padToWidth(styled, width));
+    }
+    return styled;
+  }
 
-    let styled: string;
+  private renderDiffRowContent(line: ReviewLine, prefix: string): string {
     switch (line.kind) {
       case "add":
-        styled = this.theme.fg("toolDiffAdded", raw);
-        break;
+        return `${this.theme.fg("toolDiffAdded", prefix)}${this.getHighlightedDisplayText(line)}`;
       case "remove":
-        styled = this.theme.fg("toolDiffRemoved", raw);
-        break;
+        return `${this.theme.fg("toolDiffRemoved", prefix)}${this.getHighlightedDisplayText(line)}`;
       case "context":
-        styled = this.theme.fg("toolDiffContext", raw);
-        break;
+        return `${this.theme.fg("toolDiffContext", prefix)}${this.getHighlightedDisplayText(line)}`;
       case "hunk":
-        styled = this.theme.fg("accent", raw);
-        break;
+        return this.theme.fg("accent", `${prefix}${this.getDisplayText(line)}`);
       default:
-        styled = this.theme.fg("muted", raw);
+        return this.theme.fg("muted", `${prefix}${this.getDisplayText(line)}`);
+    }
+  }
+
+  private getHighlightedDisplayText(line: ReviewLine): string {
+    const code = this.getDisplayText(line);
+    if (!code) return code;
+
+    const lang = line.filePath ? getLanguageFromPath(line.filePath) : undefined;
+    const cacheKey = `${line.id}\0${lang ?? ""}\0${code}`;
+    const cached = this.highlightedLineCache.get(cacheKey);
+    if (cached != null) return cached;
+
+    let highlighted = code;
+    try {
+      highlighted = highlightCode(code, lang)[0] ?? code;
+    } catch {
+      highlighted = code;
     }
 
+    this.highlightedLineCache.set(cacheKey, highlighted);
+    return highlighted;
+  }
+
+  private renderDiffLine(
+    line: ReviewLine,
+    index: number,
+    width: number,
+    selected: boolean,
+    selection?: SelectionBounds,
+  ): string {
+    const hasComment = this.getCommentKeysForLine(index).length > 0;
+    const commentMark = hasComment ? this.theme.fg("warning", "●") : " ";
+    const numbers = `${lineNumberCell(line.oldLineNumber)} ${lineNumberCell(line.newLineNumber)}`;
+    const prefix = `${commentMark} ${numbers} `;
+    let styled = this.renderDiffRowContent(line, prefix);
+
     styled = truncateToWidth(styled, width);
     const inSelection =
       selection != null && index >= selection.start && index <= selection.end;
     if (selected || inSelection) {
       return this.theme.bg("selectedBg", padToWidth(styled, width));
     }
-    return styled;
+    return this.applyDiffBackground(line, styled, width);
   }
 
   private renderRightPane(