fix(input): ignore stale attachment reads (#1150)

* fix(input): ignore stale attachment reads

* fix(input): centralize attachment replacement

* fix(input): route queued attachment restores through store action

---------

Co-authored-by: Isaac Sanchez <isanchez-hawkins@arize.com>
Co-authored-by: Bohdan Triapitsyn <artmore@protonmail.com>

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

@@ -0,0 +1,88 @@
+import { beforeEach, describe, expect, test } from "bun:test"
+import { useInputStore } from "./input-store"
+
+class MockFileReader {
+  result: string | ArrayBuffer | null = null
+  onload: ((this: FileReader, event: ProgressEvent<FileReader>) => unknown) | null = null
+
+  readAsDataURL() {
+    pendingReaders.push(this)
+  }
+}
+
+const pendingReaders: MockFileReader[] = []
+
+const resolveReader = (reader: MockFileReader, result: string) => {
+  reader.result = result
+  reader.onload?.call(reader as unknown as FileReader, {} as ProgressEvent<FileReader>)
+}
+
+describe("input-store attachments", () => {
+  beforeEach(() => {
+    pendingReaders.length = 0
+    globalThis.FileReader = MockFileReader as unknown as typeof FileReader
+    useInputStore.setState({
+      pendingInputText: null,
+      pendingInputMode: "replace",
+      pendingSyntheticParts: null,
+      activeEditorFile: null,
+    })
+    useInputStore.getState().setAttachedFiles([])
+  })
+
+  test("does not attach a local file that finishes reading after attachments are cleared", async () => {
+    const addPromise = useInputStore.getState().addAttachedFile(new File(["hello"], "hello.txt", { type: "text/plain" }))
+    expect(pendingReaders).toHaveLength(1)
+
+    useInputStore.getState().clearAttachedFiles()
+    resolveReader(pendingReaders[0], "data:text/plain;base64,aGVsbG8=")
+    await addPromise
+
+    expect(useInputStore.getState().attachedFiles).toEqual([])
+  })
+
+  test("does not attach a local file after attached files are replaced", async () => {
+    const addPromise = useInputStore.getState().addAttachedFile(new File(["hello"], "hello.txt", { type: "text/plain" }))
+    expect(pendingReaders).toHaveLength(1)
+
+    useInputStore.getState().setAttachedFiles([])
+    resolveReader(pendingReaders[0], "data:text/plain;base64,aGVsbG8=")
+    await addPromise
+
+    expect(useInputStore.getState().attachedFiles).toEqual([])
+  })
+
+  test("does not attach a local file after attached files are restored", async () => {
+    const addPromise = useInputStore.getState().addAttachedFile(new File(["hello"], "hello.txt", { type: "text/plain" }))
+    expect(pendingReaders).toHaveLength(1)
+
+    const restored = new File(["restored"], "restored.txt", { type: "text/plain" })
+    useInputStore.getState().setAttachedFiles([{
+      id: "restored",
+      file: restored,
+      dataUrl: "data:text/plain;base64,cmVzdG9yZWQ=",
+      mimeType: "text/plain",
+      filename: "restored.txt",
+      size: restored.size,
+      source: "local",
+    }])
+    resolveReader(pendingReaders[0], "data:text/plain;base64,aGVsbG8=")
+    await addPromise
+
+    expect(useInputStore.getState().attachedFiles.map((file) => file.filename)).toEqual(["restored.txt"])
+  })
+
+  test("does not attach a VS Code selection that finishes reading after attachments are cleared", async () => {
+    const addPromise = useInputStore.getState().addVSCodeSelectionAttachment(
+      "/workspace/hello.txt",
+      new File(["hello"], "hello.txt", { type: "text/plain" })
+    )
+    expect(pendingReaders).toHaveLength(1)
+
+    useInputStore.getState().clearAttachedFiles()
+    resolveReader(pendingReaders[0], "data:text/plain;base64,aGVsbG8=")
+    await addPromise
+
+    expect(useInputStore.getState().attachedFiles).toEqual([])
+  })
+})

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

@@ -8,6 +8,7 @@ import type { AttachedFile } from "@/stores/types/sessionTypes"
 
 const FILE_URI_PREFIX = "file://"
 const pendingVSCodeSelectionKeys = new Set<string>()
+let attachmentReadGeneration = 0
 
 const encodeFilePath = (filepath: string): string => {
   let normalized = filepath.replace(/\\/g, "/")
@@ -72,6 +73,7 @@ export type InputState = {
   consumePendingSyntheticParts: () => SyntheticContextPart[] | null
   addAttachedFile: (file: File) => Promise<void>
   removeAttachedFile: (id: string) => void
+  setAttachedFiles: (files: AttachedFile[]) => void
   clearAttachedFiles: () => void
   addVSCodeFileAttachment: (path: string, name: string, fileSize: number | null) => void
   addVSCodeSelectionAttachment: (path: string, file: File) => Promise<void>
@@ -107,11 +109,13 @@ 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) => {
       const reader = new FileReader()
       reader.onload = () => resolve(reader.result as string)
       reader.readAsDataURL(file)
     })
+    if (generation !== attachmentReadGeneration) return
     const attached: AttachedFile = {
       id,
       file,
@@ -127,7 +131,15 @@ export const useInputStore = create<InputState>()((set, get) => ({
   removeAttachedFile: (id) =>
     set((s) => ({ attachedFiles: s.attachedFiles.filter((f) => f.id !== id) })),
 
-  clearAttachedFiles: () => set({ attachedFiles: [] }),
+  setAttachedFiles: (files) => {
+    attachmentReadGeneration += 1
+    set({ attachedFiles: files })
+  },
+
+  clearAttachedFiles: () => {
+    attachmentReadGeneration += 1
+    set({ attachedFiles: [] })
+  },
 
   addVSCodeFileAttachment: (path: string, name: string, fileSize: number | null) => {
     const id = `${Date.now()}-${Math.random().toString(36).slice(2)}`
@@ -155,6 +167,7 @@ export const useInputStore = create<InputState>()((set, get) => ({
 
   addVSCodeSelectionAttachment: async (path: string, file: File) => {
     const id = `${Date.now()}-${Math.random().toString(36).slice(2)}`
+    const generation = attachmentReadGeneration
     const selectionKey = getVSCodeSelectionKey(path, file.name)
     const isDuplicate = get().attachedFiles.some(
       (f) => f.source === 'vscode' && f.vscodeSource === 'selection' && f.filename === file.name && f.vscodePath === path
@@ -171,6 +184,7 @@ export const useInputStore = create<InputState>()((set, get) => ({
     } finally {
       pendingVSCodeSelectionKeys.delete(selectionKey)
     }
+    if (generation !== attachmentReadGeneration) return
     const attached: AttachedFile = {
       id,
       file,