Harden chat assets on demand

src/app/api/sources/[sourceId]/chunks/route.test.ts

@@ -498,7 +498,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => {
       revisionKey: "job_1",
       page: 1,
       pageSize: 1,
-      assetUrlPolicy: "durable",
     })
   })
 
@@ -660,7 +659,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => {
       revisionKey: "job_result_1",
       page: 1,
       pageSize: 1,
-      assetUrlPolicy: "durable",
     })
   })
 })

src/components/workspace-citation-focus.ts

@@ -21,6 +21,7 @@ type PrefetchedChunksUpdater = (
 type WorkspaceCitationFocusInput = {
   readonly fetchChunks: (sourceId: string) => Promise<ParsedChunkView[]>
   readonly initialPrefetchedChunksBySourceId?: PrefetchedChunksBySourceId
+  readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void
   readonly onSelectSource: (sourceId: string | null) => void
   readonly selectedSourceId: string | null
   readonly sources: readonly SourceView[]
@@ -51,6 +52,7 @@ type WorkspaceCitationFocus = {
 export function useWorkspaceCitationFocus({
   fetchChunks,
   initialPrefetchedChunksBySourceId = {},
+  onRemoteSourceChunksLoaded,
   onSelectSource,
   selectedSourceId,
   sources,
@@ -88,6 +90,7 @@ export function useWorkspaceCitationFocus({
     selectedSourceId,
     sources,
     prefetchedChunksBySourceId,
+    onRemoteSourceChunksLoaded,
   })
 
   const requestChunkFocus = useCallback(

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

@@ -161,6 +161,58 @@ describe("useWorkspaceSelectedChunks", () => {
     expect(result.current.selectedChunks).toEqual([]);
   });
 
+  it("requests a source refresh after loading an unlocalized remote source", async () => {
+    const onRemoteSourceChunksLoaded = vi.fn();
+    const remoteSource: SourceView = {
+      ...readySource,
+      id: "knowhere-doc:default:doc_remote",
+      kind: "remote",
+      documentId: "doc_remote",
+      excludedFromQuery: false,
+    };
+    fetchChunkPageMock.mockResolvedValue({
+      chunks: [
+        {
+          chunkId: "chunk_1",
+          type: "text",
+          content: "Remote content",
+          sourceTitle: "Remote.pdf",
+        },
+      ],
+      pagination: {
+        page: 1,
+        pageSize: 50,
+        total: 1,
+        totalPages: 1,
+      },
+    });
+
+    const { rerender } = renderHook(
+      (input: {
+        readonly onRemoteSourceChunksLoaded: (sourceId: string) => void;
+      }) =>
+        useWorkspaceSelectedChunks({
+          selectedSourceId: "knowhere-doc:default:doc_remote",
+          sources: [remoteSource],
+          prefetchedChunksBySourceId: {},
+          onRemoteSourceChunksLoaded: input.onRemoteSourceChunksLoaded,
+        }),
+      {
+        initialProps: { onRemoteSourceChunksLoaded },
+        wrapper: createSWRWrapper,
+      },
+    );
+
+    await waitFor(() =>
+      expect(onRemoteSourceChunksLoaded).toHaveBeenCalledWith(
+        "knowhere-doc:default:doc_remote",
+      ),
+    );
+
+    rerender({ onRemoteSourceChunksLoaded });
+    expect(onRemoteSourceChunksLoaded).toHaveBeenCalledTimes(1);
+  });
+
   it("returns an empty chunk list when no source is selected", () => {
     const { result } = renderHook(
       () =>

src/components/workspace-selected-chunks.ts

@@ -1,6 +1,6 @@
 "use client"
 
-import { useMemo } from "react"
+import { useEffect, useMemo, useRef } from "react"
 import useSWRInfinite from "swr/infinite"
 
 import { workspaceClient } from "@/domains/workspace/client"
@@ -17,6 +17,7 @@ type WorkspaceSelectedChunksInput = {
   readonly selectedSourceId: string | null
   readonly sources: readonly SourceView[]
   readonly prefetchedChunksBySourceId: Readonly<Record<string, ParsedChunkView[]>>
+  readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void
 }
 
 type WorkspaceSelectedChunks = {
@@ -33,8 +34,10 @@ export function useWorkspaceSelectedChunks({
   selectedSourceId,
   sources,
   prefetchedChunksBySourceId,
+  onRemoteSourceChunksLoaded,
 }: WorkspaceSelectedChunksInput): WorkspaceSelectedChunks {
   const selectedSource = sources.find((source) => source.id === selectedSourceId)
+  const remoteSourceRefreshRequestedIdsRef = useRef<Set<string>>(new Set())
   const prefetchedSelectedChunks = selectedSourceId
     ? prefetchedChunksBySourceId[selectedSourceId]
     : undefined
@@ -100,6 +103,16 @@ export function useWorkspaceSelectedChunks({
       !selectedChunkPages &&
       isChunksLoading)
 
+  useEffect(() => {
+    const sourceId = selectedSource?.id
+    if (!sourceId || selectedSource.kind !== "remote") return
+    if (!selectedChunkPages || selectedChunkPages.length === 0) return
+    if (remoteSourceRefreshRequestedIdsRef.current.has(sourceId)) return
+
+    remoteSourceRefreshRequestedIdsRef.current.add(sourceId)
+    onRemoteSourceChunksLoaded?.(sourceId)
+  }, [onRemoteSourceChunksLoaded, selectedChunkPages, selectedSource])
+
   function handleLoadMoreChunks(): void {
     if (!hasMoreSelectedChunks || isSelectedChunksLoadingMore) return
     void setSelectedChunkPageCount(selectedChunkPageCount + 1)

src/components/workspace-shell.tsx

@@ -114,6 +114,7 @@ function WorkspaceShellContent({
     fetchChunks: workspaceClient.fetchChunks,
     initialPrefetchedChunksBySourceId:
       initialPrefetchedChunksBySourceId ?? undefined,
+    onRemoteSourceChunksLoaded: sourceWorkflow.handleSourcesRefresh,
     onSelectSource: handleCitationSourceSelected,
     selectedSourceId: sourceWorkflow.selectedSourceId,
     sources: sourceWorkflow.sources,

src/components/workspace-source-state.test.ts

@@ -83,6 +83,27 @@ describe("workspaceSourceState", () => {
     ).toBe("source_target");
   });
 
+  it("keeps a localized remote document selected after source refresh", () => {
+    const sources: readonly SourceView[] = [
+      {
+        id: "source_localized",
+        kind: "workspace",
+        title: "remote.pdf",
+        status: "ready",
+        mimeType: "application/pdf",
+        documentId: "doc_remote",
+        excludedFromQuery: false,
+      },
+    ];
+
+    expect(
+      workspaceSourceState.getResolvedSelectedSourceId(
+        sources,
+        "knowhere-doc:default:doc_remote",
+      ),
+    ).toBe("source_localized");
+  });
+
   it("applies source query exclusions without mutating the source list", () => {
     const sources: readonly SourceView[] = [
       {

src/components/workspace-source-state.ts

@@ -74,6 +74,15 @@ function getResolvedSelectedSourceId(
     return selectedSource.id
   }
 
+  const selectedDocumentId = getRemoteSourceDocumentId(selectedSourceId)
+  if (selectedDocumentId) {
+    const localizedSource = sources.find(
+      (source) =>
+        source.documentId === selectedDocumentId && isReadySource(source),
+    )
+    if (localizedSource) return localizedSource.id
+  }
+
   return getFirstReadySourceId(sources)
 }
 
@@ -134,6 +143,20 @@ function removeRecordKey<T>(
   return remaining
 }
 
+function getRemoteSourceDocumentId(sourceId: string | null): string | null {
+  if (!sourceId) return null
+
+  const parts = sourceId.split(":")
+  if (parts.length !== 3 || parts[0] !== "knowhere-doc") return null
+
+  try {
+    const documentId = decodeURIComponent(parts[2] ?? "")
+    return documentId.length > 0 ? documentId : null
+  } catch {
+    return null
+  }
+}
+
 export const workspaceSourceState: WorkspaceSourceStateModule = {
   getFirstReadySourceId,
   getInitialSelectedSourceId,

src/components/workspace-source-workflow.ts

@@ -22,6 +22,7 @@ type WorkspaceSourceWorkflow = {
   readonly handleRetrySource: (sourceId: string) => Promise<void>
   readonly handleOfficialLibrarySourceAdd: (demoSourceId: string) => Promise<boolean>
   readonly handleSelectedSourceChange: (sourceId: string | null) => void
+  readonly handleSourcesRefresh: () => void
   readonly handleSourcesMaterialized: (
     demoSourceIds: readonly string[],
     materializedSources: readonly SourceView[],
@@ -153,6 +154,10 @@ export function useWorkspaceSourceWorkflow({
     setSelectedSourceId(sourceId)
   }
 
+  function handleSourcesRefresh(): void {
+    void mutateSources()
+  }
+
   async function handleArchiveSource(sourceId: string): Promise<void> {
     setArchivingSourceIds((current) =>
       workspaceSourceState.addPendingId(current, sourceId),
@@ -237,6 +242,7 @@ export function useWorkspaceSourceWorkflow({
     handleRetrySource,
     handleOfficialLibrarySourceAdd,
     handleSelectedSourceChange,
+    handleSourcesRefresh,
     handleSourcesMaterialized,
     handleSourceUploaded,
     handleToggleIncluded,

src/domains/chat/contracts.ts

@@ -10,7 +10,7 @@ import type {
   ChatCitationView,
 } from "@/domains/chat/types"
 import type { HardenMediaAssetUrls } from "./media-asset-hardening"
-import type { LoadSourceAssetUrls } from "./media-assets"
+import type { HardenChatAssetUrl } from "./media-assets"
 
 export type RetrievalClient = {
   query(params: RetrievalQueryParams): Promise<RetrievalQueryResponse>
@@ -67,7 +67,7 @@ export type AnswerQuestionInput = {
   excludedSourceIds: readonly string[]
   retrieval: RetrievalClient
   generateAnswer: GenerateAnswer
-  loadSourceAssetUrls?: LoadSourceAssetUrls
+  hardenChatAssetUrl?: HardenChatAssetUrl
   hardenMediaAssetUrls?: HardenMediaAssetUrls
   messages: readonly ChatHistoryMessage[]
 }