fix(sync): settle failed attachment reads (#1193)

Co-authored-by: Isaac Sanchez <isanchez-hawkins@arize.com>

packages/ui/src/sync/input-store.test.ts

@@ -4,6 +4,9 @@ import { useInputStore } from "./input-store"
 class MockFileReader {
   result: string | ArrayBuffer | null = null
   onload: ((this: FileReader, event: ProgressEvent<FileReader>) => unknown) | null = null
+  onerror: ((this: FileReader, event: ProgressEvent<FileReader>) => unknown) | null = null
+  onabort: ((this: FileReader, event: ProgressEvent<FileReader>) => unknown) | null = null
+  error: DOMException | null = null
 
   readAsDataURL() {
     pendingReaders.push(this)
@@ -17,6 +20,11 @@ const resolveReader = (reader: MockFileReader, result: string) => {
   reader.onload?.call(reader as unknown as FileReader, {} as ProgressEvent<FileReader>)
 }
 
+const rejectReader = (reader: MockFileReader) => {
+  reader.error = new DOMException("read failed", "NotReadableError")
+  reader.onerror?.call(reader as unknown as FileReader, {} as ProgressEvent<FileReader>)
+}
+
 describe("input-store attachments", () => {
   beforeEach(() => {
     pendingReaders.length = 0
@@ -85,4 +93,30 @@ describe("input-store attachments", () => {
 
     expect(useInputStore.getState().attachedFiles).toEqual([])
   })
+
+  test("does not leave local file reads pending after a reader error", async () => {
+    const addPromise = useInputStore.getState().addAttachedFile(new File(["hello"], "hello.txt", { type: "text/plain" }))
+    expect(pendingReaders).toHaveLength(1)
+
+    rejectReader(pendingReaders[0])
+    await addPromise
+
+    expect(useInputStore.getState().attachedFiles).toEqual([])
+  })
+
+  test("cleans up pending VS Code selection keys after a reader error", async () => {
+    const file = new File(["hello"], "hello.txt", { type: "text/plain" })
+    const firstAdd = useInputStore.getState().addVSCodeSelectionAttachment("/workspace/hello.txt", file)
+    expect(pendingReaders).toHaveLength(1)
+
+    rejectReader(pendingReaders[0])
+    await firstAdd
+
+    const secondAdd = useInputStore.getState().addVSCodeSelectionAttachment("/workspace/hello.txt", file)
+    expect(pendingReaders).toHaveLength(2)
+    resolveReader(pendingReaders[1], "data:text/plain;base64,aGVsbG8=")
+    await secondAdd
+
+    expect(useInputStore.getState().attachedFiles.map((attached) => attached.filename)).toEqual(["hello.txt"])
+  })
 })

packages/ui/src/sync/input-store.ts

@@ -34,6 +34,14 @@ const toFileUrl = (filepath: string): string => {
 
 const getVSCodeSelectionKey = (path: string, filename: string): string => `${path}\u0000${filename}`
 
+const readFileAsDataUrl = (file: File): Promise<string> => new Promise((resolve, reject) => {
+  const reader = new FileReader()
+  reader.onload = () => resolve(reader.result as string)
+  reader.onerror = () => reject(reader.error ?? new Error("Failed to read file"))
+  reader.onabort = () => reject(new Error("File read aborted"))
+  reader.readAsDataURL(file)
+})
+
 const isSameVSCodeActiveEditorFile = (a: VSCodeActiveEditorFile | null, b: VSCodeActiveEditorFile | null): boolean => {
   if (a === b) return true
   if (!a || !b) return false
@@ -110,11 +118,12 @@ export const useInputStore = create<InputState>()((set, get) => ({
   addAttachedFile: async (file: File) => {
     const id = `${Date.now()}-${Math.random().toString(36).slice(2)}`
     const generation = attachmentReadGeneration
-    const dataUrl = await new Promise<string>((resolve) => {
+    let dataUrl: string
-      const reader = new FileReader()
+    try {
-      reader.onload = () => resolve(reader.result as string)
+      dataUrl = await readFileAsDataUrl(file)
-      reader.readAsDataURL(file)
+    } catch {
-    })
+      return
+    }
     if (generation !== attachmentReadGeneration) return
     const attached: AttachedFile = {
       id,
@@ -176,11 +185,9 @@ export const useInputStore = create<InputState>()((set, get) => ({
     pendingVSCodeSelectionKeys.add(selectionKey)
     let dataUrl: string
     try {
-      dataUrl = await new Promise<string>((resolve) => {
+      dataUrl = await readFileAsDataUrl(file)
-        const reader = new FileReader()
+    } catch {
-        reader.onload = () => resolve(reader.result as string)
+      return
-        reader.readAsDataURL(file)
-      })
     } finally {
       pendingVSCodeSelectionKeys.delete(selectionKey)
     }