Make diff stats much faster (#866)

frontend/src/components/tasks/TaskDetails/DiffTab.tsx

@@ -1,4 +1,4 @@
-import { useDiffEntries } from '@/hooks/useDiffEntries';
+import { useDiffStream } from '@/hooks/useDiffStream';
 import { useMemo, useCallback, useState, useEffect } from 'react';
 import { Loader } from '@/components/ui/loader';
 import { Button } from '@/components/ui/button';
@@ -15,7 +15,7 @@ function DiffTab({ selectedAttempt }: DiffTabProps) {
   const [loading, setLoading] = useState(true);
   const [collapsedIds, setCollapsedIds] = useState<Set<string>>(new Set());
   const [hasInitialized, setHasInitialized] = useState(false);
-  const { diffs, error } = useDiffEntries(selectedAttempt?.id ?? null, true);
+  const { diffs, error } = useDiffStream(selectedAttempt?.id ?? null, true);
   const { fileCount, added, deleted } = useDiffSummary(
     selectedAttempt?.id ?? null
   );

frontend/src/hooks/useDiffEntries.ts

@@ -1,27 +0,0 @@
-import { useMemo } from 'react';
-import { useDiffStream } from './useDiffStream';
-import type { Diff, PatchType } from 'shared/types';
-
-interface UseDiffEntriesResult {
-  diffs: Diff[];
-  isConnected: boolean;
-  error: string | null;
-}
-
-export const useDiffEntries = (
-  attemptId: string | null,
-  enabled: boolean
-): UseDiffEntriesResult => {
-  const { data, isConnected, error } = useDiffStream(attemptId, enabled);
-
-  const diffs = useMemo(() => {
-    if (!data) return [];
-    return Object.values(data.entries)
-      .filter(
-        (e): e is Extract<PatchType, { type: 'DIFF' }> => e?.type === 'DIFF'
-      )
-      .map((e) => e.content);
-  }, [data]);
-
-  return { diffs, isConnected, error };
-};

frontend/src/hooks/useDiffStream.ts

@@ -1,38 +1,60 @@
-import { useCallback } from 'react';
-import type { PatchType } from 'shared/types';
+import { useCallback, useMemo } from 'react';
+import type { Diff, PatchType } from 'shared/types';
 import { useJsonPatchWsStream } from './useJsonPatchWsStream';
 
-interface DiffState {
-  entries: Record<string, PatchType>;
+interface DiffEntries {
+  [filePath: string]: PatchType;
+}
+
+type DiffStreamEvent = {
+  entries: DiffEntries;
+};
+
+export interface UseDiffStreamOptions {
+  statsOnly?: boolean;
 }
 
 interface UseDiffStreamResult {
-  data: DiffState | undefined;
-  isConnected: boolean;
+  diffs: Diff[];
   error: string | null;
 }
 
 export const useDiffStream = (
   attemptId: string | null,
-  enabled: boolean
+  enabled: boolean,
+  options?: UseDiffStreamOptions
 ): UseDiffStreamResult => {
-  const endpoint = attemptId
-    ? `/api/task-attempts/${attemptId}/diff/ws`
-    : undefined;
+  const endpoint = (() => {
+    if (!attemptId) return undefined;
+    const query = `/api/task-attempts/${attemptId}/diff/ws`;
+    if (typeof options?.statsOnly === 'boolean') {
+      const params = new URLSearchParams();
+      params.set('stats_only', String(options.statsOnly));
+      return `${query}?${params.toString()}`;
+    } else {
+      return query;
+    }
+  })();
 
   const initialData = useCallback(
-    (): DiffState => ({
+    (): DiffStreamEvent => ({
       entries: {},
     }),
     []
   );
 
-  const { data, isConnected, error } = useJsonPatchWsStream(
+  const { data, error } = useJsonPatchWsStream<DiffStreamEvent>(
     endpoint,
     enabled && !!attemptId,
     initialData
     // No need for injectInitialEntry or deduplicatePatches for diffs
   );
 
-  return { data, isConnected, error };
+  const diffs = useMemo(() => {
+    return Object.values(data?.entries ?? {})
+      .filter((entry) => entry?.type === 'DIFF')
+      .map((entry) => entry.content);
+  }, [data?.entries]);
+
+  return { diffs, error };
 };

frontend/src/hooks/useDiffSummary.ts

@@ -1,10 +1,10 @@
-import { useDiffEntries } from '@/hooks/useDiffEntries';
-import { getHighLightLanguageFromPath } from '@/utils/extToLanguage';
-import { generateDiffFile } from '@git-diff-view/file';
+import { useDiffStream } from '@/hooks/useDiffStream';
 import { useMemo } from 'react';
 
 export function useDiffSummary(attemptId: string | null) {
-  const { diffs, error, isConnected } = useDiffEntries(attemptId, true);
+  const { diffs, error } = useDiffStream(attemptId, true, {
+    statsOnly: true,
+  });
 
   const { fileCount, added, deleted } = useMemo(() => {
     if (!attemptId || diffs.length === 0) {
@@ -13,38 +13,13 @@ export function useDiffSummary(attemptId: string | null) {
 
     return diffs.reduce(
       (acc, d) => {
-        try {
-          if (d.contentOmitted) {
-            acc.added += d.additions ?? 0;
-            acc.deleted += d.deletions ?? 0;
-            return acc;
-          }
-          const oldName = d.oldPath || d.newPath || 'old';
-          const newName = d.newPath || d.oldPath || 'new';
-          const oldContent = d.oldContent || '';
-          const newContent = d.newContent || '';
-          const oldLang = getHighLightLanguageFromPath(oldName) || 'plaintext';
-          const newLang = getHighLightLanguageFromPath(newName) || 'plaintext';
-
-          const file = generateDiffFile(
-            oldName,
-            oldContent,
-            newName,
-            newContent,
-            oldLang,
-            newLang
-          );
-          file.initRaw();
-          acc.added += file.additionLength ?? 0;
-          acc.deleted += file.deletionLength ?? 0;
-        } catch (e) {
-          console.error('Failed to compute totals for diff', e);
-        }
+        acc.added += d.additions ?? 0;
+        acc.deleted += d.deletions ?? 0;
         return acc;
       },
       { fileCount: diffs.length, added: 0, deleted: 0 }
     );
   }, [attemptId, diffs]);
 
-  return { fileCount, added, deleted, isConnected, error };
+  return { fileCount, added, deleted, error };
 }