From 891c9a682348bfabc0f4df476be909f6dab38942 Mon Sep 17 00:00:00 2001 From: Brooklyn Nicholson Date: Tue, 9 Jun 2026 18:21:10 -0500 Subject: [PATCH] fix(desktop): close eager-upload races flagged in review Two races in the drop-time eager upload: - Resurrected chip: the success path used addComposerAttachment, which re-appends when the id is gone, so a file removed mid-upload reappeared once the upload resolved. Add updateComposerAttachment (update-only; no-op when the chip was removed) and use it on both the eager success path and submit-time sync. - Duplicate upload: submit-time sync didn't join an eager upload still in flight, so drop-then-Enter could fire file.attach twice and leave a duplicate under .hermes/desktop-attachments/. Track in-flight eager uploads by id and await the pending one before deciding to re-upload, reusing its gateway ref. Tests: composer-store no-resurrect unit tests + a join-on-submit integration test asserting a single file.attach. Addresses @helix4u review on #43109. --- .../session/hooks/use-prompt-actions.test.tsx | 60 ++++++++++++++++++- .../app/session/hooks/use-prompt-actions.ts | 38 +++++++++--- apps/desktop/src/store/composer.test.ts | 43 +++++++++++++ apps/desktop/src/store/composer.ts | 19 ++++++ 4 files changed, 150 insertions(+), 10 deletions(-) create mode 100644 apps/desktop/src/store/composer.test.ts diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions.test.tsx b/apps/desktop/src/app/session/hooks/use-prompt-actions.test.tsx index cee4df7be9..145a397d87 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions.test.tsx @@ -3,9 +3,8 @@ import type { MutableRefObject } from 'react' import { useEffect } from 'react' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -import { $sessions, setSessions } from '@/store/session' -import { $connection } from '@/store/session' import { $composerAttachments, type ComposerAttachment } from '@/store/composer' +import { $connection, $sessions, setSessions } from '@/store/session' import type { SessionInfo } from '@/types/hermes' import { usePromptActions } from './use-prompt-actions' @@ -456,6 +455,63 @@ describe('usePromptActions file attachment sync', () => { }) }) +describe('usePromptActions eager-upload races', () => { + beforeEach(() => { + setSessions(() => [sessionInfo()]) + $composerAttachments.set([]) + }) + + afterEach(() => { + cleanup() + $composerAttachments.set([]) + $connection.set(null) + vi.restoreAllMocks() + }) + + it('joins an in-flight eager upload at submit instead of staging the file twice', async () => { + // Drop-then-immediately-Enter: the drop kicks off an eager file.attach; if + // submit doesn't join it, both calls stage the file and leave a duplicate + // under .hermes/desktop-attachments/. Submit must await the in-flight upload + // and reuse its gateway-side ref. + $connection.set({ mode: 'remote' } as never) + Object.defineProperty(window, 'hermesDesktop', { + configurable: true, + value: { readFileDataUrl: vi.fn(async () => 'data:application/pdf;base64,JVBERi0=') } + }) + + let releaseAttach: () => void = () => {} + const methods: string[] = [] + const requestGateway = vi.fn(async (method: string) => { + methods.push(method) + if (method === 'file.attach') { + // Block until released so submit runs while the upload is in flight. + await new Promise(resolve => { + releaseAttach = resolve + }) + return { attached: true, ref_text: '@file:.hermes/desktop-attachments/doc.pdf', uploaded: true } as never + } + return {} as never + }) + + let handle: HarnessHandle | null = null + render( (handle = h)} refreshSessions={async () => undefined} requestGateway={requestGateway} />) + await waitFor(() => expect(handle).not.toBeNull()) + + // Drop a file → the eager effect fires file.attach and blocks on it. + $composerAttachments.set([{ id: 'file:doc.pdf', kind: 'file', label: 'doc.pdf', path: '/Users/me/doc.pdf' }]) + await waitFor(() => expect(methods.filter(m => m === 'file.attach').length).toBe(1)) + + // Submit reads the store, sees the upload in flight, and joins it. + const submitting = handle!.submitText('here you go') + releaseAttach() + + expect(await submitting).toBe(true) + // Exactly one file.attach (submit reused the eager result), then the send. + expect(methods.filter(m => m === 'file.attach').length).toBe(1) + expect(methods).toContain('prompt.submit') + }) +}) + describe('usePromptActions sleep/wake session recovery', () => { const STORED_SESSION_ID = 'stored-db-xyz789' const RECOVERED_SESSION_ID = 'rt-recovered-456' diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions.ts index 21eb53f01f..63e15f0518 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions.ts @@ -25,11 +25,11 @@ import { isProviderSetupErrorMessage } from '@/lib/provider-setup-errors' import { setSessionYolo } from '@/lib/yolo-session' import { $composerAttachments, - addComposerAttachment, clearComposerAttachments, type ComposerAttachment, setComposerAttachmentUploadState, - terminalContextBlocksFromDraft + terminalContextBlocksFromDraft, + updateComposerAttachment } from '@/store/composer' import { clearNotifications, notify, notifyError } from '@/store/notifications' import { requestDesktopOnboarding } from '@/store/onboarding' @@ -310,6 +310,11 @@ export function usePromptActions({ [selectedStoredSessionIdRef, updateSessionState] ) + // In-flight drop-time eager uploads, keyed by attachment id. Submit joins + // these before re-uploading so a drop-then-immediately-Enter can't fire + // file.attach twice and stage duplicate copies on the gateway. + const eagerUploadInFlight = useRef>>(new Map()) + const syncAttachmentsForSubmit = useCallback( async ( sessionId: string, @@ -320,7 +325,21 @@ export function usePromptActions({ const remote = $connection.get()?.mode === 'remote' const synced: ComposerAttachment[] = [] - for (const attachment of attachments) { + for (const original of attachments) { + let attachment = original + + // Join a drop-time eager upload still in flight for this attachment + // before deciding anything — otherwise submit and the eager task both + // call file.attach and stage duplicate files. After it settles, take the + // store's updated copy (its gateway ref, or its failure) over the stale + // pre-upload snapshot. + const inFlight = eagerUploadInFlight.current.get(attachment.id) + + if (inFlight) { + await inFlight + attachment = $composerAttachments.get().find(item => item.id === attachment.id) ?? attachment + } + // Already-synced or pathless refs (terminal, url, etc.) pass through. // A drop-time eager upload may already have staged this one (matching // attachedSessionId) — don't re-upload it. @@ -333,8 +352,9 @@ export function usePromptActions({ if (attachment.kind === 'image' || attachment.kind === 'file') { const nextAttachment = await uploadComposerAttachment(attachment, { remote, requestGateway, sessionId }) + // Update-only: never resurrect a chip the user removed mid-upload. if (updateComposerAttachments) { - addComposerAttachment(nextAttachment) + updateComposerAttachment(nextAttachment) } synced.push(nextAttachment) @@ -367,7 +387,9 @@ export function usePromptActions({ setComposerAttachmentUploadState(attachment.id, 'uploading') try { - addComposerAttachment(await uploadComposerAttachment(attachment, { remote, requestGateway, sessionId })) + // Update-only: if the user removed the chip while this was uploading, + // don't resurrect it — just drop the staged result on the floor. + updateComposerAttachment(await uploadComposerAttachment(attachment, { remote, requestGateway, sessionId })) } catch (err) { // Leave the chip in place so submit-time sync can retry (or the user can // remove it) and flag the card; also toast so a hard failure (unreadable @@ -380,7 +402,6 @@ export function usePromptActions({ ) const composerAttachments = useStore($composerAttachments) - const eagerUploadInFlight = useRef>(new Set()) useEffect(() => { if (!activeSessionId) { @@ -399,10 +420,11 @@ export function usePromptActions({ continue } - eagerUploadInFlight.current.add(attachment.id) - void eagerlyUploadAttachment(activeSessionId, attachment).finally(() => + const task = eagerlyUploadAttachment(activeSessionId, attachment).finally(() => eagerUploadInFlight.current.delete(attachment.id) ) + + eagerUploadInFlight.current.set(attachment.id, task) } }, [activeSessionId, composerAttachments, eagerlyUploadAttachment]) diff --git a/apps/desktop/src/store/composer.test.ts b/apps/desktop/src/store/composer.test.ts new file mode 100644 index 0000000000..83f0a3feb9 --- /dev/null +++ b/apps/desktop/src/store/composer.test.ts @@ -0,0 +1,43 @@ +import { afterEach, describe, expect, it } from 'vitest' + +import { + $composerAttachments, + addComposerAttachment, + type ComposerAttachment, + removeComposerAttachment, + updateComposerAttachment +} from './composer' + +function attachment(overrides: Partial & Pick): ComposerAttachment { + return { kind: 'file', label: 'doc.pdf', ...overrides } +} + +describe('updateComposerAttachment', () => { + afterEach(() => { + $composerAttachments.set([]) + }) + + it('replaces an existing attachment in place', () => { + addComposerAttachment(attachment({ id: 'file:a', uploadState: 'uploading' })) + + const updated = updateComposerAttachment(attachment({ id: 'file:a', attachedSessionId: 'sess-1' })) + + expect(updated).toBe(true) + const current = $composerAttachments.get() + expect(current).toHaveLength(1) + expect(current[0]?.attachedSessionId).toBe('sess-1') + expect(current[0]?.uploadState).toBeUndefined() + }) + + it('does NOT resurrect an attachment the user removed mid-upload', () => { + // Drop → eager upload starts → user removes the chip → upload resolves. + // The late success must not re-add the removed attachment. + addComposerAttachment(attachment({ id: 'file:a', uploadState: 'uploading' })) + removeComposerAttachment('file:a') + + const updated = updateComposerAttachment(attachment({ id: 'file:a', attachedSessionId: 'sess-1' })) + + expect(updated).toBe(false) + expect($composerAttachments.get()).toHaveLength(0) + }) +}) diff --git a/apps/desktop/src/store/composer.ts b/apps/desktop/src/store/composer.ts index 42621e1799..6b2b58ccb8 100644 --- a/apps/desktop/src/store/composer.ts +++ b/apps/desktop/src/store/composer.ts @@ -73,6 +73,25 @@ export function removeComposerAttachment(id: string): ComposerAttachment | null return removed } +/** Replace an existing attachment in place by id. No-op (returns false) when the + * id is gone — e.g. the user removed the chip while an eager upload was still in + * flight, so a late success must NOT resurrect it. Use this instead of + * addComposerAttachment for async results that may land after a removal. */ +export function updateComposerAttachment(attachment: ComposerAttachment): boolean { + const current = $composerAttachments.get() + const index = current.findIndex(item => item.id === attachment.id) + + if (index < 0) { + return false + } + + const next = [...current] + next[index] = attachment + $composerAttachments.set(next) + + return true +} + export function clearComposerAttachments() { $composerAttachments.set([]) }