diff --git a/frontend/bun.lockb b/frontend/bun.lockb index 7204e8d51..8f298443f 100755 Binary files a/frontend/bun.lockb and b/frontend/bun.lockb differ diff --git a/frontend/package.json b/frontend/package.json index b67e5c96c..ba5e7ffe1 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -64,7 +64,8 @@ "tailwind-merge": "^2.6.0", "tailwindcss-animate": "^1.0.7", "xlsx": "^0.18.5", - "zod": "3.23.8" + "zod": "3.23.8", + "zustand": "^5.0.8" }, "devDependencies": { "@eslint/js": "^9.17.0", diff --git a/frontend/src/components/ChatBox.tsx b/frontend/src/components/ChatBox.tsx index aecde7b7b..85cb7c027 100644 --- a/frontend/src/components/ChatBox.tsx +++ b/frontend/src/components/ChatBox.tsx @@ -86,6 +86,7 @@ interface SelectedFile { uploadError?: string preview?: string // URL for image preview fileType?: FileType + uploadAborted?: boolean } export const getFileIcon = (fileType: FileType | string | undefined) => { @@ -417,6 +418,7 @@ export const ChatBox = React.forwardRef( const scrollPositionRef = useRef(0) const navigate = useNavigate() const fileInputRef = useRef(null) + const uploadControllers = useRef>(new Map()) const [showReferenceBox, setShowReferenceBox] = useState(false) const [searchMode, setSearchMode] = useState<"citations" | "global">( @@ -836,12 +838,16 @@ export const ChatBox = React.forwardRef( setSelectedFiles((prev) => prev.map((f) => files.some((file) => file.id === f.id) - ? { ...f, uploading: true, uploadError: undefined } + ? { ...f, uploading: true, uploadError: undefined, uploadAborted: false } : f, ), ) const uploadPromises = files.map(async (selectedFile) => { + // Create AbortController for this file + const abortController = new AbortController() + uploadControllers.current.set(selectedFile.id, abortController) + try { const formData = new FormData() formData.append("attachment", selectedFile.file) @@ -850,6 +856,7 @@ export const ChatBox = React.forwardRef( { method: "POST", body: formData, + signal: abortController.signal, }, ) @@ -874,6 +881,19 @@ export const ChatBox = React.forwardRef( throw new Error("No document ID returned from upload") } } catch (error) { + // Check if the error is an AbortError + if (error instanceof Error && error.name === 'AbortError') { + // Handle aborted upload - don't show error toast + setSelectedFiles((prev) => + prev.map((f) => + f.id === selectedFile.id + ? { ...f, uploading: false, uploadAborted: true } + : f, + ), + ) + return null + } + const errorMessage = error instanceof Error ? error.message : "Upload failed" setSelectedFiles((prev) => @@ -889,6 +909,8 @@ export const ChatBox = React.forwardRef( }) return null } finally { + // Clean up the controller + uploadControllers.current.delete(selectedFile.id) setUploadingFilesCount((prev) => prev - 1) } }) @@ -980,6 +1002,25 @@ export const ChatBox = React.forwardRef( const removeFile = useCallback(async (id: string) => { const fileToRemove = selectedFiles.find((f) => f.id === id) + // If the file is currently uploading, abort the upload + if (fileToRemove?.uploading) { + const controller = uploadControllers.current.get(id) + if (controller) { + controller.abort() + uploadControllers.current.delete(id) + } + + // Remove from UI immediately for aborted uploads + setSelectedFiles((prev) => { + if (fileToRemove?.preview) { + URL.revokeObjectURL(fileToRemove.preview) + } + return prev.filter((f) => f.id !== id) + }) + setUploadingFilesCount((prev) => Math.max(0, prev - 1)) + return + } + // If the file has metadata with fileId (meaning it's already uploaded), delete it from the server if (fileToRemove?.metadata?.fileId) { try { @@ -2213,6 +2254,12 @@ export const ChatBox = React.forwardRef( .map((f) => f.preview) .filter(Boolean) as string[] cleanupPreviewUrls(previewUrls) + + // Abort any ongoing uploads + uploadControllers.current.forEach((controller) => { + controller.abort() + }) + uploadControllers.current.clear() } }, []) @@ -2861,8 +2908,12 @@ export const ChatBox = React.forwardRef( @@ -2917,8 +2968,12 @@ export const ChatBox = React.forwardRef( )} @@ -2931,7 +2986,7 @@ export const ChatBox = React.forwardRef( {uploadingFilesCount > 0 && (
- Uploading files... + Uploading files... (click X to cancel)
)} {selectedFiles.length >= MAX_ATTACHMENTS && ( diff --git a/frontend/src/components/__tests__/ChatBox.test.tsx b/frontend/src/components/__tests__/ChatBox.test.tsx new file mode 100644 index 000000000..afbfe82f9 --- /dev/null +++ b/frontend/src/components/__tests__/ChatBox.test.tsx @@ -0,0 +1,106 @@ +import { renderHook, act } from '@testing-library/react' +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { ChatBox } from '../ChatBox' + +// Mock the necessary dependencies +vi.mock('@/utils/authFetch', () => ({ + authFetch: vi.fn() +})) + +vi.mock('@/hooks/use-toast', () => ({ + useToast: () => ({ + toast: { + error: vi.fn(), + success: vi.fn() + } + }) +})) + +vi.mock('@/api', () => ({ + api: { + files: { + delete: { + $post: vi.fn() + } + } + } +})) + +describe('ChatBox Upload Cancellation', () => { + let mockAuthFetch: any + + beforeEach(() => { + vi.clearAllMocks() + mockAuthFetch = vi.fn() + vi.doMock('@/utils/authFetch', () => ({ + authFetch: mockAuthFetch + })) + }) + + it('should handle upload cancellation correctly', async () => { + // Create a mock AbortController + const abortController = new AbortController() + + // Mock a file upload that can be aborted + const uploadPromise = new Promise((resolve, reject) => { + abortController.signal.addEventListener('abort', () => { + const error = new Error('Upload aborted') + error.name = 'AbortError' + reject(error) + }) + }) + + mockAuthFetch.mockReturnValue(uploadPromise) + + // Test that aborting the upload triggers the correct behavior + act(() => { + abortController.abort() + }) + + // Verify the upload was cancelled + await expect(uploadPromise).rejects.toThrow('Upload aborted') + }) + + it('should clean up upload controllers on component unmount', () => { + // This test would require more complex setup with actual component rendering + // For now, we'll just verify the concept + const controllers = new Map() + const controller1 = new AbortController() + const controller2 = new AbortController() + + controllers.set('file1', controller1) + controllers.set('file2', controller2) + + // Simulate cleanup + controllers.forEach((controller) => { + controller.abort() + }) + controllers.clear() + + expect(controllers.size).toBe(0) + expect(controller1.signal.aborted).toBe(true) + expect(controller2.signal.aborted).toBe(true) + }) + + it('should handle AbortError gracefully in upload completion', async () => { + const abortError = new Error('Request aborted') + abortError.name = 'AbortError' + + // Mock the upload function to handle AbortError + const mockUploadFile = async (file: File, signal: AbortSignal) => { + if (signal.aborted) { + throw abortError + } + return { fileId: 'test-id', fileName: file.name } + } + + const testFile = new File(['test'], 'test.txt', { type: 'text/plain' }) + const abortController = new AbortController() + + // Abort before upload + abortController.abort() + + // Verify AbortError is thrown + await expect(mockUploadFile(testFile, abortController.signal)).rejects.toThrow('Request aborted') + }) +}) \ No newline at end of file diff --git a/frontend/src/store/__tests__/useUploadProgressStore.test.ts b/frontend/src/store/__tests__/useUploadProgressStore.test.ts new file mode 100644 index 000000000..74a40766e --- /dev/null +++ b/frontend/src/store/__tests__/useUploadProgressStore.test.ts @@ -0,0 +1,118 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest' +import { useUploadProgress } from '../useUploadProgressStore' + +// Mock File API for testing +global.File = class MockFile { + name: string + size: number + type: string + lastModified: number + + constructor(bits: any[], name: string, options?: { type?: string }) { + this.name = name + this.type = options?.type || 'text/plain' + this.size = bits.reduce((acc, bit) => acc + (typeof bit === 'string' ? bit.length : 0), 0) + this.lastModified = Date.now() + } +} as any + +describe('useUploadProgressStore', () => { + beforeEach(() => { + // Reset the store before each test by getting the state and clearing it + const store = useUploadProgress.getState() + if (store.currentUpload) { + store.finishUpload(store.currentUpload.id) + } + }) + + it('should start upload with abort controller', () => { + const store = useUploadProgress.getState() + + const files = [ + { file: new File(['test'], 'test.txt'), id: 'file1' } + ] + + const result = store.startUpload('test-collection', files, 1, true) + + expect(result.uploadId).toBeDefined() + expect(result.abortController).toBeInstanceOf(AbortController) + expect(useUploadProgress.getState().currentUpload?.id).toBe(result.uploadId) + }) + + it('should cancel upload and abort controller', () => { + const store = useUploadProgress.getState() + + const files = [ + { file: new File(['test'], 'test.txt'), id: 'file1' } + ] + + const { uploadId, abortController } = store.startUpload('test-collection', files, 1, true) + + // Verify controller is not aborted initially + expect(abortController.signal.aborted).toBe(false) + + // Cancel the upload + store.cancelUpload(uploadId) + + // Verify controller is aborted and upload is removed + expect(abortController.signal.aborted).toBe(true) + expect(useUploadProgress.getState().currentUpload).toBeNull() + }) + + it('should remove progress', () => { + const store = useUploadProgress.getState() + + const files = [ + { file: new File(['test'], 'test.txt'), id: 'file1' } + ] + + const { uploadId } = store.startUpload('test-collection', files, 1, true) + + // Verify upload exists + expect(useUploadProgress.getState().currentUpload?.id).toBe(uploadId) + + // Remove progress + store.removeProgress(uploadId) + + // Verify upload is removed + expect(useUploadProgress.getState().currentUpload).toBeNull() + }) + + it('should handle file status updates', () => { + const store = useUploadProgress.getState() + + const files = [ + { file: new File(['test'], 'test.txt'), id: 'file1' } + ] + + const { uploadId } = store.startUpload('test-collection', files, 1, true) + + // Update file status to uploading + store.updateFileStatus(uploadId, 'test.txt', 'file1', 'uploading') + + let upload = useUploadProgress.getState().currentUpload + expect(upload?.files[0].status).toBe('uploading') + + // Update file status to uploaded + store.updateFileStatus(uploadId, 'test.txt', 'file1', 'uploaded') + + upload = useUploadProgress.getState().currentUpload + expect(upload?.files[0].status).toBe('uploaded') + }) + + it('should handle abort controller cleanup', () => { + const store = useUploadProgress.getState() + + const files = [ + { file: new File(['test'], 'test.txt'), id: 'file1' } + ] + + const { abortController } = store.startUpload('test-collection', files, 1, true) + + // Manually abort the controller + abortController.abort() + + // Verify the signal is aborted + expect(abortController.signal.aborted).toBe(true) + }) +}) \ No newline at end of file diff --git a/frontend/src/store/useUploadProgressStore.ts b/frontend/src/store/useUploadProgressStore.ts index 31f4d611b..f6a1c82e4 100644 --- a/frontend/src/store/useUploadProgressStore.ts +++ b/frontend/src/store/useUploadProgressStore.ts @@ -45,6 +45,7 @@ interface UploadProgressStore { ) => void finishUpload: (uploadId: string) => void cancelUpload: (uploadId: string) => void + removeProgress: (uploadId: string) => void getUploadProgress: (uploadId: string) => UploadTask | null } @@ -141,6 +142,12 @@ export const useUploadProgress = create((set, get) => ({ ) }, + removeProgress: (uploadId) => { + set((state) => + state.currentUpload?.id === uploadId ? { currentUpload: null } : state + ) + }, + getUploadProgress: (uploadId) => { const state = get() return state.currentUpload?.id === uploadId ? state.currentUpload : null diff --git a/server/api/files.ts b/server/api/files.ts index be0005ec2..13ca7a700 100644 --- a/server/api/files.ts +++ b/server/api/files.ts @@ -398,6 +398,14 @@ export const handleAttachmentUpload = async (c: Context) => { message: `Stored ${attachmentMetadata.length} attachment(s) successfully.`, }) } catch (error) { + // Handle AbortError gracefully for cancelled uploads + if (error instanceof Error && error.name === 'AbortError') { + loggerWithChild({ email }).info( + "Attachment upload was cancelled by client", + ) + throw new HTTPException(499, { message: "Upload cancelled" }) + } + loggerWithChild({ email }).error( error, "Error in attachment upload handler",