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.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,