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.
This commit is contained in:
Brooklyn Nicholson 2026-06-09 18:21:10 -05:00
parent 153060e206
commit 891c9a6823
4 changed files with 150 additions and 10 deletions

View file

@ -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<void>(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(<Harness onReady={h => (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'

View file

@ -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<Map<string, Promise<void>>>(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<Set<string>>(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])

View file

@ -0,0 +1,43 @@
import { afterEach, describe, expect, it } from 'vitest'
import {
$composerAttachments,
addComposerAttachment,
type ComposerAttachment,
removeComposerAttachment,
updateComposerAttachment
} from './composer'
function attachment(overrides: Partial<ComposerAttachment> & Pick<ComposerAttachment, 'id'>): 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)
})
})

View file

@ -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([])
}