Fix chunk image rendering after full-tree loads

src/components/chunks-panel.tsx

@@ -401,6 +401,7 @@ export function ChunksPanel({
                         hasOriginalFile ? handleChunkSelected : undefined
                       }
                       onReferenceClick={requestChunkFocus}
+                      selectedSourceFile={selectedSourceFile}
                     />
                   ))}
                 </div>
@@ -1129,6 +1130,7 @@ function VirtualChunkRow({
   measureElement,
   onChunkClick,
   onReferenceClick,
+  selectedSourceFile,
 }: {
   virtualItem: VirtualItem;
   chunk: ParsedChunkView | undefined;
@@ -1137,6 +1139,7 @@ function VirtualChunkRow({
   measureElement: (node: HTMLDivElement | null) => void;
   onChunkClick?: (chunk: ParsedChunkView) => void;
   onReferenceClick: (chunkId: string) => void;
+  selectedSourceFile: SourceOriginalFileView | null;
 }): ReactNode {
   if (!chunk) {
     return null;
@@ -1163,6 +1166,7 @@ function VirtualChunkRow({
         isOriginalPreviewAvailable={isOriginalPreviewAvailable}
         onChunkClick={onChunkClick}
         onReferenceClick={onReferenceClick}
+        sourceOriginalFile={selectedSourceFile}
       />
     </div>
   );

src/components/parsed-chunk-card.test.ts

@@ -182,4 +182,33 @@ describe("ParsedChunkCard", () => {
     expect(table.innerHTML).not.toContain("script");
     expect(table.innerHTML).not.toContain("onclick");
   });
+
+  it("renders the original image file when an image chunk has no asset URL", () => {
+    render(
+      React.createElement(ParsedChunkCard, {
+        chunk: {
+          chunkId: "image_1",
+          type: "image",
+          content: "Logo summary",
+          summary: "A black and white logo.",
+          sourceTitle: "logo.png",
+        },
+        isFocused: false,
+        onReferenceClick: vi.fn(),
+        sourceOriginalFile: {
+          url: "https://blob.example/source-uploads/logo.png",
+          mimeType: "image/png",
+        },
+      }),
+    );
+
+    const image = screen.getByRole("img", {
+      name: "A black and white logo.",
+    });
+
+    expect(image.getAttribute("src")).toBe(
+      "https://blob.example/source-uploads/logo.png",
+    );
+    expect(screen.queryByText("Image content is not available in this view.")).toBeNull();
+  });
 });

src/components/parsed-chunk-card.tsx

@@ -8,6 +8,7 @@ import { Button } from "@/components/ui/button";
 import { Card, CardContent } from "@/components/ui/card";
 import { parsedChunkCardModel } from "@/components/parsed-chunk-card-model";
 import type { ParsedChunkView } from "@/domains/chunks/types";
+import type { SourceOriginalFileView } from "@/domains/sources/types";
 import { cn } from "@/lib/utils";
 
 const keywordPanelClassName =
@@ -25,12 +26,14 @@ export function ParsedChunkCard({
   isOriginalPreviewAvailable = false,
   onChunkClick,
   onReferenceClick,
+  sourceOriginalFile = null,
 }: {
   readonly chunk: ParsedChunkView;
   readonly isFocused: boolean;
   readonly isOriginalPreviewAvailable?: boolean;
   readonly onChunkClick?: (chunk: ParsedChunkView) => void;
   readonly onReferenceClick: (chunkId: string) => void;
+  readonly sourceOriginalFile?: SourceOriginalFileView | null;
 }): ReactNode {
   if (chunk.type === "image") {
     return (
@@ -40,6 +43,7 @@ export function ParsedChunkCard({
           isFocused={isFocused}
           isOriginalPreviewAvailable={isOriginalPreviewAvailable}
           onChunkClick={onChunkClick}
+          sourceOriginalFile={sourceOriginalFile}
         />
       </ChunkCardShell>
     );
@@ -362,12 +366,16 @@ function ImageChunkCard({
   isFocused,
   isOriginalPreviewAvailable,
   onChunkClick,
+  sourceOriginalFile,
 }: {
   readonly chunk: ParsedChunkView;
   readonly isFocused: boolean;
   readonly isOriginalPreviewAvailable: boolean;
   readonly onChunkClick?: (chunk: ParsedChunkView) => void;
+  readonly sourceOriginalFile: SourceOriginalFileView | null;
 }): ReactNode {
+  const imageAssetUrl = getImageChunkAssetUrl(chunk, sourceOriginalFile);
+
   return (
     <ChunkCardFrame
       chunk={chunk}
@@ -377,11 +385,11 @@ function ImageChunkCard({
     >
       <ChunkSummaryPanel chunk={chunk} />
       <ChunkContentPanel chunk={chunk}>
-        {chunk.assetUrl ? (
+        {imageAssetUrl ? (
           <figure className="overflow-hidden rounded-lg border border-border bg-muted/30">
             {/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */}
             <img
-              src={chunk.assetUrl}
+              src={imageAssetUrl}
               alt={chunk.summary ?? "Image chunk"}
               className="max-h-[520px] w-full object-contain"
             />
@@ -407,6 +415,16 @@ function ImageChunkCard({
   );
 }
 
+function getImageChunkAssetUrl(
+  chunk: ParsedChunkView,
+  sourceOriginalFile: SourceOriginalFileView | null,
+): string | null {
+  if (chunk.assetUrl) return chunk.assetUrl;
+  if (!sourceOriginalFile?.mimeType.startsWith("image/")) return null;
+
+  return sourceOriginalFile.url;
+}
+
 function renderTextChunkContent(
   chunk: ParsedChunkView,
   onReferenceClick: (chunkId: string) => void,

src/components/workspace-selected-chunks.test.ts

@@ -1,13 +1,22 @@
 // @vitest-environment jsdom
-import { renderHook } from "@testing-library/react";
+import { renderHook, waitFor } from "@testing-library/react";
 import React from "react";
 import { SWRConfig } from "swr";
-import { describe, expect, it } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { ReactNode } from "react";
 
 import { useWorkspaceSelectedChunks } from "./workspace-selected-chunks";
+import type { ParsedChunkView } from "@/domains/chunks/types";
 import type { SourceView } from "@/domains/sources/types";
 
+const fetchChunkPageMock = vi.hoisted(() => vi.fn());
+
+vi.mock("@/domains/workspace/client", () => ({
+  workspaceClient: {
+    fetchChunkPage: fetchChunkPageMock,
+  },
+}));
+
 const readySource: SourceView = {
   id: "source_1",
   title: "lecture.pdf",
@@ -17,7 +26,20 @@ const readySource: SourceView = {
 };
 
 describe("useWorkspaceSelectedChunks", () => {
-  it("returns prefetched chunks without asking SWR to page the source", () => {
+  beforeEach(() => {
+    fetchChunkPageMock.mockReset();
+    fetchChunkPageMock.mockResolvedValue({
+      chunks: [],
+      pagination: {
+        page: 1,
+        pageSize: 50,
+        total: 0,
+        totalPages: 1,
+      },
+    });
+  });
+
+  it("returns prefetched chunks while checking the visible page for media", async () => {
     const { result } = renderHook(
       () =>
         useWorkspaceSelectedChunks({
@@ -37,6 +59,9 @@ describe("useWorkspaceSelectedChunks", () => {
       { wrapper: createSWRWrapper },
     );
 
+    await waitFor(() =>
+      expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1),
+    );
     expect(result.current.selectedSource?.id).toBe("source_1");
     expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([
       "chunk_1",
@@ -45,6 +70,65 @@ describe("useWorkspaceSelectedChunks", () => {
     expect(result.current.isSelectedChunksLoading).toBe(false);
   });
 
+  it("keeps visible page asset URLs when full-tree chunks arrive without media", async () => {
+    const pagedImageChunk: ParsedChunkView = {
+      chunkId: "image_1",
+      type: "image",
+      content: "Image summary",
+      sourceTitle: "logo.png",
+      assetUrl: "https://blob.example/chunk-assets/image-1.png",
+    };
+    const structureOnlyImageChunk: ParsedChunkView = {
+      chunkId: "image_1",
+      type: "image",
+      content: "Image summary",
+      sourceTitle: "logo.png",
+    };
+    fetchChunkPageMock.mockResolvedValue({
+      chunks: [pagedImageChunk],
+      pagination: {
+        page: 1,
+        pageSize: 50,
+        total: 1,
+        totalPages: 1,
+      },
+    });
+
+    const { result, rerender } = renderHook(
+      (input: {
+        readonly prefetchedChunksBySourceId: Readonly<
+          Record<string, ParsedChunkView[]>
+        >;
+      }) =>
+        useWorkspaceSelectedChunks({
+          selectedSourceId: "source_1",
+          sources: [readySource],
+          prefetchedChunksBySourceId: input.prefetchedChunksBySourceId,
+        }),
+      {
+        initialProps: { prefetchedChunksBySourceId: {} },
+        wrapper: createSWRWrapper,
+      },
+    );
+
+    await waitFor(() =>
+      expect(result.current.selectedChunks[0]?.assetUrl).toBe(
+        "https://blob.example/chunk-assets/image-1.png",
+      ),
+    );
+
+    rerender({
+      prefetchedChunksBySourceId: {
+        source_1: [structureOnlyImageChunk],
+      },
+    });
+
+    expect(result.current.selectedChunks[0]).toMatchObject({
+      chunkId: "image_1",
+      assetUrl: "https://blob.example/chunk-assets/image-1.png",
+    });
+  });
+
   it("returns an empty chunk list when no source is selected", () => {
     const { result } = renderHook(
       () =>

src/components/workspace-selected-chunks.ts

@@ -38,7 +38,7 @@ export function useWorkspaceSelectedChunks({
     ? prefetchedChunksBySourceId[selectedSourceId]
     : undefined
   const selectedChunkSourceId =
-    selectedSource && selectedSource.status === "ready" && !prefetchedSelectedChunks
+    selectedSource && selectedSource.status === "ready"
       ? selectedSource.id
       : null
   const {
@@ -59,20 +59,23 @@ export function useWorkspaceSelectedChunks({
       keepPreviousData: false,
     },
   )
-  const resolvedPrefetchedChunks = useMemo(
-    () =>
-      prefetchedSelectedChunks
-        ? resolveChunkConnectionTargets(prefetchedSelectedChunks)
-        : undefined,
-    [prefetchedSelectedChunks],
-  )
   const pagedSelectedChunks = useMemo(
     () =>
       resolveChunkConnectionTargets(
         (selectedChunkPages ?? []).flatMap((page) => page.chunks ?? []),
       ),
     [selectedChunkPages],
   )
+  const resolvedPrefetchedChunks = useMemo(
+    () =>
+      prefetchedSelectedChunks
+        ? mergeVisibleChunkAssetUrls(
+            resolveChunkConnectionTargets(prefetchedSelectedChunks),
+            pagedSelectedChunks,
+          )
+        : undefined,
+    [pagedSelectedChunks, prefetchedSelectedChunks],
+  )
   const selectedChunks = selectedSourceId
     ? (resolvedPrefetchedChunks ?? pagedSelectedChunks)
     : []
@@ -114,3 +117,26 @@ function fetchChunksByKey([
 ]: SourceChunksKey): Promise<SourceChunksResponse> {
   return workspaceClient.fetchChunkPage(sourceId, page)
 }
+
+function mergeVisibleChunkAssetUrls(
+  chunks: readonly ParsedChunkView[],
+  visibleChunks: readonly ParsedChunkView[],
+): ParsedChunkView[] {
+  if (visibleChunks.length === 0) return [...chunks]
+
+  const visibleChunksById = new Map(
+    visibleChunks.map((chunk) => [chunk.chunkId, chunk]),
+  )
+
+  return chunks.map((chunk) => {
+    if (chunk.assetUrl) return chunk
+
+    const visibleChunk = visibleChunksById.get(chunk.chunkId)
+    if (!visibleChunk?.assetUrl) return chunk
+
+    return {
+      ...chunk,
+      assetUrl: visibleChunk.assetUrl,
+    }
+  })
+}