feat(code): use chosen diff source for badge

Author: adboio

apps/code/src/renderer/features/code-review/components/DiffStatsBadge.tsx

@@ -1,5 +1,9 @@
 import { Tooltip } from "@components/ui/Tooltip";
-import { useGitQueries } from "@features/git-interaction/hooks/useGitQueries";
+import {
+  useBranchChangedFiles,
+  useGitQueries,
+  usePrChangedFiles,
+} from "@features/git-interaction/hooks/useGitQueries";
 import { computeDiffStats } from "@features/git-interaction/utils/diffStats";
 import { useCwd } from "@features/sidebar/hooks/useCwd";
 import { useCloudChangedFiles } from "@features/task-detail/hooks/useCloudChangedFiles";
@@ -14,22 +18,44 @@ import {
 import { useReviewNavigationStore } from "@renderer/features/code-review/stores/reviewNavigationStore";
 import type { Task } from "@shared/types";
 import { useMemo } from "react";
+import { useEffectiveDiffSource } from "../hooks/useEffectiveDiffSource";
 
 interface DiffStatsBadgeProps {
   task: Task;
 }
 
-function useChangedFileStats(task: Task) {
+interface DiffStatsValue {
+  filesChanged: number;
+  linesAdded: number;
+  linesRemoved: number;
+}
+
+function useChangedFileStats(task: Task): DiffStatsValue {
   const taskId = task.id;
   const workspace = useWorkspace(taskId);
   const isCloud =
     workspace?.mode === "cloud" || task.latest_run?.environment === "cloud";
   const repoPath = useCwd(taskId);
 
+  // Local working-tree stats — always available for local tasks; the basis
+  // for the "local" source and the fallback when other queries aren't loaded.
   const { diffStats: localDiffStats } = useGitQueries(
     isCloud ? undefined : repoPath,
   );
 
+  // Source resolution + the branch/PR queries that back the non-local
+  // sources. For cloud tasks these all run but return nothing useful — the
+  // memo below short-circuits to cloud stats regardless.
+  const { effectiveSource, repoSlug, linkedBranch, prUrl } =
+    useEffectiveDiffSource(taskId);
+  const { data: branchFiles } = useBranchChangedFiles(
+    !isCloud && effectiveSource === "branch" ? repoSlug : null,
+    !isCloud && effectiveSource === "branch" ? linkedBranch : null,
+  );
+  const { data: prFiles } = usePrChangedFiles(
+    !isCloud && effectiveSource === "pr" ? prUrl : null,
+  );
+
   const { reviewFiles } = useCloudChangedFiles(taskId, task);
 
   return useMemo(() => {
@@ -41,12 +67,25 @@ function useChangedFileStats(task: Task) {
         linesRemoved: stats.linesRemoved,
       };
     }
+    if (effectiveSource === "branch" && branchFiles) {
+      return computeDiffStats(branchFiles);
+    }
+    if (effectiveSource === "pr" && prFiles) {
+      return computeDiffStats(prFiles);
+    }
     return {
       filesChanged: localDiffStats.filesChanged,
       linesAdded: localDiffStats.linesAdded,
       linesRemoved: localDiffStats.linesRemoved,
     };
-  }, [isCloud, reviewFiles, localDiffStats]);
+  }, [
+    isCloud,
+    effectiveSource,
+    branchFiles,
+    prFiles,
+    localDiffStats,
+    reviewFiles,
+  ]);
 }
 
 export function DiffStatsBadge({ task }: DiffStatsBadgeProps) {

apps/code/src/renderer/features/code-review/components/ReviewPage.tsx

@@ -1,27 +1,21 @@
-import { useDiffViewerStore } from "@features/code-editor/stores/diffViewerStore";
 import {
   useBranchChangedFiles,
-  useGitQueries,
   usePrChangedFiles,
 } from "@features/git-interaction/hooks/useGitQueries";
-import { useLinkedBranchPrUrl } from "@features/git-interaction/hooks/useLinkedBranchPrUrl";
 import { makeFileKey } from "@features/git-interaction/utils/fileKey";
 import { usePanelLayoutStore } from "@features/panels/store/panelLayoutStore";
 import { useCwd } from "@features/sidebar/hooks/useCwd";
-import { useWorkspace } from "@features/workspace/hooks/useWorkspace";
 import type { parsePatchFiles } from "@pierre/diffs";
 import { Flex, Text } from "@radix-ui/themes";
 import { useReviewNavigationStore } from "@renderer/features/code-review/stores/reviewNavigationStore";
 import { useTRPC } from "@renderer/trpc/client";
 import type { ChangedFile, Task } from "@shared/types";
 import { useQuery } from "@tanstack/react-query";
 import { useMemo } from "react";
+import { useEffectiveDiffSource } from "../hooks/useEffectiveDiffSource";
 import { useReviewDiffs } from "../hooks/useReviewDiffs";
 import type { DiffOptions } from "../types";
-import {
-  type ResolvedDiffSource,
-  resolveDiffSource,
-} from "../utils/resolveDiffSource";
+import type { ResolvedDiffSource } from "../utils/resolveDiffSource";
 import { InteractiveFileDiff } from "./InteractiveFileDiff";
 import { LazyDiff } from "./LazyDiff";
 import { RemoteDiffList } from "./RemoteDiffList";
@@ -44,36 +38,21 @@ interface ReviewPageProps {
 export function ReviewPage({ task }: ReviewPageProps) {
   const taskId = task.id;
   const repoPath = useCwd(taskId);
-  const workspace = useWorkspace(taskId);
-  const linkedBranch = workspace?.linkedBranch ?? null;
   const openFile = usePanelLayoutStore((s) => s.openFile);
 
   const isReviewOpen = useReviewNavigationStore(
     (s) => (s.reviewModes[taskId] ?? "closed") !== "closed",
   );
 
-  const configuredSource = useDiffViewerStore(
-    (s) => s.diffSource[taskId] ?? null,
-  );
-
   const {
-    repoInfo,
-    aheadOfDefault,
-    defaultBranch,
-    changedFiles: workspaceFiles,
-  } = useGitQueries(repoPath);
-  const prUrl = useLinkedBranchPrUrl(taskId);
-  const hasLocalChanges = workspaceFiles.length > 0;
-  const branchSourceAvailable = !!linkedBranch && aheadOfDefault > 0;
-  const prSourceAvailable = !!prUrl;
-
-  const effectiveSource = resolveDiffSource({
-    configured: configuredSource,
-    hasLocalChanges,
+    effectiveSource,
+    prUrl,
     linkedBranch,
-    aheadOfDefault,
+    defaultBranch,
+    repoSlug,
+    branchSourceAvailable,
     prSourceAvailable,
-  });
+  } = useEffectiveDiffSource(taskId);
 
   const isLocalActive = isReviewOpen && effectiveSource === "local";
 
@@ -118,7 +97,7 @@ export function ReviewPage({ task }: ReviewPageProps) {
       <BranchReviewPage
         task={task}
         branch={linkedBranch as string}
-        repoInfo={repoInfo ?? undefined}
+        repoSlug={repoSlug}
         defaultBranch={defaultBranch}
         isReviewOpen={isReviewOpen}
         effectiveSource={effectiveSource}
@@ -207,7 +186,7 @@ export function ReviewPage({ task }: ReviewPageProps) {
 function BranchReviewPage({
   task,
   branch,
-  repoInfo,
+  repoSlug,
   defaultBranch,
   isReviewOpen,
   effectiveSource,
@@ -216,17 +195,14 @@ function BranchReviewPage({
 }: {
   task: Task;
   branch: string;
-  repoInfo: { organization: string; repository: string } | undefined;
+  repoSlug: string | null;
   defaultBranch: string | null;
   isReviewOpen: boolean;
   effectiveSource: ResolvedDiffSource;
   branchSourceAvailable: boolean;
   prSourceAvailable: boolean;
 }) {
   const taskId = task.id;
-  const repoSlug = repoInfo
-    ? `${repoInfo.organization}/${repoInfo.repository}`
-    : null;
 
   const { data: files = EMPTY_BRANCH_FILES, isLoading } = useBranchChangedFiles(
     isReviewOpen ? repoSlug : null,

apps/code/src/renderer/features/code-review/hooks/useEffectiveDiffSource.ts

@@ -0,0 +1,65 @@
+import type { DiffSource } from "@features/code-editor/stores/diffViewerStore";
+import { useDiffViewerStore } from "@features/code-editor/stores/diffViewerStore";
+import { useGitQueries } from "@features/git-interaction/hooks/useGitQueries";
+import { useLinkedBranchPrUrl } from "@features/git-interaction/hooks/useLinkedBranchPrUrl";
+import { useCwd } from "@features/sidebar/hooks/useCwd";
+import { useWorkspace } from "@features/workspace/hooks/useWorkspace";
+import {
+  type ResolvedDiffSource,
+  resolveDiffSource,
+} from "../utils/resolveDiffSource";
+
+export interface EffectiveDiffSource {
+  effectiveSource: ResolvedDiffSource;
+  configured: DiffSource | null;
+  prUrl: string | null;
+  linkedBranch: string | null;
+  defaultBranch: string | null;
+  repoSlug: string | null;
+  branchSourceAvailable: boolean;
+  prSourceAvailable: boolean;
+}
+
+/**
+ * Resolves which diff source should be shown for a local task, plus all the
+ * context needed to actually fetch that source. Callers (review panel, diff
+ * stats badge) share one source-of-truth so UI surfaces never disagree.
+ */
+export function useEffectiveDiffSource(taskId: string): EffectiveDiffSource {
+  const repoPath = useCwd(taskId);
+  const workspace = useWorkspace(taskId);
+  const linkedBranch = workspace?.linkedBranch ?? null;
+
+  const configured = useDiffViewerStore((s) => s.diffSource[taskId] ?? null);
+
+  const { repoInfo, aheadOfDefault, defaultBranch, changedFiles } =
+    useGitQueries(repoPath);
+  const hasLocalChanges = changedFiles.length > 0;
+  const branchSourceAvailable = !!linkedBranch && aheadOfDefault > 0;
+
+  const prUrl = useLinkedBranchPrUrl(taskId);
+  const prSourceAvailable = !!prUrl;
+
+  const repoSlug = repoInfo
+    ? `${repoInfo.organization}/${repoInfo.repository}`
+    : null;
+
+  const effectiveSource = resolveDiffSource({
+    configured,
+    hasLocalChanges,
+    linkedBranch,
+    aheadOfDefault,
+    prSourceAvailable,
+  });
+
+  return {
+    effectiveSource,
+    configured,
+    prUrl,
+    linkedBranch,
+    defaultBranch,
+    repoSlug,
+    branchSourceAvailable,
+    prSourceAvailable,
+  };
+}