fix: preserve staged and unstaged changes per file (#3104)

apps/web/components/task/task-changes-panel-state.ts

@@ -1,8 +1,9 @@
 "use client";
 
 import { useEffect, useMemo, useRef, type MutableRefObject, type RefObject } from "react";
-import { reviewFileKey, splitReviewFileKey, type ReviewFile } from "@/components/review/types";
+import { hashDiff, reviewFileKey, type ReviewFile } from "@/components/review/types";
 import type { ReviewSource } from "@/hooks/domains/session/use-review-sources";
+import type { GitChangeLayer } from "@/lib/state/slices/session-runtime/types";
 
 export function shouldDeferReviewStateForPR(
   hasSelectedPR: boolean,
@@ -104,13 +105,84 @@ type FilterVisibleFilesOpts = {
   sourceFilter: "all" | ReviewSource;
   rawPRFiles?: ReviewFile[];
   prKey?: string;
+  changeLayer?: GitChangeLayer;
 };
 
+function projectReviewFileLayer(file: ReviewFile, layer: GitChangeLayer): ReviewFile | null {
+  if (file.source !== "uncommitted") return file;
+  const facet = layer === "staged" ? file.staged_change : file.unstaged_change;
+  if (!facet) {
+    if (file.staged !== (layer === "staged")) return null;
+    return { ...file, change_layer: layer };
+  }
+  return {
+    ...file,
+    ...facet,
+    staged: layer === "staged",
+    change_layer: layer,
+  };
+}
+
+type PersistedReviewState = { reviewed: boolean; diffHash: string };
+
+function reviewStateFiles(files: ReviewFile[]): ReviewFile[] {
+  const stateFiles: ReviewFile[] = [];
+  for (const file of files) {
+    stateFiles.push(file);
+    if (
+      file.source !== "uncommitted" ||
+      file.change_layer ||
+      !file.staged_change ||
+      !file.unstaged_change
+    ) {
+      continue;
+    }
+    for (const layer of ["staged", "unstaged"] as const) {
+      const projected = projectReviewFileLayer(file, layer);
+      if (projected) stateFiles.push(projected);
+    }
+  }
+  return stateFiles;
+}
+
+export function computeChangesReviewSets(
+  files: ReviewFile[],
+  reviews: ReadonlyMap<string, PersistedReviewState>,
+): { reviewedFiles: Set<string>; staleFiles: Set<string> } {
+  const reviewedFiles = new Set<string>();
+  const staleFiles = new Set<string>();
+  for (const file of reviewStateFiles(files)) {
+    const key = reviewFileKey(file);
+    const reviewState = reviews.get(key);
+    if (!reviewState?.reviewed) continue;
+    const currentHash = hashDiff(file.diff);
+    if (reviewState.diffHash && reviewState.diffHash !== currentHash) staleFiles.add(key);
+    else reviewedFiles.add(key);
+  }
+  return { reviewedFiles, staleFiles };
+}
+
+export function reviewDiffHashForKey(files: ReviewFile[], key: string): string {
+  const file = reviewStateFiles(files).find((candidate) => reviewFileKey(candidate) === key);
+  return file ? hashDiff(file.diff) : "";
+}
+
+function projectRequestedLayer(
+  files: ReviewFile[],
+  mode: FilterVisibleFilesOpts["mode"],
+  changeLayer: GitChangeLayer | undefined,
+): ReviewFile[] {
+  if (mode !== "file" || !changeLayer) return files;
+  return files
+    .map((file) => projectReviewFileLayer(file, changeLayer))
+    .filter((file): file is ReviewFile => file !== null);
+}
+
 export function filterVisibleFiles(
   allFiles: ReviewFile[],
   opts: FilterVisibleFilesOpts,
 ): ReviewFile[] {
-  const { mode, filePath, fileRepositoryName, sourceFilter, rawPRFiles, prKey } = opts;
+  const { mode, filePath, fileRepositoryName, sourceFilter, rawPRFiles, prKey, changeLayer } = opts;
   const repositoryFilter = prKey && sourceFilter === "pr" ? undefined : fileRepositoryName;
   if (mode === "file" && filePath && sourceFilter === "pr" && rawPRFiles?.length) {
     let prFiles = rawPRFiles.filter((file) => file.path === filePath && file.source === "pr");
@@ -128,7 +200,7 @@ export function filterVisibleFiles(
     }
   }
   if (sourceFilter !== "all") files = files.filter((file) => file.source === sourceFilter);
-  return files;
+  return projectRequestedLayer(files, mode, changeLayer);
 }
 
 export function useVisibleDiffState(opts: {
@@ -139,6 +211,7 @@ export function useVisibleDiffState(opts: {
   fileRepositoryName: string | undefined;
   sourceFilter: "all" | ReviewSource;
   prKey: string | undefined;
+  changeLayer: GitChangeLayer | undefined;
   fileRefs: Map<string, RefObject<HTMLDivElement | null>>;
   reviewedFiles: Set<string>;
   staleFiles: Set<string>;
@@ -151,6 +224,7 @@ export function useVisibleDiffState(opts: {
     fileRepositoryName,
     sourceFilter,
     prKey,
+    changeLayer,
     fileRefs,
     reviewedFiles,
     staleFiles,
@@ -164,23 +238,21 @@ export function useVisibleDiffState(opts: {
         sourceFilter,
         rawPRFiles,
         prKey,
+        changeLayer,
       }),
-    [allFiles, mode, filePath, fileRepositoryName, sourceFilter, rawPRFiles, prKey],
+    [allFiles, mode, filePath, fileRepositoryName, sourceFilter, rawPRFiles, prKey, changeLayer],
   );
-  const repositoryFilter = prKey && sourceFilter === "pr" ? undefined : fileRepositoryName;
   const visibleFileRefs = useMemo(() => {
     if (mode !== "file" || !filePath) return fileRefs;
     const refs = new Map<string, RefObject<HTMLDivElement | null>>();
-    for (const [key, ref] of fileRefs.entries()) {
-      const split = splitReviewFileKey(key);
-      if (split.path !== filePath) continue;
-      if (repositoryFilter !== undefined && (split.repositoryName || "") !== repositoryFilter) {
-        continue;
-      }
-      refs.set(key, ref);
+    for (const file of visibleFiles) {
+      const visibleKey = reviewFileKey(file);
+      const rawKey = reviewFileKey({ path: file.path, repository_name: file.repository_name });
+      const ref = fileRefs.get(visibleKey) ?? fileRefs.get(rawKey);
+      if (ref) refs.set(visibleKey, ref);
     }
     return refs;
-  }, [mode, filePath, fileRefs, repositoryFilter]);
+  }, [mode, filePath, fileRefs, visibleFiles]);
   const reviewedCount = useMemo(
     () =>
       visibleFiles.reduce((count, file) => {