fix(mobile): harden multi-image attachment flow

2026-10-20 16:03:20 +00:00 · 2026-08-05 00:25:50 -07:00
parent 7497b28304
commit 294332d4a9
10 changed files with 303 additions and 68 deletions
@@ -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<PickedMobileImage>
+): Promise<PickedMobileImage[]> {
+  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
+      })
    )
  })

@@ -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<PickedMobileImage[]> {
+): AsyncGenerator<PickedMobileImage> {
  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<PickedMobileImage[]> {
+): AsyncGenerator<PickedMobileImage> {
  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<PickedMobileImage[]> {
+): AsyncIterable<PickedMobileImage> {
  if (source === 'library') {
    return pickFromLibrary(
      multiple,
@@ -170,12 +166,15 @@ export async function pickMobileImage(
  source: MobileImageSource,
  deps?: MobileImagePickerDeps
 ): Promise<PickedMobileImage | null> {
-  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<PickedMobileImage[]> {
+): AsyncIterable<PickedMobileImage> {
  return pickMobileImagesWithMode(source, true, deps)
 }
@@ -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'] }
+    })
+  })
+})
@@ -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<string, Record<string, string[]>>,
+  sessionKey: string,
+  messages: readonly NativeChatMessage[]
+): Record<string, Record<string, string[]>> {
+  const sessionPreviews = previous[sessionKey]
+  if (!sessionPreviews) {
+    return previous
+  }
+  const messageIndexById = new Map(messages.map((message, index) => [message.id, index]))
+  let nextSession: Record<string, string[]> | 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)
@@ -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<RpcClient, 'sendRequest'> & {
  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')])

@@ -13,14 +13,35 @@ export type PendingNativeChatImage = {
  readonly previewUri: string
 }

+export function appendPendingNativeChatImages(
+  current: readonly PendingNativeChatImage[],
+  uploaded: readonly Omit<PendingNativeChatImage, 'id'>[],
+  idCounter: { current: number }
+): PendingNativeChatImage[] {
+  return [
+    ...current,
+    ...uploaded.map((image) => {
+      idCounter.current += 1
+      return { id: `img-${idCounter.current}`, ...image }
+    })
+  ]
+}
+
 export type UploadNativeChatImagesDeps = {
  readonly client: Pick<RpcClient, 'sendRequest'>
  readonly getConnectionId: () => Promise<string | null>
  // Injected so this module stays free of expo/react-native imports (unit-testable).
-  readonly pickImages: (source: MobileImageSource) => Promise<PickedMobileImage[]>
+  readonly pickImages: (
+    source: MobileImageSource
+  ) =>
+    | Iterable<PickedMobileImage>
+    | AsyncIterable<PickedMobileImage>
+    | Promise<Iterable<PickedMobileImage> | AsyncIterable<PickedMobileImage>>
  // 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<PendingNativeChatImage, 'id'>) => 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<Omit<PendingNativeChatImage, 'id'>[]> {
  const picked = await pickImages(source)
-  if (picked.length === 0) {
-    return []
-  }
-  onUploadStart?.()
-  const connectionId = await getConnectionId()
  const uploaded: Omit<PendingNativeChatImage, 'id'>[] = []
-  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 <Image> 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
 }
@@ -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)
@@ -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<PendingNativeChatImage, 'id'>[] = []
+      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)
-          }
-        }
      }
    },
    [
@@ -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' }])
@@ -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({