From 294332d4a9227e98394f0065ee48f8a98e43effa Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Wed, 5 Aug 2026 00:25:50 -0700 Subject: [PATCH] fix(mobile): harden multi-image attachment flow --- .../mobile-image-source-picker.test.ts | 31 ++++++-- .../src/session/mobile-image-source-picker.ts | 33 +++++---- ...mobile-native-chat-draft-reconcile.test.ts | 72 +++++++++++++++++++ .../mobile-native-chat-draft-reconcile.ts | 57 ++++++++++++++- ...obile-native-chat-image-attachment.test.ts | 57 +++++++++++++-- .../mobile-native-chat-image-attachment.ts | 47 +++++++++--- .../session/use-mobile-native-chat-drafts.ts | 17 +++-- ...se-mobile-native-chat-image-attachments.ts | 40 +++++------ ...tive-chat-image-transcript-markers.test.ts | 14 ++++ .../native-chat-image-transcript-markers.ts | 3 +- 10 files changed, 303 insertions(+), 68 deletions(-) create mode 100644 mobile/src/session/mobile-native-chat-draft-reconcile.test.ts diff --git a/mobile/src/session/mobile-image-source-picker.test.ts b/mobile/src/session/mobile-image-source-picker.test.ts index 4b46c118297..35e3cff3b60 100644 --- a/mobile/src/session/mobile-image-source-picker.test.ts +++ b/mobile/src/session/mobile-image-source-picker.test.ts @@ -15,7 +15,8 @@ vi.mock('expo-file-system', () => ({ import { ImageLibraryPermissionError, pickMobileImage, - pickMobileImages + pickMobileImages, + type PickedMobileImage } from './mobile-image-source-picker' const granted = { granted: true } as Awaited< @@ -23,6 +24,16 @@ const granted = { granted: true } as Awaited< > const denied = { granted: false } as typeof granted +async function collectImages( + images: AsyncIterable +): Promise { + const collected: PickedMobileImage[] = [] + for await (const image of images) { + collected.push(image) + } + return collected +} + function fileFactory( bytes: Uint8Array, options?: { fileSize?: number; handleSize?: number | null; readError?: Error } @@ -101,11 +112,13 @@ describe('pickMobileImage', () => { assets: [...bytesByUri].map(([uri, bytes]) => ({ uri, fileSize: bytes.length })) }) - const result = await pickMobileImages('library', { - requestLibraryPermission: vi.fn().mockResolvedValue(granted), - launchLibrary, - createFile - }) + const result = await collectImages( + pickMobileImages('library', { + requestLibraryPermission: vi.fn().mockResolvedValue(granted), + launchLibrary, + createFile + }) + ) expect(result.map((image) => image.uri)).toEqual([ 'file:///a.jpg', @@ -113,7 +126,11 @@ describe('pickMobileImage', () => { 'file:///c.jpg' ]) expect(launchLibrary).toHaveBeenCalledWith( - expect.objectContaining({ allowsMultipleSelection: true, selectionLimit: 0 }) + expect.objectContaining({ + allowsMultipleSelection: true, + orderedSelection: true, + selectionLimit: 0 + }) ) }) diff --git a/mobile/src/session/mobile-image-source-picker.ts b/mobile/src/session/mobile-image-source-picker.ts index 91bf11e687a..062085f7270 100644 --- a/mobile/src/session/mobile-image-source-picker.ts +++ b/mobile/src/session/mobile-image-source-picker.ts @@ -83,12 +83,12 @@ async function readUriAsBase64( } } -async function pickFromLibrary( +async function* pickFromLibrary( multiple: boolean, requestPermission: typeof ImagePicker.requestMediaLibraryPermissionsAsync = ImagePicker.requestMediaLibraryPermissionsAsync, launch: typeof ImagePicker.launchImageLibraryAsync = ImagePicker.launchImageLibraryAsync, createFile: MobileImageFileFactory = defaultMobileImageFileFactory -): Promise { +): AsyncGenerator { const permission = await requestPermission() // Why: `granted` covers full + limited iOS access; only a hard denial blocks us. if (!permission.granted) { @@ -98,49 +98,45 @@ async function pickFromLibrary( mediaTypes: ['images'], base64: false, allowsMultipleSelection: multiple, - ...(multiple ? { selectionLimit: 0 } : {}), + ...(multiple ? { selectionLimit: 0, orderedSelection: true } : {}), quality: 1 }) if (result.canceled) { - return [] + return } - const picked: PickedMobileImage[] = [] for (const asset of result.assets) { if (!asset.uri) { continue } const base64 = await readUriAsBase64(asset.uri, asset.fileSize, createFile) if (base64) { - picked.push({ base64, uri: asset.uri }) + yield { base64, uri: asset.uri } } } - return picked } -async function pickFromFiles( +async function* pickFromFiles( multiple: boolean, launch: typeof DocumentPicker.getDocumentAsync = DocumentPicker.getDocumentAsync, createFile: MobileImageFileFactory = defaultMobileImageFileFactory -): Promise { +): AsyncGenerator { const result = await launch({ type: 'image/*', multiple, copyToCacheDirectory: true }) if (result.canceled) { - return [] + return } - const picked: PickedMobileImage[] = [] for (const asset of result.assets) { if (!asset.uri) { continue } const base64 = await readUriAsBase64(asset.uri, asset.size, createFile) if (base64) { - picked.push({ base64, uri: asset.uri }) + yield { base64, uri: asset.uri } } } - return picked } type MobileImagePickerDeps = { @@ -150,11 +146,11 @@ type MobileImagePickerDeps = { readonly createFile?: MobileImageFileFactory } -async function pickMobileImagesWithMode( +function pickMobileImagesWithMode( source: MobileImageSource, multiple: boolean, deps?: MobileImagePickerDeps -): Promise { +): AsyncIterable { if (source === 'library') { return pickFromLibrary( multiple, @@ -170,12 +166,15 @@ export async function pickMobileImage( source: MobileImageSource, deps?: MobileImagePickerDeps ): Promise { - return (await pickMobileImagesWithMode(source, false, deps))[0] ?? null + for await (const image of pickMobileImagesWithMode(source, false, deps)) { + return image + } + return null } export function pickMobileImages( source: MobileImageSource, deps?: MobileImagePickerDeps -): Promise { +): AsyncIterable { return pickMobileImagesWithMode(source, true, deps) } diff --git a/mobile/src/session/mobile-native-chat-draft-reconcile.test.ts b/mobile/src/session/mobile-native-chat-draft-reconcile.test.ts new file mode 100644 index 00000000000..4c0cee804c1 --- /dev/null +++ b/mobile/src/session/mobile-native-chat-draft-reconcile.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from 'vitest' +import type { NativeChatMessage } from '../../../src/shared/native-chat-types' +import { + findLandedImagePreviewEchoes, + migrateImagePreviewMessageIds, + type PendingImagePreviewEcho +} from './mobile-native-chat-draft-reconcile' + +function userText(id: string, text: string): NativeChatMessage { + return { + id, + role: 'user', + blocks: [{ type: 'text', text }], + timestamp: null, + source: 'transcript' + } +} + +function pending(id: string, images: string[], expectedOccurrence = 1): PendingImagePreviewEcho { + return { id, text: '', images, expectedOccurrence, baselineTailMessageId: null } +} + +describe('mobile native chat image preview reconciliation', () => { + it('keeps separate adjacent image-only sends independently reconcilable', () => { + const landed = findLandedImagePreviewEchoes( + [ + userText('source-a', '[Image: source: /tmp/a.png]'), + userText('source-b', '[Image: source: /tmp/b.png]') + ], + [pending('pending-a', ['file:///a.jpg']), pending('pending-b', ['file:///b.jpg'], 2)] + ) + + expect(landed).toEqual([ + { pendingId: 'pending-a', messageId: 'source-a', images: ['file:///a.jpg'] }, + { pendingId: 'pending-b', messageId: 'source-b', images: ['file:///b.jpg'] } + ]) + }) + + it('waits for a complete multi-image turn as transcript source frames stream in', () => { + const entry = pending('pending', ['file:///a.jpg', 'file:///b.jpg']) + const sourceA = userText('source-a', '[Image: source: /tmp/a.png]') + const sourceB = userText('source-b', '[Image: source: /tmp/b.png]') + + expect(findLandedImagePreviewEchoes([sourceA], [entry])).toEqual([]) + expect(findLandedImagePreviewEchoes([sourceA, sourceB], [entry])).toEqual([]) + expect( + findLandedImagePreviewEchoes( + [sourceA, sourceB, userText('prompt', '[Image #1] [Image #2]')], + [entry] + ) + ).toEqual([ + { + pendingId: 'pending', + messageId: 'prompt', + images: ['file:///a.jpg', 'file:///b.jpg'] + } + ]) + }) + + it('moves an early standalone preview to the later folded prompt id', () => { + const sessionKey = 'host\0worktree\0tab\0session' + const previous = { [sessionKey]: { source: ['file:///a.jpg'] } } + const messages = [ + userText('source', '[Image: source: /tmp/a.png]'), + userText('prompt', '[Image #1]') + ] + + expect(migrateImagePreviewMessageIds(previous, sessionKey, messages)).toEqual({ + [sessionKey]: { prompt: ['file:///a.jpg'] } + }) + }) +}) diff --git a/mobile/src/session/mobile-native-chat-draft-reconcile.ts b/mobile/src/session/mobile-native-chat-draft-reconcile.ts index 17fc5f27afa..360fd22587b 100644 --- a/mobile/src/session/mobile-native-chat-draft-reconcile.ts +++ b/mobile/src/session/mobile-native-chat-draft-reconcile.ts @@ -101,6 +101,60 @@ export function mergeLandedImagePreviewEchoes( return next } +function imagePreviewReplacementMessageId( + messages: readonly NativeChatMessage[], + sourceIndex: number +): string | null { + const source = messages[sourceIndex] + if (!source || !isImageSourceUserTurn(source)) { + return null + } + let nextIndex = sourceIndex + 1 + while ( + messages[nextIndex]?.source === source.source && + isImageSourceUserTurn(messages[nextIndex]!) + ) { + nextIndex++ + } + const prompt = messages[nextIndex] + const firstText = prompt?.blocks.find((block) => block.type === 'text') + return prompt?.role === 'user' && + prompt.source === source.source && + firstText?.type === 'text' && + stripImagePromptMarker(firstText.text) !== firstText.text + ? prompt.id + : null +} + +/** Moves previews forward when a progressive source-only transcript frame later + * folds into the marker-prefixed prompt with a different authoritative id. */ +export function migrateImagePreviewMessageIds( + previous: Record>, + sessionKey: string, + messages: readonly NativeChatMessage[] +): Record> { + const sessionPreviews = previous[sessionKey] + if (!sessionPreviews) { + return previous + } + const messageIndexById = new Map(messages.map((message, index) => [message.id, index])) + let nextSession: Record | null = null + for (const [messageId, images] of Object.entries(sessionPreviews)) { + const sourceIndex = messageIndexById.get(messageId) + if (sourceIndex === undefined) { + continue + } + const replacementId = imagePreviewReplacementMessageId(messages, sourceIndex) + if (!replacementId) { + continue + } + nextSession ??= { ...sessionPreviews } + delete nextSession[messageId] + nextSession[replacementId] = [...(nextSession[replacementId] ?? []), ...images] + } + return nextSession ? { ...previous, [sessionKey]: nextSession } : previous +} + /** Binds local preview URIs to the authoritative transcript turn that replaced * the optimistic bubble. Host paths and marker-only Codex turns cannot render * the phone-local photo without this handoff. */ @@ -125,7 +179,8 @@ export function findLandedImagePreviewEchoes( if (targetText) { return normalizedUserText(message) === targetText } - return message.blocks.length === 0 || message.blocks.some(isImageRefBlock) + const imageCount = message.blocks.filter(isImageRefBlock).length + return message.blocks.length === 0 || imageCount >= entry.images!.length }) const tailIndex = entry.baselineTailMessageId ? messageIndexById.get(entry.baselineTailMessageId) diff --git a/mobile/src/session/mobile-native-chat-image-attachment.test.ts b/mobile/src/session/mobile-native-chat-image-attachment.test.ts index b5c8c2a62c9..e61116f3caa 100644 --- a/mobile/src/session/mobile-native-chat-image-attachment.test.ts +++ b/mobile/src/session/mobile-native-chat-image-attachment.test.ts @@ -16,6 +16,10 @@ function methodNotFound(id: string): RpcResponse { } } +function failed(id: string, message: string): RpcResponse { + return { id, ok: false, error: { code: 'failed', message }, _meta: { runtimeId: 'r' } } +} + function clientWithResponses(responses: RpcResponse[]): Pick & { calls: { method: string; params: unknown }[] } { @@ -63,14 +67,22 @@ describe('uploadMobileNativeChatImages', () => { ok('save-c', '/tmp/c.png') ]) - const result = await uploadMobileNativeChatImages('library', { - client, - getConnectionId: async () => 'conn-7', - pickImages: vi.fn().mockResolvedValue([ + const order: string[] = [] + async function* pickImages() { + for (const image of [ { base64: 'AAAA', uri: 'file:///a.jpg' }, { base64: 'BBBB', uri: 'file:///b.jpg' }, { base64: 'CCCC', uri: 'file:///c.jpg' } - ]) + ]) { + order.push(`read:${image.uri}`) + yield image + } + } + const result = await uploadMobileNativeChatImages('library', { + client, + getConnectionId: async () => 'conn-7', + pickImages, + onImageUploaded: (image) => order.push(`uploaded:${image.previewUri}`) }) expect(result).toEqual([ @@ -78,6 +90,14 @@ describe('uploadMobileNativeChatImages', () => { { path: '/tmp/b.png', previewUri: 'file:///b.jpg' }, { path: '/tmp/c.png', previewUri: 'file:///c.jpg' } ]) + expect(order).toEqual([ + 'read:file:///a.jpg', + 'uploaded:file:///a.jpg', + 'read:file:///b.jpg', + 'uploaded:file:///b.jpg', + 'read:file:///c.jpg', + 'uploaded:file:///c.jpg' + ]) }) it('returns null when the picker is cancelled and uploads nothing', async () => { @@ -93,6 +113,33 @@ describe('uploadMobileNativeChatImages', () => { expect(client.calls).toEqual([]) }) + it('reports completed uploads before a later image fails', async () => { + const client = clientWithResponses([ + methodNotFound('start-a'), + ok('save-a', '/tmp/a.png'), + methodNotFound('start-b'), + failed('save-b', 'upload failed') + ]) + const onImageUploaded = vi.fn() + + await expect( + uploadMobileNativeChatImages('library', { + client, + getConnectionId: async () => null, + pickImages: vi.fn().mockResolvedValue([ + { base64: 'AAAA', uri: 'file:///a.jpg' }, + { base64: 'BBBB', uri: 'file:///b.jpg' } + ]), + onImageUploaded + }) + ).rejects.toThrow('upload failed') + expect(onImageUploaded).toHaveBeenCalledOnce() + expect(onImageUploaded).toHaveBeenCalledWith({ + path: '/tmp/a.png', + previewUri: 'file:///a.jpg' + }) + }) + it('falls back to an inline data uri for the preview when the picker omits a uri', async () => { const client = clientWithResponses([methodNotFound('start'), ok('save', '/tmp/x.png')]) diff --git a/mobile/src/session/mobile-native-chat-image-attachment.ts b/mobile/src/session/mobile-native-chat-image-attachment.ts index 7aa61259d79..fc9a962bacf 100644 --- a/mobile/src/session/mobile-native-chat-image-attachment.ts +++ b/mobile/src/session/mobile-native-chat-image-attachment.ts @@ -13,14 +13,35 @@ export type PendingNativeChatImage = { readonly previewUri: string } +export function appendPendingNativeChatImages( + current: readonly PendingNativeChatImage[], + uploaded: readonly Omit[], + idCounter: { current: number } +): PendingNativeChatImage[] { + return [ + ...current, + ...uploaded.map((image) => { + idCounter.current += 1 + return { id: `img-${idCounter.current}`, ...image } + }) + ] +} + export type UploadNativeChatImagesDeps = { readonly client: Pick readonly getConnectionId: () => Promise // Injected so this module stays free of expo/react-native imports (unit-testable). - readonly pickImages: (source: MobileImageSource) => Promise + readonly pickImages: ( + source: MobileImageSource + ) => + | Iterable + | AsyncIterable + | Promise | AsyncIterable> // Fired once the user has picked an image and the host upload is about to start — // lets the UI show the attach spinner only for the transfer, not the picker. readonly onUploadStart?: () => void + /** Retains each completed upload if a later image in the same selection fails. */ + readonly onImageUploaded?: (image: Omit) => void } /** Picks an image and uploads it to the host, returning the host path + a local @@ -30,21 +51,29 @@ export type UploadNativeChatImagesDeps = { * Returns an empty array when the user cancels the picker. */ export async function uploadMobileNativeChatImages( source: MobileImageSource, - { client, getConnectionId, pickImages, onUploadStart }: UploadNativeChatImagesDeps + { + client, + getConnectionId, + pickImages, + onUploadStart, + onImageUploaded + }: UploadNativeChatImagesDeps ): Promise[]> { const picked = await pickImages(source) - if (picked.length === 0) { - return [] - } - onUploadStart?.() - const connectionId = await getConnectionId() const uploaded: Omit[] = [] - for (const image of picked) { + let connectionId: string | null = null + for await (const image of picked) { + if (uploaded.length === 0) { + onUploadStart?.() + connectionId = await getConnectionId() + } const path = await saveMobileClipboardImageAsTempFile(client, image.base64, { connectionId }) // Prefer the picker's local URI for the thumbnail; fall back to an inline data // URI when the source omitted one (RN renders both). const previewUri = image.uri ?? `data:image/png;base64,${image.base64}` - uploaded.push({ path, previewUri }) + const result = { path, previewUri } + uploaded.push(result) + onImageUploaded?.(result) } return uploaded } diff --git a/mobile/src/session/use-mobile-native-chat-drafts.ts b/mobile/src/session/use-mobile-native-chat-drafts.ts index 88fb9002124..df81c5d0b0b 100644 --- a/mobile/src/session/use-mobile-native-chat-drafts.ts +++ b/mobile/src/session/use-mobile-native-chat-drafts.ts @@ -6,6 +6,7 @@ import { findLandedImagePreviewEchoes, findLandedUnconfirmedSends, mergeLandedImagePreviewEchoes, + migrateImagePreviewMessageIds, normalizedUserText, type UnconfirmedSend } from './mobile-native-chat-draft-reconcile' @@ -259,9 +260,7 @@ export function useMobileNativeChatDrafts(args: { const landedSet = new Set(landed) unconfirmedRef.current = unconfirmedRef.current.filter((entry) => !landedSet.has(entry)) for (const entry of landed) { - if (entry.deadline !== null) { - clearTimeout(entry.deadline) - } + clearTimeout(entry.deadline ?? undefined) } }, [messages, draftKey, pendingKey]) @@ -270,9 +269,7 @@ export function useMobileNativeChatDrafts(args: { return () => { mountedRef.current = false for (const entry of unconfirmedRef.current) { - if (entry.deadline !== null) { - clearTimeout(entry.deadline) - } + clearTimeout(entry.deadline ?? undefined) } unconfirmedRef.current = [] } @@ -282,7 +279,13 @@ export function useMobileNativeChatDrafts(args: { ? (pendingBySession[pendingKey] ?? NO_PENDING_MESSAGES) : NO_PENDING_MESSAGES useEffect(() => { - if (!pendingKey || pending.length === 0) { + if (!pendingKey) { + return + } + setImagePreviewsBySession((previous) => + migrateImagePreviewMessageIds(previous, pendingKey, messages) + ) + if (pending.length === 0) { return } const landedImagePreviews = findLandedImagePreviewEchoes(messages, pending) diff --git a/mobile/src/session/use-mobile-native-chat-image-attachments.ts b/mobile/src/session/use-mobile-native-chat-image-attachments.ts index 3060851d10a..dbcb97df527 100644 --- a/mobile/src/session/use-mobile-native-chat-image-attachments.ts +++ b/mobile/src/session/use-mobile-native-chat-image-attachments.ts @@ -9,6 +9,7 @@ import { type MobileImageSource } from './mobile-image-source-picker' import { + appendPendingNativeChatImages, uploadMobileNativeChatImages, type PendingNativeChatImage } from './mobile-native-chat-image-attachment' @@ -142,38 +143,45 @@ export function useMobileNativeChatImageAttachments({ // pick or pre-upload error never ran `onUploadStart`, so decrementing the // shared counter would clear a concurrent upload's in-flight flag early. let started = false + const uploadedImages: Omit[] = [] + let uploadError: unknown = null try { - const uploaded = await uploadMobileNativeChatImages(source, { + await uploadMobileNativeChatImages(source, { client, getConnectionId: getActiveWorktreeConnectionId, pickImages: pickMobileImages, + onImageUploaded: (image) => uploadedImages.push(image), onUploadStart: () => { started = true attachingCount.current += 1 setIsAttaching(true) } }) - // Cancelled picker: no error, no toast. - if (uploaded.length === 0) { - return + } catch (error) { + uploadError = error + } finally { + if (started) { + attachingCount.current -= 1 + if (attachingCount.current === 0) { + setIsAttaching(false) + } } - const chips = uploaded.map((image) => { - idCounter.current += 1 - return { id: `img-${idCounter.current}`, ...image } - }) + } + if (uploadedImages.length > 0) { setAttachmentsByScope((prev) => ({ ...prev, - [scope]: [...(prev[scope] ?? []), ...chips] + [scope]: appendPendingNativeChatImages(prev[scope] ?? [], uploadedImages, idCounter) })) onAttachSuccess?.() - } catch (error) { - const message = error instanceof Error ? error.message : String(error) + } + if (uploadError !== null) { + const message = uploadError instanceof Error ? uploadError.message : String(uploadError) onError?.() if (connStateRef.current !== 'connected') { showToast('Attach failed (disconnected)', 1500) return } - if (error instanceof ImageLibraryPermissionError) { + if (uploadError instanceof ImageLibraryPermissionError) { showToast('Photo permission denied', 1500) return } @@ -182,14 +190,6 @@ export function useMobileNativeChatImageAttachments({ return } showToast('Attach failed', 1500) - } finally { - if (started) { - attachingCount.current -= 1 - if (attachingCount.current <= 0) { - attachingCount.current = 0 - setIsAttaching(false) - } - } } }, [ diff --git a/src/shared/native-chat-image-transcript-markers.test.ts b/src/shared/native-chat-image-transcript-markers.test.ts index 0201f5cdab3..e4856605a74 100644 --- a/src/shared/native-chat-image-transcript-markers.test.ts +++ b/src/shared/native-chat-image-transcript-markers.test.ts @@ -67,6 +67,20 @@ describe('normalizeImageTranscriptMessages', () => { ]) }) + it('preserves adjacent standalone image turns without a prompt marker', () => { + const out = normalizeImageTranscriptMessages([ + userText('a', '[Image: source: /tmp/a.png]'), + userText('b', '[Image: source: /tmp/b.png]') + ]) + + expect(out).toHaveLength(2) + expect(out.map((message) => message.id)).toEqual(['a', 'b']) + expect(out.map((message) => message.blocks)).toEqual([ + [{ type: 'image-ref', path: '/tmp/a.png' }], + [{ type: 'image-ref', path: '/tmp/b.png' }] + ]) + }) + it('leaves ordinary user text untouched', () => { const out = normalizeImageTranscriptMessages([userText('a', 'how about this')]) expect(out[0]!.blocks).toEqual([{ type: 'text', text: 'how about this' }]) diff --git a/src/shared/native-chat-image-transcript-markers.ts b/src/shared/native-chat-image-transcript-markers.ts index b0c1b5c7246..9643cf518e0 100644 --- a/src/shared/native-chat-image-transcript-markers.ts +++ b/src/shared/native-chat-image-transcript-markers.ts @@ -84,9 +84,8 @@ export function normalizeImageTranscriptMessages( } normalized.push({ ...message, - blocks: imagePaths.map((path) => ({ type: 'image-ref' as const, path })) + blocks: [{ type: 'image-ref', path: imagePath }] }) - index = nextIndex - 1 continue } normalized.push({