fix(desktop): preserve edit-resend session and context

desktop/src/components/chat/ChatInput.test.tsx

@@ -652,6 +652,46 @@ describe('ChatInput file mentions', () => {
     expect(getComposerText()).toBe('draft before switching project')
   })
 
+  it('restores path-only files and images from an edit fallback without losing their paths', async () => {
+    render(<ChatInput compact />)
+    act(() => useChatStore.getState().queueComposerPrefill(sessionId, {
+      text: 'edited prompt',
+      attachments: [
+        { type: 'file', name: 'notes.md', path: '/repo/notes.md' },
+        { type: 'image', name: 'shot.png', path: '/repo/shot.png', mimeType: 'image/png' },
+      ],
+    }))
+    await waitFor(() => expect(getComposerText()).toBe('edited prompt'))
+    fireEvent.keyDown(getComposerElement(), { key: 'Enter' })
+    expect(mocks.wsSend).toHaveBeenCalledWith(sessionId, expect.objectContaining({
+      content: 'edited prompt',
+      attachments: [
+        expect.objectContaining({ type: 'file', path: '/repo/notes.md' }),
+        expect.objectContaining({ type: 'image', path: '/repo/shot.png', mimeType: 'image/png' }),
+      ],
+    }))
+  })
+
+  it('restores removable workspace and session references from an edit fallback', async () => {
+    render(<ChatInput compact />)
+    act(() => useChatStore.getState().queueComposerPrefill(sessionId, {
+      text: 'edited prompt',
+      attachments: [{ type: 'file', name: 'app.ts', path: '/repo/app.ts', lineStart: 3, lineEnd: 5, quote: 'for (;;) {}', note: 'stop this loop', diffSide: 'new', hunkId: 'hunk-1' }],
+      sessionReferences: [{ sessionId: 'referenced-session' }],
+    }))
+    await waitFor(() => expect(getComposerText()).toContain('edited prompt'))
+    expect(screen.getByRole('button', { name: /Remove app.ts/ })).toBeTruthy()
+    expect(getComposerElement().querySelector('[data-mention-kind="session"]')).toBeTruthy()
+    fireEvent.keyDown(getComposerElement(), { key: 'Enter' })
+    expect(mocks.wsSend).toHaveBeenCalledWith(sessionId, expect.objectContaining({
+      content: expect.stringContaining('for (;;) {}'),
+      attachments: [expect.objectContaining({ path: '/repo/app.ts', lineStart: 3, lineEnd: 5, quote: 'for (;;) {}', note: 'stop this loop' })],
+      sessionReferences: [{ sessionId: 'referenced-session' }],
+    }))
+    const message = useChatStore.getState().sessions[sessionId]!.messages.findLast((row) => row.type === 'user_text')
+    expect(message).toMatchObject({ attachments: [expect.objectContaining({ diffSide: 'new', hunkId: 'hunk-1' })] })
+  })
+
   it('restores an unsent composer draft after the composer unmounts', async () => {
     const { unmount } = render(<ChatInput compact />)
 

desktop/src/components/chat/ChatInput.tsx

@@ -2,6 +2,7 @@ import { useSideChatStore } from '@/stores/sideChatStore'
 import { openSideChat } from '@/lib/workspace/openSideChat'
 import { parseSideQuestionCommand } from './composerUtils'
 import { getSessionReferences } from '@/lib/composerMentions'
+import { normalizeSessionReferences } from '@/lib/sessionReferences'
 import { isComposerReferenceVisible, isComposerSlashCommandVisible } from '@/lib/composerCapabilityVisibility'
 import { useState, useRef, useEffect, useCallback, useMemo, useId } from 'react'
 import { useDismissable } from '@/hooks/useDismissable'
@@ -448,21 +449,49 @@ export function ChatInput({ variant = 'default', compact = false, sessionId, vis
     if (!composerPrefill || !activeTabId) return
 
     const nextAttachments = (composerPrefill.attachments ?? [])
-      .filter((attachment) => attachment.type === 'image' || attachment.data)
+      .filter((attachment) => attachment.path || attachment.data)
       .map((attachment, index) => ({
         id: `composer-prefill-${composerPrefill.nonce}-${index}`,
         name: attachment.name,
         type: attachment.type,
         mimeType: attachment.mimeType,
+        path: attachment.path,
         previewUrl: attachment.type === 'image' ? attachment.data : undefined,
         data: attachment.data,
+        isDirectory: attachment.isDirectory,
+        lineStart: attachment.lineStart,
+        lineEnd: attachment.lineEnd,
+        diffSide: attachment.diffSide,
+        hunkId: attachment.hunkId,
+        note: attachment.note,
+        quote: attachment.quote,
       }))
+    const contextAttachments = nextAttachments.filter((attachment) => attachment.type === 'file' && attachment.path && (
+      attachment.lineStart || attachment.note?.trim() || attachment.quote?.trim() || attachment.diffSide || attachment.hunkId
+    ))
+    const uploads = nextAttachments.filter((attachment) => !contextAttachments.includes(attachment))
 
     if (composerPrefill.mode === 'append') {
-      setComposerAttachments((previous) => [...previous, ...nextAttachments])
+      setComposerAttachments((previous) => [...previous, ...uploads])
     } else {
-      setComposerInput(composerPrefill.text, [])
-      setComposerAttachments(nextAttachments)
+      let restored = { text: composerPrefill.text, mentions: [] as ComposerMention[] }
+      for (const reference of normalizeSessionReferences(composerPrefill.sessionReferences)) {
+        const title = useSessionStore.getState().sessions.find((session) => session.id === reference.sessionId)?.title
+        restored = insertMentionIntoText(restored.text, restored.mentions, restored.text.length, restored.text.length, {
+          kind: 'session', id: reference.sessionId, label: title || reference.sessionId,
+          path: '', isDirectory: false,
+        })
+      }
+      setComposerInput(restored.text, restored.mentions)
+      setComposerAttachments(uploads)
+      clearWorkspaceReferences(activeTabId)
+    }
+    for (const attachment of contextAttachments) {
+      addWorkspaceReference(activeTabId, {
+        ...attachment,
+        path: attachment.path!,
+        kind: attachment.diffSide || attachment.hunkId ? 'code-comment' : attachment.lineStart ? 'code-selection' : 'file',
+      })
     }
     setPlusMenuOpen(false)
     setSlashMenuOpen(false)
@@ -480,7 +509,9 @@ export function ChatInput({ variant = 'default', compact = false, sessionId, vis
     clearComposerPrefill(activeTabId, composerPrefill.nonce)
   }, [
     activeTabId,
+    addWorkspaceReference,
     clearComposerPrefill,
+    clearWorkspaceReferences,
     composerPrefill,
     setComposerAttachments,
     setComposerInput,

desktop/src/components/chat/MessageList.test.tsx

@@ -9828,7 +9828,7 @@ describe('MessageList edit and resend', () => {
     await waitFor(() => expect(queueComposerPrefill).toHaveBeenCalledWith(ACTIVE_TAB, {
       text: 'Build a landing page',
       attachments: [],
-    }))
+    }, { restoreMissingSession: true }))
     expect(sendMessage).not.toHaveBeenCalled()
     expect(useUIStore.getState().toasts).toEqual([
       expect.objectContaining({ type: 'warning' }),
@@ -9878,9 +9878,9 @@ describe('MessageList edit and resend', () => {
 
     await waitFor(() => expect(sendMessage).toHaveBeenCalledOnce())
     const [, content, attachments, options] = sendMessage.mock.calls[0]!
-    expect(content).toBe(`${referencePrompt}\n\nHow do I stop it?`)
+    expect(content).toBe(`${referencePrompt.replace('src/app.ts:L3-L5', '/repo/src/app.ts:L3-L5')}\n\nHow do I stop it?`)
     expect(attachments).toEqual([
-      expect.objectContaining({ type: 'file', path: 'src/app.ts', lineStart: 3, lineEnd: 5, quote: 'for (;;) {}' }),
+      expect.objectContaining({ type: 'file', path: '/repo/src/app.ts', lineStart: 3, lineEnd: 5, quote: 'for (;;) {}' }),
       { type: 'image', name: 'shot.png', mimeType: 'image/png', data: 'data:image/png;base64,AAAA' },
     ])
     expect(options).toMatchObject({ displayContent: 'How do I stop it?' })

desktop/src/components/chat/useUserMessageEditResend.test.tsx

@@ -0,0 +1,180 @@
+import { act, renderHook, waitFor } from '@testing-library/react'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { sessionsApi, type SessionRewindResponse } from '@/api/sessions'
+import { useChatStore, type PerSessionState } from '@/stores/chatStore'
+import { useUIStore } from '@/stores/uiStore'
+import type { UIMessage } from '@/types/chat'
+import type { TranslationKey } from '@/i18n/locales/en'
+import { useUserMessageEditResend } from './useUserMessageEditResend'
+import type { UserMessageEditDraft } from './userMessageEdit'
+
+const messages: UIMessage[] = [
+  { id: 'shared-user', type: 'user_text', content: 'continue', timestamp: 1 },
+  { id: 'reply', type: 'assistant_text', content: 'Done', timestamp: 2 },
+]
+const draft: UserMessageEditDraft = { text: 'edited prompt', attachments: [], sessionReferences: [] }
+const target = { targetUserMessageId: 'shared-user', userMessageIndex: 0, userMessageCount: 1 }
+const preview: SessionRewindResponse = {
+  target, conversation: { messagesRemoved: 2 },
+  code: { available: false, filesChanged: [], insertions: 0, deletions: 0 },
+}
+const codePreview: SessionRewindResponse = { ...preview, code: { available: true, filesChanged: ['app.ts'], insertions: 1, deletions: 0 } }
+const translate = (key: TranslationKey) => key
+
+function makeSession(): PerSessionState {
+  return {
+    messages, chatState: 'idle', connectionState: 'connected', historyStatus: 'ready', historyHydrated: true,
+    streamingText: '', streamingToolInput: '', activeToolUseId: null, activeToolName: null, activeThinkingId: null,
+    pendingPermission: null, pendingComputerUsePermission: null,
+    tokenUsage: { input_tokens: 0, output_tokens: 0 }, streamingResponseChars: 0, elapsedSeconds: 0,
+    statusVerb: '', apiRetry: null, slashCommands: [], agentTaskNotifications: {}, elapsedTimer: null,
+  }
+}
+
+function deferred<T>() {
+  let resolve!: (value: T) => void
+  let reject!: (error: Error) => void
+  const promise = new Promise<T>((resolvePromise, rejectPromise) => { resolve = resolvePromise; reject = rejectPromise })
+  return { promise, resolve, reject }
+}
+
+function setup() {
+  const setRewindingTurnId = vi.fn()
+  return renderHook(({ sessionId }) => useUserMessageEditResend({
+    sessionId, messages,
+    turnCards: [{ target: { messageId: 'shared-user', expectedContent: 'continue' }, checkpoint: { target } }],
+    disabled: false, rewindingTurnId: null, setRewindingTurnId, t: translate,
+  }), { initialProps: { sessionId: 'session-a' } })
+}
+
+describe('useUserMessageEditResend session and recovery boundaries', () => {
+  beforeEach(() => {
+    useChatStore.setState({ ...useChatStore.getInitialState(), sessions: { 'session-a': makeSession(), 'session-b': makeSession() } }, true)
+    useUIStore.setState({ toasts: [] })
+    vi.spyOn(sessionsApi, 'rewind').mockResolvedValue(preview)
+    vi.spyOn(useChatStore.getState(), 'sendMessage').mockImplementation(() => {})
+  })
+
+  afterEach(() => { vi.restoreAllMocks() })
+
+  it.each([preview, codePreview])('ignores a pending preview after switching session instead of rewinding either session', async (response) => {
+    const pending = deferred<SessionRewindResponse>()
+    vi.mocked(sessionsApi.rewind).mockReturnValueOnce(pending.promise)
+    const { result, rerender } = setup()
+    act(() => result.current.editActionByMessageId.get('shared-user')!.onSubmit(draft))
+    rerender({ sessionId: 'session-b' })
+    await act(async () => { pending.resolve(response) })
+    expect(result.current.dialog.open).toBe(false)
+    expect(sessionsApi.rewind).toHaveBeenCalledTimes(1)
+    expect(useChatStore.getState().sendMessage).not.toHaveBeenCalled()
+  })
+
+  it('rejects a captured editor submit action after switching session', async () => {
+    const { result, rerender } = setup()
+    const submitOld = result.current.editActionByMessageId.get('shared-user')!.onSubmit
+    rerender({ sessionId: 'session-b' })
+    await act(async () => submitOld(draft))
+    expect(sessionsApi.rewind).not.toHaveBeenCalled()
+    expect(useChatStore.getState().sendMessage).not.toHaveBeenCalled()
+  })
+
+  it('ignores a stale preview even after switching back to its original session', async () => {
+    const pending = deferred<SessionRewindResponse>()
+    vi.mocked(sessionsApi.rewind).mockReturnValueOnce(pending.promise)
+    const { result, rerender } = setup()
+    act(() => result.current.editActionByMessageId.get('shared-user')!.onSubmit(draft))
+    rerender({ sessionId: 'session-b' })
+    rerender({ sessionId: 'session-a' })
+    await act(async () => { pending.resolve(codePreview) })
+    expect(result.current.dialog.open).toBe(false)
+    expect(sessionsApi.rewind).toHaveBeenCalledTimes(1)
+  })
+
+  it('hides the old dialog and rejects its captured action after switching session', async () => {
+    vi.mocked(sessionsApi.rewind).mockResolvedValue(codePreview)
+    const { result, rerender } = setup()
+    await act(async () => result.current.editActionByMessageId.get('shared-user')!.onSubmit(draft))
+    expect(result.current.dialog.open).toBe(true)
+    const confirmOld = result.current.dialog.actions[1]!.onClick
+    rerender({ sessionId: 'session-b' })
+    expect(result.current.dialog.open).toBe(false)
+    await act(async () => confirmOld())
+    expect(sessionsApi.rewind).toHaveBeenCalledTimes(1)
+  })
+
+  it('finishes an already committed rewind in its original session after navigation', async () => {
+    const pending = deferred<SessionRewindResponse>()
+    vi.mocked(sessionsApi.rewind).mockResolvedValueOnce(preview).mockReturnValueOnce(pending.promise)
+    const reload = vi.spyOn(useChatStore.getState(), 'reloadHistory').mockResolvedValue(undefined)
+    const { result, rerender } = setup()
+    act(() => result.current.editActionByMessageId.get('shared-user')!.onSubmit(draft))
+    await waitFor(() => expect(sessionsApi.rewind).toHaveBeenCalledTimes(2))
+    rerender({ sessionId: 'session-b' })
+    await act(async () => { pending.resolve(preview) })
+    expect(reload).toHaveBeenCalledWith('session-a', expect.objectContaining({ messages }), { requireApplied: true })
+    expect(useChatStore.getState().sendMessage).toHaveBeenCalledWith('session-a', draft.text, [], expect.any(Object))
+  })
+
+  it('recovers the draft when its source tab closes during the committed rewind', async () => {
+    vi.spyOn(sessionsApi, 'getFullHistory').mockResolvedValue({ messages: [] })
+    const pending = deferred<SessionRewindResponse>()
+    vi.mocked(sessionsApi.rewind).mockResolvedValueOnce(preview).mockReturnValueOnce(pending.promise)
+    const withContext: UserMessageEditDraft = {
+      ...draft,
+      attachments: [{ id: 'file', type: 'file', name: 'app.ts', path: '/repo/app.ts', lineStart: 3, quote: 'value', sendable: true }],
+      sessionReferences: [{ sessionId: 'referenced-session' }],
+    }
+    const { result, unmount } = setup()
+    act(() => result.current.editActionByMessageId.get('shared-user')!.onSubmit(withContext))
+    await waitFor(() => expect(sessionsApi.rewind).toHaveBeenCalledTimes(2))
+    act(() => useChatStore.getState().disconnectSession('session-a'))
+    unmount()
+    await act(async () => { pending.resolve(preview) })
+    expect(useChatStore.getState().sendMessage).not.toHaveBeenCalled()
+    expect(useChatStore.getState().sessions['session-a']?.composerPrefill).toMatchObject({
+      text: draft.text, attachments: withContext.attachments, sessionReferences: withContext.sessionReferences,
+    })
+    expect(useChatStore.getState().sessions['session-b']?.composerPrefill).toBeUndefined()
+  })
+
+  it('isolates drafts in branched sessions that share a transcript message id', () => {
+    const { result, rerender } = setup()
+    act(() => {
+      const action = result.current.editActionByMessageId.get('shared-user')!
+      action.onStart()
+      action.onDraftChange(draft)
+    })
+    rerender({ sessionId: 'session-b' })
+    expect(result.current.editActionByMessageId.get('shared-user')!.editing).toBe(false)
+    expect(result.current.editActionByMessageId.get('shared-user')!.getDraft().text).toBe('continue')
+    rerender({ sessionId: 'session-a' })
+    expect(result.current.editActionByMessageId.get('shared-user')!.getDraft().text).toBe(draft.text)
+  })
+
+  it('hands the complete draft to the composer when the real history reload fails', async () => {
+    vi.spyOn(sessionsApi, 'getFullHistory').mockRejectedValue(new Error('history offline'))
+    const withContext: UserMessageEditDraft = {
+      ...draft,
+      attachments: [{ id: 'file', type: 'file', name: 'app.ts', path: '/repo/app.ts', lineStart: 3, lineEnd: 5, quote: 'for (;;) {}', sendable: true }],
+      sessionReferences: [{ sessionId: 'referenced-session' }],
+    }
+    const { result } = setup()
+    await act(async () => result.current.editActionByMessageId.get('shared-user')!.onSubmit(withContext))
+    expect(useChatStore.getState().sendMessage).not.toHaveBeenCalled()
+    expect(useChatStore.getState().sessions['session-a']!.composerPrefill).toMatchObject({
+      text: draft.text, attachments: withContext.attachments, sessionReferences: withContext.sessionReferences,
+    })
+    expect(useUIStore.getState().toasts).toContainEqual(expect.objectContaining({ type: 'warning' }))
+  })
+
+  it('ignores a failed preview after the editor unmounts', async () => {
+    const pending = deferred<SessionRewindResponse>()
+    vi.mocked(sessionsApi.rewind).mockReturnValueOnce(pending.promise)
+    const { result, unmount } = setup()
+    act(() => result.current.editActionByMessageId.get('shared-user')!.onSubmit(draft))
+    unmount()
+    await act(async () => { pending.reject(new Error('late failure')) })
+    expect(useUIStore.getState().toasts).toEqual([])
+    expect(sessionsApi.rewind).toHaveBeenCalledTimes(1)
+  })
+})

desktop/src/components/chat/useUserMessageEditResend.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useMemo, useRef, useState, type ReactNode } from 'react'
+import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from 'react'
 import { sessionsApi, type SessionRewindMode, type SessionRewindResponse, type SessionTurnCheckpoint } from '../../api/sessions'
 import type { ActionDialogAction } from '@/components/ui/ActionDialog'
 import type { TranslationKey } from '../../i18n/locales/en'
@@ -23,7 +23,10 @@ export type EditableTurnCard = {
   checkpoint: { target: SessionTurnCheckpoint['target'] }
 }
 
+type EditContext = { sessionId: string | null | undefined; active: boolean }
+
 type PendingConfirm = {
+  context: EditContext
   messageId: string
   card: EditableTurnCard
   draft: UserMessageEditDraft
@@ -83,10 +86,26 @@ export function useUserMessageEditResend({
   setRewindingTurnId,
   t,
 }: Options) {
-  const [editingMessageId, setEditingMessageId] = useState<string | null>(null)
-  const [submittingMessageId, setSubmittingMessageId] = useState<string | null>(null)
+  const [editing, setEditing] = useState<{ sessionId: string; messageId: string } | null>(null)
+  const [submitting, setSubmitting] = useState<{ context: EditContext; messageId: string } | null>(null)
   const [confirm, setConfirm] = useState<PendingConfirm | null>(null)
   const draftsRef = useRef(new Map<string, UserMessageEditDraft>())
+  // A new identity also invalidates A → B → A requests and captured dialog actions.
+  const contextRef = useRef<EditContext>({ sessionId, active: true })
+  if (contextRef.current.sessionId !== sessionId) contextRef.current = { sessionId, active: true }
+  const context = contextRef.current
+  const editingMessageId = editing && editing.sessionId === sessionId ? editing.messageId : null
+  const submittingMessageId = submitting?.context === context ? submitting.messageId : null
+  const isCurrentContext = useCallback((candidate: EditContext) =>
+    candidate.active && contextRef.current === candidate, [])
+
+  useEffect(() => {
+    const current = contextRef.current
+    current.active = true
+    setRewindingTurnId(null)
+    return () => { current.active = false }
+  }, [sessionId, setRewindingTurnId])
+
   const messagesRef = useRef(messages)
   messagesRef.current = messages
 
@@ -113,24 +132,26 @@ export function useUserMessageEditResend({
     [t],
   )
 
-  const closeEditor = useCallback((messageId: string) => {
-    draftsRef.current.delete(messageId)
-    setEditingMessageId((current) => current === messageId ? null : current)
+  const closeEditor = useCallback((targetSessionId: string, messageId: string) => {
+    draftsRef.current.delete(JSON.stringify([targetSessionId, messageId]))
+    setEditing((current) => current?.sessionId === targetSessionId && current.messageId === messageId ? null : current)
   }, [])
 
   const runResend = useCallback(async (
+    requestContext: EditContext,
     messageId: string,
     card: EditableTurnCard,
     draft: UserMessageEditDraft,
     payload: UserMessageResendPayload,
     mode: SessionRewindMode,
   ) => {
-    if (!sessionId) return
+    const targetSessionId = requestContext.sessionId
+    if (!targetSessionId || !isCurrentContext(requestContext)) return
     const addToast = useUIStore.getState().addToast
     setRewindingTurnId(messageId)
     let result: SessionRewindResponse
     try {
-      result = await rewindToTurnCheckpoint(sessionId, {
+      result = await rewindToTurnCheckpoint(targetSessionId, {
         checkpointTarget: card.checkpoint.target,
         expectedContent: card.target.expectedContent,
       }, mode)
@@ -138,46 +159,55 @@ export function useUserMessageEditResend({
       // The server rewinds atomically: a failed rewind left the conversation
       // and the files as they were, so the draft stays open for another try.
       addToast({ type: 'error', message: t('chat.editResendFailed', { detail: getApiErrorMessage(error) }) })
-      setConfirm(null)
-      setRewindingTurnId(null)
+      if (isCurrentContext(requestContext)) {
+        setConfirm(null)
+        setRewindingTurnId(null)
+      }
       return
     }
 
     // From here the old turn is gone. The edit must reach either the model or
     // the composer — never be dropped with the editor.
-    closeEditor(messageId)
-    setConfirm(null)
+    closeEditor(targetSessionId, messageId)
+    if (isCurrentContext(requestContext)) setConfirm(null)
     const chatStore = useChatStore.getState()
     let sent = false
     try {
-      await chatStore.reloadHistory(sessionId)
-      if ((useChatStore.getState().sessions[sessionId]?.chatState ?? 'idle') === 'idle') {
-        useChatStore.getState().sendMessage(sessionId, payload.content, payload.attachments, payload.options)
+      const session = chatStore.sessions[targetSessionId]
+      await chatStore.reloadHistory(targetSessionId, {
+        messages: session?.messages ?? [],
+        backgroundAgentTasks: session?.backgroundAgentTasks,
+      }, { requireApplied: true })
+      if (useChatStore.getState().sessions[targetSessionId]?.chatState === 'idle') {
+        useChatStore.getState().sendMessage(targetSessionId, payload.content, payload.attachments, payload.options)
         sent = true
       }
     } catch {
       sent = false
     }
     if (!sent) {
-      useChatStore.getState().queueComposerPrefill(sessionId, {
+      useChatStore.getState().queueComposerPrefill(targetSessionId, {
         text: draft.text,
         attachments: draft.attachments.filter((attachment) => attachment.sendable),
-      })
+        ...(draft.sessionReferences.length > 0 ? { sessionReferences: draft.sessionReferences } : {}),
+      }, { restoreMissingSession: true })
       addToast({ type: 'warning', message: t('chat.editResendPrefilled') })
     } else if (mode === 'both') {
       addToast(describeRewindResult(result, mode, t))
     }
-    setRewindingTurnId(null)
-  }, [closeEditor, sessionId, setRewindingTurnId, t])
+    if (isCurrentContext(requestContext)) setRewindingTurnId(null)
+  }, [closeEditor, isCurrentContext, setRewindingTurnId, t])
 
   const submit = useCallback(async (messageId: string, draft: UserMessageEditDraft) => {
+    const requestContext = context
+    if (!isCurrentContext(requestContext)) return
     const card = editableCards.get(messageId)
     if (!sessionId || !card || disabled || submittingMessageId || rewindingTurnId) return
     const payload = buildUserMessageResendPayload(draft, { contextReferencesOnly })
     if (!payload) return
-    draftsRef.current.set(messageId, draft)
+    draftsRef.current.set(JSON.stringify([sessionId, messageId]), draft)
 
-    setSubmittingMessageId(messageId)
+    setSubmitting({ context: requestContext, messageId })
     let preview: SessionRewindResponse
     try {
       preview = await sessionsApi.rewind(sessionId, {
@@ -187,14 +217,16 @@ export function useUserMessageEditResend({
         dryRun: true,
       })
     } catch (error) {
+      if (!isCurrentContext(requestContext)) return
       useUIStore.getState().addToast({
         type: 'error',
         message: t('chat.editResendFailed', { detail: getApiErrorMessage(error) }),
       })
-      setSubmittingMessageId(null)
+      setSubmitting(null)
       return
     }
-    setSubmittingMessageId(null)
+    if (!isCurrentContext(requestContext)) return
+    setSubmitting(null)
 
     const canRestoreCode = preview.code.available &&
       preview.code.filesChanged.length > 0 &&
@@ -203,11 +235,11 @@ export function useUserMessageEditResend({
     // Nothing on disk to decide about and nothing beyond this turn to lose:
     // the edit is the whole consequence, so it goes straight through.
     if (!canRestoreCode && preview.restoreAvailable !== false && laterTurns === 0) {
-      await runResend(messageId, card, draft, payload, 'conversation')
+      await runResend(requestContext, messageId, card, draft, payload, 'conversation')
       return
     }
-    setConfirm({ messageId, card, draft, payload, preview, laterTurns, canRestoreCode })
-  }, [contextReferencesOnly, disabled, editableCards, rewindingTurnId, runResend, sessionId, submittingMessageId, t])
+    setConfirm({ context: requestContext, messageId, card, draft, payload, preview, laterTurns, canRestoreCode })
+  }, [context, contextReferencesOnly, disabled, editableCards, isCurrentContext, rewindingTurnId, runResend, sessionId, submittingMessageId, t])
 
   const editActionByMessageId = useMemo(() => {
     const result = new Map<string, UserMessageEditAction>()
@@ -224,15 +256,17 @@ export function useUserMessageEditResend({
         editing,
         submitting: submittingMessageId === messageId || rewindingTurnId === messageId,
         disabled: disabled || Boolean(rewindingTurnId && rewindingTurnId !== messageId),
-        getDraft: () => draftsRef.current.get(messageId) ?? createUserMessageEditDraft(message),
+        getDraft: () => draftsRef.current.get(JSON.stringify([sessionId, messageId])) ?? createUserMessageEditDraft(message),
         onStart: () => {
-          setEditingMessageId((current) => {
-            if (current && current !== messageId) draftsRef.current.delete(current)
-            return messageId
+          setEditing((current) => {
+            if (current && (current.sessionId !== sessionId || current.messageId !== messageId)) {
+              draftsRef.current.delete(JSON.stringify([current.sessionId, current.messageId]))
+            }
+            return { sessionId, messageId }
           })
         },
-        onCancel: () => closeEditor(messageId),
-        onDraftChange: (draft) => { draftsRef.current.set(messageId, draft) },
+        onCancel: () => closeEditor(sessionId, messageId),
+        onDraftChange: (draft) => { draftsRef.current.set(JSON.stringify([sessionId, messageId]), draft) },
         onSubmit: (draft) => { void submit(messageId, draft) },
       })
     }
@@ -242,7 +276,7 @@ export function useUserMessageEditResend({
   const dialog = useMemo<UserMessageEditDialog>(() => {
     const busy = Boolean(rewindingTurnId)
     const close = () => { if (!busy) setConfirm(null) }
-    if (!confirm) {
+    if (!confirm || !isCurrentContext(confirm.context)) {
       return { open: false, title: '', body: null, actions: [], loading: false, onClose: close }
     }
     const { preview, laterTurns, canRestoreCode } = confirm
@@ -270,7 +304,7 @@ export function useUserMessageEditResend({
       </div>
     )
     const run = (mode: SessionRewindMode) => {
-      void runResend(confirm.messageId, confirm.card, confirm.draft, confirm.payload, mode)
+      void runResend(confirm.context, confirm.messageId, confirm.card, confirm.draft, confirm.payload, mode)
     }
     const actions: ActionDialogAction[] = [
       { label: t('common.cancel'), onClick: close, variant: 'secondary' },
@@ -296,7 +330,7 @@ export function useUserMessageEditResend({
       loading: busy,
       onClose: close,
     }
-  }, [confirm, rewindingTurnId, runResend, t])
+  }, [confirm, context, isCurrentContext, rewindingTurnId, runResend, t])
 
   return { editActionByMessageId, dialog }
 }

desktop/src/components/chat/userMessageEdit.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it } from 'vitest'
-import type { UIMessage } from '../../types/chat'
+import type { UIAttachment, UIMessage } from '../../types/chat'
+import { formatWorkspaceReferencePrompt, type WorkspaceChatReference } from '../../stores/workspaceChatContextStore'
 import {
   buildUserMessageResendPayload,
   countLaterUserTurns,
@@ -37,7 +38,7 @@ describe('createUserMessageEditDraft', () => {
     expect(draft.attachments).toEqual([
       expect.objectContaining({
         type: 'file',
-        path: 'src/app.ts',
+        path: '/repo/src/app.ts',
         lineStart: 3,
         lineEnd: 5,
         note: 'why does this loop',
@@ -92,6 +93,76 @@ describe('createUserMessageEditDraft', () => {
     expect(draft.attachments).toEqual([])
   })
 
+  it.each([false, true])('preserves distinct suffix paths from composer references (loaded: %s)', (loaded) => {
+    const reference: WorkspaceChatReference = {
+      id: 'root-selection', kind: 'code-selection', name: 'config.ts',
+      path: 'config.ts', absolutePath: '/repo/config.ts',
+      lineStart: 1, lineEnd: 1, quote: 'root config',
+    }
+    const referencePrompt = formatWorkspaceReferencePrompt([reference])
+    // ChatInput prefixes uploads before the absolute workspace reference paths.
+    // Previously the display parser consumed the upload as the selection and
+    // editing permanently dropped that distinct file from the model payload.
+    const liveAttachments: UIAttachment[] = [
+      { type: 'file', name: 'config.ts', path: '/repo/sub/config.ts', mimeType: 'text/typescript' },
+      { type: 'file', name: reference.name, path: reference.path,
+        lineStart: reference.lineStart, lineEnd: reference.lineEnd, quote: reference.quote },
+    ]
+    const draft = createUserMessageEditDraft({
+      content: 'Compare these',
+      modelContent: `@"/repo/sub/config.ts" @"${reference.absolutePath}" ${referencePrompt}\n\nCompare these`,
+      // The old history display parser has already omitted the wrong prefix.
+      attachments: loaded
+        ? [{ type: 'file', name: 'config.ts', path: '/repo/config.ts' }, liveAttachments[1]!]
+        : liveAttachments,
+    })
+    const payload = buildUserMessageResendPayload(draft, labels)!
+    expect(payload.attachments.map(attachment => attachment.path))
+      .toEqual(['/repo/sub/config.ts', '/repo/config.ts'])
+    expect(payload.attachments[1]).toMatchObject({ lineStart: 1, quote: 'root config' })
+    if (!loaded) expect(payload.attachments[0]?.mimeType).toBe('text/typescript')
+  })
+
+  it('keeps separate selections and their original diff metadata on one file', () => {
+    const references: WorkspaceChatReference[] = [
+      { id: 'first', kind: 'code-comment', path: 'src/config.ts', absolutePath: '/repo/src/config.ts',
+        name: 'config.ts', lineStart: 2, diffSide: 'old', hunkId: 'hunk-first', note: ' first ', quote: ' first quote ' },
+      { id: 'second', kind: 'code-comment', path: 'src/config.ts', absolutePath: '/repo/src/config.ts',
+        name: 'config.ts', lineStart: 8, diffSide: 'new', hunkId: 'hunk-second', quote: 'second quote' },
+    ]
+    const draft = createUserMessageEditDraft({
+      content: 'Review',
+      modelContent: `@"/repo/other/src/config.ts" @"/repo/src/config.ts" @"/repo/src/config.ts" ${formatWorkspaceReferencePrompt(references)}\n\nReview`,
+      attachments: [
+        { type: 'file', name: 'config.ts', path: '/repo/other/src/config.ts' },
+        ...references.map((reference): UIAttachment => ({
+          type: 'file', name: reference.name, path: reference.path, lineStart: reference.lineStart,
+          diffSide: reference.diffSide, hunkId: reference.hunkId, note: reference.note, quote: reference.quote,
+        })),
+      ],
+    })
+    expect(draft.attachments.map(attachment => attachment.path))
+      .toEqual(['/repo/other/src/config.ts', '/repo/src/config.ts', '/repo/src/config.ts'])
+    expect(draft.attachments[1]).toMatchObject({ hunkId: 'hunk-first', lineStart: 2, diffSide: 'old' })
+    expect(draft.attachments[2]).toMatchObject({ hunkId: 'hunk-second', lineStart: 8, diffSide: 'new' })
+  })
+
+  it('merges exact original attachment metadata instead of losing image data and directory identity', () => {
+    const draft = createUserMessageEditDraft({
+      content: 'Look',
+      modelContent: '@"/repo/shot.png" @"/repo/assets" Look',
+      attachments: [
+        { type: 'image', name: 'Original screenshot', path: '/repo/shot.png', data: 'data:image/png;base64,AAAA', mimeType: 'image/png' },
+        { type: 'file', name: 'assets', path: '/repo/assets', isDirectory: true },
+      ],
+    })
+    expect(draft.attachments).toHaveLength(2)
+    expect(draft.attachments[0]).toMatchObject({ name: 'Original screenshot', data: 'data:image/png;base64,AAAA', mimeType: 'image/png' })
+    expect(draft.attachments[1]).toMatchObject({ isDirectory: true })
+    expect(buildUserMessageResendPayload(draft, labels)?.attachments[0])
+      .toEqual({ type: 'image', name: 'Original screenshot', mimeType: 'image/png', data: 'data:image/png;base64,AAAA' })
+  })
+
   it('marks an attachment with neither data nor a path as not sendable', () => {
     const draft = createUserMessageEditDraft({
       content: 'Summarize',
@@ -124,9 +195,9 @@ describe('buildUserMessageResendPayload', () => {
     })
     const payload = buildUserMessageResendPayload({ ...draft, text: 'Please delete it instead' }, labels)!
 
-    expect(payload.content).toBe(`${WORKSPACE_PROMPT}\n\nPlease delete it instead`)
+    expect(payload.content).toBe(`${WORKSPACE_PROMPT.replace('src/app.ts', '/repo/src/app.ts')}\n\nPlease delete it instead`)
     expect(payload.attachments).toEqual([
-      expect.objectContaining({ type: 'file', path: 'src/app.ts', lineStart: 3, lineEnd: 5 }),
+      expect.objectContaining({ type: 'file', path: '/repo/src/app.ts', lineStart: 3, lineEnd: 5 }),
     ])
     expect(payload.options.displayContent).toBe('Please delete it instead')
 
@@ -134,7 +205,7 @@ describe('buildUserMessageResendPayload', () => {
     // the replacement would drift.
     const reread = createUserMessageEditDraft({
       content: 'Please delete it instead',
-      modelContent: `@"src/app.ts" ${payload.content}`,
+      modelContent: `@"/repo/src/app.ts" ${payload.content}`,
     })
     expect(reread.text).toBe('Please delete it instead')
     expect(reread.attachments.map(({ id: _id, ...rest }) => rest))

desktop/src/components/chat/userMessageEdit.ts

@@ -1,4 +1,6 @@
+import { splitSessionReferenceContext } from '@/lib/sessionReferences'
 import {
+  extractLeadingFileReferences,
   extractRestoredUserDisplay,
   pathsReferToSameFile,
   stripGeneratedImageMetadataLines,
@@ -46,6 +48,77 @@ function isSendable(attachment: UIAttachment): boolean {
   return Boolean(attachment.data || attachment.path)
 }
 
+function normalizeAttachmentPath(path: string | undefined): string | undefined {
+  return path?.replace(/\\/g, '/').replace(/^\.\//, '')
+}
+
+function referenceContextKey(attachment: UIAttachment): string {
+  return JSON.stringify([
+    attachment.lineStart,
+    attachment.lineEnd ?? attachment.lineStart,
+    attachment.diffSide,
+    attachment.note?.trim() || undefined,
+    attachment.quote?.trim() || undefined,
+  ])
+}
+
+type AttachmentCandidate = { attachment: UIAttachment; paths: Array<string | undefined> }
+
+function reconstructAttachments(
+  leading: UIAttachment[],
+  workspace: UIAttachment[],
+  originals: UIAttachment[],
+): UIAttachment[] {
+  const unmatchedLeading = leading.map((attachment) => ({ attachment, paths: [attachment.path] }))
+  const references: AttachmentCandidate[] = workspace.map((attachment) => ({ attachment, paths: [attachment.path] }))
+  const findLeading = (path: string | undefined, exact: boolean) => {
+    for (let index = unmatchedLeading.length - 1; index >= 0; index -= 1) {
+      const candidatePath = unmatchedLeading[index]!.attachment.path
+      if (exact
+        ? normalizeAttachmentPath(candidatePath) === normalizeAttachmentPath(path)
+        : pathsReferToSameFile(candidatePath, path)
+      ) return index
+    }
+    return -1
+  }
+  // The composer sends upload paths first, then workspace paths in reference
+  // order. Match from the end so a root-file selection cannot consume an upload
+  // with the same suffix; retain every unmatched absolute transport path.
+  for (let index = references.length - 1; index >= 0; index -= 1) {
+    const reference = references[index]!
+    let matchingIndex = findLeading(reference.attachment.path, true)
+    if (matchingIndex < 0) matchingIndex = findLeading(reference.attachment.path, false)
+    if (matchingIndex >= 0) {
+      const [matched] = unmatchedLeading.splice(matchingIndex, 1)
+      reference.paths.push(matched!.attachment.path)
+      reference.attachment = { ...reference.attachment, path: matched!.attachment.path }
+    }
+  }
+
+  const candidates: AttachmentCandidate[] = [...unmatchedLeading, ...references]
+  const unmatchedOriginals = originals.filter((attachment) => attachment.referenceKind !== 'chat-selection')
+  for (const candidate of candidates) {
+    const matchingIndex = unmatchedOriginals.findIndex((original) =>
+      original.path && candidate.paths.some((path) =>
+        normalizeAttachmentPath(path) === normalizeAttachmentPath(original.path),
+      ) && referenceContextKey(original) === referenceContextKey(candidate.attachment),
+    )
+    if (matchingIndex < 0) continue
+    const [original] = unmatchedOriginals.splice(matchingIndex, 1)
+    const metadata = Object.fromEntries(Object.entries(original!).filter(([, value]) => value !== undefined))
+    candidate.attachment = { ...candidate.attachment, ...metadata, path: candidate.attachment.path }
+  }
+  // The display parser may already have lost a prefix on a reloaded message.
+  // Its remaining bare chip is redundant with a recovered canonical path, but
+  // distinct absolute paths and data-only uploads must survive.
+  const extraAttachments = unmatchedOriginals.filter((original) =>
+    !original.path || !candidates.some((candidate) => candidate.paths.some((path) =>
+      normalizeAttachmentPath(path) === normalizeAttachmentPath(original.path),
+    )),
+  )
+  return [...candidates.map((candidate) => candidate.attachment), ...extraAttachments]
+}
+
 /**
  * Rebuild what the user originally typed from the prompt that actually reached
  * the model. Parsing the model-facing text with the same reader history uses
@@ -60,22 +133,18 @@ export function createUserMessageEditDraft(
   const messageAttachments = message.attachments ?? []
   const hasImage = messageAttachments.some((attachment) => attachment.type === 'image')
   const source = message.modelContent ?? message.content
-  const parsed = extractRestoredUserDisplay(hasImage ? stripGeneratedImageMetadataLines(source) : source)
-  const parsedAttachments = parsed.attachments ?? []
-
-  // Message attachments add what the model text cannot carry — inline images
-  // and data-only uploads. Chat selections are already part of the text body,
-  // and anything with a path the text already referenced would be a duplicate.
-  const extraAttachments = messageAttachments.filter((attachment) =>
-    attachment.referenceKind !== 'chat-selection' &&
-    !parsedAttachments.some((candidate) => attachment.path && pathsReferToSameFile(candidate.path, attachment.path)),
-  )
-
-  const attachments = [...parsedAttachments, ...extraAttachments].map((attachment, index) => ({
-    ...attachment,
-    id: `edit-attachment-${index}`,
-    sendable: isSendable(attachment),
-  }))
+  const sanitized = hasImage ? stripGeneratedImageMetadataLines(source) : source
+  const referenceContext = splitSessionReferenceContext(sanitized)
+  // Parse transport prefixes separately: the history display reader otherwise
+  // deduplicates absolute paths against relative workspace paths by suffix.
+  const leading = extractLeadingFileReferences(referenceContext.content)
+  const parsed = extractRestoredUserDisplay(leading.content)
+  const attachments = reconstructAttachments(leading.attachments ?? [], parsed.attachments ?? [], messageAttachments)
+    .map((attachment, index) => ({
+      ...attachment,
+      id: `edit-attachment-${index}`,
+      sendable: isSendable(attachment),
+    }))
 
   const body = parsed.content.trim()
   const text = attachments.length > 0 && ATTACHMENT_ONLY_PLACEHOLDERS.has(body) ? '' : body
@@ -83,7 +152,7 @@ export function createUserMessageEditDraft(
   return {
     text,
     attachments,
-    sessionReferences: message.sessionReferences ?? parsed.sessionReferences ?? [],
+    sessionReferences: message.sessionReferences ?? referenceContext.sessionReferences,
   }
 }
 

desktop/src/stores/chatStore.test.ts

@@ -16778,3 +16778,148 @@ describe('chatStore inactive complete-page retention', () => {
     }
   })
 })
+
+// Edit-and-resend must not send on a swallowed error or an obsolete response:
+// the old turn has already been removed from the server transcript.
+describe('chatStore strict history reload', () => {
+  const sessionId = 'strict-history-reload'
+  const strict = { requireApplied: true }
+  const oldMessages: UIMessage[] = [
+    { id: 'rewound-user', type: 'user_text', content: 'Removed prompt', timestamp: 1 },
+    { id: 'rewound-reply', type: 'assistant_text', content: 'Removed reply', timestamp: 2 },
+  ]
+
+  beforeEach(() => {
+    vi.mocked(sessionsApi.getFullHistory).mockReset()
+    vi.mocked(sessionsApi.getFullHistory).mockResolvedValue({ messages: [] })
+    useChatStore.setState({
+      ...initialState,
+      sessions: { [sessionId]: makeSession({ chatState: 'idle', historyHydrated: true, messages: oldMessages }) },
+    })
+  })
+
+  it('rejects a failed fetch instead of treating stale messages as refreshed', async () => {
+    const failure = new Error('History unavailable after rewind')
+    vi.mocked(sessionsApi.getFullHistory).mockRejectedValueOnce(failure)
+    await expect(useChatStore.getState().reloadHistory(sessionId, undefined, strict)).rejects.toBe(failure)
+    expect(useChatStore.getState().sessions[sessionId]?.messages).toBe(oldMessages)
+  })
+
+  it('keeps best-effort reload failures compatible for existing callers', async () => {
+    vi.mocked(sessionsApi.getFullHistory).mockRejectedValueOnce(new Error('History unavailable'))
+    await expect(useChatStore.getState().reloadHistory(sessionId)).resolves.toBeUndefined()
+    expect(useChatStore.getState().sessions[sessionId]?.messages).toBe(oldMessages)
+  })
+
+  it('resolves once complete authoritative history replaces the removed turn', async () => {
+    await expect(useChatStore.getState().reloadHistory(sessionId, undefined, strict)).resolves.toBeUndefined()
+    expect(useChatStore.getState().sessions[sessionId]?.messages).toEqual([])
+    expect(useChatStore.getState().sessions[sessionId]?.historyStatus).toBe('ready')
+  })
+
+  it('rejects a superseded response even when the transport ignores abort', async () => {
+    let resolveFirst!: (value: { messages: MessageEntry[] }) => void
+    vi.mocked(sessionsApi.getFullHistory).mockReturnValueOnce(new Promise((resolve) => { resolveFirst = resolve }))
+    const rejected = expect(useChatStore.getState().reloadHistory(sessionId, undefined, strict)).rejects.toThrow('History reload was not applied')
+    await useChatStore.getState().reloadHistory(sessionId)
+    resolveFirst({ messages: [{ id: 'stale', type: 'user', content: 'Stale', timestamp: '2026-10-03T00:00:00Z' }] })
+    await rejected
+    expect(useChatStore.getState().sessions[sessionId]?.messages).toEqual([])
+  })
+
+  it('rejects a response for a disconnected session without recreating it', async () => {
+    let resolveHistory!: (value: { messages: MessageEntry[] }) => void
+    vi.mocked(sessionsApi.getFullHistory).mockReturnValueOnce(new Promise((resolve) => { resolveHistory = resolve }))
+    const rejected = expect(useChatStore.getState().reloadHistory(sessionId, undefined, strict)).rejects.toThrow('History reload was not applied')
+    useChatStore.getState().disconnectSession(sessionId)
+    resolveHistory({ messages: [] })
+    await rejected
+    expect(useChatStore.getState().sessions[sessionId]).toBeUndefined()
+  })
+
+  it('rejects a guarded reload when a live prompt changes the conversation', async () => {
+    let resolveHistory!: (value: { messages: MessageEntry[] }) => void
+    vi.mocked(sessionsApi.getFullHistory).mockReturnValueOnce(new Promise((resolve) => { resolveHistory = resolve }))
+    const rejected = expect(useChatStore.getState().reloadHistory(sessionId, { messages: oldMessages }, strict)).rejects.toThrow('History reload was not applied')
+    useChatStore.getState().handleServerMessage(sessionId, { type: 'user_message_replay', content: 'Competing prompt' })
+    const liveMessages = useChatStore.getState().sessions[sessionId]?.messages
+    resolveHistory({ messages: [] })
+    await rejected
+    expect(useChatStore.getState().sessions[sessionId]?.messages).toBe(liveMessages)
+    expect(liveMessages).toContainEqual(expect.objectContaining({ content: 'Competing prompt' }))
+  })
+
+  it('rejects incomplete history instead of retaining removed rows as authoritative', async () => {
+    vi.mocked(sessionsApi.getFullHistory).mockResolvedValueOnce({ messages: [], page: {
+      nextCursor: 'older', hasMore: true, historyComplete: false,
+      sourceVersion: 'version', scannedBytes: 2, omittedOversizedEntries: 1,
+    } })
+    await expect(useChatStore.getState().reloadHistory(sessionId, undefined, strict)).rejects.toThrow('History reload was not applied')
+    expect(useChatStore.getState().sessions[sessionId]?.messages).toBe(oldMessages)
+  })
+})
+
+describe('chatStore composer prefill references', () => {
+  const sessionId = 'prefill-references'
+
+  beforeEach(() => {
+    useChatStore.setState({ ...initialState, sessions: { [sessionId]: makeSession({ chatState: 'idle' }) } })
+  })
+
+  it('retains session context with an edited prompt handed to the composer', () => {
+    const sessionReferences = [{ sessionId: 'referenced-session' }]
+    useChatStore.getState().queueComposerPrefill(sessionId, { text: 'Edited prompt', sessionReferences })
+    expect(useChatStore.getState().sessions[sessionId]?.composerPrefill).toMatchObject({
+      text: 'Edited prompt', sessionReferences,
+    })
+  })
+
+  it('keeps the existing prefill shape when no session references were supplied', () => {
+    useChatStore.getState().queueComposerPrefill(sessionId, { text: 'Existing prefill' })
+    expect(useChatStore.getState().sessions[sessionId]?.composerPrefill).not.toHaveProperty('sessionReferences')
+  })
+})
+
+
+// A completed rewind must retain its edit even when the source tab was closed
+// while the API call was pending; delayed ordinary prefills must stay ignored.
+describe('chatStore closed-session edit recovery', () => {
+  const sessionId = 'closed-edit-recovery'
+  const prefill = {
+    text: 'Recovered edit',
+    attachments: [{ type: 'file' as const, name: 'app.ts', path: '/repo/app.ts' }],
+    sessionReferences: [{ sessionId: 'context-session' }],
+  }
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    useChatStore.setState({ ...initialState, sessions: {} })
+  })
+
+  it('ignores a normal delayed prefill for a missing session', () => {
+    useChatStore.getState().queueComposerPrefill(sessionId, prefill)
+    expect(useChatStore.getState().sessions[sessionId]).toBeUndefined()
+  })
+
+  it('retains an explicitly recovered edit without opening a socket or sending it', async () => {
+    const { wsManager } = await import('../api/websocket')
+    useChatStore.getState().queueComposerPrefill(sessionId, prefill, { restoreMissingSession: true })
+    expect(useChatStore.getState().sessions[sessionId]).toMatchObject({
+      connectionState: 'disconnected', chatState: 'idle', messages: [], composerPrefill: prefill,
+    })
+    expect(wsManager.connect).not.toHaveBeenCalled()
+    expect(wsManager.send).not.toHaveBeenCalled()
+    expect(sessionsApi.getFullHistory).not.toHaveBeenCalled()
+  })
+
+  it('preserves the recovered prefill when the user reopens its disconnected session', () => {
+    useChatStore.setState({ sessions: { [sessionId]: makeSession({ connectionState: 'disconnected', chatState: 'idle' }) } })
+    useChatStore.getState().queueComposerPrefill(sessionId, prefill)
+    const recovered = useChatStore.getState().sessions[sessionId]?.composerPrefill
+    useChatStore.getState().connectToSession(sessionId, { minimalBootstrap: true, prewarm: false, applyRuntimeSelection: false })
+    expect(useChatStore.getState().sessions[sessionId]?.composerPrefill).toBe(recovered)
+    expect(useChatStore.getState().sessions[sessionId]?.connectionState).toBe('connecting')
+    expect(sessionsApi.getFullHistory).not.toHaveBeenCalled()
+    expect(sendMock).not.toHaveBeenCalled()
+  })
+})

desktop/src/stores/chatStore.ts

@@ -209,6 +209,7 @@ export type PerSessionState = {
   composerPrefill?: {
     text: string
     attachments?: UIAttachment[]
+    sessionReferences?: Array<{ sessionId: string }>
     mode?: ComposerPrefillMode
     nonce: number
   } | null
@@ -442,10 +443,19 @@ type ChatStore = {
       messages: UIMessage[]
       backgroundAgentTasks?: Record<string, BackgroundAgentTask>
     },
+    /** Reject errors and skipped/incomplete reloads after a destructive rewind. */
+    options?: { requireApplied?: boolean },
   ) => Promise<void>
   queueComposerPrefill: (
     sessionId: string,
-    prefill: { text: string; attachments?: UIAttachment[]; mode?: ComposerPrefillMode },
+    prefill: {
+      text: string
+      attachments?: UIAttachment[]
+      sessionReferences?: Array<{ sessionId: string }>
+      mode?: ComposerPrefillMode
+    },
+    /** Retain an edit after its committed rewind outlives the source tab. */
+    options?: { restoreMissingSession?: boolean },
   ) => void
   clearComposerPrefill: (sessionId: string, nonce?: number) => void
   queueComposerInsertion: (
@@ -3043,6 +3053,7 @@ export const useChatStore = create<ChatStore>((setState, get) => {
           historyBootstrapDisabled: options?.minimalBootstrap === true,
           activeGoal: existing?.activeGoal ?? null,
           composerDraft: existing?.composerDraft ?? null,
+          composerPrefill: existing?.composerPrefill ?? null,
           repositoryLaunchDraft: existing?.repositoryLaunchDraft ?? null,
           queuedUserMessages: existing?.queuedUserMessages ?? [],
           backgroundAgentTasks: existing?.backgroundAgentTasks ?? {},
@@ -4072,8 +4083,11 @@ export const useChatStore = create<ChatStore>((setState, get) => {
     await loadOlderHistoryPage(sessionId, get, set)
   },
 
-  reloadHistory: async (sessionId, guard) => {
-    if (isSideChatSession(sessionId)) return
+  reloadHistory: async (sessionId, guard, options) => {
+    const skipReload = () => {
+      if (options?.requireApplied) throw new Error('History reload was not applied')
+    }
+    if (isSideChatSession(sessionId)) return skipReload()
     if (historyPageControllers.has(sessionId)) {
       historyPageControllers.get(sessionId)?.abort()
       historyPageControllers.delete(sessionId)
@@ -4092,7 +4106,7 @@ export const useChatStore = create<ChatStore>((setState, get) => {
       if (pendingLoad?.lifecycleGeneration === lifecycleGeneration) {
         await pendingLoad.promise
       }
-      if (!isCurrentHistoryLifecycle(sessionId, lifecycleGeneration) || controller.signal.aborted) return
+      if (!isCurrentHistoryLifecycle(sessionId, lifecycleGeneration) || controller.signal.aborted) return skipReload()
       historyRecoveryControllers.get(sessionId)?.abort()
       historyRecoveryControllers.delete(sessionId)
       // A reload can queue behind a cold load. Capture snapshot baselines only
@@ -4125,12 +4139,15 @@ export const useChatStore = create<ChatStore>((setState, get) => {
         sessionOwnedActivityToolUseIds(get().sessions[sessionId]),
         controller.signal,
       )
+      // A partial reload merges existing rows. It cannot prove that a rewound
+      // turn disappeared, so strict callers must keep their draft for retry.
+      if (options?.requireApplied && !historyComplete) return skipReload()
       durableHistoryRows.set(sessionId, new WeakSet(uiMessages))
 
       if (
         !isCurrentHistoryLifecycle(sessionId, lifecycleGeneration) ||
         historyReloadGenerations.get(sessionId) !== reloadGeneration
-      ) return
+      ) return skipReload()
 
       if (guard) {
         const current = get().sessions[sessionId]
@@ -4139,7 +4156,7 @@ export const useChatStore = create<ChatStore>((setState, get) => {
           current.chatState !== 'idle' ||
           (current.historyMutationEpoch ?? 0) !== requestedMutationEpoch
         ) {
-          return
+          return skipReload()
         }
       }
 
@@ -4242,7 +4259,7 @@ export const useChatStore = create<ChatStore>((setState, get) => {
         }
       })
 
-      if (!historyApplied) return
+      if (!historyApplied) return skipReload()
       if (!historyComplete && page) void recoverSessionHistory(sessionId, page.sourceVersion)
       const terminalReconnectBoundary = terminalReconnectHistoryBoundaries.get(sessionId)
       if (
@@ -4306,7 +4323,7 @@ export const useChatStore = create<ChatStore>((setState, get) => {
           requestedTasks,
         )
       }
-    } catch {
+    } catch (error) {
       // A stop failure can arrive before the task history that identifies it.
       // If that history request fails, surface the failure instead of leaving it
       // cached forever waiting for a reconciliation that may never happen.
@@ -4328,22 +4345,34 @@ export const useChatStore = create<ChatStore>((setState, get) => {
             )),
         }
       })
+      if (options?.requireApplied) throw error
     } finally {
       if (historyReloadControllers.get(sessionId) === controller) historyReloadControllers.delete(sessionId)
     }
   },
 
-  queueComposerPrefill: (sessionId, prefill) => {
-    set((state) => ({
-      sessions: updateSessionIn(state.sessions, sessionId, () => ({
-        composerPrefill: {
-          text: prefill.text,
-          attachments: prefill.attachments,
-          mode: prefill.mode,
-          nonce: Date.now(),
+  queueComposerPrefill: (sessionId, prefill, options) => {
+    set((state) => {
+      const session = state.sessions[sessionId] ?? (
+        options?.restoreMissingSession ? createDefaultSessionState() : undefined
+      )
+      if (!session) return state
+      return {
+        sessions: {
+          ...state.sessions,
+          [sessionId]: {
+            ...session,
+            composerPrefill: {
+              text: prefill.text,
+              attachments: prefill.attachments,
+              ...(prefill.sessionReferences ? { sessionReferences: prefill.sessionReferences } : {}),
+              mode: prefill.mode,
+              nonce: Date.now(),
+            },
+          },
         },
-      })),
-    }))
+      }
+    })
   },
 
   clearComposerPrefill: (sessionId, nonce) => {
@@ -7235,7 +7264,7 @@ function getReferenceName(referencePath: string): string {
   return name || referencePath
 }
 
-function extractLeadingFileReferences(text: string): {
+export function extractLeadingFileReferences(text: string): {
   content: string
   attachments?: UIAttachment[]
   modelContent?: string