fix: restore citation focus with default tree view

Author: suguanYang

src/components/chunks-panel.tsx

@@ -62,6 +62,11 @@ export type ChunksPanelProps = {
 
 type ChunkDisplayMode = "list" | "tree";
 
+type ChunkDisplayModeState = {
+  readonly handledFocusedChunkRequestId: number;
+  readonly mode: ChunkDisplayMode;
+};
+
 export function ChunksPanel({
   chunks = [],
   selectedSource = null,
@@ -86,8 +91,12 @@ export function ChunksPanel({
   const [mountedOriginalPreviewKey, setMountedOriginalPreviewKey] = useState<
     string | null
   >(null);
-  const [chunkDisplayMode, setChunkDisplayMode] =
-    useState<ChunkDisplayMode>("tree");
+  const [chunkDisplayModeState, setChunkDisplayModeState] =
+    useState<ChunkDisplayModeState>(() => ({
+      handledFocusedChunkRequestId:
+        focusedChunkId === null ? focusedChunkRequestId : -1,
+      mode: "tree",
+    }));
   const [sectionTreeZoomPercent, setSectionTreeZoomPercent] =
     useState<number>(sectionTreeDefaultZoomPercent);
   const {
@@ -141,19 +150,28 @@ export function ChunksPanel({
     selectOriginalView();
   }, [rememberOriginalPreview, selectOriginalView]);
   const handleListModeSelected = useCallback((): void => {
-    setChunkDisplayMode("list");
-  }, []);
+    setChunkDisplayModeState({
+      handledFocusedChunkRequestId: focusedChunkRequestId,
+      mode: "list",
+    });
+  }, [focusedChunkRequestId]);
   const handleTreeModeSelected = useCallback((): void => {
-    setChunkDisplayMode("tree");
-  }, []);
+    setChunkDisplayModeState({
+      handledFocusedChunkRequestId: focusedChunkRequestId,
+      mode: "tree",
+    });
+  }, [focusedChunkRequestId]);
   const handleTreeChunkFocus = useCallback(
     (chunkId: string | null): void => {
       requestChunkFocus(chunkId);
       if (chunkId !== null) {
-        setChunkDisplayMode("list");
+        setChunkDisplayModeState({
+          handledFocusedChunkRequestId: focusedChunkRequestId,
+          mode: "list",
+        });
       }
     },
-    [requestChunkFocus],
+    [focusedChunkRequestId, requestChunkFocus],
   );
   const canZoomSectionTreeOut: boolean =
     sectionTreeZoomPercent > sectionTreeMinimumZoomPercent;
@@ -198,6 +216,12 @@ export function ChunksPanel({
     },
     [],
   );
+  const chunkDisplayMode: ChunkDisplayMode =
+    focusedChunkId !== null &&
+    chunkDisplayModeState.handledFocusedChunkRequestId !==
+      focusedChunkRequestId
+      ? "list"
+      : chunkDisplayModeState.mode;
   const headerTitle = focusedChunkId ? "Referenced Chunks" : "Parsed Chunks";
   const shouldMountOriginalPreview =
     visibleView === "original" ||

src/components/workspace-citation-focus.test.ts

@@ -69,7 +69,39 @@ describe("useWorkspaceCitationFocus", () => {
     expect(result.current.pendingCitationId).toBeNull();
   });
 
-  it("clears prefetched chunks and focus when the selected source changes", () => {
+  it("clears prefetched chunks and focus when selecting a different source", () => {
+    const selectSource = vi.fn();
+    const otherSource: SourceView = {
+      id: "source_2",
+      title: "Other.pdf",
+      mimeType: "application/pdf",
+      status: "ready",
+      documentId: "document_2",
+    };
+    const { result } = renderHook(() =>
+      useWorkspaceCitationFocus({
+        fetchChunks: vi.fn(async () => []),
+        initialPrefetchedChunksBySourceId: { source_1: [prefetchedChunk] },
+        onSelectSource: selectSource,
+        selectedSourceId: "source_2",
+        sources: [readySource, otherSource],
+      }),
+      { wrapper: createSWRWrapper },
+    );
+
+    act(() => {
+      result.current.handleSourceSelected("source_1");
+    });
+
+    expect(selectSource).toHaveBeenCalledWith("source_1");
+    expect(result.current.prefetchedChunksBySourceId).toEqual({});
+    expect(result.current.focusedChunk).toEqual({
+      chunkId: null,
+      requestId: 1,
+    });
+  });
+
+  it("keeps prefetched chunks when reselecting the selected source", () => {
     const selectSource = vi.fn();
     const { result } = renderHook(() =>
       useWorkspaceCitationFocus({
@@ -87,7 +119,9 @@ describe("useWorkspaceCitationFocus", () => {
     });
 
     expect(selectSource).toHaveBeenCalledWith("source_1");
-    expect(result.current.prefetchedChunksBySourceId).toEqual({});
+    expect(result.current.prefetchedChunksBySourceId).toEqual({
+      source_1: [prefetchedChunk],
+    });
     expect(result.current.focusedChunk).toEqual({
       chunkId: null,
       requestId: 1,

src/components/workspace-citation-focus.ts

@@ -1,6 +1,6 @@
 "use client"
 
-import { useCallback, useState } from "react"
+import { useCallback, useRef, useState } from "react"
 
 import { workspaceCitationState } from "@/components/workspace-citation-state"
 import { useWorkspaceSelectedChunks } from "@/components/workspace-selected-chunks"
@@ -14,6 +14,9 @@ type FocusedChunkState = {
 }
 
 type PrefetchedChunksBySourceId = Readonly<Record<string, ParsedChunkView[]>>
+type PrefetchedChunksUpdater = (
+  current: PrefetchedChunksBySourceId,
+) => PrefetchedChunksBySourceId
 
 type WorkspaceCitationFocusInput = {
   readonly fetchChunks: (sourceId: string) => Promise<ParsedChunkView[]>
@@ -60,8 +63,15 @@ export function useWorkspaceCitationFocus({
   const [fullChunkLoadingSourceId, setFullChunkLoadingSourceId] = useState<
     string | null
   >(null)
+  const fullChunkRequestsBySourceIdRef = useRef<
+    Map<string, Promise<ParsedChunkView[]>>
+  >(new Map())
+  const fullChunkRequestedSourceIdsRef = useRef<Set<string>>(new Set())
   const [prefetchedChunksBySourceId, setPrefetchedChunksBySourceId] =
     useState<PrefetchedChunksBySourceId>(initialPrefetchedChunksBySourceId)
+  const prefetchedChunksBySourceIdRef = useRef<PrefetchedChunksBySourceId>(
+    initialPrefetchedChunksBySourceId,
+  )
   const {
     hasMoreSelectedChunks,
     handleLoadMoreChunks,
@@ -85,48 +95,79 @@ export function useWorkspaceCitationFocus({
     [],
   )
 
+  const updatePrefetchedChunksBySourceId = useCallback(
+    (updater: PrefetchedChunksUpdater): void => {
+      const next = updater(prefetchedChunksBySourceIdRef.current)
+      prefetchedChunksBySourceIdRef.current = next
+      setPrefetchedChunksBySourceId(next)
+    },
+    [],
+  )
+
   const handleSourceSelected = useCallback(
     (sourceId: string | null): void => {
       onSelectSource(sourceId)
-      if (sourceId) {
-        setPrefetchedChunksBySourceId((current) =>
+      if (sourceId && sourceId !== selectedSourceId) {
+        fullChunkRequestedSourceIdsRef.current.delete(sourceId)
+        updatePrefetchedChunksBySourceId((current) =>
           workspaceCitationState.removePrefetchedChunks(current, sourceId),
         )
       }
       requestChunkFocus(null)
     },
-    [onSelectSource, requestChunkFocus],
+    [
+      onSelectSource,
+      requestChunkFocus,
+      selectedSourceId,
+      updatePrefetchedChunksBySourceId,
+    ],
+  )
+
+  const loadAllChunksForSource = useCallback(
+    (sourceId: string): Promise<ParsedChunkView[]> => {
+      const existingRequest =
+        fullChunkRequestsBySourceIdRef.current.get(sourceId)
+      if (existingRequest) return existingRequest
+
+      setFullChunkLoadingSourceId(sourceId)
+      const request = fetchChunks(sourceId)
+        .then((chunks) => {
+          updatePrefetchedChunksBySourceId((current) =>
+            workspaceCitationState.upsertPrefetchedChunks(
+              current,
+              sourceId,
+              chunks,
+            ),
+          )
+          return chunks
+        })
+        .finally(() => {
+          fullChunkRequestsBySourceIdRef.current.delete(sourceId)
+          setFullChunkLoadingSourceId((current) =>
+            current === sourceId ? null : current,
+          )
+        })
+
+      fullChunkRequestsBySourceIdRef.current.set(sourceId, request)
+      return request
+    },
+    [fetchChunks, updatePrefetchedChunksBySourceId],
   )
 
   const handleLoadAllChunks = useCallback((): void => {
     if (
       !selectedSourceId ||
-      prefetchedChunksBySourceId[selectedSourceId] ||
-      fullChunkLoadingSourceId === selectedSourceId
+      prefetchedChunksBySourceIdRef.current[selectedSourceId] ||
+      fullChunkRequestedSourceIdsRef.current.has(selectedSourceId) ||
+      fullChunkRequestsBySourceIdRef.current.has(selectedSourceId)
     ) {
       return
     }
 
-    setFullChunkLoadingSourceId(selectedSourceId)
-    void fetchChunks(selectedSourceId)
-      .then((chunks) => {
-        setPrefetchedChunksBySourceId((current) =>
-          workspaceCitationState.upsertPrefetchedChunks(
-            current,
-            selectedSourceId,
-            chunks,
-          ),
-        )
-      })
-      .finally(() => {
-        setFullChunkLoadingSourceId((current) =>
-          current === selectedSourceId ? null : current,
-        )
-      })
+    fullChunkRequestedSourceIdsRef.current.add(selectedSourceId)
+    void loadAllChunksForSource(selectedSourceId)
   }, [
-    fetchChunks,
-    fullChunkLoadingSourceId,
-    prefetchedChunksBySourceId,
+    loadAllChunksForSource,
     selectedSourceId,
   ])
 
@@ -157,15 +198,16 @@ export function useWorkspaceCitationFocus({
         }
 
         if (!workspaceCitationState.hasExactCitationTargetHint(citation)) {
-          setPrefetchedChunksBySourceId((current) =>
+          fullChunkRequestedSourceIdsRef.current.delete(source.id)
+          updatePrefetchedChunksBySourceId((current) =>
             workspaceCitationState.removePrefetchedChunks(current, source.id),
           )
           if (selectedSourceId !== source.id) onSelectSource(source.id)
           requestChunkFocus(null)
           return
         }
 
-        const cachedChunks = prefetchedChunksBySourceId[source.id]
+        const cachedChunks = prefetchedChunksBySourceIdRef.current[source.id]
         if (cachedChunks) {
           const cachedFocusId =
             workspaceCitationState.getLoadedCitationChunkId({
@@ -175,7 +217,7 @@ export function useWorkspaceCitationFocus({
               selectedChunks: cachedChunks,
               hasMoreSelectedChunks: false,
             })
-          setPrefetchedChunksBySourceId((current) =>
+          updatePrefetchedChunksBySourceId((current) =>
             workspaceCitationState.upsertPrefetchedChunks(
               current,
               source.id,
@@ -188,14 +230,7 @@ export function useWorkspaceCitationFocus({
         }
 
         requestChunkFocus(null)
-        const chunks = await fetchChunks(source.id)
-        setPrefetchedChunksBySourceId((current) =>
-          workspaceCitationState.upsertPrefetchedChunks(
-            current,
-            source.id,
-            chunks,
-          ),
-        )
+        const chunks = await loadAllChunksForSource(source.id)
         const prefetchedChunkId =
           workspaceCitationState.getLoadedCitationChunkId({
             citation,
@@ -213,14 +248,14 @@ export function useWorkspaceCitationFocus({
       }
     },
     [
-      fetchChunks,
       hasMoreSelectedChunks,
+      loadAllChunksForSource,
       onSelectSource,
-      prefetchedChunksBySourceId,
       requestChunkFocus,
       selectedChunks,
       selectedSourceId,
       sources,
+      updatePrefetchedChunksBySourceId,
     ],
   )