diff --git a/packages/ui/src/components/chat/ChatInput.tsx b/packages/ui/src/components/chat/ChatInput.tsx index d4e167ce..e295d9ea 100644 --- a/packages/ui/src/components/chat/ChatInput.tsx +++ b/packages/ui/src/components/chat/ChatInput.tsx @@ -1646,21 +1646,21 @@ const ChatInputComponent: React.FC = ({ onOpenSettings, scrollTo if (normalized.includes('payload too large') || normalized.includes('413') || normalized.includes('entity too large')) { toast.error(t('chat.chatInput.toast.attachmentsTooLarge')); if (allAttachments.length > 0) { - useInputStore.setState({ attachedFiles: allAttachments }); + useInputStore.getState().setAttachedFiles(allAttachments); } return; } if (isSoftNetworkError) { if (allAttachments.length > 0) { - useInputStore.setState({ attachedFiles: allAttachments }); + useInputStore.getState().setAttachedFiles(allAttachments); toast.error(t('chat.chatInput.toast.sendAttachmentsFailed')); } return; } if (allAttachments.length > 0) { - useInputStore.setState({ attachedFiles: allAttachments }); + useInputStore.getState().setAttachedFiles(allAttachments); } toast.error(rawMessage || t('chat.chatInput.toast.messageSendFailed')); }); diff --git a/packages/ui/src/components/chat/QueuedMessageChips.tsx b/packages/ui/src/components/chat/QueuedMessageChips.tsx index e90d97ab..60029896 100644 --- a/packages/ui/src/components/chat/QueuedMessageChips.tsx +++ b/packages/ui/src/components/chat/QueuedMessageChips.tsx @@ -88,9 +88,7 @@ export const QueuedMessageChips = memo(({ onEditMessage }: QueuedMessageChipsPro if (popped) { if (popped.attachments && popped.attachments.length > 0) { const currentAttachments = useInputStore.getState().attachedFiles; - useInputStore.setState({ - attachedFiles: [...currentAttachments, ...popped.attachments] - }); + useInputStore.getState().setAttachedFiles([...currentAttachments, ...popped.attachments]); } onEditMessage(popped.content, popped.attachments); } diff --git a/packages/ui/src/sync/input-store.test.ts b/packages/ui/src/sync/input-store.test.ts new file mode 100644 index 00000000..4fb4aa08 --- /dev/null +++ b/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) => 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) +} + +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([]) + }) +}) diff --git a/packages/ui/src/sync/input-store.ts b/packages/ui/src/sync/input-store.ts index 284d43b7..aa58f158 100644 --- a/packages/ui/src/sync/input-store.ts +++ b/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() +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 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 @@ -107,11 +109,13 @@ export const useInputStore = create()((set, get) => ({ addAttachedFile: async (file: File) => { const id = `${Date.now()}-${Math.random().toString(36).slice(2)}` + const generation = attachmentReadGeneration const dataUrl = await new Promise((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()((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()((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()((set, get) => ({ } finally { pendingVSCodeSelectionKeys.delete(selectionKey) } + if (generation !== attachmentReadGeneration) return const attached: AttachedFile = { id, file,