fix(web): preserve queued edits across cancellation (#1324)

web/src/components/AssistantChat/QueuedMessagesBar.test.tsx

@@ -1,7 +1,457 @@
-import { describe, expect, it } from 'vitest'
+import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import type { DecryptedMessage } from '@/types/api'
-import { computeCanCancel, computeEditPendingSchedule, getQueuedMessageEditText, getQueuedMessagePreview, sortQueuedMessages } from './QueuedMessagesBar'
+import type { PendingSchedule } from '@/components/AssistantChat/ScheduleTimePicker'
+import {
+    computeCanCancel,
+    computeEditPendingSchedule,
+    getQueuedMessageEditText,
+    getQueuedMessagePreview,
+    QueuedMessagesBar,
+    sortQueuedMessages,
+} from './QueuedMessagesBar'
 import { formatScheduledTime } from '@/lib/scheduledTime'
+import { clearQueuedEditRecovery, getQueuedEditRecovery } from '@/lib/queued-edit-recovery'
+
+type DeferredCancelResult = { status: 'cancelled' | 'invoked' }
+
+const mocks = vi.hoisted(() => ({
+    composerText: '',
+    composerSetText: vi.fn(),
+    addToast: vi.fn(),
+    mutateAsync: vi.fn(),
+    resolveCancel: null as ((result: DeferredCancelResult) => void) | null,
+    rejectCancel: null as ((reason?: unknown) => void) | null,
+    saveDraft: vi.fn(),
+    messageWindowState: { messages: [] as unknown[] },
+}))
+
+vi.mock('@assistant-ui/react', () => ({
+    useAui: () => ({
+        composer: () => ({
+            getState: () => ({ text: mocks.composerText }),
+            setText: mocks.composerSetText,
+        }),
+    }),
+    useAuiState: (selector: (state: { composer: { text: string } }) => unknown) => selector({
+        composer: { text: mocks.composerText },
+    }),
+}))
+
+vi.mock('@/lib/message-window-store', () => ({
+    getMessageWindowState: () => mocks.messageWindowState,
+    subscribeMessageWindow: () => () => {},
+}))
+
+vi.mock('@/hooks/mutations/useCancelQueuedMessage', () => ({
+    useCancelQueuedMessage: () => ({
+        isPending: false,
+        variables: undefined,
+        mutateAsync: mocks.mutateAsync,
+    }),
+}))
+
+vi.mock('@/lib/composer-drafts', () => ({
+    saveDraft: mocks.saveDraft,
+}))
+
+vi.mock('@/lib/use-translation', () => ({
+    useTranslation: () => ({ t: (key: string) => key }),
+}))
+
+vi.mock('@/lib/toast-context', () => ({
+    useToast: () => ({ addToast: mocks.addToast }),
+}))
+
+function makeQueuedMessage(scheduledAt: number | null = null, id = 'server-message-id'): DecryptedMessage {
+    return {
+        id,
+        localId: `local-${id}`,
+        createdAt: 1000,
+        seq: 1,
+        scheduledAt,
+        invokedAt: null,
+        status: 'queued',
+        content: {
+            role: 'user',
+            content: { type: 'text', text: 'Queued request' },
+        },
+    } as unknown as DecryptedMessage
+}
+
+function renderQueuedMessage(
+    scheduledAt: number | null = null,
+    pendingSchedule: PendingSchedule | null = null,
+    pendingScheduleRevision = 0,
+) {
+    const onEdit = vi.fn()
+    let currentPendingScheduleRevision = pendingScheduleRevision
+    mocks.messageWindowState = { messages: [makeQueuedMessage(scheduledAt)] }
+    const view = render(
+        <QueuedMessagesBar
+            sessionId="session-1"
+            api={null}
+            pendingSchedule={pendingSchedule}
+            pendingScheduleRevision={currentPendingScheduleRevision}
+            onEdit={onEdit}
+        />
+    )
+    return {
+        onEdit,
+        unmount: view.unmount,
+        rerender: (nextPendingSchedule: PendingSchedule | null, nextPendingScheduleRevision = currentPendingScheduleRevision) => {
+            currentPendingScheduleRevision = nextPendingScheduleRevision
+            view.rerender(
+                <QueuedMessagesBar
+                    sessionId="session-1"
+                    api={null}
+                    pendingSchedule={nextPendingSchedule}
+                    pendingScheduleRevision={currentPendingScheduleRevision}
+                    onEdit={onEdit}
+                />
+            )
+        },
+    }
+}
+
+beforeEach(() => {
+    mocks.composerText = ''
+    mocks.composerSetText.mockReset()
+    mocks.addToast.mockReset()
+    mocks.mutateAsync.mockReset()
+    mocks.resolveCancel = null
+    mocks.rejectCancel = null
+    mocks.saveDraft.mockReset()
+    mocks.messageWindowState = { messages: [] }
+    clearQueuedEditRecovery('session-1')
+    mocks.mutateAsync.mockImplementation(() => new Promise<DeferredCancelResult>((resolve, reject) => {
+        mocks.resolveCancel = resolve
+        mocks.rejectCancel = reject
+    }))
+})
+
+async function resolveCancel(result: DeferredCancelResult): Promise<void> {
+    await act(async () => {
+        mocks.resolveCancel?.(result)
+        await Promise.resolve()
+    })
+}
+
+function installManualAnimationFrames() {
+    let nextHandle = 1
+    const pending = new Map<number, FrameRequestCallback>()
+    const history = new Map<number, FrameRequestCallback>()
+    vi.stubGlobal('requestAnimationFrame', (callback: FrameRequestCallback) => {
+        const handle = nextHandle++
+        pending.set(handle, callback)
+        history.set(handle, callback)
+        return handle
+    })
+    vi.stubGlobal('cancelAnimationFrame', (handle: number) => {
+        pending.delete(handle)
+    })
+
+    return {
+        flushPending() {
+            while (pending.size > 0) {
+                const [handle, callback] = pending.entries().next().value as [number, FrameRequestCallback]
+                pending.delete(handle)
+                callback(Date.now())
+            }
+        },
+        flushHistory() {
+            for (const callback of history.values()) {
+                callback(Date.now())
+            }
+        },
+    }
+}
+
+afterEach(() => {
+    vi.unstubAllGlobals()
+})
+
+describe('QueuedMessagesBar edit restore', () => {
+    it('keeps a newly typed draft and its schedule when the deferred cancel succeeds', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { onEdit } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        expect(mocks.resolveCancel).not.toBeNull()
+
+        // The user starts a new draft while DELETE /messages/:id is still pending.
+        mocks.composerText = 'New draft typed while cancelling'
+        await resolveCancel({ status: 'cancelled' })
+
+        expect(mocks.composerSetText).not.toHaveBeenCalled()
+        expect(onEdit).not.toHaveBeenCalled()
+        expect(mocks.addToast).toHaveBeenCalledWith({
+            title: 'queuedMessages.editCurrentDraftKept',
+            body: '',
+            sessionId: 'session-1',
+            url: window.location.href,
+        })
+    })
+
+    it('keeps a newly selected schedule when the composer text is unchanged', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { onEdit, rerender } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        // The user changes only the clock selection while DELETE /messages/:id is pending.
+        rerender({ type: 'preset', preset: '+30m' }, 1)
+        await resolveCancel({ status: 'cancelled' })
+
+        expect(mocks.composerSetText).not.toHaveBeenCalled()
+        expect(onEdit).not.toHaveBeenCalled()
+        expect(mocks.addToast).toHaveBeenCalledWith({
+            title: 'queuedMessages.editCurrentDraftKept',
+            body: '',
+            sessionId: 'session-1',
+            url: window.location.href,
+        })
+    })
+
+    it('keeps the current state after selecting and then clearing a schedule', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { onEdit, rerender } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        rerender({ type: 'preset', preset: '+30m' }, 1)
+        rerender(null, 2)
+        await resolveCancel({ status: 'cancelled' })
+
+        expect(mocks.composerSetText).not.toHaveBeenCalled()
+        expect(onEdit).not.toHaveBeenCalled()
+        expect(mocks.addToast).toHaveBeenCalledWith({
+            title: 'queuedMessages.editCurrentDraftKept',
+            body: '',
+            sessionId: 'session-1',
+            url: window.location.href,
+        })
+    })
+
+    it('restores both text and schedule when the composer is unchanged', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { onEdit } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        await resolveCancel({ status: 'cancelled' })
+
+        expect(mocks.composerSetText).toHaveBeenCalledWith('Queued request')
+        expect(onEdit).toHaveBeenCalledWith({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+        })
+        expect(mocks.addToast).not.toHaveBeenCalled()
+    })
+
+    it('treats structurally equal schedule props as unchanged', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { onEdit, rerender } = renderQueuedMessage(scheduledAt, { type: 'preset', preset: '+5m' })
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        rerender({ type: 'preset', preset: '+5m' })
+        await resolveCancel({ status: 'cancelled' })
+
+        expect(mocks.composerSetText).toHaveBeenCalledWith('Queued request')
+        expect(onEdit).toHaveBeenCalledWith({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+        })
+    })
+
+    it('globally disables queued operations and keeps the first edit completion', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const first = makeQueuedMessage(null, 'server-message-a')
+        const second = makeQueuedMessage(scheduledAt, 'server-message-b')
+        mocks.messageWindowState = { messages: [first, second] }
+        const onEdit = vi.fn()
+        render(
+            <QueuedMessagesBar
+                sessionId="session-1"
+                api={null}
+                pendingSchedule={null}
+                pendingScheduleRevision={0}
+                onEdit={onEdit}
+            />
+        )
+
+        const editButtons = screen.getAllByRole('button', { name: 'Edit queued message' })
+        const cancelButtons = screen.getAllByRole('button', { name: 'Cancel queued message' })
+        // QueuedMessagesBar orders immediate rows first, so the scheduled edit is second.
+        fireEvent.click(editButtons[1]!)
+
+        await waitFor(() => {
+            for (const button of [...screen.getAllByRole('button', { name: 'Edit queued message' }), ...screen.getAllByRole('button', { name: 'Cancel queued message' })]) {
+                expect(button).toBeDisabled()
+            }
+        })
+        fireEvent.click(cancelButtons[0]!)
+        fireEvent.click(editButtons[0]!)
+        expect(mocks.mutateAsync).toHaveBeenCalledTimes(1)
+
+        await resolveCancel({ status: 'cancelled' })
+        expect(mocks.composerSetText).toHaveBeenCalledWith('Queued request')
+        expect(onEdit).toHaveBeenCalledWith({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+        })
+    })
+
+    it('persists an unmounted edit result and restores it after the same session remounts', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { unmount } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        unmount()
+        await resolveCancel({ status: 'cancelled' })
+
+        expect(mocks.saveDraft).not.toHaveBeenCalled()
+        expect(getQueuedEditRecovery('session-1')).toEqual(expect.objectContaining({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+            composerTextAtEdit: '',
+            pendingScheduleAtEdit: null,
+        }))
+
+        const { onEdit } = renderQueuedMessage(scheduledAt)
+        await waitFor(() => expect(mocks.composerSetText).toHaveBeenCalledWith('Queued request'))
+        expect(onEdit).toHaveBeenCalledWith({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+        })
+        expect(getQueuedEditRecovery('session-1')).toBeNull()
+    })
+
+    it('notifies a same-session remount that returns before the edit cancellation completes', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { unmount } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        unmount()
+        const { onEdit } = renderQueuedMessage(scheduledAt)
+
+        await waitFor(() => {
+            expect(screen.getByRole('button', { name: 'Edit queued message' })).toBeDisabled()
+            expect(screen.getByRole('button', { name: 'Cancel queued message' })).toBeDisabled()
+        })
+        fireEvent.click(screen.getByRole('button', { name: 'Cancel queued message' }))
+        expect(mocks.mutateAsync).toHaveBeenCalledTimes(1)
+
+        await resolveCancel({ status: 'cancelled' })
+
+        await waitFor(() => expect(mocks.composerSetText).toHaveBeenCalledWith('Queued request'))
+        expect(onEdit).toHaveBeenCalledWith({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+        })
+        expect(getQueuedEditRecovery('session-1')).toBeNull()
+    })
+
+    it('keeps a remounted session busy until its queued-edit recovery is consumed', async () => {
+        const raf = installManualAnimationFrames()
+        const scheduledAt = Date.now() + 60_000
+        const { unmount } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        unmount()
+        const { onEdit } = renderQueuedMessage(scheduledAt)
+        await resolveCancel({ status: 'cancelled' })
+
+        // The mutation token has ended, but the unconsumed recovery keeps Q2 busy.
+        expect(screen.getByRole('button', { name: 'Edit queued message' })).toBeDisabled()
+        fireEvent.click(screen.getByRole('button', { name: 'Cancel queued message' }))
+        expect(mocks.mutateAsync).toHaveBeenCalledTimes(1)
+
+        raf.flushPending()
+        await waitFor(() => expect(screen.getByRole('button', { name: 'Edit queued message' })).not.toBeDisabled())
+        expect(onEdit).toHaveBeenCalledWith({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+        })
+    })
+
+    it('ignores overlapping and disposed recovery animation callbacks, then allows a fresh remount to consume', async () => {
+        const raf = installManualAnimationFrames()
+        const scheduledAt = Date.now() + 60_000
+        const first = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        first.unmount()
+        const second = renderQueuedMessage(scheduledAt)
+        await resolveCancel({ status: 'cancelled' })
+        second.unmount()
+
+        raf.flushHistory()
+        expect(mocks.composerSetText).not.toHaveBeenCalled()
+        expect(second.onEdit).not.toHaveBeenCalled()
+        expect(getQueuedEditRecovery('session-1')).not.toBeNull()
+
+        const third = renderQueuedMessage(scheduledAt)
+        raf.flushPending()
+        await waitFor(() => expect(mocks.composerSetText).toHaveBeenCalledWith('Queued request'))
+        expect(third.onEdit).toHaveBeenCalledWith({
+            text: 'Queued request',
+            pendingSchedule: { type: 'absolute', ms: scheduledAt },
+        })
+        expect(getQueuedEditRecovery('session-1')).toBeNull()
+    })
+
+    it('clears a conflicting recovery once and never restores it after a later composer clear', async () => {
+        const scheduledAt = Date.now() + 60_000
+        const { unmount } = renderQueuedMessage(scheduledAt)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        unmount()
+        await resolveCancel({ status: 'cancelled' })
+        mocks.composerText = 'Current draft wins'
+
+        const { onEdit, rerender } = renderQueuedMessage(scheduledAt)
+        await waitFor(() => expect(mocks.addToast).toHaveBeenCalledWith({
+            title: 'queuedMessages.editCurrentDraftKept',
+            body: '',
+            sessionId: 'session-1',
+            url: window.location.href,
+        }))
+        expect(getQueuedEditRecovery('session-1')).toBeNull()
+
+        mocks.composerText = ''
+        rerender(null, 0)
+        await new Promise((resolve) => setTimeout(resolve, 0))
+        expect(mocks.composerSetText).not.toHaveBeenCalled()
+        expect(onEdit).not.toHaveBeenCalled()
+    })
+
+    it('keeps the already-invoked toast behavior', async () => {
+        const { onEdit } = renderQueuedMessage()
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        await resolveCancel({ status: 'invoked' })
+
+        expect(mocks.composerSetText).not.toHaveBeenCalled()
+        expect(onEdit).not.toHaveBeenCalled()
+        expect(mocks.addToast).toHaveBeenCalledWith({
+            title: 'queuedMessages.editAlreadyInvoked',
+            body: '',
+            sessionId: 'session-1',
+            url: window.location.href,
+        })
+    })
+
+    it('does not restore an edit when cancel fails', async () => {
+        const { onEdit } = renderQueuedMessage(Date.now() + 60_000)
+
+        fireEvent.click(screen.getByRole('button', { name: 'Edit queued message' }))
+        await act(async () => {
+            mocks.rejectCancel?.(new Error('network failed'))
+            await Promise.resolve()
+        })
+
+        expect(mocks.composerSetText).not.toHaveBeenCalled()
+        expect(onEdit).not.toHaveBeenCalled()
+    })
+})
 
 /**
  * Unit tests for computeCanCancel — the race guard that prevents sending

web/src/components/AssistantChat/QueuedMessagesBar.tsx

@@ -1,5 +1,5 @@
-import { useAui } from '@assistant-ui/react'
-import { useCallback, useMemo, useSyncExternalStore } from 'react'
+import { useAui, useAuiState } from '@assistant-ui/react'
+import { useCallback, useEffect, useMemo, useRef, useSyncExternalStore } from 'react'
 import type { ApiClient } from '@/api/client'
 import { getMessageWindowState, subscribeMessageWindow } from '@/lib/message-window-store'
 import { isQueuedForInvocation } from '@/lib/messages'
@@ -11,6 +11,16 @@ import { useTranslation } from '@/lib/use-translation'
 import { useToast } from '@/lib/toast-context'
 import type { PendingSchedule } from '@/components/AssistantChat/ScheduleTimePicker'
 import { formatScheduledTime } from '@/lib/scheduledTime'
+import {
+    beginQueuedOperation,
+    clearQueuedEditRecovery,
+    endQueuedOperation,
+    getQueuedEditRecovery,
+    isQueuedOperationPending,
+    saveQueuedEditRecovery,
+    subscribeQueuedEditRecovery,
+    subscribeQueuedOperation,
+} from '@/lib/queued-edit-recovery'
 
 function ClockIcon() {
     return (
@@ -112,6 +122,19 @@ export function computeEditPendingSchedule(
     return { type: 'absolute', ms: scheduledAt }
 }
 
+function pendingSchedulesEqual(a: PendingSchedule | null, b: PendingSchedule | null): boolean {
+    if (a === b) return true
+    if (a === null || b === null) return false
+    switch (a.type) {
+        case 'preset':
+            return b.type === 'preset' && a.preset === b.preset
+        case 'absolute':
+            return b.type === 'absolute' && a.ms === b.ms
+    }
+    const exhaustive: never = a
+    return exhaustive
+}
+
 /**
  * Determines whether the user can cancel or edit a queued message.
  *
@@ -150,10 +173,16 @@ export function computeCanCancel({
 export function QueuedMessagesBar({
     sessionId,
     api,
+    pendingSchedule,
+    pendingScheduleRevision,
     onEdit,
 }: {
     sessionId: string
     api: ApiClient | null
+    /** Current composer schedule, used only to guard an asynchronous edit restore. */
+    pendingSchedule: PendingSchedule | null
+    /** Monotonic per-session revision; schedule selections win over an async edit restore. */
+    pendingScheduleRevision: number
     /**
      * Called when the user clicks Edit on a queued message.
      * The parent should restore `text` into the composer and `pendingSchedule` into the schedule state.
@@ -163,9 +192,108 @@ export function QueuedMessagesBar({
 }) {
     const queued = useQueuedMessages(sessionId)
     const assistantApi = useAui()
+    const composerText = useAuiState((state) => state.composer.text)
     const cancelMutation = useCancelQueuedMessage(api)
     const { t } = useTranslation()
     const { addToast } = useToast()
+    const pendingScheduleRef = useRef(pendingSchedule)
+    const pendingScheduleRevisionRef = useRef(pendingScheduleRevision)
+    const composerTextRef = useRef(composerText)
+    const onEditRef = useRef(onEdit)
+    const mountedRef = useRef(true)
+    const attemptedRecoveryIdsRef = useRef(new Set<string>())
+    // onSuccess runs after the cancel request completes, so it must read the
+    // newest schedule rather than the render that initiated the request.
+    pendingScheduleRef.current = pendingSchedule
+    pendingScheduleRevisionRef.current = pendingScheduleRevision
+    composerTextRef.current = composerText
+    onEditRef.current = onEdit
+
+    const queuedOperationPending = useSyncExternalStore(
+        useCallback((listener) => subscribeQueuedOperation(sessionId, listener), [sessionId]),
+        useCallback(() => isQueuedOperationPending(sessionId), [sessionId]),
+        () => false,
+    )
+
+    useEffect(() => {
+        mountedRef.current = true
+        return () => {
+            mountedRef.current = false
+        }
+    }, [])
+
+    const restoreQueuedEditRecovery = useCallback(() => {
+        const recovery = getQueuedEditRecovery(sessionId)
+        if (!recovery || attemptedRecoveryIdsRef.current.has(recovery.id)) return
+        attemptedRecoveryIdsRef.current.add(recovery.id)
+
+        const currentText = assistantApi.composer().getState().text
+        // A new session starts at revision 0. Any schedule interaction, even
+        // select-then-clear back to null, increments it and wins over recovery.
+        const textCompatible = currentText === recovery.composerTextAtEdit
+        const scheduleCompatible = pendingScheduleRevisionRef.current === 0
+        if (!textCompatible || !scheduleCompatible) {
+            if (mountedRef.current) {
+                addToast({
+                    title: t('queuedMessages.editCurrentDraftKept'),
+                    body: '',
+                    sessionId,
+                    url: window.location.href,
+                })
+            }
+            clearQueuedEditRecovery(sessionId)
+            return
+        }
+
+        if (recovery.text) {
+            assistantApi.composer().setText(recovery.text)
+        }
+        onEditRef.current?.({ text: recovery.text, pendingSchedule: recovery.pendingSchedule })
+        clearQueuedEditRecovery(sessionId)
+    }, [addToast, assistantApi, sessionId, t])
+
+    useEffect(() => {
+        let disposed = false
+        let generation = 0
+        const scheduledHandles = new Set<number>()
+        const scheduleFrame = (callback: FrameRequestCallback): number => {
+            if (typeof requestAnimationFrame === 'function') return requestAnimationFrame(callback)
+            return window.setTimeout(() => callback(Date.now()), 0)
+        }
+        const cancelFrame = (handle: number): void => {
+            if (typeof cancelAnimationFrame === 'function') {
+                cancelAnimationFrame(handle)
+            } else {
+                window.clearTimeout(handle)
+            }
+        }
+        const attemptAfterComposerDraft = () => {
+            const attemptGeneration = ++generation
+            const scheduleAttemptFrame = (callback: FrameRequestCallback) => {
+                let handle = 0
+                handle = scheduleFrame((timestamp) => {
+                    scheduledHandles.delete(handle)
+                    if (disposed || attemptGeneration !== generation) return
+                    callback(timestamp)
+                })
+                scheduledHandles.add(handle)
+            }
+            scheduleAttemptFrame(() => {
+                scheduleAttemptFrame(restoreQueuedEditRecovery)
+            })
+        }
+        attemptAfterComposerDraft()
+        const unsubscribe = subscribeQueuedEditRecovery(sessionId, attemptAfterComposerDraft)
+        return () => {
+            disposed = true
+            generation++
+            unsubscribe()
+            for (const handle of scheduledHandles) {
+                cancelFrame(handle)
+            }
+            scheduledHandles.clear()
+        }
+    }, [restoreQueuedEditRecovery, sessionId])
 
     if (queued.length === 0) {
         return null
@@ -192,54 +320,99 @@ export function QueuedMessagesBar({
                         const editText = getQueuedMessageEditText(preview)
                         const hasAttachments = attachmentNames.length > 0
                         const localId = msg.localId ?? msg.id
-                        const isPending = cancelMutation.isPending && cancelMutation.variables?.localId === localId
+                        const isPending = cancelMutation.isPending || queuedOperationPending
                         const canCancel = computeCanCancel({ id: msg.id, localId: msg.localId, isPending })
 
                         const handleCancel = () => {
                             if (!canCancel) return
-                            cancelMutation.mutate({
+                            const token = beginQueuedOperation(sessionId)
+                            if (!token) return
+                            void cancelMutation.mutateAsync({
                                 sessionId,
                                 messageId: msg.id,
                                 localId,
                                 snapshot: msg,
+                            }).catch(() => {
+                                // useCancelQueuedMessage restores the optimistic row and gives haptic feedback.
+                            }).finally(() => {
+                                endQueuedOperation(sessionId, token)
                             })
                         }
 
-                        const handleEdit = () => {
+                        const handleEdit = async () => {
                             if (!canCancel) return
                             // Edit = cancel + restore composer (text + schedule).
                             // Works the same for immediate-queued and future-scheduled messages.
                             const restoredPendingSchedule = computeEditPendingSchedule(msg.scheduledAt, Date.now())
+                            // The cancel request is asynchronous. Keep the exact composer text from the
+                            // click so a newer draft or schedule is never replaced when success arrives.
+                            const composerTextAtEdit = assistantApi.composer().getState().text
+                            const pendingScheduleAtEdit = pendingScheduleRef.current
+                            const pendingScheduleRevisionAtEdit = pendingScheduleRevisionRef.current
+                            const token = beginQueuedOperation(sessionId)
+                            if (!token) return
 
-                            cancelMutation.mutate(
-                                {
+                            try {
+                                const result = await cancelMutation.mutateAsync({
                                     sessionId,
                                     messageId: msg.id,
                                     localId,
                                     snapshot: msg,
-                                },
-                                {
-                                    onSuccess: (result) => {
-                                        // Race guard: if the agent already consumed this message, skip prefill
-                                        // and inform the user so they aren't confused by the row disappearing.
-                                        if (result.status === 'invoked') {
-                                            addToast({
-                                                title: t('queuedMessages.editAlreadyInvoked'),
-                                                body: '',
-                                                sessionId,
-                                                url: window.location.href,
-                                            })
-                                            return
-                                        }
-                                        // Restore text into composer
-                                        if (editText) {
-                                            assistantApi.composer().setText(editText)
-                                        }
-                                        // Restore schedule via parent callback (if provided)
-                                        onEdit?.({ text: editText, pendingSchedule: restoredPendingSchedule })
-                                    },
+                                })
+                                // Race guard: if the agent already consumed this message, skip prefill
+                                // and inform the user so they aren't confused by the row disappearing.
+                                if (result.status === 'invoked') {
+                                    if (mountedRef.current) {
+                                        addToast({
+                                            title: t('queuedMessages.editAlreadyInvoked'),
+                                            body: '',
+                                            sessionId,
+                                            url: window.location.href,
+                                        })
+                                    }
+                                    return
                                 }
-                            )
+
+                                const currentText = mountedRef.current
+                                    ? assistantApi.composer().getState().text
+                                    : composerTextRef.current
+                                const composerChanged = currentText !== composerTextAtEdit
+                                const scheduleChanged = pendingScheduleRevisionRef.current !== pendingScheduleRevisionAtEdit
+                                    || !pendingSchedulesEqual(pendingScheduleRef.current, pendingScheduleAtEdit)
+                                // Restore text and schedule as one unit. If either changed while the
+                                // cancel was pending, the user's newer composer state wins.
+                                if (composerChanged || scheduleChanged) {
+                                    if (mountedRef.current) {
+                                        addToast({
+                                            title: t('queuedMessages.editCurrentDraftKept'),
+                                            body: '',
+                                            sessionId,
+                                            url: window.location.href,
+                                        })
+                                    }
+                                    return
+                                }
+                                if (!mountedRef.current) {
+                                    // The original composer is gone. Persist both values and notify a
+                                    // same-session remount so the result is not lost or delayed until a
+                                    // later navigation cycle.
+                                    saveQueuedEditRecovery(sessionId, {
+                                        text: editText,
+                                        pendingSchedule: restoredPendingSchedule,
+                                        composerTextAtEdit,
+                                        pendingScheduleAtEdit,
+                                    })
+                                    return
+                                }
+                                if (editText) {
+                                    assistantApi.composer().setText(editText)
+                                }
+                                onEdit?.({ text: editText, pendingSchedule: restoredPendingSchedule })
+                            } catch {
+                                // useCancelQueuedMessage restores the optimistic row and gives haptic feedback.
+                            } finally {
+                                endQueuedOperation(sessionId, token)
+                            }
                         }
 
                         const canEdit = canCancel

web/src/components/SessionChat.tsx

@@ -1232,6 +1232,11 @@ function SessionChatInner(props: SessionChatProps) {
     // The ref is read at send time; resolvePendingSchedule converts it to an
     // absolute epoch-ms using Date.now() at that moment (send-time base for presets).
     const [pendingSchedule, setPendingSchedule] = useState<PendingSchedule | null>(null)
+    const [pendingScheduleRevision, setPendingScheduleRevision] = useState(0)
+    const updatePendingSchedule = useCallback((next: PendingSchedule | null) => {
+        setPendingSchedule(next)
+        setPendingScheduleRevision((revision) => revision + 1)
+    }, [])
     const pendingScheduleRef = useRef<PendingSchedule | null>(null)
     // Keep render ref in sync so onNew can snapshot at send time
     pendingScheduleRef.current = pendingSchedule
@@ -1247,12 +1252,12 @@ function SessionChatInner(props: SessionChatProps) {
         const ms = (pendingSchedule as Extract<PendingSchedule, { type: 'absolute' }>).ms
         const remaining = ms - Date.now()
         if (remaining <= 0) {
-            setPendingSchedule(null)
+            updatePendingSchedule(null)
             return
         }
-        const timer = setTimeout(() => setPendingSchedule(null), remaining)
+        const timer = setTimeout(() => updatePendingSchedule(null), remaining)
         return () => clearTimeout(timer)
-    }, [pendingSchedule])
+    }, [pendingSchedule, updatePendingSchedule])
 
     const handleSend = useCallback(async (text: string, attachments?: AttachmentMetadata[], scheduledAt?: number | null) => {
         // Route through the scratchlist-aware wrapper. When scratchlistMode
@@ -1277,10 +1282,10 @@ function SessionChatInner(props: SessionChatProps) {
             // its own send path). Schedule clear / forced scroll only
             // matter for chat sends; scratchlist adds don't have a
             // schedule and shouldn't move the chat viewport.
-            setPendingSchedule(null)
+            updatePendingSchedule(null)
             setForceScrollToken((token) => token + 1)
         }
-    }, [onSendForComposer, scratchlistMode])
+    }, [onSendForComposer, scratchlistMode, updatePendingSchedule])
 
     const attachmentAdapter = useMemo(() => {
         if (!props.session.active) {
@@ -1424,9 +1429,11 @@ function SessionChatInner(props: SessionChatProps) {
                             <QueuedMessagesBar
                                 sessionId={props.session.id}
                                 api={props.api}
+                                pendingSchedule={pendingSchedule}
+                                pendingScheduleRevision={pendingScheduleRevision}
                                 onEdit={({ pendingSchedule: restored }) => {
                                     // Restore the schedule so the clock button re-activates
-                                    setPendingSchedule(restored)
+                                    updatePendingSchedule(restored)
                                 }}
                             />
                         </div>
@@ -1437,8 +1444,8 @@ function SessionChatInner(props: SessionChatProps) {
                         resolveSessionMentionTooltip={resolveSessionMentionTooltip}
                         disabled={props.isSending}
                         pendingSchedule={pendingSchedule}
-                        onSchedule={setPendingSchedule}
-                        onClearSchedule={() => setPendingSchedule(null)}
+                        onSchedule={updatePendingSchedule}
+                        onClearSchedule={() => updatePendingSchedule(null)}
                         permissionMode={props.session.permissionMode}
                         collaborationMode={codexCollaborationModeSupported ? props.session.collaborationMode : undefined}
                         model={props.session.model}