Gate draft image ownership checks on relevant local attachments (#32098)

src/renderer/src/components/native-chat/native-chat-composer-attachments-test-probe.tsx

@@ -0,0 +1,114 @@
+import { act, createElement, useEffect, useRef, useState } from 'react'
+import { createRoot, type Root } from 'react-dom/client'
+import type { NativeChatResolvedTarget } from './native-chat-composer-target'
+import { useNativeChatComposerAttachments } from './use-native-chat-composer-attachments'
+
+type AttachmentApi = ReturnType<typeof useNativeChatComposerAttachments>
+type ProbeApi = AttachmentApi & { adoptDraft: (draft: string) => void }
+
+const target: NativeChatResolvedTarget = {
+  ptyId: 'pty-1'
+}
+
+function Probe({
+  scopeKey,
+  structured = false,
+  disabled = false,
+  isComposing,
+  onReady
+}: {
+  scopeKey: string
+  structured?: boolean
+  disabled?: boolean
+  isComposing: () => boolean
+  onReady: (api: ProbeApi) => void
+}): React.JSX.Element {
+  const [caret, setCaret] = useState(0)
+  const [draftValue, setDraftValue] = useState('')
+  const [notice, setNotice] = useState<string | null>(null)
+  const textareaRef = useRef<HTMLTextAreaElement>(null)
+  const api = useNativeChatComposerAttachments({
+    attachmentScopeKey: scopeKey,
+    allowWithoutTarget: structured,
+    imageReferences: structured,
+    caret,
+    disabled,
+    isComposing,
+    resolveTarget: () => (structured ? null : target),
+    textareaRef,
+    setCaret,
+    setDraft: (updater) => setDraftValue((previous) => updater(previous)),
+    setNotice
+  })
+  useEffect(() => {
+    onReady({ ...api, adoptDraft: setDraftValue })
+  }, [api, onReady])
+  return (
+    <div>
+      <textarea ref={textareaRef} />
+      <output data-draft>{draftValue}</output>
+      <output data-notice>{notice}</output>
+    </div>
+  )
+}
+
+export async function renderProbe(
+  scopeKey: string,
+  structured = false,
+  options: { disabled?: boolean; isComposing?: () => boolean } = {}
+): Promise<{
+  draft: () => string
+  latest: () => ProbeApi
+  notice: () => string
+  rerender: (scopeKey: string, disabled?: boolean) => Promise<void>
+  root: Root
+  textarea: () => HTMLTextAreaElement
+}> {
+  const container = document.createElement('div')
+  document.body.append(container)
+  // onReady fires on every render, so keep the freshest snapshot — reading a
+  // single captured `api` would go stale after attach/remove triggers a render.
+  let api: ProbeApi | null = null
+  const root = createRoot(container)
+  const onReady = (next: ProbeApi): void => {
+    api = next
+  }
+  const isComposing = options.isComposing ?? (() => false)
+  const render = async (nextScopeKey: string, disabled: boolean): Promise<void> => {
+    await act(async () => {
+      root.render(
+        createElement(Probe, {
+          scopeKey: nextScopeKey,
+          structured,
+          disabled,
+          isComposing,
+          onReady
+        })
+      )
+    })
+  }
+  await render(scopeKey, options.disabled ?? false)
+  if (!api) {
+    throw new Error('Probe did not render')
+  }
+  return {
+    draft: () => container.querySelector('[data-draft]')?.textContent ?? '',
+    root,
+    latest: () => {
+      if (!api) {
+        throw new Error('Probe is not mounted')
+      }
+      return api
+    },
+    notice: () => container.querySelector('[data-notice]')?.textContent ?? '',
+    rerender: (nextScopeKey: string, disabled = options.disabled ?? false) =>
+      render(nextScopeKey, disabled),
+    textarea: () => {
+      const textarea = container.querySelector('textarea')
+      if (!textarea) {
+        throw new Error('Probe textarea is not mounted')
+      }
+      return textarea
+    }
+  }
+}

src/renderer/src/components/native-chat/native-chat-composer-draft-image-check.ts

@@ -65,11 +65,18 @@ export async function findMissingNativeChatComposerDraftImages(
   return missing
 }
 
-function localStoreAuthority(scopeKey: string): string | null {
+function localStoreAuthority(scopeKey: string, checks: readonly DraftImageCheck[]): string | null {
+  if (
+    structuredAgentSessionIdOfDraftScope(scopeKey) === null ||
+    !checks.some(
+      ({ image }) => isAgentSessionAttachmentStorePath(image.path) && !image.connectionId
+    )
+  ) {
+    return null
+  }
   const stored = readNativeChatComposerDraft(scopeKey).owner
   const current = resolveNativeChatComposerDraftOwner(scopeKey)
-  return structuredAgentSessionIdOfDraftScope(scopeKey) !== null &&
-    (stored ?? current)?.executionHostId === 'local' &&
+  return (stored ?? current)?.executionHostId === 'local' &&
     (!current || current.executionHostId === 'local')
     ? JSON.stringify([stored, current])
     : null
@@ -105,7 +112,7 @@ function attemptIsCurrent(scopeKey: string, entry: DraftImageCheck): boolean {
     readNativeChatComposerDraft(scopeKey).images.some(
       (image) => restoredImageKey(image) === restoredImageKey(entry.image)
     ) &&
-    entry.attempt.authority === checkAuthority(entry.image, localStoreAuthority(scopeKey))
+    entry.attempt.authority === checkAuthority(entry.image, localStoreAuthority(scopeKey, [entry]))
   )
 }
 
@@ -135,7 +142,7 @@ async function runCheck(scopeKey: string, entry: DraftImageCheck): Promise<void>
 export async function verifyRestoredNativeChatComposerDraftImages(scopeKey: string): Promise<void> {
   await Promise.all(
     (restoredImageChecks.get(scopeKey) ?? NO_IMAGE_CHECKS).map(async (entry) => {
-      const authority = checkAuthority(entry.image, localStoreAuthority(scopeKey))
+      const authority = checkAuthority(entry.image, localStoreAuthority(scopeKey, [entry]))
       if (entry.attempt && entry.attempt.authority !== authority) {
         replaceCheck(scopeKey, entry, { image: entry.image })
       } else if (!entry.attempt && authority !== null) {
@@ -192,7 +199,7 @@ export function useRestoredNativeChatComposerDraftImageCheck(
     subscribe,
     () => restoredImageChecks.get(scopeKey) ?? NO_IMAGE_CHECKS
   )
-  const local = useAppStore(() => localStoreAuthority(scopeKey))
+  const local = useAppStore(() => localStoreAuthority(scopeKey, checks))
   const pending = useMemo(
     () =>
       new Set(

src/renderer/src/components/native-chat/use-native-chat-composer-attachments.test.tsx

@@ -1,15 +1,13 @@
 // @vitest-environment happy-dom
 import { afterEach, describe, expect, it, vi } from 'vitest'
-import { act, createElement, useEffect, useRef, useState } from 'react'
-import { createRoot, type Root } from 'react-dom/client'
+import { act } from 'react'
+import { renderProbe } from './native-chat-composer-attachments-test-probe'
 import {
   appendNativeChatAttachmentCache,
   clearNativeChatAttachmentCacheForTests,
-  readNativeChatAttachmentCache,
-  useNativeChatComposerAttachments
+  readNativeChatAttachmentCache
 } from './use-native-chat-composer-attachments'
 import * as draftCache from './native-chat-draft-cache'
-import type { NativeChatResolvedTarget } from './native-chat-composer-target'
 import { nativeChatPendingAttachmentSnapshot } from './native-chat-pending-attachment-cache'
 import { readNativeChatDraftCache } from './native-chat-draft-cache'
 import { NATIVE_FILE_DROP_MAX_PATHS } from '../../../../shared/native-file-drop'
@@ -31,113 +29,27 @@ vi.mock('../../../../shared/remote-runtime-pty-id', async (importOriginal) => ({
   isRemoteRuntimePtyId: () => runtimeTarget.remote
 }))
 
-type AttachmentApi = ReturnType<typeof useNativeChatComposerAttachments>
-type ProbeApi = AttachmentApi & { adoptDraft: (draft: string) => void }
-
-const target: NativeChatResolvedTarget = {
-  ptyId: 'pty-1'
-}
-
-function Probe({
-  scopeKey,
-  structured = false,
-  disabled = false,
-  isComposing,
-  onReady
-}: {
-  scopeKey: string
-  structured?: boolean
-  disabled?: boolean
-  isComposing: () => boolean
-  onReady: (api: ProbeApi) => void
-}): React.JSX.Element {
-  const [caret, setCaret] = useState(0)
-  const [draftValue, setDraftValue] = useState('')
-  const [notice, setNotice] = useState<string | null>(null)
-  const textareaRef = useRef<HTMLTextAreaElement>(null)
-  const api = useNativeChatComposerAttachments({
-    attachmentScopeKey: scopeKey,
-    allowWithoutTarget: structured,
-    imageReferences: structured,
-    caret,
-    disabled,
-    isComposing,
-    resolveTarget: () => (structured ? null : target),
-    textareaRef,
-    setCaret,
-    setDraft: (updater) => setDraftValue((previous) => updater(previous)),
-    setNotice
-  })
-  useEffect(() => {
-    onReady({ ...api, adoptDraft: setDraftValue })
-  }, [api, onReady])
-  return (
-    <div>
-      <textarea ref={textareaRef} />
-      <output data-draft>{draftValue}</output>
-      <output data-notice>{notice}</output>
-    </div>
-  )
-}
-
-async function renderProbe(
-  scopeKey: string,
-  structured = false,
-  options: { disabled?: boolean; isComposing?: () => boolean } = {}
-): Promise<{
-  draft: () => string
-  latest: () => ProbeApi
-  notice: () => string
-  rerender: (scopeKey: string, disabled?: boolean) => Promise<void>
-  root: Root
-  textarea: () => HTMLTextAreaElement
-}> {
-  const container = document.createElement('div')
-  document.body.append(container)
-  // onReady fires on every render, so keep the freshest snapshot — reading a
-  // single captured `api` would go stale after attach/remove triggers a render.
-  let api: ProbeApi | null = null
-  const root = createRoot(container)
-  const onReady = (next: ProbeApi): void => {
-    api = next
-  }
-  const isComposing = options.isComposing ?? (() => false)
-  const render = async (nextScopeKey: string, disabled: boolean): Promise<void> => {
-    await act(async () => {
-      root.render(
-        createElement(Probe, {
-          scopeKey: nextScopeKey,
-          structured,
-          disabled,
-          isComposing,
-          onReady
-        })
-      )
-    })
-  }
-  await render(scopeKey, options.disabled ?? false)
-  if (!api) {
-    throw new Error('Probe did not render')
-  }
-  return {
-    draft: () => container.querySelector('[data-draft]')?.textContent ?? '',
-    root,
-    latest: () => {
-      if (!api) {
-        throw new Error('Probe is not mounted')
+function publishUnrelatedStatuses(): void {
+  for (let tick = 1; tick <= 24; tick += 1) {
+    useAppStore.setState((state) => ({
+      runtimeStatusByEnvironmentId: new Map(state.runtimeStatusByEnvironmentId).set('other-host', {
+        status: null,
+        checkedAt: tick
+      })
+    }))
+    useAppStore.setState((state) => ({
+      agentStatusByPaneKey: {
+        ...state.agentStatusByPaneKey,
+        'other-tab:other-pane': {
+          paneKey: 'other-tab:other-pane',
+          state: 'working',
+          prompt: 'background work',
+          updatedAt: tick,
+          stateStartedAt: 1,
+          stateHistory: []
+        }
       }
-      return api
-    },
-    notice: () => container.querySelector('[data-notice]')?.textContent ?? '',
-    rerender: (nextScopeKey: string, disabled = options.disabled ?? false) =>
-      render(nextScopeKey, disabled),
-    textarea: () => {
-      const textarea = container.querySelector('textarea')
-      if (!textarea) {
-        throw new Error('Probe textarea is not mounted')
-      }
-      return textarea
-    }
+    }))
   }
 }
 
@@ -156,6 +68,147 @@ describe('useNativeChatComposerAttachments', () => {
     )
   })
 
+  it.each<{
+    name: string
+    scope: string
+    images: draftStorage.NativeChatComposerDraftImage[]
+    check: 'none' | 'path' | 'paste'
+  }>([
+    { name: 'empty checks', scope: 'agent-session:empty', images: [], check: 'none' },
+    {
+      name: 'ordinary paths',
+      scope: 'agent-session:path',
+      images: [{ id: 'path', path: '/repo/shot.png' }],
+      check: 'path'
+    },
+    {
+      name: 'kept local pastes',
+      scope: 'agent-session:paste',
+      images: [{ id: 'paste', path: '/state/native-chat-pastes/orca-paste-1-a.png' }],
+      check: 'paste'
+    },
+    {
+      name: 'SSH attachment-store paths',
+      scope: 'agent-session:ssh',
+      images: [
+        {
+          id: 'ssh',
+          path: '/state/agent-session-attachments/upload/shot.png',
+          connectionId: 'ssh-host'
+        }
+      ],
+      check: 'none'
+    },
+    {
+      name: 'ineligible terminal scopes',
+      scope: 'tab:leaf',
+      images: [{ id: 'terminal', path: '/state/agent-session-attachments/upload/shot.png' }],
+      check: 'none'
+    }
+  ])(
+    'does no ownership resolution for $name across status publications',
+    async ({ scope, images, check }) => {
+      const storage = draftStorage.createMemoryNativeChatComposerDraftStorage()
+      const owner: draftStorage.NativeChatComposerDraftOwner = {
+        workspaceId: 'folder-workspace',
+        executionHostId: 'local'
+      }
+      storage.drafts.set(scope, { text: '', images, owner, savedAt: 1 })
+      draftStorage.setNativeChatComposerDraftStorageForTests(storage)
+      const resolveOwner = vi.fn(() => owner)
+      draftStore.setNativeChatComposerDraftOwnerResolver(resolveOwner)
+      const validation = Promise.withResolvers<boolean>()
+      const pathExists = vi.fn(() => validation.promise)
+      const restoreNativeChatPastes = vi.fn(async () => {
+        const exists = await validation.promise
+        return images.map(({ path }) => ({ path, kept: true, exists }))
+      })
+      const originalApi = window.api
+      const originalState = useAppStore.getState()
+      let probe: Awaited<ReturnType<typeof renderProbe>> | undefined
+      Object.assign(window, { api: { fs: { pathExists }, ui: { restoreNativeChatPastes } } })
+      try {
+        await draftStore.hydrateNativeChatComposerDrafts()
+        probe = await renderProbe(scope, true)
+        expect(pathExists).toHaveBeenCalledTimes(check === 'path' ? 1 : 0)
+        expect(restoreNativeChatPastes).toHaveBeenCalledTimes(check === 'paste' ? 1 : 0)
+        const callsAtMount = resolveOwner.mock.calls.length
+        act(publishUnrelatedStatuses)
+        expect(resolveOwner.mock.calls.length - callsAtMount).toBe(0)
+        expect(resolveOwner).not.toHaveBeenCalled()
+        await act(async () => validation.resolve(false))
+        expect(probe.latest().imageAttachments).toMatchObject(
+          check === 'none'
+            ? images
+            : [{ id: images[0].id, path: '', unavailableName: expect.any(String) }]
+        )
+        act(publishUnrelatedStatuses)
+        expect(resolveOwner).not.toHaveBeenCalled()
+        expect(pathExists).toHaveBeenCalledTimes(check === 'path' ? 1 : 0)
+        expect(restoreNativeChatPastes).toHaveBeenCalledTimes(check === 'paste' ? 1 : 0)
+      } finally {
+        act(() => probe?.root.unmount())
+        useAppStore.setState({
+          runtimeStatusByEnvironmentId: originalState.runtimeStatusByEnvironmentId,
+          agentStatusByPaneKey: originalState.agentStatusByPaneKey
+        })
+        Object.assign(window, { api: originalApi })
+      }
+    }
+  )
+
+  it('ignores a late missing-file result from a previous host and retries when local again', async () => {
+    const scope = 'agent-session:host-changed'
+    const image = { id: 'old', path: '/state/agent-session-attachments/upload/shot.png' }
+    const storedOwner: draftStorage.NativeChatComposerDraftOwner = {
+      workspaceId: 'folder-workspace',
+      executionHostId: 'local'
+    }
+    const storage = draftStorage.createMemoryNativeChatComposerDraftStorage()
+    storage.drafts.set(scope, { text: '', images: [image], owner: storedOwner, savedAt: 1 })
+    draftStorage.setNativeChatComposerDraftStorageForTests(storage)
+    let owner = storedOwner
+    const resolveOwner = vi.fn(() => owner)
+    draftStore.setNativeChatComposerDraftOwnerResolver(resolveOwner)
+    const old = Promise.withResolvers<boolean>()
+    const next = Promise.withResolvers<boolean>()
+    const pathExists = vi.fn().mockReturnValueOnce(old.promise).mockReturnValueOnce(next.promise)
+    const originalApi = window.api
+    Object.assign(window, { api: { fs: { pathExists } } })
+    let probe: Awaited<ReturnType<typeof renderProbe>> | undefined
+    const originalState = useAppStore.getState()
+    try {
+      await draftStore.hydrateNativeChatComposerDrafts()
+      probe = await renderProbe(scope, true)
+      expect(nativeChatImagesHoldSend(probe.latest().imageAttachments)).toBe(true)
+      expect(pathExists).toHaveBeenCalledTimes(1)
+      owner = { ...storedOwner, executionHostId: 'runtime:remote' }
+      await act(async () => old.resolve(false))
+      expect(draftStore.readNativeChatComposerDraft(scope).images).toEqual([image])
+      await act(async () => useAppStore.setState({}))
+      expect(nativeChatImagesHoldSend(probe.latest().imageAttachments)).toBe(false)
+      owner = storedOwner
+      await act(async () => useAppStore.setState({}))
+      expect(pathExists).toHaveBeenCalledTimes(2)
+      expect(nativeChatImagesHoldSend(probe.latest().imageAttachments)).toBe(true)
+      await act(async () => next.resolve(false))
+      expect(probe.latest().imageAttachments).toMatchObject([
+        { id: image.id, path: '', unavailableName: 'shot.png' }
+      ])
+      expect(draftStore.readNativeChatComposerDraft(scope).owner).toEqual(storedOwner)
+      resolveOwner.mockClear()
+      act(publishUnrelatedStatuses)
+      expect(resolveOwner).not.toHaveBeenCalled()
+    } finally {
+      act(() => probe?.root.unmount())
+      useAppStore.setState({
+        runtimeStatusByEnvironmentId: originalState.runtimeStatusByEnvironmentId,
+        agentStatusByPaneKey: originalState.agentStatusByPaneKey
+      })
+      Object.assign(window, { api: originalApi })
+    }
+  })
+
   it('checks a restored store image when its owner becomes local, holding Send and preserving a newer edit', async () => {
     const scope = 'agent-session:restored'
     const image = { id: 'old', path: '/state/agent-session-attachments/upload/shot.png' }
@@ -210,42 +263,54 @@ describe('useNativeChatComposerAttachments', () => {
     }
   })
 
-  it('marks a restored image Not kept when a later sweep removes its upload, composer shown or not', async () => {
-    const scope = 'agent-session:swept-later'
-    const image = { id: 'old', path: '/state/agent-session-attachments/upload/shot.png' }
-    const other = { id: 'other', path: '/state/agent-session-attachments/upload-2/b.png' }
-    const storage = draftStorage.createMemoryNativeChatComposerDraftStorage()
-    storage.drafts.set(scope, { text: 'hi', images: [image, other], savedAt: 1 })
-    draftStorage.setNativeChatComposerDraftStorageForTests(storage)
-    draftStore.setNativeChatComposerDraftOwnerResolver(() => ({
-      workspaceId: 'wt',
-      executionHostId: 'local'
-    }))
-    let present = true
-    const pathExists = vi.fn(async () => present)
-    const originalApi = window.api
-    Object.assign(window, { api: { fs: { pathExists }, ui: {} } })
-    try {
-      await draftStore.hydrateNativeChatComposerDrafts()
-      const probe = await renderProbe(scope, true)
-      await act(async () => {})
-      expect(nativeChatImagesHoldSend(probe.latest().imageAttachments)).toBe(false)
-      act(() => probe.root.unmount())
-      present = false
-      await act(async () =>
-        recheckNativeChatComposerDraftImagesIn('/state/agent-session-attachments/upload')
-      )
-      expect(draftStore.readNativeChatComposerDraft(scope).images).toMatchObject([
-        { id: 'old', path: '', unavailableName: 'shot.png' },
-        other
-      ])
-      const reopened = await renderProbe(scope, true)
-      expect(nativeChatImagesHoldSend(reopened.latest().imageAttachments)).toBe(true)
-      act(() => reopened.root.unmount())
-    } finally {
-      Object.assign(window, { api: originalApi })
+  it.each([true, false])(
+    'rechecks a removed upload after restore settles (composer mounted: %s)',
+    async (mounted) => {
+      const scope = 'agent-session:swept-later'
+      const image = { id: 'old', path: '/state/agent-session-attachments/upload/shot.png' }
+      const other = { id: 'other', path: '/state/agent-session-attachments/upload-2/b.png' }
+      const storage = draftStorage.createMemoryNativeChatComposerDraftStorage()
+      storage.drafts.set(scope, { text: 'hi', images: [image, other], savedAt: 1 })
+      draftStorage.setNativeChatComposerDraftStorageForTests(storage)
+      draftStore.setNativeChatComposerDraftOwnerResolver(() => ({
+        workspaceId: 'wt',
+        executionHostId: 'local'
+      }))
+      let present = true
+      const pathExists = vi.fn(async () => present)
+      const originalApi = window.api
+      Object.assign(window, { api: { fs: { pathExists }, ui: {} } })
+      try {
+        await draftStore.hydrateNativeChatComposerDrafts()
+        const probe = await renderProbe(scope, true)
+        await act(async () => {})
+        expect(nativeChatImagesHoldSend(probe.latest().imageAttachments)).toBe(false)
+        if (!mounted) {
+          act(() => probe.root.unmount())
+        }
+        present = false
+        await act(async () =>
+          recheckNativeChatComposerDraftImagesIn('/state/agent-session-attachments/upload')
+        )
+        expect(draftStore.readNativeChatComposerDraft(scope).images).toMatchObject([
+          { id: 'old', path: '', unavailableName: 'shot.png' },
+          other
+        ])
+        if (mounted) {
+          expect(probe.latest().imageAttachments).toMatchObject([
+            { id: 'old', path: '', unavailableName: 'shot.png' },
+            other
+          ])
+          act(() => probe.root.unmount())
+        }
+        const reopened = await renderProbe(scope, true)
+        expect(nativeChatImagesHoldSend(reopened.latest().imageAttachments)).toBe(true)
+        act(() => reopened.root.unmount())
+      } finally {
+        Object.assign(window, { api: originalApi })
+      }
     }
-  })
+  )
 
   it('holds attached images as chips (deferred to submit) and restores them on remount', async () => {
     const first = await renderProbe('pty-1')